Coverage for model_workflow/utils/selections.py: 87%

1from typing import List, Optional 

2 

3from model_workflow.utils.auxiliar import ranger 

4 

5# A selection is a list of atom indices from a structure 

6class Selection: 

7 

8 def __init__ (self, atom_indices : Optional[ List[int] ] = None): 

9 self.atom_indices = atom_indices if atom_indices != None else [] 

10 

11 def __repr__ (self): 

12 return f'<Selection ({len(self.atom_indices)} atoms)>' 

13 

14 def __hash__ (self): 

15 return hash(tuple(self.atom_indices)) 

16 

17 def __len__ (self): 

18 return len(self.atom_indices) 

19 

20 # Return true if there is at least one atom index on the selection and false otherwise 

21 def __bool__ (self): 

22 return len(self.atom_indices) > 0 

23 

24 # Two selections are equal if they have the same atom indices 

25 def __eq__ (self, other): 

26 if not isinstance(other, self.__class__): 

27 return False 

28 return set(self.atom_indices) == set(other.atom_indices) 

29 

30 # Return a new selection with atom indices from both self and the other selection 

31 def __add__ (self, other): 

32 return self.merge(other) 

33 

34 # Return a new selection with self atom indices except for those atom indices in other 

35 def __sub__ (self, other): 

36 return self.substract(other) 

37 

38 # Return a new selection with the intersection of both selections 

39 def __and__ (self, other): 

40 return self.intersection(other) 

41 

42 # Return a new selection with atom indices from both self and the other selection (same as add) 

43 def __or__ (self, other): 

44 return self.merge(other) 

45 

46 # Return a new selection made of self and other selection atom indices 

47 def merge (self, other : Optional['Selection']) -> 'Selection': 

48 if not other: 

49 return self 

50 unique_atom_indices = list(set( self.atom_indices + other.atom_indices )) 

51 return Selection(unique_atom_indices) 

52 

53 # Return a new selection made of self and not other selection atom indices  

54 def substract (self, other : Optional['Selection']) -> 'Selection': 

55 if not other: 

56 return self 

57 remaining_atom_indices = [ atom_index for atom_index in self.atom_indices if atom_index not in other.atom_indices ] 

58 return Selection(remaining_atom_indices) 

59 

60 # Return a new selection with the intersection of both selections 

61 def intersection (self, other : Optional['Selection']) -> 'Selection': 

62 if not other: 

63 return Selection() 

64 self_atom_indices = set(self.atom_indices) 

65 other_atom_indices = set(other.atom_indices) 

66 intersection_atom_indices = list(self_atom_indices.intersection(other_atom_indices)) 

67 return Selection(intersection_atom_indices) 

68 

69 def to_mdanalysis (self) -> str: 

70 # Make sure it is not an empty selection 

71 if not self: raise ValueError('Trying to get MDAnalysis selection from an empty selection') 

72 return 'index ' + ' '.join([ str(index) for index in self.atom_indices ]) 

73 

74 def to_pytraj (self) -> str: 

75 # Make sure it is not an empty selection 

76 if not self: raise ValueError('Trying to get PyTraj selection from an empty selection') 

77 # NEVER FORGET: Pytraj counts atoms starting at 1, not at 0 

78 indices = [ index + 1 for index in self.atom_indices ] 

79 # Make ranges for atoms in a row 

80 return '@' + ranger(indices) 

81 

82 def to_ngl (self) -> str: 

83 return '@' + ','.join([ str(index) for index in self.atom_indices ]) 

84 

85 # Get a string made of all indexes separated by underscores 

86 # This string can be then passed as a bash argument and easily parsed by other programms 

87 # Indices can start from 0 or from 1 

88 def to_bash (self, one_start : bool = False) -> str: 

89 # Make sure it is not an empty selection 

90 if not self: raise ValueError('Trying to get Bash selection from an empty selection') 

91 if one_start: 

92 return '_'.join([ str(index + 1) for index in self.atom_indices ]) 

93 else: 

94 return '_'.join([ str(index) for index in self.atom_indices ]) 

95 

96 # Produce a vmd selection in tcl format 

97 def to_vmd (self) -> str: 

98 # Make sure it is not an empty selection 

99 if not self: raise ValueError('Trying to get VMD selection from an empty selection') 

100 return 'index ' + ' '.join([ str(index) for index in self.atom_indices ]) 

101 

102 # Produce the content of gromacs ndx file 

103 def to_ndx (self, selection_name : str = 'Selection') -> str: 

104 # Make sure it is not an empty selection 

105 if not self: raise ValueError('Trying to get NDX selection from an empty selection') 

106 # Add a header 

107 # WARNING: Sometimes (and only sometimes) if the selection name includes white spaces there is an error 

108 # WARNING: Gromacs may get the first word only as the name of the selection, so we must purge white spaces 

109 # WARNING: However, if a call to this function passes a selection name it will expect it to be in the index file 

110 # WARNING: For this reason we must kill here the proccess and warn the user 

111 if ' ' in selection_name: 

112 raise ValueError(f'A Gromacs index file selection name must never include white spaces: {selection_name}') 

113 content = '[ ' + selection_name + ' ]\n' 

114 count = 0 

115 for index in self.atom_indices: 

116 # Add a breakline each 15 indices 

117 count += 1 

118 if count == 15: 

119 content += '\n' 

120 count = 0 

121 # Add a space between indices 

122 # Atom indices go from 0 to n-1 

123 # Add +1 to the index since gromacs counts from 1 to n 

124 content += str(index + 1) + ' ' 

125 content += '\n' 

126 return content 

127 

128 # Create a gromacs ndx file 

129 def to_ndx_file (self, selection_name : str = 'Selection', output_filepath : str = 'index.ndx'): 

130 index_content = self.to_ndx(selection_name) 

131 with open(output_filepath, 'w') as file: 

132 file.write(index_content)