From 795a7b6187d67e2b3bc34023f83e5be3200c3156 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Fri, 14 Feb 2025 16:55:02 +0000 Subject: [PATCH] Save selection as array of bool values --- .../Framework/AtomSelector/atom_selection.py | 8 ++++++-- .../Framework/AtomSelector/general_selection.py | 8 ++++++-- .../Framework/AtomSelector/group_selection.py | 8 ++++++-- .../Framework/AtomSelector/molecule_selection.py | 8 ++++++-- .../MDANSE/Framework/AtomSelector/selector.py | 10 ++++++++-- .../Configurators/AtomSelectionConfigurator.py | 9 ++++++--- .../AtomTransmutationConfigurator.py | 4 +--- .../Configurators/PartialChargeConfigurator.py | 4 +--- MDANSE/Src/MDANSE/Framework/Jobs/IJob.py | 16 ++++++++++++++++ .../Src/MDANSE_GUI/Widgets/SelectionWidgets.py | 8 +++----- 10 files changed, 59 insertions(+), 24 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py index 60b7ab462..6351dd61b 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py @@ -20,7 +20,9 @@ from MDANSE.MolecularDynamics.Trajectory import Trajectory -def select_atoms(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Set[int]: +def select_atoms( + trajectory: Trajectory, **function_parameters: Dict[str, Any] +) -> Set[int]: """Selects all the atoms in the trajectory. Parameters @@ -46,7 +48,9 @@ def select_atoms(trajectory: Trajectory, **function_parameters: Dict[str, Any]) if index_list is not None: selection = selection.union(indices.intersection(index_list)) if index_range is not None: - selection = selection.union(indices.intersection(range(index_range[0], index_range[1]))) + selection = selection.union( + indices.intersection(range(index_range[0], index_range[1])) + ) if index_slice is not None: selection = selection.union( indices.intersection(range(index_slice[0], index_slice[1], index_slice[2])) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py index b947a01c3..8c404e894 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py @@ -19,7 +19,9 @@ from MDANSE.MolecularDynamics.Trajectory import Trajectory -def select_all(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Set[int]: +def select_all( + trajectory: Trajectory, **function_parameters: Dict[str, Any] +) -> Set[int]: """Selects all the atoms in the trajectory. Parameters @@ -37,7 +39,9 @@ def select_all(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> return set(range(len(trajectory.chemical_system.atom_list))) -def select_none(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Set[int]: +def select_none( + trajectory: Trajectory, **function_parameters: Dict[str, Any] +) -> Set[int]: """Returns an empty selection. Parameters diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py index ad94a9994..49615d553 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py @@ -21,7 +21,9 @@ from MDANSE.MolecularDynamics.Trajectory import Trajectory -def select_labels(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Set[int]: +def select_labels( + trajectory: Trajectory, **function_parameters: Dict[str, Any] +) -> Set[int]: """Selects all the atoms in the trajectory. Parameters @@ -45,7 +47,9 @@ def select_labels(trajectory: Trajectory, **function_parameters: Dict[str, Any]) return selection -def select_pattern(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Set[int]: +def select_pattern( + trajectory: Trajectory, **function_parameters: Dict[str, Any] +) -> Set[int]: """Selects all the atoms in the trajectory. Parameters diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py index b77e47651..0513d25c6 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py @@ -21,7 +21,9 @@ from MDANSE.MolecularDynamics.Trajectory import Trajectory -def select_molecules(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Set[int]: +def select_molecules( + trajectory: Trajectory, **function_parameters: Dict[str, Any] +) -> Set[int]: """Selects all the atoms in the trajectory. Parameters @@ -41,5 +43,7 @@ def select_molecules(trajectory: Trajectory, **function_parameters: Dict[str, An molecule_names = function_parameters.get("molecule_names", None) for molecule in molecule_names: if molecule in system._clusters: - selection = selection.union(reduce(list.__add__, system._clusters[molecule])) + selection = selection.union( + reduce(list.__add__, system._clusters[molecule]) + ) return selection diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index 57be31fd2..d722c8b52 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -113,9 +113,15 @@ def validate_selection_string( selection = selection.difference(temp_selection) else: selection = temp_selection - if len(selection.difference(current_selection)) > 0 and operation_type == "union": + if ( + len(selection.difference(current_selection)) > 0 + and operation_type == "union" + ): return True - elif len(current_selection.difference(selection)) > 0 and operation_type != "union": + elif ( + len(current_selection.difference(selection)) > 0 + and operation_type != "union" + ): return True return False diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py index 9c2c5aa90..d6da23792 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py @@ -33,7 +33,7 @@ class AtomSelectionConfigurator(IConfigurator): The defaults selection setting. """ - _default = '{}' + _default = "{}" def configure(self, value: str) -> None: """Configure an input value. @@ -43,6 +43,8 @@ def configure(self, value: str) -> None: value : str The selection setting in a json readable format. """ + self._original_input = value + trajConfig = self._configurable[self._dependencies["trajectory"]] self.selector = ReusableSelection() @@ -67,6 +69,7 @@ def configure(self, value: str) -> None: self["flatten_indices"] = sorted(list(indices)) atoms = trajConfig["instance"].chemical_system.atom_list + self["total_number_of_atoms"] = len(atoms) selectedAtoms = [atoms[idx] for idx in self["flatten_indices"]] self["selection_length"] = len(self["flatten_indices"]) @@ -93,7 +96,7 @@ def get_natoms(self) -> dict[str, int]: A dictionary of the number of atom per element. """ names, counts = np.unique(self["names"], return_counts=True) - nAtomsPerElement = {names[n]:counts[n] for n in range(len(names))} + nAtomsPerElement = {names[n]: counts[n] for n in range(len(names))} return nAtomsPerElement @@ -132,7 +135,7 @@ def get_information(self) -> str: return "\n".join(info) + "\n" - def get_selector(self) -> 'ReusableSelection': + def get_selector(self) -> "ReusableSelection": """ Returns ------- diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py index 986258192..4bbc17ee4 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py @@ -41,9 +41,7 @@ def __init__(self, trajectory: Trajectory) -> None: self._new_map = {} self._current_trajectory = trajectory - def apply_transmutation( - self, selection_string: str, symbol: str - ) -> None: + def apply_transmutation(self, selection_string: str, symbol: str) -> None: """With the selection dictionary update selector and then update the transmutation map. diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py index 6518c1c87..3f66c7f89 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py @@ -44,9 +44,7 @@ def __init__(self, trajectory: Trajectory) -> None: self._original_map[at_num] = 0.0 self._new_map = {} - def update_charges( - self, selection_string: str, charge: float - ) -> None: + def update_charges(self, selection_string: str, charge: float) -> None: """With the selection dictionary update the selector and then update the partial charge map. diff --git a/MDANSE/Src/MDANSE/Framework/Jobs/IJob.py b/MDANSE/Src/MDANSE/Framework/Jobs/IJob.py index e0fc7ba37..474a812ad 100644 --- a/MDANSE/Src/MDANSE/Framework/Jobs/IJob.py +++ b/MDANSE/Src/MDANSE/Framework/Jobs/IJob.py @@ -169,6 +169,22 @@ def initialize(self): ) except KeyError: LOG.error("IJob did not find 'write_logs' in output_files") + if "atom_selection" in self.configuration: + try: + array_length = self.configuration["atom_selection"][ + "total_number_of_atoms" + ] + except KeyError: + LOG.warning( + "Job could not find total number of atoms in atom selection." + ) + else: + valid_indices = self.configuration["atom_selection"]["flatten_indices"] + self._outputData.add( + "selected_atoms", + "LineOutputVariable", + [index in valid_indices for index in range(array_length)], + ) @abc.abstractmethod def run_step(self, index): diff --git a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py index 91bd9446d..317394272 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py @@ -95,10 +95,8 @@ def __init__( self.atom_names = [] if trajectory: self.atom_types = list(np.unique(trajectory.chemical_system.atom_list)) - if trajectory.chemical_system._atom_names: - self.atom_names = list( - np.unique(trajectory.chemical_system._atom_names) - ) + if trajectory.chemical_system.name_list: + self.atom_names = list(np.unique(trajectory.chemical_system.name_list)) self.selection_types = [] self.selection_keyword = "" if self.atom_types: @@ -127,7 +125,7 @@ def switch_mode(self, new_mode: str): self.selection_field.addItems(self.atom_types) self.selection_keyword = "atom_types" elif new_mode == "name": - self.selection_field.addItems(self.atom_types) + self.selection_field.addItems(self.atom_names) self.selection_keyword = "atom_names" def parameter_dictionary(self):