Coverage for model_workflow/tools/get_atoms_count.py: 100%

1# Count different type of atoms and residues in the structure 

2def get_atoms_count (structure : 'Structure') -> tuple: 

3 

4 # Number of system atoms and residues 

5 system_atoms = len(structure.atoms) 

6 system_residues = len(structure.residues) 

7 # Number of protein atoms and residues 

8 protein_selection = structure.select_protein() 

9 protein_atoms = len(protein_selection) 

10 protein_residues = len(structure.get_selection_residue_indices(protein_selection)) 

11 # Number of nucleic atoms and residues 

12 nucleic_selection = structure.select_nucleic() 

13 nucleic_atoms = len(nucleic_selection) 

14 nucleic_residues = len(structure.get_selection_residue_indices(nucleic_selection)) 

15 # Number of lipid atoms and residues 

16 lipids_selection = structure.select_lipids() 

17 lipid_atoms = len(lipids_selection) 

18 lipid_residues = len(structure.get_selection_residue_indices(lipids_selection)) 

19 # Number of carbohydrates atoms and residues 

20 carbohydrates_selection = structure.select_carbohydrates() 

21 carbohydrates_atoms = len(carbohydrates_selection) 

22 carbohydrates_residues = len(structure.get_selection_residue_indices(carbohydrates_selection)) 

23 # Number of solvent atoms and residues 

24 solvent_selection = structure.select_water() 

25 solvent_atoms = len(solvent_selection) 

26 solvent_residues = len(structure.get_selection_residue_indices(solvent_selection)) 

27 # Number of counter ions 

28 counter_cations = len(structure.select_counter_ions(charge='+')) 

29 counter_anions = len(structure.select_counter_ions(charge='-')) 

30 counter_ions = counter_cations + counter_anions 

31 

32 # Display a summary of atom and residue counts 

33 print('Atom and residue counts:') 

34 print(f' System atoms: {system_atoms}') 

35 print(f' System residues: {system_residues}') 

36 print(f' Protein atoms: {protein_atoms}') 

37 print(f' Protein residues: {protein_residues}') 

38 print(f' Nucleic atoms: {nucleic_atoms}') 

39 print(f' Nucleic residues: {nucleic_residues}') 

40 print(f' Lipid atoms: {lipid_atoms}') 

41 print(f' Lipid residues: {lipid_residues}') 

42 print(f' Carbohydrate atoms: {carbohydrates_atoms}') 

43 print(f' Carbohydrate residues: {carbohydrates_residues}') 

44 print(f' Solvent atoms: {solvent_atoms}') 

45 print(f' Solvent residues: {solvent_residues}') 

46 print(f' Counter cations: {counter_cations}') 

47 print(f' Counter anions: {counter_anions}') 

48 print(f' Counter ions: {counter_ions}') 

49 

50 # Return the counts 

51 return (system_atoms, system_residues, protein_atoms, protein_residues, 

52 nucleic_atoms, nucleic_residues, lipid_atoms, lipid_residues, 

53 carbohydrates_atoms, carbohydrates_residues, solvent_atoms, solvent_residues, 

54 counter_cations, counter_anions, counter_ions)