Coverage for src / molecular_simulations / utils / amber_utils.py: 100%

1import MDAnalysis as mda 

2import string 

3 

4def assign_chainids(u: mda.Universe, 

5 terminus_selection: str='name OXT') -> mda.Universe: 

6 if not hasattr(u.atoms, 'chainIDs'): 

7 u.add_TopologyAttr('chainIDs') 

8 

9 termini_atoms = u.select_atoms(terminus_selection) 

10 termini_resindices = set(termini_atoms.resindices) 

11 

12 def get_chain_label(index): 

13 if index < 26: 

14 return string.ascii_uppercase[index] 

15 first = string.ascii_uppercase[(index // 26) - 1] 

16 second = string.ascii_uppercase[index % 26] 

17 return first + second 

18 

19 chain_index = 0 

20 for residue in u.residues: 

21 residue.atoms.chainIDs = get_chain_label(chain_index) 

22 

23 if residue.resindex in termini_resindices: 

24 chain_index += 1 

25 

26 return u