Coverage for mddb_workflow/utils/selections.py: 86%
79 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-10-29 15:48 +0000
1from typing import List, Optional
3from mddb_workflow.utils.auxiliar import ranger
5# A selection is a list of atom indices from a structure
6class Selection:
8 def __init__ (self, atom_indices : Optional[ list[int] ] = None):
9 self.atom_indices = atom_indices if atom_indices != None else []
11 def __repr__ (self):
12 return f'<Selection ({len(self.atom_indices)} atoms)>'
14 def __hash__ (self):
15 return hash(tuple(self.atom_indices))
17 def __len__ (self):
18 return len(self.atom_indices)
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
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)
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)
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)
38 # Return a new selection with the intersection of both selections
39 def __and__ (self, other):
40 return self.intersection(other)
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)
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)
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)
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)
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 ])
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)
82 def to_ngl (self) -> str:
83 return '@' + ','.join([ str(index) for index in self.atom_indices ])
85 def to_list (self) -> list[int]:
86 """Return a copy of the atom indices as a list."""
87 return self.atom_indices.copy()
89 # Get a string made of all indexes separated by underscores
90 # This string can be then passed as a bash argument and easily parsed by other programms
91 # Indices can start from 0 or from 1
92 def to_bash (self, one_start : bool = False) -> str:
93 # Make sure it is not an empty selection
94 if not self: raise ValueError('Trying to get Bash selection from an empty selection')
95 if one_start:
96 return '_'.join([ str(index + 1) for index in self.atom_indices ])
97 else:
98 return '_'.join([ str(index) for index in self.atom_indices ])
100 # Produce a vmd selection in tcl format
101 def to_vmd (self) -> str:
102 # Make sure it is not an empty selection
103 if not self: raise ValueError('Trying to get VMD selection from an empty selection')
104 return 'index ' + ' '.join([ str(index) for index in self.atom_indices ])
106 # Produce the content of gromacs ndx file
107 def to_ndx (self, selection_name : str = 'Selection') -> str:
108 # Make sure it is not an empty selection
109 if not self: raise ValueError('Trying to get NDX selection from an empty selection')
110 # Add a header
111 # WARNING: Sometimes (and only sometimes) if the selection name includes white spaces there is an error
112 # WARNING: Gromacs may get the first word only as the name of the selection, so we must purge white spaces
113 # WARNING: However, if a call to this function passes a selection name it will expect it to be in the index file
114 # WARNING: For this reason we must kill here the proccess and warn the user
115 if ' ' in selection_name:
116 raise ValueError(f'A Gromacs index file selection name must never include white spaces: {selection_name}')
117 content = '[ ' + selection_name + ' ]\n'
118 count = 0
119 for index in self.atom_indices:
120 # Add a breakline each 15 indices
121 count += 1
122 if count == 15:
123 content += '\n'
124 count = 0
125 # Add a space between indices
126 # Atom indices go from 0 to n-1
127 # Add +1 to the index since gromacs counts from 1 to n
128 content += str(index + 1) + ' '
129 content += '\n'
130 return content
132 # Create a gromacs ndx file
133 def to_ndx_file (self, selection_name : str = 'Selection', output_filepath : str = 'index.ndx'):
134 index_content = self.to_ndx(selection_name)
135 with open(output_filepath, 'w') as file:
136 file.write(index_content)