From a32a0f9ac8d70c2ae1c79405031096a18023500a Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Tue, 11 Feb 2025 14:17:51 +0000 Subject: [PATCH 01/37] Add new selection functions --- .../Framework/AtomSelector/atom_selection.py | 66 ++++ .../AtomSelector/general_selection.py | 59 +++ .../Framework/AtomSelector/group_selection.py | 72 ++++ .../AtomSelector/molecule_selection.py | 49 +++ .../MDANSE/Framework/AtomSelector/selector.py | 371 +++--------------- 5 files changed, 303 insertions(+), 314 deletions(-) create mode 100644 MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py create mode 100644 MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py create mode 100644 MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py create mode 100644 MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py new file mode 100644 index 000000000..97a8ac2ca --- /dev/null +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py @@ -0,0 +1,66 @@ +# This file is part of MDANSE. +# +# MDANSE is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . +# + +from typing import Union, Dict, Any, Set + +from MDANSE.Chemistry.ChemicalSystem import ChemicalSystem +from MDANSE.MolecularDynamics.Trajectory import Trajectory + + +def select_atoms( + trajectory: Trajectory, **function_parameters: Dict[str, Any] +) -> Set[int]: + """Selects all the atoms in the trajectory. + + Parameters + ---------- + selection : Set[int] + A set of atom indices + trajectory : Trajectory + A trajectory instance to which the selection is applied + + Returns + ------- + Set[int] + Set of all the atom indices + """ + selection = set() + system = trajectory.chemical_system + element_list = system.atom_list + name_list = system.name_list + indices = set(range(len(element_list))) + index_list = function_parameters.get("index_list", None) + index_range = function_parameters.get("index_range", None) + index_slice = function_parameters.get("index_range", None) + 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])) + ) + if index_slice is not None: + selection = selection.union( + indices.intersection(range(index_slice[0], index_slice[1], index_slice[2])) + ) + atom_types = function_parameters.get("atom_types", None) + if atom_types: + new_indices = [index for index in indices if element_list[index] in atom_types] + selection = selection.union(new_indices) + atom_names = function_parameters.get("atom_names", None) + if atom_names: + new_indices = [index for index in indices if name_list[index] in atom_names] + selection = selection.union(new_indices) + return selection diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py new file mode 100644 index 000000000..a58cd06fe --- /dev/null +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py @@ -0,0 +1,59 @@ +# This file is part of MDANSE. +# +# MDANSE is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . +# + +from typing import Union, Dict, Any, Set +from MDANSE.Chemistry.ChemicalSystem import ChemicalSystem +from MDANSE.MolecularDynamics.Trajectory import Trajectory + + +def select_all( + trajectory: Trajectory, **function_parameters: Dict[str, Any] +) -> Set[int]: + """Selects all the atoms in the trajectory. + + Parameters + ---------- + selection : Set[int] + A set of atom indices + trajectory : Trajectory + A trajectory instance to which the selection is applied + + Returns + ------- + Set[int] + Set of all the atom indices + """ + return set(range(len(trajectory.chemical_system.atom_list))) + + +def select_none( + trajectory: Trajectory, **function_parameters: Dict[str, Any] +) -> Set[int]: + """Returns an empty selection. + + Parameters + ---------- + selection : Set[int] + A set of atom indices + trajectory : Trajectory + A trajectory instance to which the selection is applied + + Returns + ------- + Set[int] + An empty set. + """ + return set() diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py new file mode 100644 index 000000000..04ab009c7 --- /dev/null +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py @@ -0,0 +1,72 @@ +# This file is part of MDANSE. +# +# MDANSE is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . +# + +from typing import Union, Dict, Any, Set +from functools import reduce + +from MDANSE.Chemistry.ChemicalSystem import ChemicalSystem +from MDANSE.MolecularDynamics.Trajectory import Trajectory + + +def select_labels( + trajectory: Trajectory, **function_parameters: Dict[str, Any] +) -> Set[int]: + """Selects all the atoms in the trajectory. + + Parameters + ---------- + selection : Set[int] + A set of atom indices + trajectory : Trajectory + A trajectory instance to which the selection is applied + + Returns + ------- + Set[int] + Set of all the atom indices + """ + selection = set() + system = trajectory.chemical_system + atom_labels = function_parameters.get("atom_labels", None) + for label in atom_labels: + if label in system._labels: + selection = selection.union(reduce(list.__add__, system._labels[label])) + return selection + + +def select_pattern( + trajectory: Trajectory, **function_parameters: Dict[str, Any] +) -> Set[int]: + """Selects all the atoms in the trajectory. + + Parameters + ---------- + selection : Set[int] + A set of atom indices + trajectory : Trajectory + A trajectory instance to which the selection is applied + + Returns + ------- + Set[int] + Set of all the atom indices + """ + selection = set() + system = trajectory.chemical_system + pattern = function_parameters.get("rdkit_pattern", None) + if pattern: + selection = system.get_substructure_matches(pattern) + return selection diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py new file mode 100644 index 000000000..0513d25c6 --- /dev/null +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py @@ -0,0 +1,49 @@ +# This file is part of MDANSE. +# +# MDANSE is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . +# + +from typing import Union, Dict, Any, Set +from functools import reduce + +from MDANSE.Chemistry.ChemicalSystem import ChemicalSystem +from MDANSE.MolecularDynamics.Trajectory import Trajectory + + +def select_molecules( + trajectory: Trajectory, **function_parameters: Dict[str, Any] +) -> Set[int]: + """Selects all the atoms in the trajectory. + + Parameters + ---------- + selection : Set[int] + A set of atom indices + trajectory : Trajectory + A trajectory instance to which the selection is applied + + Returns + ------- + Set[int] + Set of all the atom indices + """ + selection = set() + system = trajectory.chemical_system + 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]) + ) + return selection diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index f98d9aafd..6bd689490 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -15,347 +15,90 @@ # import json import copy -from typing import Union +from typing import Union, Dict, Any, Set from MDANSE.Chemistry.ChemicalSystem import ChemicalSystem from MDANSE.MolecularDynamics.Trajectory import Trajectory from MDANSE.Framework.AtomSelector.all_selector import select_all -from MDANSE.Framework.AtomSelector.atom_selectors import * -from MDANSE.Framework.AtomSelector.group_selectors import * -from MDANSE.Framework.AtomSelector.molecule_selectors import * +from MDANSE.Framework.AtomSelector.general_selection import select_all, select_none -class Selector: - """Used to get the indices of a subset of atoms of a chemical system. +function_lookup = { + function.__name__: function for function in [select_all, select_none] +} - Attributes - ---------- - _default : dict[str, bool | dict] - The default settings. - _funcs : dict[str, Callable] - A dictionary of the functions. - _kwarg_keys : dict[str, str] - A dictionary of the function arg keys. - """ - - _default = { - "all": True, - "dummy": False, - "hs_on_heteroatom": False, - "primary_amine": False, - "hydroxy": False, - "methyl": False, - "phosphate": False, - "sulphate": False, - "thiol": False, - "water": False, - # e.g. {"S": True} - "hs_on_element": {}, - "element": {}, - "name": {}, - "fullname": {}, - # e.g. {1: True} - "index": {}, - } - - _funcs = { - "all": select_all, - "dummy": select_dummy, - "hs_on_heteroatom": select_hs_on_heteroatom, - "primary_amine": select_primary_amine, - "hydroxy": select_hydroxy, - "methyl": select_methyl, - "phosphate": select_phosphate, - "sulphate": select_sulphate, - "thiol": select_thiol, - "water": select_water, - "hs_on_element": select_hs_on_element, - "element": select_element, - "name": select_atom_name, - "fullname": select_atom_fullname, - "index": select_index, - } - _kwarg_keys = { - "hs_on_element": "symbol", - "element": "symbol", - "name": "name", - "fullname": "fullname", - "index": "index", - } +class ReusableSelection: + """A reusable sequence of operations which, when applied + to a trajectory, returns a set of atom indices based + on the specified criteria. + """ - def __init__(self, trajectory: Trajectory) -> None: + def __init__(self) -> None: """ Parameters ---------- trajectory: Trajectory The chemical system to apply the selection to. """ - system = trajectory.chemical_system - self.system = system - self.trajectory = trajectory - self.all_idxs = set(system._atom_indices) - self.settings = copy.deepcopy(self._default) - - symbols = set(system.atom_list) - # all possible values for the system - self._kwarg_vals = { - "element": symbols, - "hs_on_element": set( - [ - symbol - for symbol in symbols - if select_hs_on_element(trajectory, symbol, check_exists=True) - ] - ), - "name": set(system.atom_list), - "fullname": set(system.name_list), - "index": self.all_idxs, - } - - # figure out if a match exists for the selector function - self.match_exists = self.create_default_settings() - for k0, v0 in self.match_exists.items(): - if isinstance(v0, dict): - for k1 in v0.keys(): - self.match_exists[k0][k1] = True - else: - self.match_exists[k0] = self._funcs[k0]( - self.trajectory, check_exists=True - ) - - self.settings = self.create_default_settings() - - def create_default_settings(self) -> dict[str, Union[bool, dict]]: - """Create a new settings dictionary with default settings. - - Returns - ------- - dict[str, Union[bool, dict]] - A settings dictionary. - """ - settings = copy.deepcopy(self._default) - for k, vs in self._kwarg_vals.items(): - for v in sorted(vs): - settings[k][v] = False - return settings - - def reset_settings(self) -> None: - """Resets the settings back to the defaults.""" - self.settings = self.create_default_settings() - - def update_settings( - self, settings: dict[str, Union[bool, dict]], reset_first: bool = False - ) -> None: - """Updates the selection settings. - - Parameters - ---------- - settings : dict[str, bool | dict] - The selection settings. - reset_first : bool, optional - Resets the settings to the default before loading. - - Raises - ------ - ValueError - Raises a ValueError if the inputted settings are not valid. - """ - if not self.check_valid_setting(settings): - raise ValueError( - f"Settings are not valid for the given chemical system - {settings}." - ) - - if reset_first: - self.reset_settings() - - for k0, v0 in settings.items(): - if isinstance(self.settings[k0], dict): - for k1, v1 in v0.items(): - self.settings[k0][k1] = v1 - else: - self.settings[k0] = v0 - - def get_idxs(self) -> set[int]: - """The atom indices after applying the selection to the system. - - Returns - ------- - set[int] - The atoms indices. - """ - idxs = set([]) - - for k, v in self.settings.items(): - - if isinstance(v, dict): - args = [{self._kwarg_keys[k]: i} for i in v.keys()] - switches = v.values() + self.reset() + + def reset(self): + self.system = None + self.trajectory = None + self.all_idxs = set() + self.operations = {} + + def set_selection( + self, number: Union[int, None] = None, function_parameters: Dict[str, Any] = {} + ): + if number is None: + number = len(self.operations) + self.operations[number] = function_parameters + + def select_in_trajectory(self, trajectory: Trajectory) -> Set[int]: + selection = set() + self.all_idxs = set(range(len(trajectory.chemical_system.atom_list))) + sequence = sorted([int(x) for x in self.operations.keys()]) + if len(sequence) == 0: + return self.all_idxs + for number in sequence: + function_parameters = self.operations[number] + function_name = function_parameters.pop("function_name", "select_all") + if function_name == "invert_selection": + selection = self.all_idxs.difference(selection) else: - args = [{}] - switches = [v] - - for arg, switch in zip(args, switches): - if not switch: - continue - - idxs.update(self._funcs[k](self.trajectory, **arg)) - - return idxs - - def update_with_idxs(self, idxs: set[int]) -> None: - """Using the inputted idxs change the selection setting so - that it would return the same idxs with get_idxs. It will - switch off the setting if idxs is not a superset of the - selection for that setting. - - Parameters - ---------- - idxs : set[int] - With the indices of the atom selection. - """ - new_settings = self.create_default_settings() - new_settings["all"] = False - - added = set([]) - for k, v in self.settings.items(): - - if k == "index": - continue - - if isinstance(v, dict): - args = [{self._kwarg_keys[k]: i} for i in v.keys()] - switches = v.values() - else: - args = [{}] - switches = [v] - - for arg, switch in zip(args, switches): - if not switch: - continue - - selection = self._funcs[k](self.trajectory, **arg) - if not idxs.issuperset(selection): - continue - - added.update(selection) - if isinstance(v, dict): - new_settings[k][arg[self._kwarg_keys[k]]] = True + operation_type = function_parameters.pop("operation_type", "union") + function = function_lookup[function_name] + temp_selection = function(trajectory, **function_parameters) + if operation_type == "union": + selection = selection.union(temp_selection) + elif operation_type == "intersection": + selection = selection.intersection(temp_selection) else: - new_settings[k] = True - - for idx in idxs - added: - new_settings["index"][idx] = True - - self.settings = new_settings + selection = temp_selection + return selection - def settings_to_json(self) -> str: - """Return the minimal json string required to achieve the same - settings with the settings_from_json method. + def convert_to_json(self) -> str: + """For the purpose of storing the selection independent of the + trajectory it is acting on, this method encodes the sequence + of selection operations as a string. Returns ------- str - A JSON string. + All the operations of this selection, encoded as string """ - minimal_dict = {} - for k0, v0 in self.settings.items(): - if isinstance(v0, bool) and (k0 == "all" or k0 != "all" and v0): - minimal_dict[k0] = v0 - elif isinstance(v0, dict): - sub_list = [] - for k1, v1 in v0.items(): - if v1: - sub_list.append(k1) - if sub_list: - minimal_dict[k0] = sorted(sub_list) - return json.dumps(minimal_dict) + return json.dumps(self.operations) - def json_to_settings(self, json_string: str) -> dict[str, Union[bool, dict]]: - """Loads the json string and converts to a settings. + def read_from_json(self, json_string: str): + """_summary_ Parameters ---------- json_string : str - The JSON string of settings. - - Returns - ------- - dict[str, Union[bool, dict]] - The selection settings. + A sequence of selection operations, encoded as a JSON string """ json_setting = json.loads(json_string) - settings = {} for k0, v0 in json_setting.items(): - if isinstance(v0, bool): - settings[k0] = v0 - elif isinstance(v0, list): - sub_dict = {} - for k1 in v0: - sub_dict[k1] = True - if sub_dict: - settings[k0] = sub_dict - return settings - - def load_from_json(self, json_string: str) -> None: - """Load the selection settings from a JSON string. - - Parameters - ---------- - json_string : str - The JSON string of settings. - """ - self.update_settings(self.json_to_settings(json_string), reset_first=True) - - def check_valid_setting(self, settings: dict[str, Union[bool, dict]]) -> bool: - """Checks that the input settings are valid. - - Parameters - ---------- - settings : dict[str, bool | dict] - The selection settings. - - Returns - ------- - bool - True if settings are valid. - """ - setting_keys = self._default.keys() - dict_setting_keys = self._kwarg_keys.keys() - for k0, v0 in settings.items(): - - if k0 not in setting_keys: - return False - - if k0 not in dict_setting_keys: - if not isinstance(v0, bool): - return False - - if k0 in dict_setting_keys: - if not isinstance(v0, dict): - return False - for k1, v1 in v0.items(): - if k1 not in self._kwarg_vals[k0]: - return False - if not isinstance(v1, bool): - return False - - return True - - def check_valid_json_settings(self, json_string: str) -> bool: - """Checks that the input JSON setting string is valid. - - Parameters - ---------- - json_string : str - The JSON string of settings. - - Returns - ------- - bool - True if settings are valid. - """ - try: - settings = self.json_to_settings(json_string) - except ValueError: - return False - return self.check_valid_setting(settings) + if isinstance(v0, dict): + self.append_operation(k0, v0) From 881d4c627edcfa20c304b0947fd775209a3b3a04 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Wed, 12 Feb 2025 08:47:08 +0000 Subject: [PATCH 02/37] Add first unit tests for selection --- .../MDANSE/Framework/AtomSelector/__init__.py | 1 - .../Framework/AtomSelector/all_selector.py | 12 +--- .../Framework/AtomSelector/atom_selection.py | 10 +-- .../Framework/AtomSelector/atom_selectors.py | 18 ++--- .../AtomSelector/general_selection.py | 8 +-- .../Framework/AtomSelector/group_selection.py | 8 +-- .../Framework/AtomSelector/group_selectors.py | 20 ++---- .../AtomSelector/molecule_selection.py | 8 +-- .../AtomSelector/molecule_selectors.py | 4 +- .../MDANSE/Framework/AtomSelector/selector.py | 4 +- .../AtomSelectionConfigurator.py | 3 +- .../AtomTransmutationConfigurator.py | 1 - .../PartialChargeConfigurator.py | 2 - MDANSE/Tests/UnitTests/test_selection.py | 70 +++++++++++++++++++ .../InputWidgets/AtomSelectionWidget.py | 3 +- 15 files changed, 94 insertions(+), 78 deletions(-) create mode 100644 MDANSE/Tests/UnitTests/test_selection.py diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/__init__.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/__init__.py index 1ce9f3c79..268c9cc94 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/__init__.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/__init__.py @@ -13,4 +13,3 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . # -from .selector import Selector diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/all_selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/all_selector.py index 8b1813e10..cf1a452c8 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/all_selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/all_selector.py @@ -17,9 +17,7 @@ from MDANSE.MolecularDynamics.Trajectory import Trajectory -def select_all( - trajectory: Trajectory, check_exists: bool = False -) -> Union[set[int], bool]: +def select_all(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: """Selects all atoms in the chemical system except for the dummy atoms. @@ -45,10 +43,4 @@ def select_all( for atm in system._unique_elements: if trajectory.get_atom_property(atm, "dummy"): dummy_list.append(atm) - return set( - [ - index - for index in system._atom_indices - if atom_list[index] not in dummy_list - ] - ) + return set([index for index in system._atom_indices if atom_list[index] not in dummy_list]) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py index 97a8ac2ca..60b7ab462 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py @@ -20,9 +20,7 @@ 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 @@ -44,13 +42,11 @@ def select_atoms( indices = set(range(len(element_list))) index_list = function_parameters.get("index_list", None) index_range = function_parameters.get("index_range", None) - index_slice = function_parameters.get("index_range", None) + index_slice = function_parameters.get("index_slice", None) 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/atom_selectors.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selectors.py index 1b7a05481..1b0ba6211 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selectors.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selectors.py @@ -55,9 +55,7 @@ def select_element( return system.get_substructure_matches(pattern) -def select_dummy( - trajectory: Trajectory, check_exists: bool = False -) -> Union[set[int], bool]: +def select_dummy(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: """Selects all dummy atoms in the chemical system. Parameters @@ -86,11 +84,7 @@ def select_dummy( if trajectory.get_atom_property(atm, "dummy"): dummy_list.append(atm) return set( - [ - index - for index, element in enumerate(system.atom_list) - if element in dummy_list - ] + [index for index, element in enumerate(system.atom_list) if element in dummy_list] ) @@ -119,9 +113,7 @@ def select_atom_name( return True return False else: - return set( - [index for index, element in enumerate(system.atom_list) if element == name] - ) + return set([index for index, element in enumerate(system.atom_list) if element == name]) def select_atom_fullname( @@ -149,9 +141,7 @@ def select_atom_fullname( return True return False else: - return set( - [index for index, name in enumerate(system.name_list) if name == fullname] - ) + return set([index for index, name in enumerate(system.name_list) if name == fullname]) def select_hs_on_element( diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py index a58cd06fe..39216c61b 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py @@ -19,9 +19,7 @@ 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 @@ -39,9 +37,7 @@ def select_all( 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 04ab009c7..6b1ecf42f 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py @@ -21,9 +21,7 @@ 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 @@ -47,9 +45,7 @@ def select_labels( 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/group_selectors.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selectors.py index 0bc1e8567..3f6c330ca 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selectors.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selectors.py @@ -52,9 +52,7 @@ def select_primary_amine( return system.get_substructure_matches(pattern) -def select_hydroxy( - trajectory: Trajectory, check_exists: bool = False -) -> Union[set[int], bool]: +def select_hydroxy(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: """Selects the O and H atoms of all hydroxy groups including water. Parameters @@ -77,9 +75,7 @@ def select_hydroxy( return system.get_substructure_matches(pattern) -def select_methyl( - trajectory: Trajectory, check_exists: bool = False -) -> Union[set[int], bool]: +def select_methyl(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: """Selects the C and H atoms of all methyl groups. Parameters @@ -102,9 +98,7 @@ def select_methyl( return system.get_substructure_matches(pattern) -def select_phosphate( - trajectory: Trajectory, check_exists: bool = False -) -> Union[set[int], bool]: +def select_phosphate(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: """Selects the P and O atoms of all phosphate groups. Parameters @@ -127,9 +121,7 @@ def select_phosphate( return system.get_substructure_matches(pattern) -def select_sulphate( - trajectory: Trajectory, check_exists: bool = False -) -> Union[set[int], bool]: +def select_sulphate(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: """Selects the S and O atoms of all sulphate groups. Parameters @@ -152,9 +144,7 @@ def select_sulphate( return system.get_substructure_matches(pattern) -def select_thiol( - trajectory: Trajectory, check_exists: bool = False -) -> Union[set[int], bool]: +def select_thiol(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: """Selects the S and H atoms of all thiol groups. Parameters diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py index 0513d25c6..b77e47651 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py @@ -21,9 +21,7 @@ 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 @@ -43,7 +41,5 @@ def select_molecules( 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/molecule_selectors.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selectors.py index 366f59983..d793fe35a 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selectors.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selectors.py @@ -22,9 +22,7 @@ ] -def select_water( - trajectory: Trajectory, check_exists: bool = False -) -> Union[set[int], bool]: +def select_water(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: """Selects the O and H atoms of all water molecules. Parameters diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index 8930a9283..1385bae8b 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -22,9 +22,7 @@ from MDANSE.Framework.AtomSelector.general_selection import select_all, select_none -function_lookup = { - function.__name__: function for function in [select_all, select_none] -} +function_lookup = {function.__name__: function for function in [select_all, select_none]} class ReusableSelection: diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py index 5ed11f7e8..6f72767b6 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py @@ -15,7 +15,6 @@ # from MDANSE.Framework.Configurators.IConfigurator import IConfigurator -from MDANSE.Framework.AtomSelector import Selector class AtomSelectionConfigurator(IConfigurator): @@ -132,7 +131,7 @@ def get_information(self) -> str: return "\n".join(info) + "\n" - def get_selector(self) -> Selector: + def get_selector(self): """ Returns ------- diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py index ddd2ce27d..033dd597d 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py @@ -19,7 +19,6 @@ from MDANSE.Framework.Configurators.IConfigurator import IConfigurator from MDANSE.Chemistry import ATOMS_DATABASE from MDANSE.MolecularDynamics.Trajectory import Trajectory -from MDANSE.Framework.AtomSelector import Selector class AtomTransmuter: diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py index a85bd46a9..2ca82f85c 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py @@ -17,7 +17,6 @@ import json from MDANSE.Framework.Configurators.IConfigurator import IConfigurator -from MDANSE.Framework.AtomSelector import Selector from MDANSE.MolecularDynamics.Trajectory import Trajectory @@ -35,7 +34,6 @@ def __init__(self, trajectory: Trajectory) -> None: """ system = trajectory.chemical_system charges = trajectory.charges(0) - self.selector = Selector(trajectory) self._original_map = {} for at_num, at in enumerate(system.atom_list): try: diff --git a/MDANSE/Tests/UnitTests/test_selection.py b/MDANSE/Tests/UnitTests/test_selection.py new file mode 100644 index 000000000..cf1a53d33 --- /dev/null +++ b/MDANSE/Tests/UnitTests/test_selection.py @@ -0,0 +1,70 @@ +import tempfile +import os +from os import path + +import h5py +import numpy as np +import pytest + +from MDANSE.Framework.AtomSelector.general_selection import select_all, select_none +from MDANSE.Framework.AtomSelector.atom_selection import select_atoms +from MDANSE.Framework.InputData.HDFTrajectoryInputData import HDFTrajectoryInputData + + +short_traj = os.path.join( + os.path.dirname(os.path.realpath(__file__)), + "Converted", + "short_trajectory_after_changes.mdt", +) +mdmc_traj = os.path.join( + os.path.dirname(os.path.realpath(__file__)), + "Converted", + "Ar_mdmc_h5md.h5", +) +com_traj = os.path.join( + os.path.dirname(os.path.realpath(__file__)), + "Converted", + "com_trajectory.mdt", +) + +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) +def test_select_all(trajectory): + traj_object = HDFTrajectoryInputData(trajectory) + n_atoms = len(traj_object.chemical_system.atom_list) + selection = select_all(traj_object.trajectory) + assert len(selection) == n_atoms + + +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) +def test_select_none(trajectory): + traj_object = HDFTrajectoryInputData(trajectory) + selection = select_none(traj_object.trajectory) + assert len(selection) == 0 + + +def test_select_atoms_selects_by_element(): + traj_object = HDFTrajectoryInputData(short_traj) + s_selection = select_atoms(traj_object.trajectory, atom_types=['S']) + assert len(s_selection) == 208 + cu_selection = select_atoms(traj_object.trajectory, atom_types=['Cu']) + assert len(cu_selection) == 208 + sb_selection = select_atoms(traj_object.trajectory, atom_types=['Sb']) + assert len(sb_selection) == 64 + cusbs_selection = select_atoms(traj_object.trajectory, atom_types=['Cu','Sb', 'S']) + assert len(cusbs_selection) == 480 + + +def test_select_atoms_selects_by_range(): + traj_object = HDFTrajectoryInputData(short_traj) + range_selection = select_atoms(traj_object.trajectory, index_range=[15,35]) + assert len(range_selection) == 20 + overshoot_selection = select_atoms(traj_object.trajectory, index_range=[470,510]) + assert len(overshoot_selection) == 10 + + +def test_select_atoms_selects_by_slice(): + traj_object = HDFTrajectoryInputData(short_traj) + range_selection = select_atoms(traj_object.trajectory, index_slice=[150,350, 10]) + assert len(range_selection) == 20 + overshoot_selection = select_atoms(traj_object.trajectory, index_slice=[470,510,5]) + assert len(overshoot_selection) == 2 diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index c7c6c784d..ac23908e9 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -26,7 +26,6 @@ QPlainTextEdit, QWidget, ) -from MDANSE.Framework.AtomSelector import Selector from MDANSE_GUI.InputWidgets.WidgetBase import WidgetBase from MDANSE_GUI.Tabs.Visualisers.View3D import View3D from MDANSE_GUI.MolecularViewer.MolecularViewer import MolecularViewerWithPicking @@ -67,7 +66,7 @@ class SelectionHelper(QDialog): def __init__( self, - selector: Selector, + selector, traj_data: tuple[str, HDFTrajectoryInputData], field: QLineEdit, parent, From f21f07eca40d4c8f2a01badca1be5c135f772dc8 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Wed, 12 Feb 2025 09:16:34 +0000 Subject: [PATCH 03/37] Add tests for molecule and group selection --- .../AtomSelector/general_selection.py | 6 +++ MDANSE/Tests/UnitTests/test_selection.py | 49 +++++++++++++++++-- 2 files changed, 50 insertions(+), 5 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py index 39216c61b..b947a01c3 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py @@ -53,3 +53,9 @@ def select_none(trajectory: Trajectory, **function_parameters: Dict[str, Any]) - An empty set. """ return set() + + +def invert_selection(trajectory: Trajectory, selection: Set[int]) -> Set[int]: + all_indices = select_all(trajectory) + inverted = all_indices.difference(selection) + return inverted diff --git a/MDANSE/Tests/UnitTests/test_selection.py b/MDANSE/Tests/UnitTests/test_selection.py index cf1a53d33..650cae3d8 100644 --- a/MDANSE/Tests/UnitTests/test_selection.py +++ b/MDANSE/Tests/UnitTests/test_selection.py @@ -1,13 +1,12 @@ -import tempfile + import os -from os import path -import h5py -import numpy as np import pytest -from MDANSE.Framework.AtomSelector.general_selection import select_all, select_none +from MDANSE.Framework.AtomSelector.general_selection import select_all, select_none, invert_selection from MDANSE.Framework.AtomSelector.atom_selection import select_atoms +from MDANSE.Framework.AtomSelector.molecule_selection import select_molecules +from MDANSE.Framework.AtomSelector.group_selection import select_labels, select_pattern from MDANSE.Framework.InputData.HDFTrajectoryInputData import HDFTrajectoryInputData @@ -26,6 +25,12 @@ "Converted", "com_trajectory.mdt", ) +traj_2vb1 = os.path.join( + os.path.dirname(os.path.realpath(__file__)), + "Converted", + "2vb1.mdt" +) + @pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) def test_select_all(trajectory): @@ -42,6 +47,15 @@ def test_select_none(trajectory): assert len(selection) == 0 +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) +def test_inverted_none_is_all(trajectory): + traj_object = HDFTrajectoryInputData(trajectory) + none_selection = select_none(traj_object.trajectory) + all_selection = select_all(traj_object.trajectory) + inverted_none = invert_selection(traj_object.trajectory, none_selection) + assert all_selection == inverted_none + + def test_select_atoms_selects_by_element(): traj_object = HDFTrajectoryInputData(short_traj) s_selection = select_atoms(traj_object.trajectory, atom_types=['S']) @@ -68,3 +82,28 @@ def test_select_atoms_selects_by_slice(): assert len(range_selection) == 20 overshoot_selection = select_atoms(traj_object.trajectory, index_slice=[470,510,5]) assert len(overshoot_selection) == 2 + + +def test_select_molecules_selects_water(): + traj_object = HDFTrajectoryInputData(traj_2vb1) + water_selection = select_molecules(traj_object.trajectory, molecule_names = ['H2 O1']) + assert len(water_selection) == 28746 + + +def test_select_molecules_selects_protein(): + traj_object = HDFTrajectoryInputData(traj_2vb1) + protein_selecton = select_molecules(traj_object.trajectory, molecule_names = ['C613 H959 N193 O185 S10']) + assert len(protein_selecton) == 613 + 959 + 193 + 185 + 10 + + +def test_select_molecules_inverted_selects_ions(): + traj_object = HDFTrajectoryInputData(traj_2vb1) + all_molecules_selection = select_molecules(traj_object.trajectory, molecule_names = ['C613 H959 N193 O185 S10', 'H2 O1']) + non_molecules_selection = invert_selection(traj_object.trajectory, all_molecules_selection) + assert all([traj_object.chemical_system.atom_list[index] in ['Na', 'Cl'] for index in non_molecules_selection]) + + +def test_select_pattern_selects_water(): + traj_object = HDFTrajectoryInputData(traj_2vb1) + water_selection = select_pattern(traj_object.trajectory, rdkit_pattern="[#8X2;H2](~[H])~[H]") + assert len(water_selection) == 28746 From 438a46207309925f142b7c1458e9d6c63038c400 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Wed, 12 Feb 2025 09:51:35 +0000 Subject: [PATCH 04/37] Add tests for selection saving --- .../MDANSE/Framework/AtomSelector/selector.py | 37 +++- .../UnitTests/test_reusable_selection.py | 176 ++++++++++++++++++ 2 files changed, 205 insertions(+), 8 deletions(-) create mode 100644 MDANSE/Tests/UnitTests/test_reusable_selection.py diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index 1385bae8b..82eef73d0 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -13,16 +13,32 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . # -import copy + import json from typing import Union, Dict, Any, Set -from MDANSE.Chemistry.ChemicalSystem import ChemicalSystem from MDANSE.MolecularDynamics.Trajectory import Trajectory -from MDANSE.Framework.AtomSelector.all_selector import select_all -from MDANSE.Framework.AtomSelector.general_selection import select_all, select_none +from MDANSE.Framework.AtomSelector.general_selection import ( + select_all, + select_none, + invert_selection, +) +from MDANSE.Framework.AtomSelector.atom_selection import select_atoms +from MDANSE.Framework.AtomSelector.molecule_selection import select_molecules +from MDANSE.Framework.AtomSelector.group_selection import select_labels, select_pattern -function_lookup = {function.__name__: function for function in [select_all, select_none]} +function_lookup = { + function.__name__: function + for function in [ + select_all, + select_none, + invert_selection, + select_atoms, + select_molecules, + select_labels, + select_pattern, + ] +} class ReusableSelection: @@ -51,6 +67,11 @@ def set_selection( ): if number is None: number = len(self.operations) + else: + try: + number = int(number) + except TypeError: + number = len(self.operations) self.operations[number] = function_parameters def select_in_trajectory(self, trajectory: Trajectory) -> Set[int]: @@ -61,11 +82,11 @@ def select_in_trajectory(self, trajectory: Trajectory) -> Set[int]: return self.all_idxs for number in sequence: function_parameters = self.operations[number] - function_name = function_parameters.pop("function_name", "select_all") + function_name = function_parameters.get("function_name", "select_all") if function_name == "invert_selection": selection = self.all_idxs.difference(selection) else: - operation_type = function_parameters.pop("operation_type", "union") + operation_type = function_parameters.get("operation_type", "union") function = function_lookup[function_name] temp_selection = function(trajectory, **function_parameters) if operation_type == "union": @@ -99,4 +120,4 @@ def read_from_json(self, json_string: str): json_setting = json.loads(json_string) for k0, v0 in json_setting.items(): if isinstance(v0, dict): - self.append_operation(k0, v0) + self.set_selection(k0, v0) diff --git a/MDANSE/Tests/UnitTests/test_reusable_selection.py b/MDANSE/Tests/UnitTests/test_reusable_selection.py new file mode 100644 index 000000000..ef30169e3 --- /dev/null +++ b/MDANSE/Tests/UnitTests/test_reusable_selection.py @@ -0,0 +1,176 @@ + +import os + +import pytest + +from MDANSE.Framework.AtomSelector.selector import ReusableSelection +from MDANSE.Framework.InputData.HDFTrajectoryInputData import HDFTrajectoryInputData + + +short_traj = os.path.join( + os.path.dirname(os.path.realpath(__file__)), + "Converted", + "short_trajectory_after_changes.mdt", +) +mdmc_traj = os.path.join( + os.path.dirname(os.path.realpath(__file__)), + "Converted", + "Ar_mdmc_h5md.h5", +) +com_traj = os.path.join( + os.path.dirname(os.path.realpath(__file__)), + "Converted", + "com_trajectory.mdt", +) +traj_2vb1 = os.path.join( + os.path.dirname(os.path.realpath(__file__)), + "Converted", + "2vb1.mdt" +) + + +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) +def test_select_all(trajectory): + traj_object = HDFTrajectoryInputData(trajectory) + n_atoms = len(traj_object.chemical_system.atom_list) + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_all'}) + selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + assert len(selection) == n_atoms + + +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) +def test_select_none(trajectory): + traj_object = HDFTrajectoryInputData(trajectory) + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_none'}) + selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + assert len(selection) == 0 + + +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) +def test_inverted_all_is_none(trajectory): + traj_object = HDFTrajectoryInputData(trajectory) + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_all'}) + reusable_selection.set_selection(None, {'function_name': 'invert_selection'}) + selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + assert len(selection) == 0 + + +def test_json_saving_is_reversible(): + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_all'}) + reusable_selection.set_selection(None, {'function_name': 'invert_selection'}) + json_string = reusable_selection.convert_to_json() + another_selection = ReusableSelection() + another_selection.read_from_json(json_string) + json_string_2 = another_selection.convert_to_json() + assert json_string == json_string_2 + + +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) +def test_selection_from_json_is_the_same_as_from_runtime(trajectory): + traj_object = HDFTrajectoryInputData(trajectory) + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_all'}) + reusable_selection.set_selection(None, {'function_name': 'invert_selection'}) + selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + json_string = reusable_selection.convert_to_json() + another_selection = ReusableSelection() + another_selection.read_from_json(json_string) + selection2 = another_selection.select_in_trajectory(traj_object.trajectory) + print(f"original: {reusable_selection.operations}") + print(f"another: {another_selection.operations}") + assert selection == selection2 + +def test_select_atoms_selects_by_element(): + traj_object = HDFTrajectoryInputData(short_traj) + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'atom_types': ['S']}) + s_selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + assert len(s_selection) == 208 + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'atom_types': ['Cu']}) + cu_selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + assert len(cu_selection) == 208 + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'atom_types': ['Sb']}) + sb_selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + assert len(sb_selection) == 64 + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'atom_types': ['S', 'Sb', 'Cu']}) + json_string = reusable_selection.convert_to_json() + another_selection = ReusableSelection() + another_selection.read_from_json(json_string) + cusbs_selection = another_selection.select_in_trajectory(traj_object.trajectory) + assert len(cusbs_selection) == 480 + + +def test_select_atoms_selects_by_range(): + traj_object = HDFTrajectoryInputData(short_traj) + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'index_range': [15,35]}) + range_selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + assert len(range_selection) == 20 + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'index_range': [470,510]}) + overshoot_selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + assert len(overshoot_selection) == 10 + + +def test_select_atoms_selects_by_slice(): + traj_object = HDFTrajectoryInputData(short_traj) + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'index_slice': [150,350,10]}) + range_selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + assert len(range_selection) == 20 + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'index_slice': [470,510,5]}) + json_string = reusable_selection.convert_to_json() + another_selection = ReusableSelection() + another_selection.read_from_json(json_string) + overshoot_selection = another_selection.select_in_trajectory(traj_object.trajectory) + assert len(overshoot_selection) == 2 + + +def test_select_molecules_selects_water(): + traj_object = HDFTrajectoryInputData(traj_2vb1) + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_molecules', 'molecule_names': ['H2 O1']}) + water_selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + assert len(water_selection) == 28746 + + +def test_select_molecules_selects_protein(): + traj_object = HDFTrajectoryInputData(traj_2vb1) + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_molecules', 'molecule_names': ['C613 H959 N193 O185 S10']}) + json_string = reusable_selection.convert_to_json() + another_selection = ReusableSelection() + another_selection.read_from_json(json_string) + protein_selection = another_selection.select_in_trajectory(traj_object.trajectory) + assert len(protein_selection) == 613 + 959 + 193 + 185 + 10 + + +def test_select_molecules_inverted_selects_ions(): + traj_object = HDFTrajectoryInputData(traj_2vb1) + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_molecules', 'molecule_names': ['C613 H959 N193 O185 S10', 'H2 O1']}) + reusable_selection.set_selection(None, {'function_name': 'invert_selection'}) + json_string = reusable_selection.convert_to_json() + another_selection = ReusableSelection() + another_selection.read_from_json(json_string) + non_molecules_selection = another_selection.select_in_trajectory(traj_object.trajectory) + assert all([traj_object.chemical_system.atom_list[index] in ['Na', 'Cl'] for index in non_molecules_selection]) + + +def test_select_pattern_selects_water(): + traj_object = HDFTrajectoryInputData(traj_2vb1) + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_pattern', 'rdkit_pattern': "[#8X2;H2](~[H])~[H]"}) + json_string = reusable_selection.convert_to_json() + another_selection = ReusableSelection() + another_selection.read_from_json(json_string) + water_selection = another_selection.select_in_trajectory(traj_object.trajectory) + assert len(water_selection) == 28746 From 4f10e0f188249e472c9d96504bd30106753ae954 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Wed, 12 Feb 2025 11:01:37 +0000 Subject: [PATCH 05/37] Update configurators to use the new selection mechanism --- .../Framework/AtomSelector/all_selector.py | 46 --- .../Framework/AtomSelector/atom_selectors.py | 226 ------------ .../Framework/AtomSelector/group_selectors.py | 167 --------- .../AtomSelector/molecule_selectors.py | 45 --- .../Tests/UnitTests/AtomSelector/__init__.py | 0 .../AtomSelector/test_all_selector.py | 20 -- .../AtomSelector/test_atom_selectors.py | 90 ----- .../AtomSelector/test_group_selectors.py | 83 ----- .../AtomSelector/test_molecule_selectors.py | 29 -- .../UnitTests/AtomSelector/test_selector.py | 325 ------------------ 10 files changed, 1031 deletions(-) delete mode 100644 MDANSE/Src/MDANSE/Framework/AtomSelector/all_selector.py delete mode 100644 MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selectors.py delete mode 100644 MDANSE/Src/MDANSE/Framework/AtomSelector/group_selectors.py delete mode 100644 MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selectors.py delete mode 100644 MDANSE/Tests/UnitTests/AtomSelector/__init__.py delete mode 100644 MDANSE/Tests/UnitTests/AtomSelector/test_all_selector.py delete mode 100644 MDANSE/Tests/UnitTests/AtomSelector/test_atom_selectors.py delete mode 100644 MDANSE/Tests/UnitTests/AtomSelector/test_group_selectors.py delete mode 100644 MDANSE/Tests/UnitTests/AtomSelector/test_molecule_selectors.py delete mode 100644 MDANSE/Tests/UnitTests/AtomSelector/test_selector.py diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/all_selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/all_selector.py deleted file mode 100644 index cf1a452c8..000000000 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/all_selector.py +++ /dev/null @@ -1,46 +0,0 @@ -# This file is part of MDANSE. -# -# MDANSE is free software: you can redistribute it and/or modify -# it under the terms of the GNU General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU General Public License -# along with this program. If not, see . -# -from typing import Union -from MDANSE.MolecularDynamics.Trajectory import Trajectory - - -def select_all(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: - """Selects all atoms in the chemical system except for the dummy - atoms. - - Parameters - ---------- - system : ChemicalSystem - The MDANSE chemical system. - check_exists : bool, optional - Check if a match exists. - - Returns - ------- - Union[set[int], bool] - All atom indices except for dummy atoms or a bool if checking - match. - """ - system = trajectory.chemical_system - if check_exists: - return True - else: - dummy_list = [] - atom_list = system.atom_list - for atm in system._unique_elements: - if trajectory.get_atom_property(atm, "dummy"): - dummy_list.append(atm) - return set([index for index in system._atom_indices if atom_list[index] not in dummy_list]) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selectors.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selectors.py deleted file mode 100644 index 1b0ba6211..000000000 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selectors.py +++ /dev/null @@ -1,226 +0,0 @@ -# This file is part of MDANSE. -# -# MDANSE is free software: you can redistribute it and/or modify -# it under the terms of the GNU General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU General Public License -# along with this program. If not, see . -# -from typing import Union -from MDANSE.MolecularDynamics.Trajectory import Trajectory - - -__all__ = [ - "select_element", - "select_dummy", - "select_atom_name", - "select_atom_fullname", - "select_hs_on_element", - "select_hs_on_heteroatom", - "select_index", -] - - -def select_element( - trajectory: Trajectory, symbol: str, check_exists: bool = False -) -> Union[set[int], bool]: - """Selects all atoms for the input element. - - Parameters - ---------- - system : ChemicalSystem - The MDANSE chemical system. - symbol : str - Symbol of the element. - check_exists : bool, optional - Check if a match exists. - - Returns - ------- - Union[set[int], bool] - The atom indices of the matched atoms. - """ - system = trajectory.chemical_system - pattern = f"[#{trajectory.get_atom_property(symbol, 'atomic_number')}]" - if check_exists: - return system.has_substructure_match(pattern) - else: - return system.get_substructure_matches(pattern) - - -def select_dummy(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: - """Selects all dummy atoms in the chemical system. - - Parameters - ---------- - system : ChemicalSystem - The MDANSE chemical system. - check_exists : bool, optional - Check if a match exists. - - Returns - ------- - Union[set[int], bool] - All dummy atom indices or a bool if checking match. - """ - system = trajectory.chemical_system - dummy_list = ["Du", "dummy"] - if check_exists: - for atm in system.atom_list: - if atm in dummy_list: - return True - elif trajectory.get_atom_property(atm, "dummy"): - return True - return False - else: - for atm in system._unique_elements: - if trajectory.get_atom_property(atm, "dummy"): - dummy_list.append(atm) - return set( - [index for index, element in enumerate(system.atom_list) if element in dummy_list] - ) - - -def select_atom_name( - trajectory: Trajectory, name: str, check_exists: bool = False -) -> Union[set[int], bool]: - """Selects all atoms with the input name in the chemical system. - - Parameters - ---------- - system : ChemicalSystem - The MDANSE chemical system. - name : str - The name of the atom to match. - check_exists : bool, optional - Check if a match exists. - - Returns - ------- - Union[set[int], bool] - All atom indices or a bool if checking match. - """ - system = trajectory.chemical_system - if check_exists: - if name in system.atom_list: - return True - return False - else: - return set([index for index, element in enumerate(system.atom_list) if element == name]) - - -def select_atom_fullname( - trajectory: Trajectory, fullname: str, check_exists: bool = False -) -> Union[set[int], bool]: - """Selects all atoms with the input fullname in the chemical system. - - Parameters - ---------- - system : ChemicalSystem - The MDANSE chemical system. - fullname : str - The fullname of the atom to match. - check_exists : bool, optional - Check if a match exists. - - Returns - ------- - Union[set[int], bool] - All atom indices or a bool if checking match. - """ - system = trajectory.chemical_system - if check_exists: - if fullname in system.name_list: - return True - return False - else: - return set([index for index, name in enumerate(system.name_list) if name == fullname]) - - -def select_hs_on_element( - trajectory: Trajectory, symbol: str, check_exists: bool = False -) -> Union[set[int], bool]: - """Selects all H atoms bonded to the input element. - - Parameters - ---------- - system : ChemicalSystem - The MDANSE chemical system. - symbol : str - Symbol of the element that the H atoms are bonded to. - check_exists : bool, optional - Check if a match exists. - - Returns - ------- - Union[set[int], bool] - The atom indices of the matched atoms. - """ - system = trajectory.chemical_system - num = trajectory.get_atom_property(symbol, "atomic_number") - if check_exists: - return system.has_substructure_match(f"[#{num}]~[H]") - else: - xh_matches = system.get_substructure_matches(f"[#{num}]~[H]") - x_matches = system.get_substructure_matches(f"[#{num}]") - return xh_matches - x_matches - - -def select_hs_on_heteroatom( - trajectory: Trajectory, check_exists: bool = False -) -> Union[set[int], bool]: - """Selects all H atoms bonded to any atom except carbon and - hydrogen. - - Parameters - ---------- - system : ChemicalSystem - The MDANSE chemical system. - check_exists : bool, optional - Check if a match exists. - - Returns - ------- - Union[set[int], bool] - The atom indices of the matched atoms. - """ - system = trajectory.chemical_system - if check_exists: - return system.has_substructure_match("[!#6&!#1]~[H]") - else: - xh_matches = system.get_substructure_matches("[!#6&!#1]~[H]") - x_matches = system.get_substructure_matches("[!#6&!#1]") - return xh_matches - x_matches - - -def select_index( - trajectory: Trajectory, index: Union[int, str], check_exists: bool = False -) -> Union[set[int], bool]: - """Selects atom with index - just returns the set with the - index in it. - - Parameters - ---------- - system : ChemicalSystem - The MDANSE chemical system. - index : int or str - The index to select. - check_exists : bool, optional - Check if a match exists. - - Returns - ------- - Union[set[int], bool] - The index in a set or a bool if checking match. - """ - if check_exists: - return True - else: - return {int(index)} diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selectors.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selectors.py deleted file mode 100644 index 3f6c330ca..000000000 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selectors.py +++ /dev/null @@ -1,167 +0,0 @@ -# This file is part of MDANSE. -# -# MDANSE is free software: you can redistribute it and/or modify -# it under the terms of the GNU General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU General Public License -# along with this program. If not, see . -# -from typing import Union -from MDANSE.MolecularDynamics.Trajectory import Trajectory - - -__all__ = [ - "select_primary_amine", - "select_hydroxy", - "select_methyl", - "select_phosphate", - "select_sulphate", - "select_thiol", -] - - -def select_primary_amine( - trajectory: Trajectory, check_exists: bool = False -) -> Union[set[int], bool]: - """Selects the N and H atoms of all primary amines. - - Parameters - ---------- - system : ChemicalSystem - The MDANSE chemical system. - check_exists : bool, optional - Check if a match exists. - - Returns - ------- - Union[set[int], bool] - The atom indices of the matched atoms or a bool if checking match. - """ - system = trajectory.chemical_system - pattern = "[#7X3;H2;!$([#7][#6X3][!#6]);!$([#7][#6X2][!#6])](~[H])~[H]" - if check_exists: - return system.has_substructure_match(pattern) - else: - return system.get_substructure_matches(pattern) - - -def select_hydroxy(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: - """Selects the O and H atoms of all hydroxy groups including water. - - Parameters - ---------- - system : ChemicalSystem - The MDANSE chemical system. - check_exists : bool, optional - Check if a match exists. - - Returns - ------- - Union[set[int], bool] - The atom indices of the matched atoms or a bool if checking match. - """ - system = trajectory.chemical_system - pattern = "[#8;H1,H2]~[H]" - if check_exists: - return system.has_substructure_match(pattern) - else: - return system.get_substructure_matches(pattern) - - -def select_methyl(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: - """Selects the C and H atoms of all methyl groups. - - Parameters - ---------- - system : ChemicalSystem - The MDANSE chemical system. - check_exists : bool, optional - Check if a match exists. - - Returns - ------- - Union[set[int], bool] - The atom indices of the matched atoms or a bool if checking match. - """ - system = trajectory.chemical_system - pattern = "[#6;H3](~[H])(~[H])~[H]" - if check_exists: - return system.has_substructure_match(pattern) - else: - return system.get_substructure_matches(pattern) - - -def select_phosphate(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: - """Selects the P and O atoms of all phosphate groups. - - Parameters - ---------- - system : ChemicalSystem - The MDANSE chemical system. - check_exists : bool, optional - Check if a match exists. - - Returns - ------- - set[int] - The atom indices of the matched atoms or a bool if checking match. - """ - system = trajectory.chemical_system - pattern = "[#15X4](~[#8])(~[#8])(~[#8])~[#8]" - if check_exists: - return system.has_substructure_match(pattern) - else: - return system.get_substructure_matches(pattern) - - -def select_sulphate(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: - """Selects the S and O atoms of all sulphate groups. - - Parameters - ---------- - system : ChemicalSystem - The MDANSE chemical system. - check_exists : bool, optional - Check if a match exists. - - Returns - ------- - Union[set[int], bool] - The atom indices of the matched atoms or a bool if checking match. - """ - system = trajectory.chemical_system - pattern = "[#16X4](~[#8])(~[#8])(~[#8])~[#8]" - if check_exists: - return system.has_substructure_match(pattern) - else: - return system.get_substructure_matches(pattern) - - -def select_thiol(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: - """Selects the S and H atoms of all thiol groups. - - Parameters - ---------- - system : ChemicalSystem - The MDANSE chemical system. - check_exists : bool, optional - Check if a match exists. - - Returns - ------- - Union[set[int], bool] - The atom indices of the matched atoms or a bool if checking match. - """ - system = trajectory.chemical_system - pattern = "[#16X2;H1]~[H]" - if check_exists: - return system.has_substructure_match(pattern) - else: - return system.get_substructure_matches(pattern) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selectors.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selectors.py deleted file mode 100644 index d793fe35a..000000000 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selectors.py +++ /dev/null @@ -1,45 +0,0 @@ -# This file is part of MDANSE. -# -# MDANSE is free software: you can redistribute it and/or modify -# it under the terms of the GNU General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU General Public License -# along with this program. If not, see . -# -from typing import Union -from MDANSE.MolecularDynamics.Trajectory import Trajectory - - -__all__ = [ - "select_water", -] - - -def select_water(trajectory: Trajectory, check_exists: bool = False) -> Union[set[int], bool]: - """Selects the O and H atoms of all water molecules. - - Parameters - ---------- - system : ChemicalSystem - The MDANSE chemical system. - check_exists : bool, optional - Check if a match exists. - - Returns - ------- - Union[set[int], bool] - The atom indices of the matched atoms or a bool if checking match. - """ - system = trajectory.chemical_system - pattern = "[#8X2;H2](~[H])~[H]" - if check_exists: - return system.has_substructure_match(pattern) - else: - return system.get_substructure_matches(pattern) diff --git a/MDANSE/Tests/UnitTests/AtomSelector/__init__.py b/MDANSE/Tests/UnitTests/AtomSelector/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/MDANSE/Tests/UnitTests/AtomSelector/test_all_selector.py b/MDANSE/Tests/UnitTests/AtomSelector/test_all_selector.py deleted file mode 100644 index 22eed92b8..000000000 --- a/MDANSE/Tests/UnitTests/AtomSelector/test_all_selector.py +++ /dev/null @@ -1,20 +0,0 @@ -import os -import pytest -from MDANSE.Framework.InputData.HDFTrajectoryInputData import HDFTrajectoryInputData -from MDANSE.Framework.AtomSelector.all_selector import select_all - - -traj_2vb1 = os.path.join( - os.path.dirname(os.path.realpath(__file__)), "..", "Converted", "2vb1.mdt" -) - - -@pytest.fixture(scope="module") -def protein_trajectory(): - protein_trajectory = HDFTrajectoryInputData(traj_2vb1) - return protein_trajectory.trajectory - - -def test_select_all_returns_correct_number_of_atoms_matches(protein_trajectory): - selection = select_all(protein_trajectory) - assert len(selection) == 30714 diff --git a/MDANSE/Tests/UnitTests/AtomSelector/test_atom_selectors.py b/MDANSE/Tests/UnitTests/AtomSelector/test_atom_selectors.py deleted file mode 100644 index 6903bbd3a..000000000 --- a/MDANSE/Tests/UnitTests/AtomSelector/test_atom_selectors.py +++ /dev/null @@ -1,90 +0,0 @@ -import os -import pytest -from MDANSE.Framework.InputData.HDFTrajectoryInputData import HDFTrajectoryInputData -from MDANSE.Framework.AtomSelector.atom_selectors import ( - select_element, - select_hs_on_heteroatom, - select_hs_on_element, - select_dummy, -) - - -traj_2vb1 = os.path.join( - os.path.dirname(os.path.realpath(__file__)), "..", "Converted", "2vb1.mdt" -) - - -@pytest.fixture(scope="module") -def protein_trajectory(): - protein_trajectory = HDFTrajectoryInputData(traj_2vb1) - return protein_trajectory.trajectory - - -def test_select_element_returns_true_as_match_exist( - protein_trajectory, -): - exists = select_element(protein_trajectory, "S", check_exists=True) - assert exists - - -def test_select_element_returns_false_as_match_does_not_exist( - protein_trajectory, -): - exists = select_element(protein_trajectory, "Si", check_exists=True) - assert not exists - - -def test_select_element_returns_correct_number_of_atom_matches( - protein_trajectory, -): - selection = select_element(protein_trajectory, "S") - assert len(selection) == 10 - - -def test_select_hs_on_carbon_returns_correct_number_of_atom_matches( - protein_trajectory, -): - selection = select_hs_on_element(protein_trajectory, "C") - assert len(selection) == 696 - - -def test_select_hs_on_nitrogen_returns_correct_number_of_atom_matches( - protein_trajectory, -): - selection = select_hs_on_element(protein_trajectory, "N") - assert len(selection) == 243 - - -def test_select_hs_on_oxygen_returns_correct_number_of_atom_matches( - protein_trajectory, -): - selection = select_hs_on_element(protein_trajectory, "O") - assert len(selection) == 19184 - - -def test_select_hs_on_sulfur_returns_correct_number_of_atom_matches( - protein_trajectory, -): - selection = select_hs_on_element(protein_trajectory, "S") - assert len(selection) == 0 - - -def test_select_hs_on_silicon_returns_correct_number_of_atom_matches( - protein_trajectory, -): - selection = select_hs_on_element(protein_trajectory, "Si") - assert len(selection) == 0 - - -def test_select_hs_on_heteroatom_returns_correct_number_of_atom_matches( - protein_trajectory, -): - selection = select_hs_on_heteroatom(protein_trajectory) - assert len(selection) == 19427 - - -def test_select_dummy_returns_correct_number_of_atom_matches( - protein_trajectory, -): - selection = select_dummy(protein_trajectory) - assert len(selection) == 0 diff --git a/MDANSE/Tests/UnitTests/AtomSelector/test_group_selectors.py b/MDANSE/Tests/UnitTests/AtomSelector/test_group_selectors.py deleted file mode 100644 index 322697cc3..000000000 --- a/MDANSE/Tests/UnitTests/AtomSelector/test_group_selectors.py +++ /dev/null @@ -1,83 +0,0 @@ -import os -import pytest -from MDANSE.Framework.InputData.HDFTrajectoryInputData import HDFTrajectoryInputData -from MDANSE.Framework.AtomSelector.group_selectors import ( - select_primary_amine, - select_hydroxy, - select_methyl, - select_phosphate, - select_sulphate, - select_thiol, -) - - -traj_2vb1 = os.path.join( - os.path.dirname(os.path.realpath(__file__)), "..", "Converted", "2vb1.mdt" -) -traj_1gip = os.path.join( - os.path.dirname(os.path.realpath(__file__)), "..", "Converted", "1gip.mdt" -) - - -@pytest.fixture(scope="module") -def protein_trajectory(): - protein_trajectory = HDFTrajectoryInputData(traj_2vb1) - return protein_trajectory.trajectory - - -@pytest.fixture(scope="module") -def nucleic_acid_chemical_system(): - protein_trajectory = HDFTrajectoryInputData(traj_1gip) - return protein_trajectory.trajectory - - -def test_select_primary_amine_returns_true_as_match_exists( - protein_trajectory, -): - exists = select_primary_amine(protein_trajectory, check_exists=True) - assert exists - - -def test_select_sulphate_returns_false_as_match_does_not_exist( - nucleic_acid_chemical_system, -): - exists = select_sulphate(nucleic_acid_chemical_system, check_exists=True) - assert not exists - - -def test_select_primary_amine_returns_correct_number_of_atom_matches( - protein_trajectory, -): - selection = select_primary_amine(protein_trajectory) - assert len(selection) == 117 - - -def test_select_hydroxy_returns_correct_number_of_atom_matches( - protein_trajectory, -): - selection = select_hydroxy(protein_trajectory) - assert len(selection) == 28786 - - -def test_select_methyl_returns_correct_number_of_atom_matches(protein_trajectory): - selection = select_methyl(protein_trajectory) - assert len(selection) == 244 - - -def test_select_phosphate_returns_correct_number_of_atom_matches( - nucleic_acid_chemical_system, -): - selection = select_phosphate(nucleic_acid_chemical_system) - assert len(selection) == 110 - - -def test_select_sulphate_returns_correct_number_of_atom_matches( - nucleic_acid_chemical_system, -): - selection = select_sulphate(nucleic_acid_chemical_system) - assert len(selection) == 0 - - -def test_select_thiol_returns_correct_number_of_atoms_matches(protein_trajectory): - selection = select_thiol(protein_trajectory) - assert len(selection) == 0 diff --git a/MDANSE/Tests/UnitTests/AtomSelector/test_molecule_selectors.py b/MDANSE/Tests/UnitTests/AtomSelector/test_molecule_selectors.py deleted file mode 100644 index d540fa6fd..000000000 --- a/MDANSE/Tests/UnitTests/AtomSelector/test_molecule_selectors.py +++ /dev/null @@ -1,29 +0,0 @@ -import os -import pytest -from MDANSE.Framework.InputData.HDFTrajectoryInputData import HDFTrajectoryInputData -from MDANSE.Framework.AtomSelector.molecule_selectors import select_water - - -traj_2vb1 = os.path.join( - os.path.dirname(os.path.realpath(__file__)), "..", "Converted", "2vb1.mdt" -) - - -@pytest.fixture(scope="module") -def protein_trajectory(): - protein_trajectory = HDFTrajectoryInputData(traj_2vb1) - return protein_trajectory.trajectory - - -def test_select_water_returns_true_as_match_exists( - protein_trajectory, -): - exists = select_water(protein_trajectory, check_exists=True) - assert exists - - -def test_select_water_returns_correct_number_of_atom_matches( - protein_trajectory, -): - selection = select_water(protein_trajectory) - assert len(selection) == 28746 diff --git a/MDANSE/Tests/UnitTests/AtomSelector/test_selector.py b/MDANSE/Tests/UnitTests/AtomSelector/test_selector.py deleted file mode 100644 index a853a036f..000000000 --- a/MDANSE/Tests/UnitTests/AtomSelector/test_selector.py +++ /dev/null @@ -1,325 +0,0 @@ -import os -import pytest -from MDANSE.Framework.InputData.HDFTrajectoryInputData import HDFTrajectoryInputData -from MDANSE.Framework.AtomSelector.selector import Selector - - -traj_2vb1 = os.path.join( - os.path.dirname(os.path.realpath(__file__)), "..", "Converted", "2vb1.mdt" -) - - -@pytest.fixture(scope="module") -def protein_trajectory(): - protein_trajectory = HDFTrajectoryInputData(traj_2vb1) - return protein_trajectory.trajectory - - -def test_selector_returns_all_atom_idxs(protein_trajectory): - selector = Selector(protein_trajectory) - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 30714 - - -def test_selector_returns_all_atom_idxs_with_all_and_sulfurs_selected( - protein_trajectory, -): - selector = Selector(protein_trajectory) - selector.settings["all"] = True - selector.settings["element"] = {"S": True} - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 30714 - - -def test_selector_returns_correct_number_of_atom_idxs_when_sulfur_atoms_are_selected( - protein_trajectory, -): - selector = Selector(protein_trajectory) - selector.settings["all"] = False - selector.settings["element"] = {"S": True} - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 10 - - -def test_selector_returns_correct_number_of_atom_idxs_when_sulfur_atoms_are_selected_when_get_idxs_is_called_twice( - protein_trajectory, -): - selector = Selector(protein_trajectory) - selector.settings["all"] = False - selector.settings["element"] = {"S": True} - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 10 - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 10 - - -def test_selector_returns_correct_number_of_atom_idxs_when_waters_are_selected( - protein_trajectory, -): - selector = Selector(protein_trajectory) - selector.settings["all"] = False - selector.settings["water"] = True - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 28746 - - -def test_selector_returns_correct_number_of_atom_idxs_when_water_is_turned_on_and_off( - protein_trajectory, -): - selector = Selector(protein_trajectory) - selector.settings["all"] = False - selector.settings["water"] = True - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 28746 - selector.settings["water"] = False - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 0 - - -def test_selector_returns_correct_number_of_atom_idxs_when_waters_and_sulfurs_are_selected( - protein_trajectory, -): - selector = Selector(protein_trajectory) - selector.settings["all"] = False - selector.settings["water"] = True - selector.settings["element"] = {"S": True} - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 28746 + 10 - - -def test_selector_returns_correct_number_of_atom_idxs_when_waters_and_sulfurs_are_selected_with_settings_loaded_as_a_dict( - protein_trajectory, -): - selector = Selector(protein_trajectory) - selector.update_settings({"all": False, "element": {"S": True}, "water": True}) - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 28746 + 10 - - -def test_selector_json_dump_0(protein_trajectory): - selector = Selector(protein_trajectory) - selector.update_settings({"all": False, "element": {"S": True}}) - json_dump = selector.settings_to_json() - assert json_dump == '{"all": false, "element": ["S"]}' - - -def test_selector_json_dump_1(protein_trajectory): - selector = Selector(protein_trajectory) - selector.update_settings({"all": False, "element": {"S": True}, "water": True}) - json_dump = selector.settings_to_json() - assert json_dump == '{"all": false, "water": true, "element": ["S"]}' - - -def test_selector_json_dump_2(protein_trajectory): - selector = Selector(protein_trajectory) - selector.update_settings({"all": False, "water": True}) - json_dump = selector.settings_to_json() - assert json_dump == '{"all": false, "water": true}' - - -def test_selector_json_dump_3(protein_trajectory): - selector = Selector(protein_trajectory) - selector.update_settings( - {"all": False, "element": {"S": True, "H": True}, "water": True} - ) - json_dump = selector.settings_to_json() - assert json_dump == '{"all": false, "water": true, "element": ["H", "S"]}' - - -def test_selector_json_dump_4(protein_trajectory): - selector = Selector(protein_trajectory) - selector.update_settings( - { - "all": False, - "element": {"S": True, "H": True}, - "water": True, - "index": {0: True, 1: True}, - } - ) - json_dump = selector.settings_to_json() - assert ( - json_dump - == '{"all": false, "water": true, "element": ["H", "S"], "index": [0, 1]}' - ) - - -def test_selector_json_dump_with_second_update(protein_trajectory): - selector = Selector(protein_trajectory) - selector.update_settings({"all": False}) - selector.update_settings({"element": {"S": True, "O": True}, "water": True}) - json_dump = selector.settings_to_json() - assert json_dump == '{"all": false, "water": true, "element": ["O", "S"]}' - - -def test_selector_json_dump_with_third_update(protein_trajectory): - selector = Selector(protein_trajectory) - selector.update_settings({"all": False}) - selector.update_settings({"element": {"S": True, "O": True}, "water": True}) - selector.update_settings({"element": {"S": False}}) - json_dump = selector.settings_to_json() - assert json_dump == '{"all": false, "water": true, "element": ["O"]}' - - -def test_selector_json_dump_with_fourth_update(protein_trajectory): - selector = Selector(protein_trajectory) - selector.update_settings({"all": False}) - selector.update_settings({"element": {"S": True, "O": True}, "water": True}) - selector.update_settings({"element": {"S": False}}) - selector.update_settings({"water": False}) - json_dump = selector.settings_to_json() - assert json_dump == '{"all": false, "element": ["O"]}' - - -def test_selector_returns_correct_number_of_atom_idxs_after_setting_settings_again_with_reset_first( - protein_trajectory, -): - selector = Selector(protein_trajectory) - selector.update_settings({"all": False, "element": {"S": True}, "water": True}) - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 28746 + 10 - - selector.update_settings( - { - "all": False, - "element": {"S": True}, - }, - reset_first=True, - ) - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 10 - - -def test_selector_json_dump_and_load_0(protein_trajectory): - selector = Selector(protein_trajectory) - selector.update_settings({"all": False, "index": {0: True, 1: True}}) - json_dump = selector.settings_to_json() - assert json_dump == '{"all": false, "index": [0, 1]}' - selector.load_from_json(json_dump) - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 2 - - -def test_selector_json_dump_and_load_1(protein_trajectory): - selector = Selector(protein_trajectory) - selector.update_settings({"all": False, "element": {"S": True}, "water": True}) - json_dump = selector.settings_to_json() - assert json_dump == '{"all": false, "water": true, "element": ["S"]}' - selector.load_from_json(json_dump) - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 28746 + 10 - - -def test_selector_returns_correct_number_of_atom_idxs_when_indexes_0_and_1_are_selected( - protein_trajectory, -): - selector = Selector(protein_trajectory) - selector.update_settings( - { - "all": False, - "index": {0: True, 1: True}, - } - ) - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 2 - - -def test_selector_returns_true_with_correct_setting_check(protein_trajectory): - selector = Selector(protein_trajectory) - assert selector.check_valid_setting( - { - "all": False, - "index": {0: True, 1: True}, - } - ) - - -def test_selector_returns_false_with_incorrect_setting_check_0(protein_trajectory): - selector = Selector(protein_trajectory) - assert not selector.check_valid_setting( - { - "alle": False, - "index": {0: True, 1: True}, - } - ) - - -def test_selector_returns_false_with_incorrect_setting_check_1(protein_trajectory): - selector = Selector(protein_trajectory) - assert not selector.check_valid_setting( - { - "all": False, - "index": {-1: True, 1: True}, - } - ) - - -def test_selector_returns_false_with_incorrect_setting_check_2(protein_trajectory): - selector = Selector(protein_trajectory) - assert not selector.check_valid_setting( - { - "all": False, - "index": {0: True, 1: True}, - "element": {"Ss": True}, - } - ) - - -def test_selector_returns_true_with_correct_json_setting_0(protein_trajectory): - selector = Selector(protein_trajectory) - assert selector.check_valid_json_settings( - '{"all": false, "water": true, "element": {"S": true}}' - ) - - -def test_selector_returns_true_with_correct_json_setting_1(protein_trajectory): - selector = Selector(protein_trajectory) - assert selector.check_valid_json_settings('{"all": false, "index": [0, 1]}') - - -def test_selector_returns_false_with_incorrect_json_setting_0(protein_trajectory): - selector = Selector(protein_trajectory) - assert not selector.check_valid_json_settings( - '{all: false, "water": true, "element": {"S": true}}' - ) - - -def test_selector_returns_false_with_incorrect_json_setting_1(protein_trajectory): - selector = Selector(protein_trajectory) - assert not selector.check_valid_json_settings('{"all": false, "index": [0, "1"]}') - - -def test_selector_returns_false_with_incorrect_json_setting_2(protein_trajectory): - selector = Selector(protein_trajectory) - assert not selector.check_valid_json_settings('{"all": False, "index": ["0", "1"]}') - - -@pytest.mark.xfail(reason="see docstring") -def test_selector_with_atom_fullname(protein_trajectory): - """For the moment, full names of atoms are not implemented - in the ChemicalSystem.""" - selector = Selector(protein_trajectory) - selector.update_settings( - { - "all": False, - "fullname": {"...LYS1.N": True, "...VAL2.O": True}, - } - ) - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 2 - - -@pytest.mark.xfail(reason="see docstring") -def test_selector_with_atom_name(protein_trajectory): - """At the moment the oxygen in water has the same - atom name as the oxygen in the protein. - We will have to decide if this is acceptable. - """ - selector = Selector(protein_trajectory) - selector.update_settings( - { - "all": False, - "name": {"N": True, "O": True}, - } - ) - atm_idxs = selector.get_idxs() - assert len(atm_idxs) == 258 From f4042e044ecf9f959a0ecbea4006b4ba063f4224 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Wed, 12 Feb 2025 11:01:42 +0000 Subject: [PATCH 06/37] Update configurators to use the new selection mechanism --- .../AtomSelectionConfigurator.py | 38 +++++++++---------- .../AtomTransmutationConfigurator.py | 16 ++++---- .../PartialChargeConfigurator.py | 10 +++-- .../UnitTests/test_reusable_selection.py | 28 ++++++++++++++ .../InputWidgets/AtomSelectionWidget.py | 14 +++---- 5 files changed, 69 insertions(+), 37 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py index 6f72767b6..8b8b3ab3b 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py @@ -14,7 +14,12 @@ # along with this program. If not, see . # +from json import JSONDecodeError + +import numpy as np + from MDANSE.Framework.Configurators.IConfigurator import IConfigurator +from MDANSE.Framework.AtomSelector.selector import ReusableSelection class AtomSelectionConfigurator(IConfigurator): @@ -28,7 +33,7 @@ class AtomSelectionConfigurator(IConfigurator): The defaults selection setting. """ - _default = '{"all": true}' + _default = '{}' def configure(self, value: str) -> None: """Configure an input value. @@ -39,6 +44,7 @@ def configure(self, value: str) -> None: The selection setting in a json readable format. """ trajConfig = self._configurable[self._dependencies["trajectory"]] + self.selector = ReusableSelection() if value is None: value = self._default @@ -47,20 +53,19 @@ def configure(self, value: str) -> None: self.error_status = "Invalid input value." return - selector = Selector(trajConfig["instance"]) - if not selector.check_valid_json_settings(value): + try: + self.selector.read_from_json(value) + except JSONDecodeError: self.error_status = "Invalid JSON string." return self["value"] = value - selector.load_from_json(value) - indices = selector.get_idxs() + self.selector.load_from_json(value) + indices = self.selector.select_in_trajectory(trajConfig["instance"]) self["flatten_indices"] = sorted(list(indices)) - trajConfig = self._configurable[self._dependencies["trajectory"]] - atoms = trajConfig["instance"].chemical_system.atom_list selectedAtoms = [atoms[idx] for idx in self["flatten_indices"]] @@ -87,12 +92,8 @@ def get_natoms(self) -> dict[str, int]: dict A dictionary of the number of atom per element. """ - nAtomsPerElement = {} - for v in self["names"]: - if v in nAtomsPerElement: - nAtomsPerElement[v] += 1 - else: - nAtomsPerElement[v] = 1 + names, counts = np.unique(self["names"], return_counts=True) + nAtomsPerElement = {names[n]:counts[n] for n in range(len(names))} return nAtomsPerElement @@ -131,14 +132,11 @@ def get_information(self) -> str: return "\n".join(info) + "\n" - def get_selector(self): + def get_selector(self) -> 'ReusableSelection': """ Returns ------- - Selector - The atom selector object initialised with the trajectories - chemical system. + ReusableSelection + the instance of the class which selects atoms in a trajectory """ - traj_config = self._configurable[self._dependencies["trajectory"]] - selector = Selector(traj_config["instance"]) - return selector + return self.selector diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py index 033dd597d..7aaa414cc 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py @@ -18,6 +18,7 @@ from MDANSE.Framework.Configurators.IConfigurator import IConfigurator from MDANSE.Chemistry import ATOMS_DATABASE +from MDANSE.Framework.AtomSelector.selector import ReusableSelection from MDANSE.MolecularDynamics.Trajectory import Trajectory @@ -33,31 +34,32 @@ def __init__(self, trajectory: Trajectory) -> None: system : ChemicalSystem The chemical system object. """ - self.selector = Selector(trajectory) + self.selector = ReusableSelection() self._original_map = {} for number, element in enumerate(trajectory.chemical_system.atom_list): self._original_map[number] = element self._new_map = {} + self._current_trajectory = trajectory def apply_transmutation( - self, selection_dict: dict[str, Union[bool, dict]], symbol: str + self, selection_string: str, symbol: str ) -> None: """With the selection dictionary update selector and then update the transmutation map. Parameters ---------- - selection_dict: dict[str, Union[bool, dict]] - The selection setting to get the indices to map the inputted - symbol. + selection_string: str + the JSON string of the selection operation to use. symbol: str The element to map the selected atoms to. """ if symbol not in ATOMS_DATABASE: raise ValueError(f"{symbol} not found in the atom database.") - self.selector.update_settings(selection_dict, reset_first=True) - for idx in self.selector.get_idxs(): + self.selector.read_from_json(selection_string) + indices = self.selector.select_in_trajectory(self._current_trajectory) + for idx in indices: self._new_map[idx] = symbol def get_setting(self) -> dict[int, str]: diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py index 2ca82f85c..544d6f7f4 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py @@ -17,6 +17,7 @@ import json from MDANSE.Framework.Configurators.IConfigurator import IConfigurator +from MDANSE.Framework.AtomSelector.selector import ReusableSelection from MDANSE.MolecularDynamics.Trajectory import Trajectory @@ -34,6 +35,7 @@ def __init__(self, trajectory: Trajectory) -> None: """ system = trajectory.chemical_system charges = trajectory.charges(0) + self._current_trajectory = trajectory self._original_map = {} for at_num, at in enumerate(system.atom_list): try: @@ -43,7 +45,7 @@ def __init__(self, trajectory: Trajectory) -> None: self._new_map = {} def update_charges( - self, selection_dict: dict[str, Union[bool, dict]], charge: float + self, selection_string: str, charge: float ) -> None: """With the selection dictionary update the selector and then update the partial charge map. @@ -56,8 +58,10 @@ def update_charges( charge: float The partial charge to map the selected atoms to. """ - self.selector.update_settings(selection_dict, reset_first=True) - for idx in self.selector.get_idxs(): + selector = ReusableSelection() + selector.read_from_json(selection_string) + indices = selector.select_in_trajectory(self._current_trajectory) + for idx in indices: self._new_map[idx] = charge def get_full_setting(self) -> dict[int, float]: diff --git a/MDANSE/Tests/UnitTests/test_reusable_selection.py b/MDANSE/Tests/UnitTests/test_reusable_selection.py index ef30169e3..6e83f78f4 100644 --- a/MDANSE/Tests/UnitTests/test_reusable_selection.py +++ b/MDANSE/Tests/UnitTests/test_reusable_selection.py @@ -39,6 +39,16 @@ def test_select_all(trajectory): assert len(selection) == n_atoms +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) +def test_empty_json_string_selects_all(trajectory): + traj_object = HDFTrajectoryInputData(trajectory) + n_atoms = len(traj_object.chemical_system.atom_list) + reusable_selection = ReusableSelection() + reusable_selection.read_from_json('{}') + selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + assert len(selection) == n_atoms + + @pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) def test_select_none(trajectory): traj_object = HDFTrajectoryInputData(trajectory) @@ -174,3 +184,21 @@ def test_select_pattern_selects_water(): another_selection.read_from_json(json_string) water_selection = another_selection.select_in_trajectory(traj_object.trajectory) assert len(water_selection) == 28746 + + +def test_selection_with_multiple_steps(): + """This tests if the ReusableSelection can select oxygen only in + the water molecules. It combines two steps: + 1. water is selected using rdkit pattern matching + 2. oxygen is selected using simple atom type matching; intersection of the selections is applied + The selection is then saved to a JSON string, loaded from the string and applied to the trajectory. + """ + traj_object = HDFTrajectoryInputData(traj_2vb1) + reusable_selection = ReusableSelection() + reusable_selection.set_selection(None, {'function_name': 'select_pattern', 'rdkit_pattern': "[#8X2;H2](~[H])~[H]"}) + reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'atom_types': ['O'], 'operation_type': 'intersection'}) + json_string = reusable_selection.convert_to_json() + another_selection = ReusableSelection() + another_selection.read_from_json(json_string) + water_oxygen_selection = another_selection.select_in_trajectory(traj_object.trajectory) + assert len(water_oxygen_selection) == int(28746/3) diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index ac23908e9..fb07409ea 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -30,6 +30,7 @@ from MDANSE_GUI.Tabs.Visualisers.View3D import View3D from MDANSE_GUI.MolecularViewer.MolecularViewer import MolecularViewerWithPicking from MDANSE.Framework.InputData.HDFTrajectoryInputData import HDFTrajectoryInputData +from MDANSE.Framework.AtomSelector.selector import ReusableSelection from .CheckableComboBox import CheckableComboBox @@ -66,7 +67,6 @@ class SelectionHelper(QDialog): def __init__( self, - selector, traj_data: tuple[str, HDFTrajectoryInputData], field: QLineEdit, parent, @@ -76,9 +76,6 @@ def __init__( """ Parameters ---------- - selector : Selector - The MDANSE selector initialized with the current chemical - system. traj_data : tuple[str, HDFTrajectoryInputData] A tuple of the trajectory data used to load the 3D viewer. field : QLineEdit @@ -88,10 +85,13 @@ def __init__( super().__init__(parent, *args, **kwargs) self.setWindowTitle(self._helper_title) - self.selector = selector + self.selector = ReusableSelection() + self.trajectory = traj_data[1].trajectory + self.system = self.trajectory.chemical_system self._field = field - self.settings = self.selector.settings - self.atm_full_names = self.selector.system.name_list + self.atm_full_names = self.system.name_list + self.molecule_names = [str(x) for x in self.system._clusters.keys()] + self.labels = [str(x) for x in self.system._labels.keys()] self.selection_textbox = QPlainTextEdit() self.selection_textbox.setReadOnly(True) From d2fa0cdade06dd99c45aab11e390d753ef28bf0f Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Wed, 12 Feb 2025 11:59:40 +0000 Subject: [PATCH 07/37] Consistently rename load_from_json method --- .../MDANSE/Framework/AtomSelector/selector.py | 2 +- .../Configurators/AtomSelectionConfigurator.py | 2 +- .../AtomTransmutationConfigurator.py | 2 +- .../Configurators/PartialChargeConfigurator.py | 2 +- .../Tests/UnitTests/test_reusable_selection.py | 18 +++++++++--------- 5 files changed, 13 insertions(+), 13 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index 82eef73d0..c41344a0f 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -109,7 +109,7 @@ def convert_to_json(self) -> str: """ return json.dumps(self.operations) - def read_from_json(self, json_string: str): + def load_from_json(self, json_string: str): """_summary_ Parameters diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py index 8b8b3ab3b..9c2c5aa90 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py @@ -54,7 +54,7 @@ def configure(self, value: str) -> None: return try: - self.selector.read_from_json(value) + self.selector.load_from_json(value) except JSONDecodeError: self.error_status = "Invalid JSON string." return diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py index 7aaa414cc..986258192 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py @@ -57,7 +57,7 @@ def apply_transmutation( if symbol not in ATOMS_DATABASE: raise ValueError(f"{symbol} not found in the atom database.") - self.selector.read_from_json(selection_string) + self.selector.load_from_json(selection_string) indices = self.selector.select_in_trajectory(self._current_trajectory) for idx in indices: self._new_map[idx] = symbol diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py index 544d6f7f4..6518c1c87 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/PartialChargeConfigurator.py @@ -59,7 +59,7 @@ def update_charges( The partial charge to map the selected atoms to. """ selector = ReusableSelection() - selector.read_from_json(selection_string) + selector.load_from_json(selection_string) indices = selector.select_in_trajectory(self._current_trajectory) for idx in indices: self._new_map[idx] = charge diff --git a/MDANSE/Tests/UnitTests/test_reusable_selection.py b/MDANSE/Tests/UnitTests/test_reusable_selection.py index 6e83f78f4..74a36b6b6 100644 --- a/MDANSE/Tests/UnitTests/test_reusable_selection.py +++ b/MDANSE/Tests/UnitTests/test_reusable_selection.py @@ -44,7 +44,7 @@ def test_empty_json_string_selects_all(trajectory): traj_object = HDFTrajectoryInputData(trajectory) n_atoms = len(traj_object.chemical_system.atom_list) reusable_selection = ReusableSelection() - reusable_selection.read_from_json('{}') + reusable_selection.load_from_json('{}') selection = reusable_selection.select_in_trajectory(traj_object.trajectory) assert len(selection) == n_atoms @@ -74,7 +74,7 @@ def test_json_saving_is_reversible(): reusable_selection.set_selection(None, {'function_name': 'invert_selection'}) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() - another_selection.read_from_json(json_string) + another_selection.load_from_json(json_string) json_string_2 = another_selection.convert_to_json() assert json_string == json_string_2 @@ -88,7 +88,7 @@ def test_selection_from_json_is_the_same_as_from_runtime(trajectory): selection = reusable_selection.select_in_trajectory(traj_object.trajectory) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() - another_selection.read_from_json(json_string) + another_selection.load_from_json(json_string) selection2 = another_selection.select_in_trajectory(traj_object.trajectory) print(f"original: {reusable_selection.operations}") print(f"another: {another_selection.operations}") @@ -112,7 +112,7 @@ def test_select_atoms_selects_by_element(): reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'atom_types': ['S', 'Sb', 'Cu']}) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() - another_selection.read_from_json(json_string) + another_selection.load_from_json(json_string) cusbs_selection = another_selection.select_in_trajectory(traj_object.trajectory) assert len(cusbs_selection) == 480 @@ -139,7 +139,7 @@ def test_select_atoms_selects_by_slice(): reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'index_slice': [470,510,5]}) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() - another_selection.read_from_json(json_string) + another_selection.load_from_json(json_string) overshoot_selection = another_selection.select_in_trajectory(traj_object.trajectory) assert len(overshoot_selection) == 2 @@ -158,7 +158,7 @@ def test_select_molecules_selects_protein(): reusable_selection.set_selection(None, {'function_name': 'select_molecules', 'molecule_names': ['C613 H959 N193 O185 S10']}) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() - another_selection.read_from_json(json_string) + another_selection.load_from_json(json_string) protein_selection = another_selection.select_in_trajectory(traj_object.trajectory) assert len(protein_selection) == 613 + 959 + 193 + 185 + 10 @@ -170,7 +170,7 @@ def test_select_molecules_inverted_selects_ions(): reusable_selection.set_selection(None, {'function_name': 'invert_selection'}) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() - another_selection.read_from_json(json_string) + another_selection.load_from_json(json_string) non_molecules_selection = another_selection.select_in_trajectory(traj_object.trajectory) assert all([traj_object.chemical_system.atom_list[index] in ['Na', 'Cl'] for index in non_molecules_selection]) @@ -181,7 +181,7 @@ def test_select_pattern_selects_water(): reusable_selection.set_selection(None, {'function_name': 'select_pattern', 'rdkit_pattern': "[#8X2;H2](~[H])~[H]"}) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() - another_selection.read_from_json(json_string) + another_selection.load_from_json(json_string) water_selection = another_selection.select_in_trajectory(traj_object.trajectory) assert len(water_selection) == 28746 @@ -199,6 +199,6 @@ def test_selection_with_multiple_steps(): reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'atom_types': ['O'], 'operation_type': 'intersection'}) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() - another_selection.read_from_json(json_string) + another_selection.load_from_json(json_string) water_oxygen_selection = another_selection.select_in_trajectory(traj_object.trajectory) assert len(water_oxygen_selection) == int(28746/3) From 10263c878aed6fd56852ef1800f732d36537a328 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Wed, 12 Feb 2025 12:24:19 +0000 Subject: [PATCH 08/37] Make all tests pass --- .../AtomTransmutation/test_transmutation.py | 24 ++++++++----------- .../UnitTests/TrajectoryEditor/test_editor.py | 2 +- 2 files changed, 11 insertions(+), 15 deletions(-) diff --git a/MDANSE/Tests/UnitTests/AtomTransmutation/test_transmutation.py b/MDANSE/Tests/UnitTests/AtomTransmutation/test_transmutation.py index 60e92cc42..eb3eaa671 100644 --- a/MDANSE/Tests/UnitTests/AtomTransmutation/test_transmutation.py +++ b/MDANSE/Tests/UnitTests/AtomTransmutation/test_transmutation.py @@ -29,7 +29,7 @@ def test_atom_transmutation_return_dict_with_transmutations_with_incorrect_eleme atm_transmuter = AtomTransmuter(protein_trajectory) with pytest.raises(ValueError): atm_transmuter.apply_transmutation( - {"all": False, "element": {"S": True}}, "CCC" + '{"0": {"function_name": "select_atoms", "atom_types": ["S"]}}', "CCC" ) @@ -37,7 +37,7 @@ def test_atom_transmutation_return_dict_with_transmutations_with_s_element_trans protein_trajectory, ): atm_transmuter = AtomTransmuter(protein_trajectory) - atm_transmuter.apply_transmutation({"all": False, "element": {"S": True}}, "C") + atm_transmuter.apply_transmutation('{"0": {"function_name": "select_atoms", "atom_types": ["S"]}}', "C") mapping = atm_transmuter.get_setting() assert mapping == { 98: "C", @@ -57,8 +57,8 @@ def test_atom_transmutation_return_dict_with_transmutations_with_s_element_trans protein_trajectory, ): atm_transmuter = AtomTransmuter(protein_trajectory) - atm_transmuter.apply_transmutation({"all": False, "element": {"S": True}}, "C") - atm_transmuter.apply_transmutation({"all": False, "index": {98: True}}, "N") + atm_transmuter.apply_transmutation('{"0": {"function_name": "select_atoms", "atom_types": ["S"]}}', "C") + atm_transmuter.apply_transmutation('{"0": {"function_name": "select_atoms", "index_list": [98]}}', "N") mapping = atm_transmuter.get_setting() assert mapping == { 98: "N", @@ -78,8 +78,8 @@ def test_atom_transmutation_return_dict_with_transmutations_with_s_element_trans protein_trajectory, ): atm_transmuter = AtomTransmuter(protein_trajectory) - atm_transmuter.apply_transmutation({"all": False, "element": {"S": True}}, "C") - atm_transmuter.apply_transmutation({"all": False, "index": {98: True}}, "S") + atm_transmuter.apply_transmutation('{"0": {"function_name": "select_atoms", "atom_types": ["S"]}}', "C") + atm_transmuter.apply_transmutation('{"0": {"function_name": "select_atoms", "index_list": [98]}}', "S") mapping = atm_transmuter.get_setting() assert mapping == { 175: "C", @@ -98,10 +98,8 @@ def test_atom_transmutation_return_dict_with_transmutations_with_s_element_trans protein_trajectory, ): atm_transmuter = AtomTransmuter(protein_trajectory) - atm_transmuter.apply_transmutation({"all": False, "element": {"S": True}}, "C") - atm_transmuter.apply_transmutation( - {"all": False, "index": {98: True, 99: True}}, "S" - ) + atm_transmuter.apply_transmutation('{"0": {"function_name": "select_atoms", "atom_types": ["S"]}}', "C") + atm_transmuter.apply_transmutation('{"0": {"function_name": "select_atoms", "index_list": [98, 99]}}', "S") mapping = atm_transmuter.get_setting() assert mapping == { 99: "S", @@ -119,10 +117,8 @@ def test_atom_transmutation_return_dict_with_transmutations_with_s_element_trans def test_atom_transmutation_return_empty_dict_after_reset(protein_trajectory): atm_transmuter = AtomTransmuter(protein_trajectory) - atm_transmuter.apply_transmutation({"all": False, "element": {"S": True}}, "C") - atm_transmuter.apply_transmutation( - {"all": False, "index": {98: True, 99: True}}, "S" - ) + atm_transmuter.apply_transmutation('{"0": {"function_name": "select_atoms", "atom_types": ["S"]}}', "C") + atm_transmuter.apply_transmutation('{"0": {"function_name": "select_atoms", "index_list": [98, 99]}}', "S") atm_transmuter.reset_setting() mapping = atm_transmuter.get_setting() assert mapping == {} diff --git a/MDANSE/Tests/UnitTests/TrajectoryEditor/test_editor.py b/MDANSE/Tests/UnitTests/TrajectoryEditor/test_editor.py index dfa27ed05..e609a7ff4 100644 --- a/MDANSE/Tests/UnitTests/TrajectoryEditor/test_editor.py +++ b/MDANSE/Tests/UnitTests/TrajectoryEditor/test_editor.py @@ -112,7 +112,7 @@ def test_editor_atoms(): parameters["output_files"] = (temp_name, 64, 128, "gzip", "INFO") parameters["trajectory"] = short_traj parameters["frames"] = (0, 501, 1) - parameters["atom_selection"] = '{"all": false, "element": ["H"]}' + parameters["atom_selection"] = '{"0": {"function_name": "select_atoms", "atom_types": ["H"]}}' temp = IJob.create("TrajectoryEditor") temp.run(parameters, status=True) assert path.exists(temp_name + ".mdt") From 2abad5352a11e4e151760fac9c10f56886b6824c Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Thu, 13 Feb 2025 08:10:53 +0000 Subject: [PATCH 09/37] Add a selection mechanism via widgets --- .../MDANSE/Framework/AtomSelector/selector.py | 45 ++++ .../InputWidgets/AtomSelectionWidget.py | 226 +++++++++--------- .../MDANSE_GUI/Widgets/SelectionWidgets.py | 78 ++++++ 3 files changed, 238 insertions(+), 111 deletions(-) create mode 100644 MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index c41344a0f..131b687ee 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -74,6 +74,49 @@ def set_selection( number = len(self.operations) self.operations[number] = function_parameters + def validate_selection_string(self, json_string: str, trajectory: Trajectory, current_selection: Set[int]) -> bool: + """Checks if the selection operation encoded in the input JSON string + will add any new atoms to the current selection on the given trajectory. + + Parameters + ---------- + json_string : str + new selection operation in a JSON string + trajectory : Trajectory + a trajectory instance for which current_selection is defined + current_selection : Set[int] + set of currently selected atom indices + + Returns + ------- + bool + True if the selection adds atoms, False otherwise + """ + function_parameters = json.loads(json_string) + if len(self.operations) == 0: + return True + function_name = function_parameters.get("function_name", "select_all") + if function_name == "invert_selection": + selection = invert_selection(trajectory, current_selection) + return True + else: + operation_type = function_parameters.get("operation_type", "union") + function = function_lookup[function_name] + temp_selection = function(trajectory, **function_parameters) + if operation_type == "union": + selection = selection.union(temp_selection) + elif operation_type == "intersection": + selection = selection.intersection(temp_selection) + elif operation_type == "difference": + selection = selection.difference(temp_selection) + else: + selection = temp_selection + if len(selection.difference(current_selection)) > 0 and operation_type == "union": + return True + elif len(current_selection.difference(selection)) > 0 and operation_type != "union": + return True + return False + def select_in_trajectory(self, trajectory: Trajectory) -> Set[int]: selection = set() self.all_idxs = set(range(len(trajectory.chemical_system.atom_list))) @@ -93,6 +136,8 @@ def select_in_trajectory(self, trajectory: Trajectory) -> Set[int]: selection = selection.union(temp_selection) elif operation_type == "intersection": selection = selection.intersection(temp_selection) + elif operation_type == "difference": + selection = selection.difference(temp_selection) else: selection = temp_selection return selection diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index fb07409ea..eea75da4a 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -13,25 +13,88 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . # -from qtpy.QtCore import Qt, Slot + +from typing import Set +import json + +from qtpy.QtCore import Slot, Signal +from qtpy.QtGui import QStandardItemModel, QStandardItem from qtpy.QtWidgets import ( QLineEdit, QPushButton, QDialog, - QCheckBox, QVBoxLayout, QHBoxLayout, QGroupBox, QLabel, QPlainTextEdit, QWidget, + QListView, ) from MDANSE_GUI.InputWidgets.WidgetBase import WidgetBase from MDANSE_GUI.Tabs.Visualisers.View3D import View3D from MDANSE_GUI.MolecularViewer.MolecularViewer import MolecularViewerWithPicking from MDANSE.Framework.InputData.HDFTrajectoryInputData import HDFTrajectoryInputData from MDANSE.Framework.AtomSelector.selector import ReusableSelection -from .CheckableComboBox import CheckableComboBox +from MDANSE_GUI.Widgets.SelectionWidgets import AllAtomSelection + + +VALID_SELECTION = "Valid selection" +USELESS_SELECTION = "Selection did not change. This operation is not needed." +MALFORMED_SELECTION = "This is not a valid JSON string." + + +class SelectionModel(QStandardItemModel): + selection_changed = Signal() + + def __init__(self, trajectory): + super().__init__(None) + self._trajectory = trajectory + self._selection = ReusableSelection() + self._current_selection = set() + + def rebuild_selection(self, last_operation: str): + self._selection = ReusableSelection() + self._current_selection = set() + for row in range(self.rowCount()): + index = self.index(row, 0) + item = self.itemFromIndex(index) + json_string = item.text() + self._selection.load_from_json(json_string) + self._current_selection = self._selection.select_in_trajectory(self._trajectory) + if last_operation: + try: + valid = self._selection.validate_selection_string( + last_operation, self._trajectory, self._current_selection + ) + except json.JSONDecodeError: + return MALFORMED_SELECTION + if valid: + self._selection.load_from_json(json_string) + return VALID_SELECTION + else: + return USELESS_SELECTION + + def current_selection(self, last_operation: str = "") -> Set[int]: + self.rebuild_selection(last_operation) + return self._selection.select_in_trajectory(self._trajectory) + + def current_steps(self) -> str: + result = {} + for row in range(self.rowCount()): + index = self.index(row, 0) + item = self.itemFromIndex(index) + json_string = item.text() + python_object = json.loads(json_string) + result[row] = python_object + return json.dumps(result) + + @Slot(str) + def accept_from_widget(self, json_string: str): + new_item = QStandardItem(json_string) + new_item.setEditable(False) + self.appendRow(new_item) + self.selection_changed.emit() class SelectionHelper(QDialog): @@ -85,9 +148,9 @@ def __init__( super().__init__(parent, *args, **kwargs) self.setWindowTitle(self._helper_title) - self.selector = ReusableSelection() self.trajectory = traj_data[1].trajectory self.system = self.trajectory.chemical_system + self.selection_model = SelectionModel(self.trajectory) self._field = field self.atm_full_names = self.system.name_list self.molecule_names = [str(x) for x in self.system._clusters.keys()] @@ -114,7 +177,6 @@ def __init__( helper_layout.addLayout(layout) self.setLayout(helper_layout) - self.update_others() self.all_selection = True self.selected = set([]) @@ -180,81 +242,41 @@ def left_widgets(self) -> list[QWidget]: List of QWidgets to add to the left layout from create_layouts. """ - match_exists = self.selector.match_exists select = QGroupBox("selection") select_layout = QVBoxLayout() - self.check_boxes = [] - self.combo_boxes = [] - - for k, v in self.settings.items(): - - if isinstance(v, bool): - check_layout = QHBoxLayout() - checkbox = QCheckBox() - checkbox.setChecked(v) - checkbox.setLayoutDirection(Qt.RightToLeft) - label = QLabel(self._cbox_text[k]) - checkbox.setObjectName(k) - checkbox.stateChanged.connect(self.update_others) - if not match_exists[k]: - checkbox.setEnabled(False) - label.setStyleSheet("color: grey;") - self.check_boxes.append(checkbox) - check_layout.addWidget(label) - check_layout.addWidget(checkbox) - select_layout.addLayout(check_layout) - - elif isinstance(v, dict): - combo_layout = QHBoxLayout() - combo = CheckableComboBox() - items = [str(i) for i in v.keys() if match_exists[k][i]] - # we blocksignals here as there can be some - # performance issues with a large number of items - combo.model().blockSignals(True) - combo.addItems(items) - combo.model().blockSignals(False) - combo.setObjectName(k) - combo.model().dataChanged.connect(self.update_others) - label = QLabel(self._cbox_text[k]) - if len(items) == 0: - combo.setEnabled(False) - label.setStyleSheet("color: grey;") - self.combo_boxes.append(combo) - combo_layout.addWidget(label) - combo_layout.addWidget(combo) - select_layout.addLayout(combo_layout) + self.selection_widgets = [AllAtomSelection(self)] + + for widget in self.selection_widgets: + select_layout.addWidget(widget) + widget.new_selection.connect(self.selection_model.accept_from_widget) invert_layout = QHBoxLayout() - label = QLabel("Invert selection:") + label = QLabel("Current selection:") + self.selection_line = QLineEdit("", self) apply = QPushButton("Apply") - apply.clicked.connect(self.invert_selection) + apply.clicked.connect(self.append_selection) invert_layout.addWidget(label) + invert_layout.addWidget(self.selection_line) invert_layout.addWidget(apply) select_layout.addLayout(invert_layout) select.setLayout(select_layout) - return [select] - def update_others(self) -> None: - """Using the checkbox and combobox widgets: update the settings, - get the selection and update the textedit box with details of - the current selection and the 3d view to match the selection. - """ - for check_box in self.check_boxes: - self.settings[check_box.objectName()] = check_box.isChecked() - for combo_box in self.combo_boxes: - for i in range(combo_box.n_items): - txt = combo_box.text[i] - if combo_box.objectName() == "index": - key = int(txt) - else: - key = txt - self.settings[combo_box.objectName()][key] = combo_box.checked[i] - - self.selector.update_settings(self.settings) - self.selected = self.selector.get_idxs() + preview = QGroupBox("Selection operations") + preview_layout = QVBoxLayout() + preview.setLayout(preview_layout) + self.selection_operations_view = QListView(preview) + self.selection_operations_view.setDragEnabled(True) + self.selection_operations_view.setModel(self.selection_model) + self.selection_model.selection_changed.connect(self.recalculate_selection) + preview_layout.addWidget(self.selection_operations_view) + return [select, preview] + + @Slot() + def recalculate_selection(self): + self.selected = self.selection_model.current_selection() self.view_3d._viewer.change_picked(self.selected) self.update_selection_textbox() @@ -267,46 +289,34 @@ def update_from_3d_view(self, selection: set[int]) -> None: selection : set[int] Selection indexes from the 3d view. """ - self.selector.update_with_idxs(selection) - self.settings = self.selector.settings - self.update_selection_widgets() - self.selected = self.selector.get_idxs() self.update_selection_textbox() - def invert_selection(self): - """Inverts the selection.""" - self.selected = self.selector.all_idxs - self.selected - self.selector.update_with_idxs(self.selected) - self.settings = self.selector.settings - self.update_selection_widgets() + @Slot(str) + def update_operation(self, new_json_string: str): + self.selection_line.setText(new_json_string) self.view_3d._viewer.change_picked(self.selected) self.update_selection_textbox() - def update_selection_widgets(self) -> None: - """Updates the selection widgets so that it matches the full - setting. - """ - for check_box in self.check_boxes: - check_box.blockSignals(True) - if self.settings[check_box.objectName()]: - check_box.setCheckState(Qt.Checked) - else: - check_box.setCheckState(Qt.Unchecked) - check_box.blockSignals(False) - for combo_box in self.combo_boxes: - combo_box.model().blockSignals(True) - for i in range(combo_box.n_items): - txt = combo_box.text[i] - if combo_box.objectName() == "index": - key = int(txt) - else: - key = txt - combo_box.set_item_checked_state( - i, self.settings[combo_box.objectName()][key] - ) - combo_box.update_all_selected() - combo_box.update_line_edit() - combo_box.model().blockSignals(False) + @Slot() + def append_selection(self): + self.selection_line.setStyleSheet("") + self.selection_line.setToolTip("") + selection_text = self.selection_line.text() + validation = self.selection_model.rebuild_selection(selection_text) + if validation == MALFORMED_SELECTION: + self.selection_line.setStyleSheet( + "QWidget#InputWidget { background-color:rgb(180,20,180); font-weight: bold }" + ) + self.selection_line.setToolTip(validation) + elif validation == USELESS_SELECTION: + self.selection_line.setStyleSheet( + "QWidget#InputWidget { background-color:rgb(180,20,180); font-weight: bold }" + ) + self.selection_line.setToolTip(validation) + elif validation == VALID_SELECTION: + self.selection_model.appendRow(QStandardItem(selection_text)) + self.view_3d._viewer.change_picked(self.selected) + self.update_selection_textbox() def update_selection_textbox(self) -> None: """Update the selection textbox.""" @@ -320,16 +330,11 @@ def apply(self) -> None: """Set the field of the AtomSelectionWidget to the currently chosen setting in this widget. """ - self.selector.update_settings(self.settings) - self._field.setText(self.selector.settings_to_json()) + self._field.setText(self.selection_model.current_steps()) def reset(self) -> None: """Resets the helper to the default state.""" - self.selector.reset_settings() - self.selector.settings["all"] = self.all_selection - self.settings = self.selector.settings - self.update_selection_widgets() - self.selected = self.selector.get_idxs() + self.selection_model.clear() self.view_3d._viewer.change_picked(self.selected) self.update_selection_textbox() @@ -376,8 +381,7 @@ def create_helper( SelectionHelper Create and return the selection helper QDialog. """ - selector = self._configurator.get_selector() - return SelectionHelper(selector, traj_data, self._field, self._base) + return SelectionHelper(traj_data, self._field, self._base) @Slot() def helper_dialog(self) -> None: diff --git a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py new file mode 100644 index 000000000..8c051a0da --- /dev/null +++ b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py @@ -0,0 +1,78 @@ +# This file is part of MDANSE_GUI. +# +# MDANSE_GUI is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . +# + +from typing import Dict, Any +import json + +from qtpy.QtCore import Signal, Slot +from qtpy.QtWidgets import QGroupBox, QHBoxLayout, QPushButton, QComboBox + +from MDANSE.Framework.AtomSelector.general_selection import ( + select_all, + select_none, + invert_selection, +) + + +class BasicSelectionWidget(QGroupBox): + new_selection = Signal(str) + + def __init__(self, parent=None, widget_label="Atom selection widget"): + super().__init__(parent) + layout = QHBoxLayout() + self.setLayout(layout) + self.setTitle(widget_label) + self.add_specific_widgets() + self.add_standard_widgets() + + def parameter_dictionary(self) -> Dict[str, Any]: + return {} + + def add_specific_widgets(self): + return None + + def add_standard_widgets(self): + self.mode_box = QComboBox(self) + self.mode_box.setEditable(False) + self.mode_box.addItems( + ["Add (union)", "Filter (intersection)", "Remove (difference)"] + ) + self.commit_button = QPushButton("Apply", self) + layout = self.layout() + layout.addWidget(self.mode_box) + layout.addWidget(self.commit_button) + self.commit_button.clicked.connect(self.create_selection) + + def get_mode(self) -> str: + if self.mode_box.currentIndex == 0: + return "difference" + elif self.mode_box.currentIndex == 1: + return "intersection" + else: + return "union" + + def create_selection(self): + funtion_parameters = self.parameter_dictionary() + funtion_parameters["operation_type"] = self.get_mode() + self.new_selection.emit(json.dumps(funtion_parameters)) + + +class AllAtomSelection(BasicSelectionWidget): + def __init__(self, parent=None, widget_label="ALL ATOMS"): + super().__init__(parent, widget_label) + + def parameter_dictionary(self): + return {"function_name": "select_all"} From 5371440322137dab1bf4a14c40d61dd4ff8f7e31 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Thu, 13 Feb 2025 09:46:43 +0000 Subject: [PATCH 10/37] Add atom selection widgets --- .../MDANSE/Framework/AtomSelector/selector.py | 4 +- .../InputWidgets/AtomSelectionWidget.py | 15 ++- .../InputWidgets/CheckableComboBox.py | 6 + .../MDANSE_GUI/Widgets/SelectionWidgets.py | 112 +++++++++++++++++- 4 files changed, 129 insertions(+), 8 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index 131b687ee..57be31fd2 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -74,7 +74,9 @@ def set_selection( number = len(self.operations) self.operations[number] = function_parameters - def validate_selection_string(self, json_string: str, trajectory: Trajectory, current_selection: Set[int]) -> bool: + def validate_selection_string( + self, json_string: str, trajectory: Trajectory, current_selection: Set[int] + ) -> bool: """Checks if the selection operation encoded in the input JSON string will add any new atoms to the current selection on the given trajectory. diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index eea75da4a..e8f356d2c 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -36,7 +36,11 @@ from MDANSE_GUI.MolecularViewer.MolecularViewer import MolecularViewerWithPicking from MDANSE.Framework.InputData.HDFTrajectoryInputData import HDFTrajectoryInputData from MDANSE.Framework.AtomSelector.selector import ReusableSelection -from MDANSE_GUI.Widgets.SelectionWidgets import AllAtomSelection +from MDANSE_GUI.Widgets.SelectionWidgets import ( + AllAtomSelection, + AtomSelection, + IndexSelection, +) VALID_SELECTION = "Valid selection" @@ -246,7 +250,11 @@ def left_widgets(self) -> list[QWidget]: select = QGroupBox("selection") select_layout = QVBoxLayout() - self.selection_widgets = [AllAtomSelection(self)] + self.selection_widgets = [ + AllAtomSelection(self), + AtomSelection(self, self.trajectory), + IndexSelection(self), + ] for widget in self.selection_widgets: select_layout.addWidget(widget) @@ -335,8 +343,7 @@ def apply(self) -> None: def reset(self) -> None: """Resets the helper to the default state.""" self.selection_model.clear() - self.view_3d._viewer.change_picked(self.selected) - self.update_selection_textbox() + self.recalculate_selection() class AtomSelectionWidget(WidgetBase): diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/CheckableComboBox.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/CheckableComboBox.py index 327fca0b3..bc21c7f22 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/CheckableComboBox.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/CheckableComboBox.py @@ -42,6 +42,12 @@ def __init__(self, *args, **kwargs): self.addItem("select all", underline=True) self.lineEdit().setText("") + def clear(self): + result = super().clear() + self.addItem("select all", underline=True) + self.lineEdit().setText("") + return result + def eventFilter(self, a0: Union[QObject, None], a1: Union[QEvent, None]) -> bool: """Updates the check state of the items and the lineEdit. diff --git a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py index 8c051a0da..9d1454490 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py @@ -17,9 +17,19 @@ from typing import Dict, Any import json +import numpy as np from qtpy.QtCore import Signal, Slot -from qtpy.QtWidgets import QGroupBox, QHBoxLayout, QPushButton, QComboBox +from qtpy.QtWidgets import ( + QGroupBox, + QHBoxLayout, + QPushButton, + QComboBox, + QLabel, + QLineEdit, +) +from MDANSE_GUI.InputWidgets.CheckableComboBox import CheckableComboBox +from MDANSE.MolecularDynamics.Trajectory import Trajectory from MDANSE.Framework.AtomSelector.general_selection import ( select_all, select_none, @@ -57,9 +67,9 @@ def add_standard_widgets(self): self.commit_button.clicked.connect(self.create_selection) def get_mode(self) -> str: - if self.mode_box.currentIndex == 0: + if self.mode_box.currentIndex() == 2: return "difference" - elif self.mode_box.currentIndex == 1: + elif self.mode_box.currentIndex() == 1: return "intersection" else: return "union" @@ -74,5 +84,101 @@ class AllAtomSelection(BasicSelectionWidget): def __init__(self, parent=None, widget_label="ALL ATOMS"): super().__init__(parent, widget_label) + def add_specific_widgets(self): + layout = self.layout() + layout.addWidget(QLabel("Add/remove ALL atoms")) + def parameter_dictionary(self): return {"function_name": "select_all"} + + +class AtomSelection(BasicSelectionWidget): + def __init__( + self, parent=None, trajectory: Trajectory = None, widget_label="Select atoms" + ): + self.atom_types = [] + 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) + ) + self.selection_types = [] + self.selection_keyword = "" + if self.atom_types: + self.selection_types += ["type"] + if self.atom_names: + self.selection_types += ["name"] + super().__init__(parent, widget_label) + + def add_specific_widgets(self): + layout = self.layout() + layout.addWidget(QLabel("Select atoms by atom")) + self.selection_type_combo = QComboBox(self) + self.selection_type_combo.addItems(self.selection_types) + self.selection_type_combo.setEditable(False) + layout.addWidget(self.selection_type_combo) + self.selection_field = CheckableComboBox(self) + layout.addWidget(self.selection_field) + self.selection_type_combo.currentTextChanged.connect(self.switch_mode) + self.selection_type_combo.setCurrentText(self.selection_types[0]) + self.switch_mode(self.selection_types[0]) + + @Slot(str) + def switch_mode(self, new_mode: str): + self.selection_field.clear() + if new_mode == "type": + 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_keyword = "atom_names" + + def parameter_dictionary(self): + function_parameters = {"function_name": "select_atoms"} + selection = self.selection_field.checked_values() + function_parameters[self.selection_keyword] = selection + return function_parameters + + +class IndexSelection(BasicSelectionWidget): + def __init__(self, parent=None, widget_label="Select atoms"): + super().__init__(parent, widget_label) + self.selection_keyword = "index_list" + + def add_specific_widgets(self): + layout = self.layout() + layout.addWidget(QLabel("Select atoms by index")) + self.selection_type_combo = QComboBox(self) + self.selection_type_combo.addItems( + [ + "list", + "range", + "slice", + ] + ) + self.selection_type_combo.setEditable(False) + layout.addWidget(self.selection_type_combo) + self.selection_field = QLineEdit(self) + layout.addWidget(self.selection_field) + self.selection_type_combo.currentTextChanged.connect(self.switch_mode) + + @Slot(str) + def switch_mode(self, new_mode: str): + self.selection_field.setText("") + if new_mode == "list": + self.selection_field.setPlaceholderText("0,1,2") + self.selection_keyword = "index_list" + if new_mode == "range": + self.selection_field.setPlaceholderText("0-20") + self.selection_keyword = "index_range" + if new_mode == "slice": + self.selection_field.setPlaceholderText("first:last:step") + self.selection_keyword = "index_slice" + + def parameter_dictionary(self): + function_parameters = {"function_name": "select_atoms"} + selection = self.selection_field.text() + function_parameters[self.selection_keyword] = selection + return function_parameters From d8604f992c8aa2fe7b5160df8e51591dfb6398b9 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Fri, 14 Feb 2025 09:40:27 +0000 Subject: [PATCH 11/37] Add string parsing to index range selection widget --- .../Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py | 4 +++- MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py | 5 ++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index e8f356d2c..858ff7be9 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -60,11 +60,13 @@ def __init__(self, trajectory): def rebuild_selection(self, last_operation: str): self._selection = ReusableSelection() self._current_selection = set() + total_dict = {} for row in range(self.rowCount()): index = self.index(row, 0) item = self.itemFromIndex(index) json_string = item.text() - self._selection.load_from_json(json_string) + total_dict[row] = json.loads(json_string) + self._selection.load_from_json(json.dumps(total_dict)) self._current_selection = self._selection.select_in_trajectory(self._trajectory) if last_operation: try: diff --git a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py index 9d1454490..9d7e734b5 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py @@ -170,15 +170,18 @@ def switch_mode(self, new_mode: str): if new_mode == "list": self.selection_field.setPlaceholderText("0,1,2") self.selection_keyword = "index_list" + self.selection_separator = ',' if new_mode == "range": self.selection_field.setPlaceholderText("0-20") self.selection_keyword = "index_range" + self.selection_separator = '-' if new_mode == "slice": self.selection_field.setPlaceholderText("first:last:step") self.selection_keyword = "index_slice" + self.selection_separator = ':' def parameter_dictionary(self): function_parameters = {"function_name": "select_atoms"} selection = self.selection_field.text() - function_parameters[self.selection_keyword] = selection + function_parameters[self.selection_keyword] = [int(x) for x in selection.split(self.selection_separator)] return function_parameters From edaedc773ed7e52cedbdf4e736792cb589204a66 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Fri, 14 Feb 2025 14:19:06 +0000 Subject: [PATCH 12/37] Add new selectors and change layout in AtomSelectionWidget --- .../Framework/AtomSelector/group_selection.py | 2 +- .../InputWidgets/AtomSelectionWidget.py | 57 ++++----- .../InputWidgets/CheckableComboBox.py | 4 + .../MDANSE_GUI/Widgets/SelectionWidgets.py | 109 ++++++++++++++++-- 4 files changed, 129 insertions(+), 43 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py index 6b1ecf42f..ad94a9994 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py @@ -41,7 +41,7 @@ def select_labels(trajectory: Trajectory, **function_parameters: Dict[str, Any]) atom_labels = function_parameters.get("atom_labels", None) for label in atom_labels: if label in system._labels: - selection = selection.union(reduce(list.__add__, system._labels[label])) + selection = selection.union(system._labels[label]) return selection diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index 858ff7be9..b8473290a 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -30,6 +30,8 @@ QPlainTextEdit, QWidget, QListView, + QAbstractItemView, + QScrollArea, ) from MDANSE_GUI.InputWidgets.WidgetBase import WidgetBase from MDANSE_GUI.Tabs.Visualisers.View3D import View3D @@ -40,6 +42,9 @@ AllAtomSelection, AtomSelection, IndexSelection, + MoleculeSelection, + PatternSelection, + LabelSelection, ) @@ -110,29 +115,9 @@ class SelectionHelper(QDialog): ---------- _helper_title : str The title of the helper dialog window. - _cbox_text : dict - The dictionary that maps the selector settings to text used in - the helper dialog. """ _helper_title = "Atom selection helper" - _cbox_text = { - "all": "All atoms (excl. dummy atoms):", - "dummy": "All dummy atoms:", - "hs_on_heteroatom": "Hs on heteroatoms:", - "primary_amine": "Primary amine groups:", - "hydroxy": "Hydroxy groups:", - "methyl": "Methyl groups:", - "phosphate": "Phosphate groups:", - "sulphate": "Sulphate groups:", - "thiol": "Thiol groups:", - "water": "Water molecules:", - "hs_on_element": "Hs on elements:", - "element": "Elements:", - "name": "Atom name:", - "fullname": "Atom fullname:", - "index": "Indexes:", - } def __init__( self, @@ -176,16 +161,19 @@ def __init__( for button in self.create_buttons(): bottom.addWidget(button) - layouts[-1].addLayout(bottom) - helper_layout = QHBoxLayout() - for layout in layouts: - helper_layout.addLayout(layout) + sub_layout = QVBoxLayout() + helper_layout.addLayout(layouts[0]) + helper_layout.addLayout(sub_layout) + for layout in layouts[1:]: + sub_layout.addLayout(layout) + sub_layout.addLayout(bottom) self.setLayout(helper_layout) self.all_selection = True self.selected = set([]) + self.reset() def closeEvent(self, a0): """Hide the window instead of closing. Some issues occur in the @@ -224,7 +212,7 @@ def create_layouts(self) -> list[QVBoxLayout]: for widget in self.left_widgets(): left.addWidget(widget) - right = QVBoxLayout() + right = QHBoxLayout() for widget in self.right_widgets(): right.addWidget(widget) @@ -238,7 +226,7 @@ def right_widgets(self) -> list[QWidget]: List of QWidgets to add to the right layout from create_layouts. """ - return [self.selection_textbox] + return [self.selection_operations_view, self.selection_textbox] def left_widgets(self) -> list[QWidget]: """ @@ -251,11 +239,16 @@ def left_widgets(self) -> list[QWidget]: select = QGroupBox("selection") select_layout = QVBoxLayout() + scroll_area = QScrollArea() + scroll_area.setLayout(select_layout) self.selection_widgets = [ AllAtomSelection(self), AtomSelection(self, self.trajectory), IndexSelection(self), + MoleculeSelection(self, self.trajectory), + PatternSelection(self), + LabelSelection(self, self.trajectory), ] for widget in self.selection_widgets: @@ -274,15 +267,15 @@ def left_widgets(self) -> list[QWidget]: select.setLayout(select_layout) - preview = QGroupBox("Selection operations") - preview_layout = QVBoxLayout() - preview.setLayout(preview_layout) - self.selection_operations_view = QListView(preview) + self.selection_operations_view = QListView(self) self.selection_operations_view.setDragEnabled(True) + self.selection_operations_view.setAcceptDrops(True) + self.selection_operations_view.setDragDropMode( + QAbstractItemView.DragDropMode.InternalMove + ) self.selection_operations_view.setModel(self.selection_model) self.selection_model.selection_changed.connect(self.recalculate_selection) - preview_layout.addWidget(self.selection_operations_view) - return [select, preview] + return [select] @Slot() def recalculate_selection(self): diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/CheckableComboBox.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/CheckableComboBox.py index bc21c7f22..9a621d33a 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/CheckableComboBox.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/CheckableComboBox.py @@ -44,6 +44,10 @@ def __init__(self, *args, **kwargs): def clear(self): result = super().clear() + self.items = [] + self.checked = [] + self.text = [] + self.select_all_item = None self.addItem("select all", underline=True) self.lineEdit().setText("") return result diff --git a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py index 9d7e734b5..91bd9446d 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py @@ -30,11 +30,6 @@ from MDANSE_GUI.InputWidgets.CheckableComboBox import CheckableComboBox from MDANSE.MolecularDynamics.Trajectory import Trajectory -from MDANSE.Framework.AtomSelector.general_selection import ( - select_all, - select_none, - invert_selection, -) class BasicSelectionWidget(QGroupBox): @@ -143,7 +138,7 @@ def parameter_dictionary(self): class IndexSelection(BasicSelectionWidget): - def __init__(self, parent=None, widget_label="Select atoms"): + def __init__(self, parent=None, widget_label="Index selection"): super().__init__(parent, widget_label) self.selection_keyword = "index_list" @@ -170,18 +165,112 @@ def switch_mode(self, new_mode: str): if new_mode == "list": self.selection_field.setPlaceholderText("0,1,2") self.selection_keyword = "index_list" - self.selection_separator = ',' + self.selection_separator = "," if new_mode == "range": self.selection_field.setPlaceholderText("0-20") self.selection_keyword = "index_range" - self.selection_separator = '-' + self.selection_separator = "-" if new_mode == "slice": self.selection_field.setPlaceholderText("first:last:step") self.selection_keyword = "index_slice" - self.selection_separator = ':' + self.selection_separator = ":" def parameter_dictionary(self): function_parameters = {"function_name": "select_atoms"} selection = self.selection_field.text() - function_parameters[self.selection_keyword] = [int(x) for x in selection.split(self.selection_separator)] + function_parameters[self.selection_keyword] = [ + int(x) for x in selection.split(self.selection_separator) + ] + return function_parameters + + +class MoleculeSelection(BasicSelectionWidget): + def __init__( + self, + parent=None, + trajectory: Trajectory = None, + widget_label="Select molecules", + ): + self.molecule_names = [] + if trajectory: + self.molecule_names = list(trajectory.chemical_system._clusters.keys()) + super().__init__(parent, widget_label) + + def add_specific_widgets(self): + layout = self.layout() + layout.addWidget(QLabel("Select molecules named: ")) + self.selection_field = CheckableComboBox(self) + layout.addWidget(self.selection_field) + self.selection_field.addItems(self.molecule_names) + + def parameter_dictionary(self): + function_parameters = {"function_name": "select_molecules"} + selection = self.selection_field.checked_values() + function_parameters["molecule_names"] = selection + return function_parameters + + +class LabelSelection(BasicSelectionWidget): + def __init__( + self, + parent=None, + trajectory: Trajectory = None, + widget_label="Select by label", + ): + self.labels = [] + if trajectory: + self.labels = list(trajectory.chemical_system._labels.keys()) + super().__init__(parent, widget_label) + + def add_specific_widgets(self): + layout = self.layout() + layout.addWidget(QLabel("Select atoms with label: ")) + self.selection_field = CheckableComboBox(self) + layout.addWidget(self.selection_field) + self.selection_field.addItems(self.labels) + + def parameter_dictionary(self): + function_parameters = {"function_name": "select_labels"} + selection = self.selection_field.checked_values() + function_parameters["atom_labels"] = selection + return function_parameters + + +class PatternSelection(BasicSelectionWidget): + def __init__( + self, + parent=None, + widget_label="SMARTS pattern matching", + ): + self.pattern_dictionary = { + "primary amine": "[#7X3;H2;!$([#7][#6X3][!#6]);!$([#7][#6X2][!#6])](~[H])~[H]", + "hydroxy": "[#8;H1,H2]~[H]", + "methyl": "[#6;H3](~[H])(~[H])~[H]", + "phosphate": "[#15X4](~[#8])(~[#8])(~[#8])~[#8]", + "sulphate": "[#16X4](~[#8])(~[#8])(~[#8])~[#8]", + "thiol": "[#16X2;H1]~[H]", + } + super().__init__(parent, widget_label) + + def add_specific_widgets(self): + layout = self.layout() + layout.addWidget(QLabel("Pick a group")) + self.selection_field = QComboBox(self) + layout.addWidget(self.selection_field) + self.selection_field.addItems(self.pattern_dictionary.keys()) + layout.addWidget(QLabel("pattern:")) + self.input_field = QLineEdit("", self) + self.input_field.setPlaceholderText("can be edited") + layout.addWidget(self.input_field) + self.selection_field.currentTextChanged.connect(self.update_string) + + @Slot(str) + def update_string(self, key_string: str): + if key_string in self.pattern_dictionary: + self.input_field.setText(self.pattern_dictionary[key_string]) + + def parameter_dictionary(self): + function_parameters = {"function_name": "select_pattern"} + selection = self.input_field.text() + function_parameters["rdkit_pattern"] = selection return function_parameters From 795a7b6187d67e2b3bc34023f83e5be3200c3156 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Fri, 14 Feb 2025 16:55:02 +0000 Subject: [PATCH 13/37] 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): From 16c56f95818e3af9ecb91022041196aee9e19ef6 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak <108934199+MBartkowiakSTFC@users.noreply.github.com> Date: Wed, 19 Feb 2025 15:57:42 +0000 Subject: [PATCH 14/37] Apply suggestions from code review Co-authored-by: Jacob Wilkins <46597752+oerc0122@users.noreply.github.com> Signed-off-by: Maciej Bartkowiak <108934199+MBartkowiakSTFC@users.noreply.github.com> --- .../Framework/AtomSelector/atom_selection.py | 28 ++++++++----------- .../AtomSelector/general_selection.py | 2 +- .../Framework/AtomSelector/group_selection.py | 3 +- .../AtomSelector/molecule_selection.py | 6 +--- 4 files changed, 15 insertions(+), 24 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py index 6351dd61b..d1e13669c 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py @@ -42,25 +42,21 @@ def select_atoms( element_list = system.atom_list name_list = system.name_list indices = set(range(len(element_list))) - index_list = function_parameters.get("index_list", None) - index_range = function_parameters.get("index_range", None) - index_slice = function_parameters.get("index_slice", None) + index_list = function_parameters.get("index_list") + index_range = function_parameters.get("index_range") + index_slice = function_parameters.get("index_slice") if index_list is not None: - selection = selection.union(indices.intersection(index_list)) + selection |= indices & index_list if index_range is not None: - selection = selection.union( - indices.intersection(range(index_range[0], index_range[1])) - ) + selection |= indices & set(range(*index_range)) if index_slice is not None: - selection = selection.union( - indices.intersection(range(index_slice[0], index_slice[1], index_slice[2])) - ) - atom_types = function_parameters.get("atom_types", None) + selection |= indices & set(range(*index_slice)) + atom_types = function_parameters.get("atom_types", ()) if atom_types: - new_indices = [index for index in indices if element_list[index] in atom_types] - selection = selection.union(new_indices) - atom_names = function_parameters.get("atom_names", None) + new_indices = {index for index in indices if element_list[index] in atom_types} + selection |= new_indices + atom_names = function_parameters.get("atom_names", ()) if atom_names: - new_indices = [index for index in indices if name_list[index] in atom_names] - selection = selection.union(new_indices) + new_indices = {index for index in indices if name_list[index] in atom_names} + selection |= new_indices return selection diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py index 8c404e894..fb0c57623 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py @@ -61,5 +61,5 @@ def select_none( def invert_selection(trajectory: Trajectory, selection: Set[int]) -> Set[int]: all_indices = select_all(trajectory) - inverted = all_indices.difference(selection) + inverted = all_indices - selection return inverted diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py index 49615d553..f68c2818f 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py @@ -66,7 +66,6 @@ def select_pattern( """ selection = set() system = trajectory.chemical_system - pattern = function_parameters.get("rdkit_pattern", None) - if pattern: + if pattern := function_parameters.get("rdkit_pattern"): selection = system.get_substructure_matches(pattern) return selection diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py index 0513d25c6..c727b40f9 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py @@ -41,9 +41,5 @@ def select_molecules( selection = set() system = trajectory.chemical_system 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 = {cluster for cluster in system._clusters.get(molecule, ()) for molecule in molecule_names} return selection From e1f971be4418df383096450615de591521dd4299 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Wed, 19 Feb 2025 16:10:24 +0000 Subject: [PATCH 15/37] Adjust code after review suggestions --- .../MDANSE/Framework/AtomSelector/atom_selection.py | 8 +++----- .../Framework/AtomSelector/general_selection.py | 8 ++------ .../MDANSE/Framework/AtomSelector/group_selection.py | 8 ++------ .../Framework/AtomSelector/molecule_selection.py | 11 +++++++---- MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py | 10 ++-------- 5 files changed, 16 insertions(+), 29 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py index d1e13669c..f575c2f44 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py @@ -20,9 +20,7 @@ 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 @@ -47,9 +45,9 @@ def select_atoms( index_slice = function_parameters.get("index_slice") if index_list is not None: selection |= indices & index_list - if index_range is not None: + elif index_range is not None: selection |= indices & set(range(*index_range)) - if index_slice is not None: + elif index_slice is not None: selection |= indices & set(range(*index_slice)) atom_types = function_parameters.get("atom_types", ()) if atom_types: diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py index fb0c57623..dd37ec3ce 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py @@ -19,9 +19,7 @@ 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 @@ -39,9 +37,7 @@ def select_all( 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 f68c2818f..7c34bfe95 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py @@ -21,9 +21,7 @@ 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 @@ -47,9 +45,7 @@ def select_labels( 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 c727b40f9..5c8df8d46 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py @@ -21,9 +21,7 @@ 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 +39,10 @@ def select_molecules( selection = set() system = trajectory.chemical_system molecule_names = function_parameters.get("molecule_names", None) - selection = {cluster for cluster in system._clusters.get(molecule, ()) for molecule in molecule_names} + selection = { + index + for molecule in molecule_names + for cluster in system._clusters.get(molecule, ()) + for index in cluster + } return selection diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index d722c8b52..57be31fd2 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -113,15 +113,9 @@ 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 From 727fe532f2a287439ba94e16ca834f6bd95b43a0 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Wed, 19 Feb 2025 16:29:50 +0000 Subject: [PATCH 16/37] Correct the docstrings of selection functions --- .../Framework/AtomSelector/atom_selection.py | 19 ++++++++++++----- .../AtomSelector/general_selection.py | 21 ++++++++++++++----- .../Framework/AtomSelector/group_selection.py | 16 ++++++++------ .../AtomSelector/molecule_selection.py | 6 +++--- 4 files changed, 43 insertions(+), 19 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py index f575c2f44..792dc4bd3 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py @@ -21,14 +21,23 @@ def select_atoms(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Set[int]: - """Selects all the atoms in the trajectory. + """Selects specific atoms in the trajectory. These can be selected based + on indices, atom type or trajectory-specific atom name. + The atom type is normally the chemical element, while + the atom name can be more specific and depend on the + force field used. Parameters ---------- - selection : Set[int] - A set of atom indices trajectory : Trajectory A trajectory instance to which the selection is applied + function_parameters : Dict[str, Any] + may include any combination of the following + "index_list" : List[int] + "index_range" : a start, stop pair of indices for range(start, stop) + "index_slice" : s start, stop, step sequence of indices for range(start, stop, step) + "atom_types" : List[str] + "atom_names" : List[str] Returns ------- @@ -45,9 +54,9 @@ def select_atoms(trajectory: Trajectory, **function_parameters: Dict[str, Any]) index_slice = function_parameters.get("index_slice") if index_list is not None: selection |= indices & index_list - elif index_range is not None: + if index_range is not None: selection |= indices & set(range(*index_range)) - elif index_slice is not None: + if index_slice is not None: selection |= indices & set(range(*index_slice)) atom_types = function_parameters.get("atom_types", ()) if atom_types: diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py index dd37ec3ce..778c54fb2 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py @@ -24,8 +24,6 @@ def select_all(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Parameters ---------- - selection : Set[int] - A set of atom indices trajectory : Trajectory A trajectory instance to which the selection is applied @@ -42,10 +40,8 @@ def select_none(trajectory: Trajectory, **function_parameters: Dict[str, Any]) - Parameters ---------- - selection : Set[int] - A set of atom indices trajectory : Trajectory - A trajectory instance to which the selection is applied + A trajectory instance, ignored in this selection Returns ------- @@ -56,6 +52,21 @@ def select_none(trajectory: Trajectory, **function_parameters: Dict[str, Any]) - def invert_selection(trajectory: Trajectory, selection: Set[int]) -> Set[int]: + """Returns a set of all the indices that are present in the trajectory + and were not included in the input selection. + + Parameters + ---------- + trajectory : Trajectory + a trajectory containing atoms to be selected + selection : Set[int] + set of indices to be excluded from the set of all indices + + Returns + ------- + Set[int] + set of all the indices in the trajectory which were not in the selection + """ all_indices = select_all(trajectory) inverted = all_indices - selection return inverted diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py index 7c34bfe95..ee453dec4 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py @@ -22,14 +22,15 @@ def select_labels(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Set[int]: - """Selects all the atoms in the trajectory. + """Selects atoms with a specific label in the trajectory. + A residue name can be read as a label by MDANSE. Parameters ---------- - selection : Set[int] - A set of atom indices trajectory : Trajectory A trajectory instance to which the selection is applied + function_parameters : Dict[str, Any] + should include a list of string labels under key "atom_labels" Returns ------- @@ -46,14 +47,17 @@ def select_labels(trajectory: Trajectory, **function_parameters: Dict[str, Any]) def select_pattern(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Set[int]: - """Selects all the atoms in the trajectory. + """Selects atoms according to the SMARTS string given as input. + This will only work if molecules and bonds have been detected in the system. + If the bond information was not read from the input trajectory on conversion, + it can still be determined in a TrajectoryEditor run. Parameters ---------- - selection : Set[int] - A set of atom indices trajectory : Trajectory A trajectory instance to which the selection is applied + function_parameters : Dict[str, Any] + should include a SMARTS string under key "rdkit_pattern" Returns ------- diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py index 5c8df8d46..499071bf4 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py @@ -22,14 +22,14 @@ def select_molecules(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Set[int]: - """Selects all the atoms in the trajectory. + """Selects all the atoms belonging to the specified molecule types. Parameters ---------- - selection : Set[int] - A set of atom indices trajectory : Trajectory A trajectory instance to which the selection is applied + function_parameters : Dict[str, Any] + should include a list of str molecule names under key "atom_labels" Returns ------- From f51741904bef521e5e95a43754740cd0dc6c2562 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Thu, 20 Feb 2025 12:55:14 +0000 Subject: [PATCH 17/37] Add inversion to the GUI --- .../AtomSelector/spatial_selection.py | 79 +++++++++++++++++++ .../InputWidgets/AtomSelectionWidget.py | 20 ++--- .../MDANSE_GUI/Widgets/SelectionWidgets.py | 6 ++ 3 files changed, 96 insertions(+), 9 deletions(-) create mode 100644 MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py new file mode 100644 index 000000000..5736e5038 --- /dev/null +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py @@ -0,0 +1,79 @@ +# This file is part of MDANSE. +# +# MDANSE is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . +# + +from typing import Union, Dict, Any, Set + +import numpy as np +from scipy.spatial import KDTree + +from MDANSE.Chemistry.ChemicalSystem import ChemicalSystem +from MDANSE.MolecularDynamics.Trajectory import Trajectory + + +def select_positions(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Set[int]: + """Selects atoms based on their positions at a specified frame number. + + + Parameters + ---------- + trajectory : Trajectory + A trajectory instance to which the selection is applied + function_parameters : Dict[str, Any] + may include + "frame_number" : int + "position_minimum" : np.ndarray[float] + "position_maximum" : np.ndarray[float] + + Returns + ------- + Set[int] + Set of all the atom indices + """ + coordinates = trajectory.coordinates(function_parameters.get("frame_number", 0)) + lower_limits = trajectory.coordinates( + function_parameters.get("position_minimum", 3 * [-np.inf]) + ) + upper_limits = trajectory.coordinates(function_parameters.get("position_maximum", 3 * [np.inf])) + mask1 = np.all(coordinates > lower_limits.reshape((1, 3)), axis=1) + mask2 = np.all(coordinates < upper_limits.reshape((1, 3)), axis=1) + return set(np.where(np.logical_and(mask1, mask2))[0]) + + +def select_sphere(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Set[int]: + """Selects atoms within a distance from a fixed point in space, + based on coordinates at a specific frame number. + + Parameters + ---------- + trajectory : Trajectory + A trajectory instance to which the selection is applied + function_parameters : Dict[str, Any] + may include + "frame_number" : int + "sphere_centre" : np.ndarray[float] + "sphere_radius" : float + + Returns + ------- + Set[int] + Set of all the atom indices + """ + coordinates = trajectory.coordinates(function_parameters.get("frame_number", 0)) + sphere_centre = function_parameters.get("sphere_centre", np.zeros(3)) + sphere_radius = function_parameters.get("sphere_radius", 1.0) + kdtree = KDTree(coordinates) + indices = kdtree.query_ball_point(sphere_centre, sphere_radius) + return set(indices) diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index b8473290a..b0b329de1 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -16,6 +16,7 @@ from typing import Set import json +from enum import Enum from qtpy.QtCore import Slot, Signal from qtpy.QtGui import QStandardItemModel, QStandardItem @@ -48,9 +49,10 @@ ) -VALID_SELECTION = "Valid selection" -USELESS_SELECTION = "Selection did not change. This operation is not needed." -MALFORMED_SELECTION = "This is not a valid JSON string." +class SelectionValidity(Enum): + VALID_SELECTION = "Valid selection" + USELESS_SELECTION = "Selection did not change. This operation is not needed." + MALFORMED_SELECTION = "This is not a valid JSON string." class SelectionModel(QStandardItemModel): @@ -79,12 +81,12 @@ def rebuild_selection(self, last_operation: str): last_operation, self._trajectory, self._current_selection ) except json.JSONDecodeError: - return MALFORMED_SELECTION + return SelectionValidity.MALFORMED_SELECTION if valid: self._selection.load_from_json(json_string) - return VALID_SELECTION + return SelectionValidity.VALID_SELECTION else: - return USELESS_SELECTION + return SelectionValidity.USELESS_SELECTION def current_selection(self, last_operation: str = "") -> Set[int]: self.rebuild_selection(last_operation) @@ -306,17 +308,17 @@ def append_selection(self): self.selection_line.setToolTip("") selection_text = self.selection_line.text() validation = self.selection_model.rebuild_selection(selection_text) - if validation == MALFORMED_SELECTION: + if validation == SelectionValidity.MALFORMED_SELECTION: self.selection_line.setStyleSheet( "QWidget#InputWidget { background-color:rgb(180,20,180); font-weight: bold }" ) self.selection_line.setToolTip(validation) - elif validation == USELESS_SELECTION: + elif validation == SelectionValidity.USELESS_SELECTION: self.selection_line.setStyleSheet( "QWidget#InputWidget { background-color:rgb(180,20,180); font-weight: bold }" ) self.selection_line.setToolTip(validation) - elif validation == VALID_SELECTION: + elif validation == SelectionValidity.VALID_SELECTION: self.selection_model.appendRow(QStandardItem(selection_text)) self.view_3d._viewer.change_picked(self.selected) self.update_selection_textbox() diff --git a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py index 317394272..64166a5fd 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py @@ -81,8 +81,14 @@ def __init__(self, parent=None, widget_label="ALL ATOMS"): def add_specific_widgets(self): layout = self.layout() + inversion_button = QPushButton("INVERT selection", self) + inversion_button.clicked.connect(self.invert_selection) + layout.addWidget(inversion_button) layout.addWidget(QLabel("Add/remove ALL atoms")) + def invert_selection(self): + self.new_selection.emit(json.dumps({"function_name": "invert_selection"})) + def parameter_dictionary(self): return {"function_name": "select_all"} From d20fa6421aaa8971a83ebd979c598708edd6e7b4 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Thu, 20 Feb 2025 14:37:22 +0000 Subject: [PATCH 18/37] Add position selection widget --- .../MDANSE/Framework/AtomSelector/selector.py | 3 + .../AtomSelector/spatial_selection.py | 2 +- .../InputWidgets/AtomSelectionWidget.py | 7 +- .../Src/MDANSE_GUI/InputWidgets/WidgetBase.py | 1 - .../MDANSE_GUI/Widgets/SelectionWidgets.py | 115 +++++++++++++++++- 5 files changed, 123 insertions(+), 5 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index 57be31fd2..75c5deb17 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -25,6 +25,7 @@ from MDANSE.Framework.AtomSelector.atom_selection import select_atoms from MDANSE.Framework.AtomSelector.molecule_selection import select_molecules from MDANSE.Framework.AtomSelector.group_selection import select_labels, select_pattern +from MDANSE.Framework.AtomSelector.spatial_selection import select_positions, select_sphere function_lookup = { @@ -37,6 +38,8 @@ select_molecules, select_labels, select_pattern, + select_positions, + select_sphere, ] } diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py index 5736e5038..0e4b7846c 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py @@ -25,7 +25,7 @@ def select_positions(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Set[int]: """Selects atoms based on their positions at a specified frame number. - + Lower and upper limits of x, y and z coordinates can be given as input. Parameters ---------- diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index b0b329de1..5ec7c53c8 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -46,6 +46,7 @@ MoleculeSelection, PatternSelection, LabelSelection, + PositionSelection, ) @@ -242,7 +243,8 @@ def left_widgets(self) -> list[QWidget]: select = QGroupBox("selection") select_layout = QVBoxLayout() scroll_area = QScrollArea() - scroll_area.setLayout(select_layout) + scroll_area.setWidget(select) + scroll_area.setWidgetResizable(True) self.selection_widgets = [ AllAtomSelection(self), @@ -251,6 +253,7 @@ def left_widgets(self) -> list[QWidget]: MoleculeSelection(self, self.trajectory), PatternSelection(self), LabelSelection(self, self.trajectory), + PositionSelection(self, self.trajectory, self.view_3d._viewer), ] for widget in self.selection_widgets: @@ -277,7 +280,7 @@ def left_widgets(self) -> list[QWidget]: ) self.selection_operations_view.setModel(self.selection_model) self.selection_model.selection_changed.connect(self.recalculate_selection) - return [select] + return [scroll_area] @Slot() def recalculate_selection(self): diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/WidgetBase.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/WidgetBase.py index 8591de9f5..86be1bc07 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/WidgetBase.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/WidgetBase.py @@ -29,7 +29,6 @@ class WidgetBase(QObject): - valid_changed = Signal() value_updated = Signal() value_changed = Signal() diff --git a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py index 64166a5fd..415b304e3 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py @@ -14,11 +14,12 @@ # along with this program. If not, see . # -from typing import Dict, Any +from typing import Dict, Any, TYPE_CHECKING, Tuple import json import numpy as np from qtpy.QtCore import Signal, Slot +from qtpy.QtGui import QValidator from qtpy.QtWidgets import ( QGroupBox, QHBoxLayout, @@ -31,6 +32,58 @@ from MDANSE_GUI.InputWidgets.CheckableComboBox import CheckableComboBox from MDANSE.MolecularDynamics.Trajectory import Trajectory +if TYPE_CHECKING: + from MDANSE_GUI.MolecularViewer.MolecularViewer import MolecularViewer + + +class XYZValidator(QValidator): + """A custom validator for a QLineEdit. + It is intended to limit the input to a string + of 3 comma-separated float numbers. + + Additional checks are necessary later in the code, + since the validator cannot exclude the cases of + 1 or 2 comma-separated values, since they are + a preliminary step when typing in 3 numbers. + """ + + def validate(self, input_string: str, position: int) -> Tuple[int, str]: + """Implementation of the virtual method of QValidator. + It takes in the string from a QLineEdit and the cursor position, + and an enum value of the validator state. Widgets will reject + inputs which change the state to Invalid. + + Parameters + ---------- + input_string : str + current contents of a text input field + position : int + position of the cursor in the text input field + + Returns + ------- + Tuple[int,str] + a tuple of (validator state, input string, cursor position) + """ + state = QValidator.State.Intermediate + comma_count = input_string.count(",") + if len(input_string) > 0: + try: + values = [int(x) for x in input_string.split(",")] + except (TypeError, ValueError): + if input_string[-1] == "," and comma_count < 3: + state = QValidator.State.Intermediate + else: + state = QValidator.State.Invalid + else: + if len(values) > 3: + state = QValidator.State.Invalid + elif len(values) == 3: + state = QValidator.State.Acceptable + else: + state = QValidator.State.Intermediate + return state, input_string, position + class BasicSelectionWidget(QGroupBox): new_selection = Signal(str) @@ -278,3 +331,63 @@ def parameter_dictionary(self): selection = self.input_field.text() function_parameters["rdkit_pattern"] = selection return function_parameters + + +class PositionSelection(BasicSelectionWidget): + def __init__( + self, + parent=None, + trajectory: Trajectory = None, + molecular_viewer: "MolecularViewer" = None, + widget_label="ALL ATOMS", + ): + self._viewer = molecular_viewer + self._lower_limit = np.zeros(3) + self._upper_limit = np.linalg.norm(trajectory.unit_cell(0), axis=1) + self._current_lower_limit = self._lower_limit.copy() + self._current_upper_limit = self._upper_limit.copy() + super().__init__(parent, widget_label) + + def add_specific_widgets(self): + layout = self.layout() + layout.addWidget(QLabel("Lower limits")) + self._lower_limit_input = QLineEdit( + ",".join([str(round(x, 3)) for x in self._lower_limit]) + ) + layout.addWidget(self._lower_limit_input) + layout.addWidget(QLabel("Upper limits")) + self._upper_limit_input = QLineEdit( + ",".join([str(round(x, 3)) for x in self._upper_limit]) + ) + layout.addWidget(self._upper_limit_input) + for field in [self._lower_limit_input, self._upper_limit_input]: + field.setValidator(XYZValidator(self)) + field.textChanged.connect(self.check_inputs) + + @Slot() + def check_inputs(self): + enable = True + try: + self._current_lower_limit = [ + float(x) for x in self._lower_limit_input.text().split(",") + ] + self._current_upper_limit = [ + float(x) for x in self._upper_limit_input.text().split(",") + ] + except (TypeError, ValueError): + enable = False + else: + if ( + len(self._current_lower_limit) != 3 + or len(self._current_upper_limit) != 3 + ): + enable = False + self.commit_button.setEnabled(enable) + + def parameter_dictionary(self): + return { + "function_name": "select_positions", + "frame_number": self._viewer._current_frame, + "lower_limits": self._current_lower_limit, + "upper_limits": self._current_upper_limit, + } From aaf0b68ba2c936c2d7189d1f17d0f258d84908ca Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Thu, 20 Feb 2025 16:38:00 +0000 Subject: [PATCH 19/37] Add cube and sphere selection to the GUI --- .../AtomSelector/spatial_selection.py | 8 +- MDANSE/Tests/UnitTests/test_selection.py | 79 ++++++++++++------- .../InputWidgets/AtomSelectionWidget.py | 2 + .../MDANSE_GUI/Widgets/SelectionWidgets.py | 64 +++++++++++++-- 4 files changed, 114 insertions(+), 39 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py index 0e4b7846c..ad04678a3 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py @@ -43,10 +43,8 @@ def select_positions(trajectory: Trajectory, **function_parameters: Dict[str, An Set of all the atom indices """ coordinates = trajectory.coordinates(function_parameters.get("frame_number", 0)) - lower_limits = trajectory.coordinates( - function_parameters.get("position_minimum", 3 * [-np.inf]) - ) - upper_limits = trajectory.coordinates(function_parameters.get("position_maximum", 3 * [np.inf])) + lower_limits = np.array(function_parameters.get("position_minimum", 3 * [-np.inf])) + upper_limits = np.array(function_parameters.get("position_maximum", 3 * [np.inf])) mask1 = np.all(coordinates > lower_limits.reshape((1, 3)), axis=1) mask2 = np.all(coordinates < upper_limits.reshape((1, 3)), axis=1) return set(np.where(np.logical_and(mask1, mask2))[0]) @@ -72,7 +70,7 @@ def select_sphere(trajectory: Trajectory, **function_parameters: Dict[str, Any]) Set of all the atom indices """ coordinates = trajectory.coordinates(function_parameters.get("frame_number", 0)) - sphere_centre = function_parameters.get("sphere_centre", np.zeros(3)) + sphere_centre = np.array(function_parameters.get("sphere_centre", 3 * [0.0])) sphere_radius = function_parameters.get("sphere_radius", 1.0) kdtree = KDTree(coordinates) indices = kdtree.query_ball_point(sphere_centre, sphere_radius) diff --git a/MDANSE/Tests/UnitTests/test_selection.py b/MDANSE/Tests/UnitTests/test_selection.py index 650cae3d8..5749d27a9 100644 --- a/MDANSE/Tests/UnitTests/test_selection.py +++ b/MDANSE/Tests/UnitTests/test_selection.py @@ -2,11 +2,13 @@ import os import pytest +import numpy as np from MDANSE.Framework.AtomSelector.general_selection import select_all, select_none, invert_selection from MDANSE.Framework.AtomSelector.atom_selection import select_atoms from MDANSE.Framework.AtomSelector.molecule_selection import select_molecules from MDANSE.Framework.AtomSelector.group_selection import select_labels, select_pattern +from MDANSE.Framework.AtomSelector.spatial_selection import select_positions, select_sphere from MDANSE.Framework.InputData.HDFTrajectoryInputData import HDFTrajectoryInputData @@ -32,6 +34,20 @@ ) +@pytest.fixture(scope="module") +def short_trajectory(): + traj_object = HDFTrajectoryInputData(short_traj) + yield traj_object.trajectory + traj_object.close() + + +@pytest.fixture(scope="module") +def gromacs_trajectory(): + traj_object = HDFTrajectoryInputData(traj_2vb1) + yield traj_object.trajectory + traj_object.close() + + @pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) def test_select_all(trajectory): traj_object = HDFTrajectoryInputData(trajectory) @@ -56,54 +72,61 @@ def test_inverted_none_is_all(trajectory): assert all_selection == inverted_none -def test_select_atoms_selects_by_element(): - traj_object = HDFTrajectoryInputData(short_traj) - s_selection = select_atoms(traj_object.trajectory, atom_types=['S']) +def test_select_atoms_selects_by_element(short_trajectory): + s_selection = select_atoms(short_trajectory, atom_types=['S']) assert len(s_selection) == 208 - cu_selection = select_atoms(traj_object.trajectory, atom_types=['Cu']) + cu_selection = select_atoms(short_trajectory, atom_types=['Cu']) assert len(cu_selection) == 208 - sb_selection = select_atoms(traj_object.trajectory, atom_types=['Sb']) + sb_selection = select_atoms(short_trajectory, atom_types=['Sb']) assert len(sb_selection) == 64 - cusbs_selection = select_atoms(traj_object.trajectory, atom_types=['Cu','Sb', 'S']) + cusbs_selection = select_atoms(short_trajectory, atom_types=['Cu','Sb', 'S']) assert len(cusbs_selection) == 480 -def test_select_atoms_selects_by_range(): - traj_object = HDFTrajectoryInputData(short_traj) - range_selection = select_atoms(traj_object.trajectory, index_range=[15,35]) +def test_select_atoms_selects_by_range(short_trajectory): + range_selection = select_atoms(short_trajectory, index_range=[15,35]) assert len(range_selection) == 20 - overshoot_selection = select_atoms(traj_object.trajectory, index_range=[470,510]) + overshoot_selection = select_atoms(short_trajectory, index_range=[470,510]) assert len(overshoot_selection) == 10 -def test_select_atoms_selects_by_slice(): - traj_object = HDFTrajectoryInputData(short_traj) - range_selection = select_atoms(traj_object.trajectory, index_slice=[150,350, 10]) +def test_select_atoms_selects_by_slice(short_trajectory): + range_selection = select_atoms(short_trajectory, index_slice=[150,350, 10]) assert len(range_selection) == 20 - overshoot_selection = select_atoms(traj_object.trajectory, index_slice=[470,510,5]) + overshoot_selection = select_atoms(short_trajectory, index_slice=[470,510,5]) assert len(overshoot_selection) == 2 -def test_select_molecules_selects_water(): - traj_object = HDFTrajectoryInputData(traj_2vb1) - water_selection = select_molecules(traj_object.trajectory, molecule_names = ['H2 O1']) +def test_select_molecules_selects_water(gromacs_trajectory): + water_selection = select_molecules(gromacs_trajectory, molecule_names = ['H2 O1']) assert len(water_selection) == 28746 -def test_select_molecules_selects_protein(): - traj_object = HDFTrajectoryInputData(traj_2vb1) - protein_selecton = select_molecules(traj_object.trajectory, molecule_names = ['C613 H959 N193 O185 S10']) +def test_select_molecules_selects_protein(gromacs_trajectory): + protein_selecton = select_molecules(gromacs_trajectory, molecule_names = ['C613 H959 N193 O185 S10']) assert len(protein_selecton) == 613 + 959 + 193 + 185 + 10 -def test_select_molecules_inverted_selects_ions(): - traj_object = HDFTrajectoryInputData(traj_2vb1) - all_molecules_selection = select_molecules(traj_object.trajectory, molecule_names = ['C613 H959 N193 O185 S10', 'H2 O1']) - non_molecules_selection = invert_selection(traj_object.trajectory, all_molecules_selection) - assert all([traj_object.chemical_system.atom_list[index] in ['Na', 'Cl'] for index in non_molecules_selection]) +def test_select_molecules_inverted_selects_ions(gromacs_trajectory): + all_molecules_selection = select_molecules(gromacs_trajectory, molecule_names = ['C613 H959 N193 O185 S10', 'H2 O1']) + non_molecules_selection = invert_selection(gromacs_trajectory, all_molecules_selection) + assert all([gromacs_trajectory.chemical_system.atom_list[index] in ['Na', 'Cl'] for index in non_molecules_selection]) -def test_select_pattern_selects_water(): - traj_object = HDFTrajectoryInputData(traj_2vb1) - water_selection = select_pattern(traj_object.trajectory, rdkit_pattern="[#8X2;H2](~[H])~[H]") +def test_select_pattern_selects_water(gromacs_trajectory): + water_selection = select_pattern(gromacs_trajectory, rdkit_pattern="[#8X2;H2](~[H])~[H]") assert len(water_selection) == 28746 + + +def test_select_positions(short_trajectory): + cube_selection = select_positions(short_trajectory, + position_minimum = 0.5*np.ones(3), + position_maximum = 0.7*np.ones(3)) + assert cube_selection == {59} + + +def test_select_sphere(short_trajectory): + sphere_selection = select_sphere(short_trajectory, + sphere_centre = 0.6*np.ones(3), + sphere_radius = 0.2) + assert sphere_selection == {19, 17, 59, 15} diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index 5ec7c53c8..99679ff97 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -47,6 +47,7 @@ PatternSelection, LabelSelection, PositionSelection, + SphereSelection, ) @@ -254,6 +255,7 @@ def left_widgets(self) -> list[QWidget]: PatternSelection(self), LabelSelection(self, self.trajectory), PositionSelection(self, self.trajectory, self.view_3d._viewer), + SphereSelection(self, self.trajectory, self.view_3d._viewer), ] for widget in self.selection_widgets: diff --git a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py index 415b304e3..0b87d07e7 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py @@ -19,7 +19,7 @@ import numpy as np from qtpy.QtCore import Signal, Slot -from qtpy.QtGui import QValidator +from qtpy.QtGui import QValidator, QDoubleValidator from qtpy.QtWidgets import ( QGroupBox, QHBoxLayout, @@ -69,7 +69,7 @@ def validate(self, input_string: str, position: int) -> Tuple[int, str]: comma_count = input_string.count(",") if len(input_string) > 0: try: - values = [int(x) for x in input_string.split(",")] + values = [float(x) for x in input_string.split(",")] except (TypeError, ValueError): if input_string[-1] == "," and comma_count < 3: state = QValidator.State.Intermediate @@ -339,11 +339,11 @@ def __init__( parent=None, trajectory: Trajectory = None, molecular_viewer: "MolecularViewer" = None, - widget_label="ALL ATOMS", + widget_label="Select by position", ): self._viewer = molecular_viewer self._lower_limit = np.zeros(3) - self._upper_limit = np.linalg.norm(trajectory.unit_cell(0), axis=1) + self._upper_limit = np.linalg.norm(trajectory.unit_cell(0)._unit_cell, axis=1) self._current_lower_limit = self._lower_limit.copy() self._current_upper_limit = self._upper_limit.copy() super().__init__(parent, widget_label) @@ -388,6 +388,58 @@ def parameter_dictionary(self): return { "function_name": "select_positions", "frame_number": self._viewer._current_frame, - "lower_limits": self._current_lower_limit, - "upper_limits": self._current_upper_limit, + "position_minimum": list(self._current_lower_limit), + "position_maximum": list(self._current_upper_limit), + } + + +class SphereSelection(BasicSelectionWidget): + def __init__( + self, + parent=None, + trajectory: Trajectory = None, + molecular_viewer: "MolecularViewer" = None, + widget_label="Select in a sphere", + ): + self._viewer = molecular_viewer + self._current_sphere_centre = np.diag(trajectory.unit_cell(0)._unit_cell) * 0.5 + self._current_sphere_radius = np.min(self._current_sphere_centre) + super().__init__(parent, widget_label) + + def add_specific_widgets(self): + layout = self.layout() + layout.addWidget(QLabel("Sphere centre")) + self._sphere_centre_input = QLineEdit( + ",".join([str(round(x, 3)) for x in self._current_sphere_centre]), self + ) + layout.addWidget(self._sphere_centre_input) + layout.addWidget(QLabel("Sphere radius (nm)")) + self._sphere_radius_input = QLineEdit("0.5", self) + layout.addWidget(self._sphere_radius_input) + self._sphere_centre_input.setValidator(XYZValidator()) + self._sphere_centre_input.textChanged.connect(self.check_inputs) + self._sphere_radius_input.setValidator(QDoubleValidator()) + self._sphere_radius_input.textChanged.connect(self.check_inputs) + + @Slot() + def check_inputs(self): + enable = True + try: + self._current_sphere_centre = [ + float(x) for x in self._sphere_centre_input.text().split(",") + ] + self._current_sphere_radius = float(self._sphere_radius_input.text()) + except (TypeError, ValueError): + enable = False + else: + if len(self._current_sphere_centre) != 3: + enable = False + self.commit_button.setEnabled(enable) + + def parameter_dictionary(self): + return { + "function_name": "select_sphere", + "frame_number": self._viewer._current_frame, + "sphere_centre": list(self._current_sphere_centre), + "sphere_radius": self._current_sphere_radius, } From 47b2a8e6c59b01e8f036f11799490fd1769aeebc Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Thu, 20 Feb 2025 16:51:19 +0000 Subject: [PATCH 20/37] Format code using black --- .../Framework/AtomSelector/atom_selection.py | 4 +++- .../Framework/AtomSelector/general_selection.py | 8 ++++++-- .../Framework/AtomSelector/group_selection.py | 8 ++++++-- .../Framework/AtomSelector/molecule_selection.py | 4 +++- .../Src/MDANSE/Framework/AtomSelector/selector.py | 15 ++++++++++++--- .../Framework/AtomSelector/spatial_selection.py | 8 ++++++-- 6 files changed, 36 insertions(+), 11 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py index 792dc4bd3..eb61f9687 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 specific atoms in the trajectory. These can be selected based on indices, atom type or trajectory-specific atom name. The atom type is normally the chemical element, while diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py index 778c54fb2..9e87a6534 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 @@ -35,7 +37,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 ee453dec4..2781fa197 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 atoms with a specific label in the trajectory. A residue name can be read as a label by MDANSE. @@ -46,7 +48,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 atoms according to the SMARTS string given as input. This will only work if molecules and bonds have been detected in the system. If the bond information was not read from the input trajectory on conversion, diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py index 499071bf4..a636b49db 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 belonging to the specified molecule types. Parameters diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index 75c5deb17..24c12e852 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -25,7 +25,10 @@ from MDANSE.Framework.AtomSelector.atom_selection import select_atoms from MDANSE.Framework.AtomSelector.molecule_selection import select_molecules from MDANSE.Framework.AtomSelector.group_selection import select_labels, select_pattern -from MDANSE.Framework.AtomSelector.spatial_selection import select_positions, select_sphere +from MDANSE.Framework.AtomSelector.spatial_selection import ( + select_positions, + select_sphere, +) function_lookup = { @@ -116,9 +119,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/AtomSelector/spatial_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py index ad04678a3..ec8928bbb 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py @@ -23,7 +23,9 @@ from MDANSE.MolecularDynamics.Trajectory import Trajectory -def select_positions(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Set[int]: +def select_positions( + trajectory: Trajectory, **function_parameters: Dict[str, Any] +) -> Set[int]: """Selects atoms based on their positions at a specified frame number. Lower and upper limits of x, y and z coordinates can be given as input. @@ -50,7 +52,9 @@ def select_positions(trajectory: Trajectory, **function_parameters: Dict[str, An return set(np.where(np.logical_and(mask1, mask2))[0]) -def select_sphere(trajectory: Trajectory, **function_parameters: Dict[str, Any]) -> Set[int]: +def select_sphere( + trajectory: Trajectory, **function_parameters: Dict[str, Any] +) -> Set[int]: """Selects atoms within a distance from a fixed point in space, based on coordinates at a specific frame number. From e1ffb960b1a2ed51c44dec471215e835f19491ce Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Thu, 20 Feb 2025 16:54:41 +0000 Subject: [PATCH 21/37] Convert list to set in select_atoms --- MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py index eb61f9687..186142356 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py @@ -55,7 +55,7 @@ def select_atoms( index_range = function_parameters.get("index_range") index_slice = function_parameters.get("index_slice") if index_list is not None: - selection |= indices & index_list + selection |= indices & set(index_list) if index_range is not None: selection |= indices & set(range(*index_range)) if index_slice is not None: From b540e8fdc233728704240aaa979e221c29df9a80 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Thu, 27 Feb 2025 16:59:08 +0000 Subject: [PATCH 22/37] Correct atom transmutation and partial charge widgets --- .../MDANSE_GUI/InputWidgets/AtomTransmutationWidget.py | 9 +++------ .../Src/MDANSE_GUI/InputWidgets/PartialChargeWidget.py | 9 ++++----- 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomTransmutationWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomTransmutationWidget.py index 901c846a5..305ef038b 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomTransmutationWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomTransmutationWidget.py @@ -69,16 +69,13 @@ def __init__( self.transmutation_textbox.setReadOnly(True) self.transmutation_combo = QComboBox() self.transmutation_combo.addItems(ATOMS_DATABASE.atoms) - self.transmuter.selector.settings["all"] = False super().__init__( - transmuter.selector, traj_data, field, parent, *args, **kwargs, ) - self.all_selection = False self.update_transmutation_textbox() def right_widgets(self) -> list[QWidget]: @@ -125,8 +122,9 @@ def apply_transmutation(self) -> None: transmutation and update the transmutation textbox with the new transmutation setting. """ + selection_string = self.selection_model.current_steps() self.transmuter.apply_transmutation( - self.settings, self.transmutation_combo.currentText() + selection_string, self.transmutation_combo.currentText() ) self.update_transmutation_textbox() @@ -137,9 +135,8 @@ def update_transmutation_textbox(self) -> None: map = self.transmuter.get_setting() text = [f"Number of atoms transmuted:\n{len(map)}\n\nTransmuted atoms:\n"] - atoms = self.selector.system.atom_list for idx, symbol in map.items(): - text.append(f"{idx} {atoms[idx]} -> {symbol}\n") + text.append(f"{idx} {self.atm_full_names[idx]} -> {symbol}\n") self.transmutation_textbox.setText("".join(text)) diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/PartialChargeWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/PartialChargeWidget.py index 4d357189a..2f6d97711 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/PartialChargeWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/PartialChargeWidget.py @@ -68,8 +68,7 @@ def __init__( self.charge_textbox.setReadOnly(True) self.charge_qline = QLineEdit() self.charge_qline.setValidator(QDoubleValidator()) - self.mapper.selector.settings["all"] = False - super().__init__(mapper.selector, traj_data, field, parent, *args, **kwargs) + super().__init__(traj_data, field, parent, *args, **kwargs) self.all_selection = False self.update_charge_textbox() @@ -122,7 +121,8 @@ def apply_charges(self) -> None: except ValueError: # probably an empty QLineEdit box return - self.mapper.update_charges(self.settings, charge) + selection_string = self.selection_model.current_steps() + self.mapper.update_charges(selection_string, charge) self.update_charge_textbox() def update_charge_textbox(self) -> None: @@ -132,9 +132,8 @@ def update_charge_textbox(self) -> None: map = self.mapper.get_full_setting() text = ["Partial charge mapping:\n"] - atoms = self.selector.system.atom_list for idx, charge in map.items(): - text.append(f"{idx} ({atoms[idx]}) -> {charge}\n") + text.append(f"{idx} ({self.atm_full_names[idx]}) -> {charge}\n") self.charge_textbox.setText("".join(text)) From d128109c4b68ba8bc6fcbc3892018499ce4fda14 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak <108934199+MBartkowiakSTFC@users.noreply.github.com> Date: Fri, 28 Feb 2025 09:28:21 +0000 Subject: [PATCH 23/37] Apply suggestions from code review Co-authored-by: Jacob Wilkins <46597752+oerc0122@users.noreply.github.com> Signed-off-by: Maciej Bartkowiak <108934199+MBartkowiakSTFC@users.noreply.github.com> --- .../Src/MDANSE/Framework/AtomSelector/group_selection.py | 7 ++----- MDANSE/Src/MDANSE/Framework/Jobs/IJob.py | 8 +++----- .../Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py | 4 ++-- 3 files changed, 7 insertions(+), 12 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py index 2781fa197..1522877f4 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py @@ -39,12 +39,9 @@ def select_labels( Set[int] Set of all the atom indices """ - selection = set() system = trajectory.chemical_system - atom_labels = function_parameters.get("atom_labels", None) - for label in atom_labels: - if label in system._labels: - selection = selection.union(system._labels[label]) + atom_labels = function_parameters.get("atom_labels", ()) + selection = {system._labels[label] for label in atom_labels if label in system._labels} return selection diff --git a/MDANSE/Src/MDANSE/Framework/Jobs/IJob.py b/MDANSE/Src/MDANSE/Framework/Jobs/IJob.py index 94b47a345..edb27ce83 100644 --- a/MDANSE/Src/MDANSE/Framework/Jobs/IJob.py +++ b/MDANSE/Src/MDANSE/Framework/Jobs/IJob.py @@ -169,17 +169,15 @@ def initialize(self): ) except KeyError: LOG.error("IJob did not find 'write_logs' in output_files") - if "atom_selection" in self.configuration: + if selection := self.configuration.get("atom_selection"): try: - array_length = self.configuration["atom_selection"][ - "total_number_of_atoms" - ] + array_length = 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"] + valid_indices = selection["flatten_indices"] self._outputData.add( "selected_atoms", "LineOutputVariable", diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index 99679ff97..65ef227d5 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -148,8 +148,8 @@ def __init__( self.selection_model = SelectionModel(self.trajectory) self._field = field self.atm_full_names = self.system.name_list - self.molecule_names = [str(x) for x in self.system._clusters.keys()] - self.labels = [str(x) for x in self.system._labels.keys()] + self.molecule_names = list(map(str, self.system._clusters)) + self.labels = list(map(str, self.system._labels)) self.selection_textbox = QPlainTextEdit() self.selection_textbox.setReadOnly(True) From 437de9ac1f8f95660769d03a5a6748a30ef150b0 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Fri, 28 Feb 2025 09:29:43 +0000 Subject: [PATCH 24/37] Implement changes from code review --- MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py | 4 +++- .../Src/MDANSE/Framework/AtomSelector/molecule_selection.py | 2 +- .../Framework/Configurators/AtomSelectionConfigurator.py | 5 ++--- .../Framework/Configurators/AtomTransmutationConfigurator.py | 3 +-- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py index 1522877f4..3b7ce593c 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py @@ -41,7 +41,9 @@ def select_labels( """ system = trajectory.chemical_system atom_labels = function_parameters.get("atom_labels", ()) - selection = {system._labels[label] for label in atom_labels if label in system._labels} + selection = { + system._labels[label] for label in atom_labels if label in system._labels + } return selection diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py index a636b49db..47fb539d9 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py @@ -40,7 +40,7 @@ def select_molecules( """ selection = set() system = trajectory.chemical_system - molecule_names = function_parameters.get("molecule_names", None) + molecule_names = function_parameters.get("molecule_names", []) selection = { index for molecule in molecule_names diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py index d6da23792..04825504e 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py @@ -15,6 +15,7 @@ # from json import JSONDecodeError +from collections import Counter import numpy as np @@ -95,9 +96,7 @@ def get_natoms(self) -> dict[str, int]: dict 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 = Counter(self["names"]) return nAtomsPerElement def get_total_natoms(self) -> int: diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py index b9ae9c136..4bd55dd94 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py @@ -57,8 +57,7 @@ def apply_transmutation(self, selection_string: str, symbol: str) -> None: self.selector.load_from_json(selection_string) indices = self.selector.select_in_trajectory(self._current_trajectory) - for idx in indices: - self._new_map[idx] = symbol + self._new_map.update(dict.fromkeys(indices, symbol)) def get_setting(self) -> dict[int, str]: """ From 634eab40104b4f8d12596c500aa9003579d37e93 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Fri, 28 Feb 2025 11:40:29 +0000 Subject: [PATCH 25/37] Replace some constants with StrEnum --- .../InputWidgets/AtomSelectionWidget.py | 4 +-- .../MDANSE_GUI/Widgets/SelectionWidgets.py | 29 +++++++++---------- 2 files changed, 15 insertions(+), 18 deletions(-) diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index 65ef227d5..5cc53defe 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -16,7 +16,7 @@ from typing import Set import json -from enum import Enum +from enum import StrEnum from qtpy.QtCore import Slot, Signal from qtpy.QtGui import QStandardItemModel, QStandardItem @@ -51,7 +51,7 @@ ) -class SelectionValidity(Enum): +class SelectionValidity(StrEnum): VALID_SELECTION = "Valid selection" USELESS_SELECTION = "Selection did not change. This operation is not needed." MALFORMED_SELECTION = "This is not a valid JSON string." diff --git a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py index 0b87d07e7..9687e83ae 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py @@ -16,6 +16,7 @@ from typing import Dict, Any, TYPE_CHECKING, Tuple import json +from enum import StrEnum import numpy as np from qtpy.QtCore import Signal, Slot @@ -36,6 +37,12 @@ from MDANSE_GUI.MolecularViewer.MolecularViewer import MolecularViewer +class IndexSelectionMode(StrEnum): + LIST = "list" + RANGE = "range" + SLICE = "slice" + + class XYZValidator(QValidator): """A custom validator for a QLineEdit. It is intended to limit the input to a string @@ -108,6 +115,7 @@ def add_standard_widgets(self): self.mode_box.addItems( ["Add (union)", "Filter (intersection)", "Remove (difference)"] ) + self._mode_box_values = ["union", "intersection", "difference"] self.commit_button = QPushButton("Apply", self) layout = self.layout() layout.addWidget(self.mode_box) @@ -115,12 +123,7 @@ def add_standard_widgets(self): self.commit_button.clicked.connect(self.create_selection) def get_mode(self) -> str: - if self.mode_box.currentIndex() == 2: - return "difference" - elif self.mode_box.currentIndex() == 1: - return "intersection" - else: - return "union" + return self._mode_box_values[self.mode_box.currentIndex()] def create_selection(self): funtion_parameters = self.parameter_dictionary() @@ -203,13 +206,7 @@ def add_specific_widgets(self): layout = self.layout() layout.addWidget(QLabel("Select atoms by index")) self.selection_type_combo = QComboBox(self) - self.selection_type_combo.addItems( - [ - "list", - "range", - "slice", - ] - ) + self.selection_type_combo.addItems(str(mode) for mode in IndexSelectionMode) self.selection_type_combo.setEditable(False) layout.addWidget(self.selection_type_combo) self.selection_field = QLineEdit(self) @@ -219,15 +216,15 @@ def add_specific_widgets(self): @Slot(str) def switch_mode(self, new_mode: str): self.selection_field.setText("") - if new_mode == "list": + if new_mode == IndexSelectionMode.LIST: self.selection_field.setPlaceholderText("0,1,2") self.selection_keyword = "index_list" self.selection_separator = "," - if new_mode == "range": + elif new_mode == IndexSelectionMode.RANGE: self.selection_field.setPlaceholderText("0-20") self.selection_keyword = "index_range" self.selection_separator = "-" - if new_mode == "slice": + elif new_mode == IndexSelectionMode.SLICE: self.selection_field.setPlaceholderText("first:last:step") self.selection_keyword = "index_slice" self.selection_separator = ":" From 52538231646caa42599f26c6b257feb4a05f1499 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Fri, 28 Feb 2025 18:08:14 +0000 Subject: [PATCH 26/37] Rewrite selection tests using parametrize --- .../UnitTests/test_reusable_selection.py | 149 +++++++++--------- MDANSE/Tests/UnitTests/test_selection.py | 131 ++++++++------- .../InputWidgets/AtomTransmutationWidget.py | 8 +- 3 files changed, 150 insertions(+), 138 deletions(-) diff --git a/MDANSE/Tests/UnitTests/test_reusable_selection.py b/MDANSE/Tests/UnitTests/test_reusable_selection.py index 74a36b6b6..ecdf9bc4e 100644 --- a/MDANSE/Tests/UnitTests/test_reusable_selection.py +++ b/MDANSE/Tests/UnitTests/test_reusable_selection.py @@ -28,43 +28,43 @@ "2vb1.mdt" ) +@pytest.fixture(scope='module') +def trajectory(request): + return HDFTrajectoryInputData(request.param) -@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) + +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj], indirect=True) def test_select_all(trajectory): - traj_object = HDFTrajectoryInputData(trajectory) - n_atoms = len(traj_object.chemical_system.atom_list) + n_atoms = len(trajectory.chemical_system.atom_list) reusable_selection = ReusableSelection() reusable_selection.set_selection(None, {'function_name': 'select_all'}) - selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + selection = reusable_selection.select_in_trajectory(trajectory.trajectory) assert len(selection) == n_atoms -@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj], indirect=True) def test_empty_json_string_selects_all(trajectory): - traj_object = HDFTrajectoryInputData(trajectory) - n_atoms = len(traj_object.chemical_system.atom_list) + n_atoms = len(trajectory.chemical_system.atom_list) reusable_selection = ReusableSelection() reusable_selection.load_from_json('{}') - selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + selection = reusable_selection.select_in_trajectory(trajectory.trajectory) assert len(selection) == n_atoms -@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj], indirect=True) def test_select_none(trajectory): - traj_object = HDFTrajectoryInputData(trajectory) reusable_selection = ReusableSelection() reusable_selection.set_selection(None, {'function_name': 'select_none'}) - selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + selection = reusable_selection.select_in_trajectory(trajectory.trajectory) assert len(selection) == 0 -@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj], indirect=True) def test_inverted_all_is_none(trajectory): - traj_object = HDFTrajectoryInputData(trajectory) reusable_selection = ReusableSelection() reusable_selection.set_selection(None, {'function_name': 'select_all'}) reusable_selection.set_selection(None, {'function_name': 'invert_selection'}) - selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + selection = reusable_selection.select_in_trajectory(trajectory.trajectory) assert len(selection) == 0 @@ -79,126 +79,123 @@ def test_json_saving_is_reversible(): assert json_string == json_string_2 -@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj], indirect=True) def test_selection_from_json_is_the_same_as_from_runtime(trajectory): - traj_object = HDFTrajectoryInputData(trajectory) reusable_selection = ReusableSelection() reusable_selection.set_selection(None, {'function_name': 'select_all'}) reusable_selection.set_selection(None, {'function_name': 'invert_selection'}) - selection = reusable_selection.select_in_trajectory(traj_object.trajectory) + selection = reusable_selection.select_in_trajectory(trajectory.trajectory) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() another_selection.load_from_json(json_string) - selection2 = another_selection.select_in_trajectory(traj_object.trajectory) + selection2 = another_selection.select_in_trajectory(trajectory.trajectory) print(f"original: {reusable_selection.operations}") print(f"another: {another_selection.operations}") assert selection == selection2 -def test_select_atoms_selects_by_element(): - traj_object = HDFTrajectoryInputData(short_traj) - reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'atom_types': ['S']}) - s_selection = reusable_selection.select_in_trajectory(traj_object.trajectory) - assert len(s_selection) == 208 - reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'atom_types': ['Cu']}) - cu_selection = reusable_selection.select_in_trajectory(traj_object.trajectory) - assert len(cu_selection) == 208 - reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'atom_types': ['Sb']}) - sb_selection = reusable_selection.select_in_trajectory(traj_object.trajectory) - assert len(sb_selection) == 64 + +@pytest.mark.parametrize("trajectory", [short_traj], indirect=True) +@pytest.mark.parametrize("element, expected", ( + (["Cu"], 208), + (["S"], 208), + (["Sb"], 64), + (["S", "Sb", "Cu"], 480) +)) +def test_select_atoms_selects_by_element(trajectory, element, expected): reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'atom_types': ['S', 'Sb', 'Cu']}) + reusable_selection.set_selection(None, {"function_name": "select_atoms", + "atom_types": element}) + selection = reusable_selection.select_in_trajectory(trajectory) + assert len(selection) == expected json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() another_selection.load_from_json(json_string) - cusbs_selection = another_selection.select_in_trajectory(traj_object.trajectory) - assert len(cusbs_selection) == 480 + loaded_selection = another_selection.select_in_trajectory(trajectory) + assert len(loaded_selection) == expected -def test_select_atoms_selects_by_range(): - traj_object = HDFTrajectoryInputData(short_traj) - reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'index_range': [15,35]}) - range_selection = reusable_selection.select_in_trajectory(traj_object.trajectory) - assert len(range_selection) == 20 +@pytest.mark.parametrize("trajectory", [short_traj], indirect=True) +@pytest.mark.parametrize("index_range, expected", ( + ([15,35], 20), + ([470,510], 10), +)) +def test_select_atoms_selects_by_range(trajectory, index_range, expected): reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'index_range': [470,510]}) - overshoot_selection = reusable_selection.select_in_trajectory(traj_object.trajectory) - assert len(overshoot_selection) == 10 + reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'index_range': index_range}) + range_selection = reusable_selection.select_in_trajectory(trajectory.trajectory) + assert len(range_selection) == expected -def test_select_atoms_selects_by_slice(): - traj_object = HDFTrajectoryInputData(short_traj) +@pytest.mark.parametrize("trajectory", [short_traj], indirect=True) +@pytest.mark.parametrize("index_slice, expected", ( + ([150,350,10], 20), + ([470,510,5], 2), +)) +def test_select_atoms_selects_by_slice(trajectory, index_slice, expected): reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'index_slice': [150,350,10]}) - range_selection = reusable_selection.select_in_trajectory(traj_object.trajectory) - assert len(range_selection) == 20 - reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'index_slice': [470,510,5]}) + reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'index_slice': index_slice}) + range_selection = reusable_selection.select_in_trajectory(trajectory.trajectory) + assert len(range_selection) == expected json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() another_selection.load_from_json(json_string) - overshoot_selection = another_selection.select_in_trajectory(traj_object.trajectory) - assert len(overshoot_selection) == 2 - - -def test_select_molecules_selects_water(): - traj_object = HDFTrajectoryInputData(traj_2vb1) - reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_molecules', 'molecule_names': ['H2 O1']}) - water_selection = reusable_selection.select_in_trajectory(traj_object.trajectory) - assert len(water_selection) == 28746 + overshoot_selection = another_selection.select_in_trajectory(trajectory.trajectory) + assert len(overshoot_selection) == expected -def test_select_molecules_selects_protein(): - traj_object = HDFTrajectoryInputData(traj_2vb1) +@pytest.mark.parametrize("trajectory", [traj_2vb1], indirect=True) +@pytest.mark.parametrize("molecule_names, expected", ( + (['H2 O1'], 28746), + (['C613 H959 N193 O185 S10'], 1960), +)) +def test_select_molecules_selects_water(trajectory, molecule_names, expected): reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_molecules', 'molecule_names': ['C613 H959 N193 O185 S10']}) + reusable_selection.set_selection(None, {'function_name': 'select_molecules', 'molecule_names': molecule_names}) + first_selection = reusable_selection.select_in_trajectory(trajectory.trajectory) + assert len(first_selection) == expected json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() another_selection.load_from_json(json_string) - protein_selection = another_selection.select_in_trajectory(traj_object.trajectory) - assert len(protein_selection) == 613 + 959 + 193 + 185 + 10 + second_selection = another_selection.select_in_trajectory(trajectory.trajectory) + assert len(second_selection) == expected -def test_select_molecules_inverted_selects_ions(): - traj_object = HDFTrajectoryInputData(traj_2vb1) +@pytest.mark.parametrize("trajectory", [traj_2vb1], indirect=True) +def test_select_molecules_inverted_selects_ions(trajectory): reusable_selection = ReusableSelection() reusable_selection.set_selection(None, {'function_name': 'select_molecules', 'molecule_names': ['C613 H959 N193 O185 S10', 'H2 O1']}) reusable_selection.set_selection(None, {'function_name': 'invert_selection'}) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() another_selection.load_from_json(json_string) - non_molecules_selection = another_selection.select_in_trajectory(traj_object.trajectory) - assert all([traj_object.chemical_system.atom_list[index] in ['Na', 'Cl'] for index in non_molecules_selection]) + non_molecules_selection = another_selection.select_in_trajectory(trajectory.trajectory) + assert all([trajectory.chemical_system.atom_list[index] in ['Na', 'Cl'] for index in non_molecules_selection]) -def test_select_pattern_selects_water(): - traj_object = HDFTrajectoryInputData(traj_2vb1) +@pytest.mark.parametrize("trajectory", [traj_2vb1], indirect=True) +def test_select_pattern_selects_water(trajectory): reusable_selection = ReusableSelection() reusable_selection.set_selection(None, {'function_name': 'select_pattern', 'rdkit_pattern': "[#8X2;H2](~[H])~[H]"}) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() another_selection.load_from_json(json_string) - water_selection = another_selection.select_in_trajectory(traj_object.trajectory) + water_selection = another_selection.select_in_trajectory(trajectory.trajectory) assert len(water_selection) == 28746 -def test_selection_with_multiple_steps(): +@pytest.mark.parametrize("trajectory", [traj_2vb1], indirect=True) +def test_selection_with_multiple_steps(trajectory): """This tests if the ReusableSelection can select oxygen only in the water molecules. It combines two steps: 1. water is selected using rdkit pattern matching 2. oxygen is selected using simple atom type matching; intersection of the selections is applied The selection is then saved to a JSON string, loaded from the string and applied to the trajectory. """ - traj_object = HDFTrajectoryInputData(traj_2vb1) reusable_selection = ReusableSelection() reusable_selection.set_selection(None, {'function_name': 'select_pattern', 'rdkit_pattern': "[#8X2;H2](~[H])~[H]"}) reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'atom_types': ['O'], 'operation_type': 'intersection'}) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() another_selection.load_from_json(json_string) - water_oxygen_selection = another_selection.select_in_trajectory(traj_object.trajectory) + water_oxygen_selection = another_selection.select_in_trajectory(trajectory.trajectory) assert len(water_oxygen_selection) == int(28746/3) diff --git a/MDANSE/Tests/UnitTests/test_selection.py b/MDANSE/Tests/UnitTests/test_selection.py index 5749d27a9..6e8a89a78 100644 --- a/MDANSE/Tests/UnitTests/test_selection.py +++ b/MDANSE/Tests/UnitTests/test_selection.py @@ -34,6 +34,11 @@ ) +@pytest.fixture(scope='module') +def trajectory(request): + return HDFTrajectoryInputData(request.param) + + @pytest.fixture(scope="module") def short_trajectory(): traj_object = HDFTrajectoryInputData(short_traj) @@ -48,85 +53,93 @@ def gromacs_trajectory(): traj_object.close() -@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj], indirect=True) def test_select_all(trajectory): - traj_object = HDFTrajectoryInputData(trajectory) - n_atoms = len(traj_object.chemical_system.atom_list) - selection = select_all(traj_object.trajectory) + n_atoms = len(trajectory.chemical_system.atom_list) + selection = select_all(trajectory.trajectory) assert len(selection) == n_atoms -@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj], indirect=True) def test_select_none(trajectory): - traj_object = HDFTrajectoryInputData(trajectory) - selection = select_none(traj_object.trajectory) + selection = select_none(trajectory.trajectory) assert len(selection) == 0 -@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj]) +@pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj], indirect=True) def test_inverted_none_is_all(trajectory): - traj_object = HDFTrajectoryInputData(trajectory) - none_selection = select_none(traj_object.trajectory) - all_selection = select_all(traj_object.trajectory) - inverted_none = invert_selection(traj_object.trajectory, none_selection) + none_selection = select_none(trajectory.trajectory) + all_selection = select_all(trajectory.trajectory) + inverted_none = invert_selection(trajectory.trajectory, none_selection) assert all_selection == inverted_none -def test_select_atoms_selects_by_element(short_trajectory): - s_selection = select_atoms(short_trajectory, atom_types=['S']) - assert len(s_selection) == 208 - cu_selection = select_atoms(short_trajectory, atom_types=['Cu']) - assert len(cu_selection) == 208 - sb_selection = select_atoms(short_trajectory, atom_types=['Sb']) - assert len(sb_selection) == 64 - cusbs_selection = select_atoms(short_trajectory, atom_types=['Cu','Sb', 'S']) - assert len(cusbs_selection) == 480 - - -def test_select_atoms_selects_by_range(short_trajectory): - range_selection = select_atoms(short_trajectory, index_range=[15,35]) - assert len(range_selection) == 20 - overshoot_selection = select_atoms(short_trajectory, index_range=[470,510]) - assert len(overshoot_selection) == 10 - - -def test_select_atoms_selects_by_slice(short_trajectory): - range_selection = select_atoms(short_trajectory, index_slice=[150,350, 10]) - assert len(range_selection) == 20 - overshoot_selection = select_atoms(short_trajectory, index_slice=[470,510,5]) - assert len(overshoot_selection) == 2 - - -def test_select_molecules_selects_water(gromacs_trajectory): - water_selection = select_molecules(gromacs_trajectory, molecule_names = ['H2 O1']) - assert len(water_selection) == 28746 - - -def test_select_molecules_selects_protein(gromacs_trajectory): - protein_selecton = select_molecules(gromacs_trajectory, molecule_names = ['C613 H959 N193 O185 S10']) - assert len(protein_selecton) == 613 + 959 + 193 + 185 + 10 - - -def test_select_molecules_inverted_selects_ions(gromacs_trajectory): - all_molecules_selection = select_molecules(gromacs_trajectory, molecule_names = ['C613 H959 N193 O185 S10', 'H2 O1']) - non_molecules_selection = invert_selection(gromacs_trajectory, all_molecules_selection) - assert all([gromacs_trajectory.chemical_system.atom_list[index] in ['Na', 'Cl'] for index in non_molecules_selection]) - - -def test_select_pattern_selects_water(gromacs_trajectory): - water_selection = select_pattern(gromacs_trajectory, rdkit_pattern="[#8X2;H2](~[H])~[H]") +@pytest.mark.parametrize("trajectory", [short_traj], indirect=True) +@pytest.mark.parametrize("element, expected", ( + (["Cu"], 208), + (["S"], 208), + (["Sb"], 64), + (["S", "Sb", "Cu"], 480) +)) +def test_select_atoms_selects_by_element(trajectory, element, expected): + s_selection = select_atoms(trajectory.trajectory, atom_types=element) + assert len(s_selection) == expected + + +@pytest.mark.parametrize("trajectory", [short_traj], indirect=True) +@pytest.mark.parametrize("index_range, expected", ( + ([15,35], 20), + ([470,510], 10), +)) +def test_select_atoms_selects_by_range(trajectory, index_range, expected): + range_selection = select_atoms(trajectory.trajectory, index_range=index_range) + assert len(range_selection) == expected + + +@pytest.mark.parametrize("trajectory", [short_traj], indirect=True) +@pytest.mark.parametrize("index_slice, expected", ( + ([150,350,10], 20), + ([470,510,5], 2), +)) +def test_select_atoms_selects_by_slice(trajectory, index_slice, expected): + range_selection = select_atoms(trajectory.trajectory, index_slice=index_slice) + assert len(range_selection) == expected + + +@pytest.mark.parametrize("trajectory", [traj_2vb1], indirect=True) +@pytest.mark.parametrize("molecule_names, expected", ( + (['H2 O1'], 28746), + (['C613 H959 N193 O185 S10'], 1960), +)) +def test_select_molecules(trajectory, molecule_names, expected): + water_selection = select_molecules(trajectory.trajectory, molecule_names = molecule_names) + assert len(water_selection) == expected + + +@pytest.mark.parametrize("trajectory", [traj_2vb1], indirect=True) +def test_select_molecules_inverted_selects_ions(trajectory): + all_molecules_selection = select_molecules(trajectory.trajectory, molecule_names = ['C613 H959 N193 O185 S10', 'H2 O1']) + non_molecules_selection = invert_selection(trajectory.trajectory, all_molecules_selection) + assert all([trajectory.chemical_system.atom_list[index] in ['Na', 'Cl'] for index in non_molecules_selection]) + + +@pytest.mark.parametrize("trajectory", [traj_2vb1], indirect=True) +def test_select_pattern_selects_water(trajectory): + water_selection = select_pattern(trajectory.trajectory, rdkit_pattern="[#8X2;H2](~[H])~[H]") assert len(water_selection) == 28746 -def test_select_positions(short_trajectory): - cube_selection = select_positions(short_trajectory, +@pytest.mark.parametrize("trajectory", [short_traj], indirect=True) +def test_select_positions(trajectory): + cube_selection = select_positions(trajectory.trajectory, position_minimum = 0.5*np.ones(3), position_maximum = 0.7*np.ones(3)) assert cube_selection == {59} -def test_select_sphere(short_trajectory): - sphere_selection = select_sphere(short_trajectory, +@pytest.mark.parametrize("trajectory", [short_traj], indirect=True) +def test_select_sphere(trajectory): + sphere_selection = select_sphere(trajectory.trajectory, sphere_centre = 0.6*np.ones(3), sphere_radius = 0.2) assert sphere_selection == {19, 17, 59, 15} diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomTransmutationWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomTransmutationWidget.py index 305ef038b..486dabe90 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomTransmutationWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomTransmutationWidget.py @@ -132,10 +132,12 @@ def update_transmutation_textbox(self) -> None: """Update the transmutation textbox with the current transmuter setting information. """ - map = self.transmuter.get_setting() + substitutions = self.transmuter.get_setting() - text = [f"Number of atoms transmuted:\n{len(map)}\n\nTransmuted atoms:\n"] - for idx, symbol in map.items(): + text = [ + f"Number of atoms transmuted:\n{len(substitutions)}\n\nTransmuted atoms:\n" + ] + for idx, symbol in substitutions.items(): text.append(f"{idx} {self.atm_full_names[idx]} -> {symbol}\n") self.transmutation_textbox.setText("".join(text)) From d117d6ef1856e95ec2c7fc7536446f74c4557a86 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak <108934199+MBartkowiakSTFC@users.noreply.github.com> Date: Mon, 3 Mar 2025 14:24:12 +0000 Subject: [PATCH 27/37] Apply suggestions from code review Co-authored-by: Jacob Wilkins <46597752+oerc0122@users.noreply.github.com> Signed-off-by: Maciej Bartkowiak <108934199+MBartkowiakSTFC@users.noreply.github.com> --- MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py | 6 +++--- .../Framework/AtomSelector/spatial_selection.py | 6 +++--- .../Src/MDANSE_GUI/Widgets/SelectionWidgets.py | 12 ++++++++---- 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index 24c12e852..3c56bc1ce 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -101,7 +101,7 @@ def validate_selection_string( True if the selection adds atoms, False otherwise """ function_parameters = json.loads(json_string) - if len(self.operations) == 0: + if not self.operations: return True function_name = function_parameters.get("function_name", "select_all") if function_name == "invert_selection": @@ -134,8 +134,8 @@ def validate_selection_string( def select_in_trajectory(self, trajectory: Trajectory) -> Set[int]: selection = set() self.all_idxs = set(range(len(trajectory.chemical_system.atom_list))) - sequence = sorted([int(x) for x in self.operations.keys()]) - if len(sequence) == 0: + sequence = sorted(map(int, self.operations)) + if not sequence: return self.all_idxs for number in sequence: function_parameters = self.operations[number] diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py index ec8928bbb..03b42d8db 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py @@ -47,9 +47,9 @@ def select_positions( coordinates = trajectory.coordinates(function_parameters.get("frame_number", 0)) lower_limits = np.array(function_parameters.get("position_minimum", 3 * [-np.inf])) upper_limits = np.array(function_parameters.get("position_maximum", 3 * [np.inf])) - mask1 = np.all(coordinates > lower_limits.reshape((1, 3)), axis=1) - mask2 = np.all(coordinates < upper_limits.reshape((1, 3)), axis=1) - return set(np.where(np.logical_and(mask1, mask2))[0]) + valid = np.where(((coordinates > lower_limts) & + (coordinates < upper_limits)).all(axis=1)) + return set(valid) def select_sphere( diff --git a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py index 9687e83ae..588dcd206 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py @@ -69,16 +69,20 @@ def validate(self, input_string: str, position: int) -> Tuple[int, str]: Returns ------- - Tuple[int,str] - a tuple of (validator state, input string, cursor position) + int + Validator state. + str + Original input string. + int + Cursor position. """ state = QValidator.State.Intermediate comma_count = input_string.count(",") - if len(input_string) > 0: + if input_string: try: values = [float(x) for x in input_string.split(",")] except (TypeError, ValueError): - if input_string[-1] == "," and comma_count < 3: + if input_string.endswith(",") and comma_count < 3: state = QValidator.State.Intermediate else: state = QValidator.State.Invalid From baac644e633f2b973444a8200b874207947d92f9 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Tue, 4 Mar 2025 15:05:01 +0000 Subject: [PATCH 28/37] Respond to review, correct selection functions --- .../Framework/AtomSelector/atom_selection.py | 37 +++++----- .../AtomSelector/general_selection.py | 17 ++--- .../Framework/AtomSelector/group_selection.py | 24 +++---- .../AtomSelector/molecule_selection.py | 13 ++-- .../AtomSelector/spatial_selection.py | 68 ++++++++++++------- .../Src/MDANSE/Trajectory/H5MDTrajectory.py | 1 + MDANSE/Tests/UnitTests/test_selection.py | 13 ++-- 7 files changed, 96 insertions(+), 77 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py index 186142356..169d8b99e 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py @@ -14,14 +14,20 @@ # along with this program. If not, see . # -from typing import Union, Dict, Any, Set +from typing import Set, Sequence -from MDANSE.Chemistry.ChemicalSystem import ChemicalSystem from MDANSE.MolecularDynamics.Trajectory import Trajectory def select_atoms( - trajectory: Trajectory, **function_parameters: Dict[str, Any] + trajectory: Trajectory, + *, + index_list: Sequence[int] = None, + index_range: Sequence[int] = None, + index_slice: Sequence[int] = None, + atom_types: Sequence[str] = (), + atom_names: Sequence[str] = (), + **kwargs: str, ) -> Set[int]: """Selects specific atoms in the trajectory. These can be selected based on indices, atom type or trajectory-specific atom name. @@ -33,38 +39,37 @@ def select_atoms( ---------- trajectory : Trajectory A trajectory instance to which the selection is applied - function_parameters : Dict[str, Any] - may include any combination of the following - "index_list" : List[int] - "index_range" : a start, stop pair of indices for range(start, stop) - "index_slice" : s start, stop, step sequence of indices for range(start, stop, step) - "atom_types" : List[str] - "atom_names" : List[str] + index_list : Sequence[int] + a list of indices to be selected + index_range : Sequence[int] + a pair of (first, last+1) indices defining a range + index_slice : Sequence[int] + a sequence of (first, last+1, step) indices defining a slice + atom_types : Sequence[str] + a list of atom types (i.e. chemical elements) to be selected, given as string + atom_names : Sequence[str] + a list of atom names (as used by the MD engine, force field, etc.) to be selected Returns ------- Set[int] - Set of all the atom indices + A set of indices which have been selected """ + selection = set() system = trajectory.chemical_system element_list = system.atom_list name_list = system.name_list indices = set(range(len(element_list))) - index_list = function_parameters.get("index_list") - index_range = function_parameters.get("index_range") - index_slice = function_parameters.get("index_slice") if index_list is not None: selection |= indices & set(index_list) if index_range is not None: selection |= indices & set(range(*index_range)) if index_slice is not None: selection |= indices & set(range(*index_slice)) - atom_types = function_parameters.get("atom_types", ()) if atom_types: new_indices = {index for index in indices if element_list[index] in atom_types} selection |= new_indices - atom_names = function_parameters.get("atom_names", ()) if atom_names: new_indices = {index for index in indices if name_list[index] in atom_names} selection |= new_indices diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py index 9e87a6534..2b7ac3ae3 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py @@ -14,14 +14,11 @@ # along with this program. If not, see . # -from typing import Union, Dict, Any, Set -from MDANSE.Chemistry.ChemicalSystem import ChemicalSystem +from typing import Set from MDANSE.MolecularDynamics.Trajectory import Trajectory -def select_all( - trajectory: Trajectory, **function_parameters: Dict[str, Any] -) -> Set[int]: +def select_all(trajectory: Trajectory, **kwargs: str) -> Set[int]: """Selects all the atoms in the trajectory. Parameters @@ -37,14 +34,12 @@ def select_all( 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, **kwargs: str) -> Set[int]: """Returns an empty selection. Parameters ---------- - trajectory : Trajectory + _trajectory : Trajectory A trajectory instance, ignored in this selection Returns @@ -55,7 +50,9 @@ def select_none( return set() -def invert_selection(trajectory: Trajectory, selection: Set[int]) -> Set[int]: +def invert_selection( + trajectory: Trajectory, selection: Set[int], **kwargs: str +) -> Set[int]: """Returns a set of all the indices that are present in the trajectory and were not included in the input selection. diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py index 3b7ce593c..8be60687f 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py @@ -14,15 +14,14 @@ # along with this program. If not, see . # -from typing import Union, Dict, Any, Set +from typing import Set, Sequence from functools import reduce -from MDANSE.Chemistry.ChemicalSystem import ChemicalSystem from MDANSE.MolecularDynamics.Trajectory import Trajectory def select_labels( - trajectory: Trajectory, **function_parameters: Dict[str, Any] + trajectory: Trajectory, atom_labels: Sequence[str] = (), **kwargs: str ) -> Set[int]: """Selects atoms with a specific label in the trajectory. A residue name can be read as a label by MDANSE. @@ -31,16 +30,15 @@ def select_labels( ---------- trajectory : Trajectory A trajectory instance to which the selection is applied - function_parameters : Dict[str, Any] - should include a list of string labels under key "atom_labels" + atom_labels : Sequence[str] + a list of string labels (e.g. residue names) by which to select atoms Returns ------- Set[int] - Set of all the atom indices + Set of atom indices corresponding to the selected labels """ system = trajectory.chemical_system - atom_labels = function_parameters.get("atom_labels", ()) selection = { system._labels[label] for label in atom_labels if label in system._labels } @@ -48,7 +46,7 @@ def select_labels( def select_pattern( - trajectory: Trajectory, **function_parameters: Dict[str, Any] + trajectory: Trajectory, rdkit_pattern: str = "", **kwargs: str ) -> Set[int]: """Selects atoms according to the SMARTS string given as input. This will only work if molecules and bonds have been detected in the system. @@ -59,16 +57,16 @@ def select_pattern( ---------- trajectory : Trajectory A trajectory instance to which the selection is applied - function_parameters : Dict[str, Any] - should include a SMARTS string under key "rdkit_pattern" + rdkit_pattern : str + a SMARTS string to be matched Returns ------- Set[int] - Set of all the atom indices + Set of atom indices matched by rdkit """ selection = set() system = trajectory.chemical_system - if pattern := function_parameters.get("rdkit_pattern"): - selection = system.get_substructure_matches(pattern) + if rdkit_pattern: + selection = system.get_substructure_matches(rdkit_pattern) return selection diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py index 47fb539d9..4040abefa 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py @@ -14,15 +14,13 @@ # along with this program. If not, see . # -from typing import Union, Dict, Any, Set -from functools import reduce +from typing import Set, Sequence -from MDANSE.Chemistry.ChemicalSystem import ChemicalSystem from MDANSE.MolecularDynamics.Trajectory import Trajectory def select_molecules( - trajectory: Trajectory, **function_parameters: Dict[str, Any] + trajectory: Trajectory, molecule_names: Sequence[str] = (), **kwargs: str ) -> Set[int]: """Selects all the atoms belonging to the specified molecule types. @@ -30,17 +28,16 @@ def select_molecules( ---------- trajectory : Trajectory A trajectory instance to which the selection is applied - function_parameters : Dict[str, Any] - should include a list of str molecule names under key "atom_labels" + molecule_names : Sequence[str] + a list of molecule names (str) which are keys of ChemicalSystem._clusters Returns ------- Set[int] - Set of all the atom indices + Set of indices of atoms belonging to molecules from molecule_names """ selection = set() system = trajectory.chemical_system - molecule_names = function_parameters.get("molecule_names", []) selection = { index for molecule in molecule_names diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py index 03b42d8db..7cd63e9ca 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py @@ -14,17 +14,21 @@ # along with this program. If not, see . # -from typing import Union, Dict, Any, Set +from typing import Set, Sequence, Union import numpy as np from scipy.spatial import KDTree -from MDANSE.Chemistry.ChemicalSystem import ChemicalSystem from MDANSE.MolecularDynamics.Trajectory import Trajectory def select_positions( - trajectory: Trajectory, **function_parameters: Dict[str, Any] + trajectory: Trajectory, + *, + frame_number: int = 0, + position_minimum: Union[Sequence[float], None] = None, + position_maximum: Union[Sequence[float], None] = None, + **kwargs: str, ) -> Set[int]: """Selects atoms based on their positions at a specified frame number. Lower and upper limits of x, y and z coordinates can be given as input. @@ -32,28 +36,41 @@ def select_positions( Parameters ---------- trajectory : Trajectory - A trajectory instance to which the selection is applied - function_parameters : Dict[str, Any] - may include - "frame_number" : int - "position_minimum" : np.ndarray[float] - "position_maximum" : np.ndarray[float] + a trajectory instance in which the atoms are being selected + frame_number : int, optional + trajectory frame at which to check the coordinates, by default 0 + position_minimum : Sequence[float], optional + (x, y, z) lower limits of coordinates to be selected, by default None + position_maximum : Sequence[float], optional + (x, y, z) upper limits of coordinates to be selected, by default None Returns ------- Set[int] - Set of all the atom indices + _description_ """ - coordinates = trajectory.coordinates(function_parameters.get("frame_number", 0)) - lower_limits = np.array(function_parameters.get("position_minimum", 3 * [-np.inf])) - upper_limits = np.array(function_parameters.get("position_maximum", 3 * [np.inf])) - valid = np.where(((coordinates > lower_limts) & - (coordinates < upper_limits)).all(axis=1)) - return set(valid) + coordinates = trajectory.coordinates(frame_number) + if position_minimum is None: + lower_limits = np.array(3 * [-np.inf]) + else: + lower_limits = np.array(position_minimum) + if position_maximum is None: + upper_limits = np.array(3 * [np.inf]) + else: + upper_limits = np.array(position_maximum) + valid = np.where( + ((coordinates > lower_limits) & (coordinates < upper_limits)).all(axis=1) + ) + return set(valid[0]) def select_sphere( - trajectory: Trajectory, **function_parameters: Dict[str, Any] + trajectory: Trajectory, + *, + frame_number: int = 0, + sphere_centre: Sequence[float], + sphere_radius: float, + **kwargs: str, ) -> Set[int]: """Selects atoms within a distance from a fixed point in space, based on coordinates at a specific frame number. @@ -62,20 +79,19 @@ def select_sphere( ---------- trajectory : Trajectory A trajectory instance to which the selection is applied - function_parameters : Dict[str, Any] - may include - "frame_number" : int - "sphere_centre" : np.ndarray[float] - "sphere_radius" : float + frame_number : int, optional + trajectory frame at which to check the coordinates, by default 0 + sphere_centre : Sequence[float] + (x, y, z) coordinates of the centre of the selection + sphere_radius : float + distance from the centre within which to select atoms Returns ------- Set[int] - Set of all the atom indices + set of indices of atoms inside the sphere """ - coordinates = trajectory.coordinates(function_parameters.get("frame_number", 0)) - sphere_centre = np.array(function_parameters.get("sphere_centre", 3 * [0.0])) - sphere_radius = function_parameters.get("sphere_radius", 1.0) + coordinates = trajectory.coordinates(frame_number) kdtree = KDTree(coordinates) indices = kdtree.query_ball_point(sphere_centre, sphere_radius) return set(indices) diff --git a/MDANSE/Src/MDANSE/Trajectory/H5MDTrajectory.py b/MDANSE/Src/MDANSE/Trajectory/H5MDTrajectory.py index cea09067b..2287303a2 100644 --- a/MDANSE/Src/MDANSE/Trajectory/H5MDTrajectory.py +++ b/MDANSE/Src/MDANSE/Trajectory/H5MDTrajectory.py @@ -127,6 +127,7 @@ def __getitem__(self, frame): :return: the configuration :rtype: dict of ndarray """ + grp = self._h5_file["/particles/all/position/value"] try: pos_unit = grp.attrs["unit"] diff --git a/MDANSE/Tests/UnitTests/test_selection.py b/MDANSE/Tests/UnitTests/test_selection.py index 6e8a89a78..de1cb9ba1 100644 --- a/MDANSE/Tests/UnitTests/test_selection.py +++ b/MDANSE/Tests/UnitTests/test_selection.py @@ -130,11 +130,16 @@ def test_select_pattern_selects_water(trajectory): @pytest.mark.parametrize("trajectory", [short_traj], indirect=True) -def test_select_positions(trajectory): +@pytest.mark.parametrize("lower_limit, upper_limit, expected", ( + ([0.5, 0.5, 0.5], [0.7, 0.7, 0.7], {59}), + (None, [0.2, 0.2, 0.2], {46}), + ([1.8, 1.8, 1.8], None, {453}), +)) +def test_select_positions(trajectory, lower_limit, upper_limit, expected): cube_selection = select_positions(trajectory.trajectory, - position_minimum = 0.5*np.ones(3), - position_maximum = 0.7*np.ones(3)) - assert cube_selection == {59} + position_minimum = lower_limit, + position_maximum = upper_limit) + assert cube_selection == expected @pytest.mark.parametrize("trajectory", [short_traj], indirect=True) From 05e0bd8d31835138dda4f8ae0e82ee2373246464 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Tue, 4 Mar 2025 15:58:24 +0000 Subject: [PATCH 29/37] Improve docstrings, add all_indices property --- MDANSE/Src/MDANSE/Chemistry/ChemicalSystem.py | 7 ++++- .../Framework/AtomSelector/atom_selection.py | 4 +-- .../AtomSelector/general_selection.py | 10 +++---- .../Framework/AtomSelector/group_selection.py | 4 +-- .../AtomSelector/molecule_selection.py | 2 +- .../MDANSE/Framework/AtomSelector/selector.py | 28 +++++++++++++++++-- .../AtomSelector/spatial_selection.py | 4 +-- .../InputWidgets/AtomSelectionWidget.py | 2 +- .../MDANSE_GUI/Widgets/SelectionWidgets.py | 2 +- 9 files changed, 46 insertions(+), 17 deletions(-) diff --git a/MDANSE/Src/MDANSE/Chemistry/ChemicalSystem.py b/MDANSE/Src/MDANSE/Chemistry/ChemicalSystem.py index 7cf476dde..0f22ad5b7 100644 --- a/MDANSE/Src/MDANSE/Chemistry/ChemicalSystem.py +++ b/MDANSE/Src/MDANSE/Chemistry/ChemicalSystem.py @@ -15,7 +15,7 @@ # from __future__ import annotations -from typing import List, Tuple, Dict, Any +from typing import List, Tuple, Dict, Any, Set import copy from functools import reduce @@ -229,6 +229,11 @@ def number_of_atoms(self) -> int: """The number of non-ghost atoms in the ChemicalSystem.""" return self._total_number_of_atoms + @property + def all_indices(self) -> Set[int]: + """The number of non-ghost atoms in the ChemicalSystem.""" + return set(self._atom_indices) + @property def total_number_of_atoms(self) -> int: """The number of all atoms in the ChemicalSystem, including ghost ones.""" diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py index 169d8b99e..0e0c32052 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py @@ -27,7 +27,7 @@ def select_atoms( index_slice: Sequence[int] = None, atom_types: Sequence[str] = (), atom_names: Sequence[str] = (), - **kwargs: str, + **_kwargs: str, ) -> Set[int]: """Selects specific atoms in the trajectory. These can be selected based on indices, atom type or trajectory-specific atom name. @@ -60,7 +60,7 @@ def select_atoms( system = trajectory.chemical_system element_list = system.atom_list name_list = system.name_list - indices = set(range(len(element_list))) + indices = system.all_indices if index_list is not None: selection |= indices & set(index_list) if index_range is not None: diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py index 2b7ac3ae3..fb2174e1d 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py @@ -18,7 +18,7 @@ from MDANSE.MolecularDynamics.Trajectory import Trajectory -def select_all(trajectory: Trajectory, **kwargs: str) -> Set[int]: +def select_all(trajectory: Trajectory, **_kwargs: str) -> Set[int]: """Selects all the atoms in the trajectory. Parameters @@ -31,10 +31,10 @@ def select_all(trajectory: Trajectory, **kwargs: str) -> Set[int]: Set[int] Set of all the atom indices """ - return set(range(len(trajectory.chemical_system.atom_list))) + return trajectory.chemical_system.all_indices -def select_none(_trajectory: Trajectory, **kwargs: str) -> Set[int]: +def select_none(_trajectory: Trajectory, **_kwargs: str) -> Set[int]: """Returns an empty selection. Parameters @@ -51,7 +51,7 @@ def select_none(_trajectory: Trajectory, **kwargs: str) -> Set[int]: def invert_selection( - trajectory: Trajectory, selection: Set[int], **kwargs: str + trajectory: Trajectory, selection: Set[int], **_kwargs: str ) -> Set[int]: """Returns a set of all the indices that are present in the trajectory and were not included in the input selection. @@ -66,7 +66,7 @@ def invert_selection( Returns ------- Set[int] - set of all the indices in the trajectory which were not in the selection + set of all the indices in the trajectory which were not in the input selection """ all_indices = select_all(trajectory) inverted = all_indices - selection diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py index 8be60687f..6c4e77704 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py @@ -21,7 +21,7 @@ def select_labels( - trajectory: Trajectory, atom_labels: Sequence[str] = (), **kwargs: str + trajectory: Trajectory, atom_labels: Sequence[str] = (), **_kwargs: str ) -> Set[int]: """Selects atoms with a specific label in the trajectory. A residue name can be read as a label by MDANSE. @@ -46,7 +46,7 @@ def select_labels( def select_pattern( - trajectory: Trajectory, rdkit_pattern: str = "", **kwargs: str + trajectory: Trajectory, rdkit_pattern: str = "", **_kwargs: str ) -> Set[int]: """Selects atoms according to the SMARTS string given as input. This will only work if molecules and bonds have been detected in the system. diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py index 4040abefa..4d8c388b9 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/molecule_selection.py @@ -20,7 +20,7 @@ def select_molecules( - trajectory: Trajectory, molecule_names: Sequence[str] = (), **kwargs: str + trajectory: Trajectory, molecule_names: Sequence[str] = (), **_kwargs: str ) -> Set[int]: """Selects all the atoms belonging to the specified molecule types. diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index 3c56bc1ce..a5cbfd083 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -63,6 +63,7 @@ def __init__(self) -> None: self.reset() def reset(self): + """Initialises the attributes to an empty list of operations.""" self.system = None self.trajectory = None self.all_idxs = set() @@ -71,6 +72,15 @@ def reset(self): def set_selection( self, number: Union[int, None] = None, function_parameters: Dict[str, Any] = {} ): + """Appends a new selection operation, or overwrites an existing one. + + Parameters + ---------- + number : Union[int, None], optional + the position of the new selection in the sequence of operations, by default None + function_parameters : Dict[str, Any], optional + the dictionary of keyword arguments defining a selection operation, by default {} + """ if number is None: number = len(self.operations) else: @@ -132,8 +142,21 @@ def validate_selection_string( return False def select_in_trajectory(self, trajectory: Trajectory) -> Set[int]: + """Applies all the selection operations in sequence to the + input trajectory, and returns the resulting set of indices. + + Parameters + ---------- + trajectory : Trajectory + trajectory object in which the atoms will be selected + + Returns + ------- + Set[int] + set of atom indices that have been selected in the input trajectory + """ selection = set() - self.all_idxs = set(range(len(trajectory.chemical_system.atom_list))) + self.all_idxs = trajectory.chemical_system.all_indices sequence = sorted(map(int, self.operations)) if not sequence: return self.all_idxs @@ -169,7 +192,8 @@ def convert_to_json(self) -> str: return json.dumps(self.operations) def load_from_json(self, json_string: str): - """_summary_ + """Loads the atom selection operations from a JSON string. + Adds the operations to the selection sequence. Parameters ---------- diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py index 7cd63e9ca..274a7770b 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py @@ -28,7 +28,7 @@ def select_positions( frame_number: int = 0, position_minimum: Union[Sequence[float], None] = None, position_maximum: Union[Sequence[float], None] = None, - **kwargs: str, + **_kwargs: str, ) -> Set[int]: """Selects atoms based on their positions at a specified frame number. Lower and upper limits of x, y and z coordinates can be given as input. @@ -70,7 +70,7 @@ def select_sphere( frame_number: int = 0, sphere_centre: Sequence[float], sphere_radius: float, - **kwargs: str, + **_kwargs: str, ) -> Set[int]: """Selects atoms within a distance from a fixed point in space, based on coordinates at a specific frame number. diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index 5cc53defe..e2ec1ae94 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -148,7 +148,7 @@ def __init__( self.selection_model = SelectionModel(self.trajectory) self._field = field self.atm_full_names = self.system.name_list - self.molecule_names = list(map(str, self.system._clusters)) + self.molecule_names = self.system.unique_molecules() self.labels = list(map(str, self.system._labels)) self.selection_textbox = QPlainTextEdit() diff --git a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py index 588dcd206..a2bf234a9 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py @@ -251,7 +251,7 @@ def __init__( ): self.molecule_names = [] if trajectory: - self.molecule_names = list(trajectory.chemical_system._clusters.keys()) + self.molecule_names = trajectory.chemical_system.unique_molecules() super().__init__(parent, widget_label) def add_specific_widgets(self): From 4eab2038e2f3bd59bb57c8467a93e801cb0cc54d Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Wed, 5 Mar 2025 15:27:21 +0000 Subject: [PATCH 30/37] Format the code according to linter --- .../Framework/AtomSelector/atom_selection.py | 19 ++-- .../AtomSelector/general_selection.py | 25 +++-- .../Framework/AtomSelector/group_selection.py | 30 ++--- .../MDANSE/Framework/AtomSelector/selector.py | 104 ++++++++++-------- .../AtomSelector/spatial_selection.py | 22 ++-- .../AtomSelectionConfigurator.py | 53 +++++---- 6 files changed, 148 insertions(+), 105 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py index 0e0c32052..10280df9a 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py @@ -14,7 +14,8 @@ # along with this program. If not, see . # -from typing import Set, Sequence +from collections.abc import Sequence +from typing import Optional from MDANSE.MolecularDynamics.Trajectory import Trajectory @@ -22,14 +23,16 @@ def select_atoms( trajectory: Trajectory, *, - index_list: Sequence[int] = None, - index_range: Sequence[int] = None, - index_slice: Sequence[int] = None, + index_list: Optional[Sequence[int]] = None, + index_range: Optional[Sequence[int]] = None, + index_slice: Optional[Sequence[int]] = None, atom_types: Sequence[str] = (), atom_names: Sequence[str] = (), **_kwargs: str, -) -> Set[int]: - """Selects specific atoms in the trajectory. These can be selected based +) -> set[int]: + """Select specific atoms in the trajectory. + + Atoms can be selected based on indices, atom type or trajectory-specific atom name. The atom type is normally the chemical element, while the atom name can be more specific and depend on the @@ -52,10 +55,10 @@ def select_atoms( Returns ------- - Set[int] + set[int] A set of indices which have been selected - """ + """ selection = set() system = trajectory.chemical_system element_list = system.atom_list diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py index fb2174e1d..ea5b5ff18 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/general_selection.py @@ -14,12 +14,11 @@ # along with this program. If not, see . # -from typing import Set from MDANSE.MolecularDynamics.Trajectory import Trajectory -def select_all(trajectory: Trajectory, **_kwargs: str) -> Set[int]: - """Selects all the atoms in the trajectory. +def select_all(trajectory: Trajectory, **_kwargs: str) -> set[int]: + """Select all the atoms in the trajectory. Parameters ---------- @@ -30,12 +29,13 @@ def select_all(trajectory: Trajectory, **_kwargs: str) -> Set[int]: ------- Set[int] Set of all the atom indices + """ return trajectory.chemical_system.all_indices -def select_none(_trajectory: Trajectory, **_kwargs: str) -> Set[int]: - """Returns an empty selection. +def select_none(_trajectory: Trajectory, **_kwargs: str) -> set[int]: + """Return an empty selection. Parameters ---------- @@ -46,14 +46,19 @@ def select_none(_trajectory: Trajectory, **_kwargs: str) -> Set[int]: ------- Set[int] An empty set. + """ return set() def invert_selection( - trajectory: Trajectory, selection: Set[int], **_kwargs: str -) -> Set[int]: - """Returns a set of all the indices that are present in the trajectory + trajectory: Trajectory, + selection: set[int], + **_kwargs: str, +) -> set[int]: + """Invert the current selection for the input trajectory. + + Return a set of all the indices that are present in the trajectory and were not included in the input selection. Parameters @@ -67,7 +72,7 @@ def invert_selection( ------- Set[int] set of all the indices in the trajectory which were not in the input selection + """ all_indices = select_all(trajectory) - inverted = all_indices - selection - return inverted + return all_indices - selection diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py index 6c4e77704..8a5937668 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py @@ -14,17 +14,19 @@ # along with this program. If not, see . # -from typing import Set, Sequence -from functools import reduce +from collections.abc import Sequence from MDANSE.MolecularDynamics.Trajectory import Trajectory def select_labels( - trajectory: Trajectory, atom_labels: Sequence[str] = (), **_kwargs: str -) -> Set[int]: - """Selects atoms with a specific label in the trajectory. - A residue name can be read as a label by MDANSE. + trajectory: Trajectory, + atom_labels: Sequence[str] = (), + **_kwargs: str, +) -> set[int]: + """Select atoms with a specific label in the trajectory. + + A residue name can be used as a label by MDANSE. Parameters ---------- @@ -37,18 +39,19 @@ def select_labels( ------- Set[int] Set of atom indices corresponding to the selected labels + """ system = trajectory.chemical_system - selection = { - system._labels[label] for label in atom_labels if label in system._labels - } - return selection + return {system._labels[label] for label in atom_labels if label in system._labels} def select_pattern( - trajectory: Trajectory, rdkit_pattern: str = "", **_kwargs: str -) -> Set[int]: - """Selects atoms according to the SMARTS string given as input. + trajectory: Trajectory, + rdkit_pattern: str = "", + **_kwargs: str, +) -> set[int]: + """Select atoms according to the SMARTS string given as input. + This will only work if molecules and bonds have been detected in the system. If the bond information was not read from the input trajectory on conversion, it can still be determined in a TrajectoryEditor run. @@ -64,6 +67,7 @@ def select_pattern( ------- Set[int] Set of atom indices matched by rdkit + """ selection = set() system = trajectory.chemical_system diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index a5cbfd083..740fcb551 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -15,21 +15,21 @@ # import json -from typing import Union, Dict, Any, Set -from MDANSE.MolecularDynamics.Trajectory import Trajectory +from typing import Any, Union + +from MDANSE.Framework.AtomSelector.atom_selection import select_atoms from MDANSE.Framework.AtomSelector.general_selection import ( + invert_selection, select_all, select_none, - invert_selection, ) -from MDANSE.Framework.AtomSelector.atom_selection import select_atoms -from MDANSE.Framework.AtomSelector.molecule_selection import select_molecules from MDANSE.Framework.AtomSelector.group_selection import select_labels, select_pattern +from MDANSE.Framework.AtomSelector.molecule_selection import select_molecules from MDANSE.Framework.AtomSelector.spatial_selection import ( select_positions, select_sphere, ) - +from MDANSE.MolecularDynamics.Trajectory import Trajectory function_lookup = { function.__name__: function @@ -48,38 +48,47 @@ class ReusableSelection: - """A reusable sequence of operations which, when applied + """Stores an applies atom selection operations. + + A reusable sequence of operations which, when applied to a trajectory, returns a set of atom indices based on the specified criteria. + """ def __init__(self) -> None: - """ + """Create an empty selection. + Parameters ---------- trajectory: Trajectory The chemical system to apply the selection to. + """ self.reset() def reset(self): - """Initialises the attributes to an empty list of operations.""" + """Initialise the attributes to an empty list of operations.""" self.system = None self.trajectory = None self.all_idxs = set() self.operations = {} def set_selection( - self, number: Union[int, None] = None, function_parameters: Dict[str, Any] = {} + self, + *, + number: Union[int, None] = None, + function_parameters: dict[str, Any], ): - """Appends a new selection operation, or overwrites an existing one. + """Append a new selection operation, or overwrite an existing one. Parameters ---------- number : Union[int, None], optional - the position of the new selection in the sequence of operations, by default None + the position of the new selection in the sequence of operations function_parameters : Dict[str, Any], optional - the dictionary of keyword arguments defining a selection operation, by default {} + the dictionary of keyword arguments defining a selection operation + """ if number is None: number = len(self.operations) @@ -91,9 +100,14 @@ def set_selection( self.operations[number] = function_parameters def validate_selection_string( - self, json_string: str, trajectory: Trajectory, current_selection: Set[int] + self, + json_string: str, + trajectory: Trajectory, + current_selection: set[int], ) -> bool: - """Checks if the selection operation encoded in the input JSON string + """Check if the new selection string changes the current selection. + + Checks if the selection operation encoded in the input JSON string will add any new atoms to the current selection on the given trajectory. Parameters @@ -109,6 +123,7 @@ def validate_selection_string( ------- bool True if the selection adds atoms, False otherwise + """ function_parameters = json.loads(json_string) if not self.operations: @@ -117,32 +132,28 @@ def validate_selection_string( if function_name == "invert_selection": selection = invert_selection(trajectory, current_selection) return True + operation_type = function_parameters.get("operation_type", "union") + function = function_lookup[function_name] + temp_selection = function(trajectory, **function_parameters) + if operation_type == "union": + selection = selection.union(temp_selection) + elif operation_type == "intersection": + selection = selection.intersection(temp_selection) + elif operation_type == "difference": + selection = selection.difference(temp_selection) else: - operation_type = function_parameters.get("operation_type", "union") - function = function_lookup[function_name] - temp_selection = function(trajectory, **function_parameters) - if operation_type == "union": - selection = selection.union(temp_selection) - elif operation_type == "intersection": - selection = selection.intersection(temp_selection) - elif operation_type == "difference": - selection = selection.difference(temp_selection) - else: - selection = temp_selection - if ( - len(selection.difference(current_selection)) > 0 - and operation_type == "union" - ): - return True - elif ( - len(current_selection.difference(selection)) > 0 - and operation_type != "union" - ): - return True - return False - - def select_in_trajectory(self, trajectory: Trajectory) -> Set[int]: - """Applies all the selection operations in sequence to the + selection = temp_selection + return bool( + (len(selection.difference(current_selection)) > 0 + and operation_type == "union") + or (len(current_selection.difference(selection)) > 0 + and operation_type != "union") + ) + + def select_in_trajectory(self, trajectory: Trajectory) -> set[int]: + """Select atoms in the input trajectory. + + Applies all the selection operations in sequence to the input trajectory, and returns the resulting set of indices. Parameters @@ -152,8 +163,9 @@ def select_in_trajectory(self, trajectory: Trajectory) -> Set[int]: Returns ------- - Set[int] + set[int] set of atom indices that have been selected in the input trajectory + """ selection = set() self.all_idxs = trajectory.chemical_system.all_indices @@ -180,7 +192,9 @@ def select_in_trajectory(self, trajectory: Trajectory) -> Set[int]: return selection def convert_to_json(self) -> str: - """For the purpose of storing the selection independent of the + """Output all the operations as a JSON string. + + For the purpose of storing the selection independent of the trajectory it is acting on, this method encodes the sequence of selection operations as a string. @@ -188,17 +202,21 @@ def convert_to_json(self) -> str: ------- str All the operations of this selection, encoded as string + """ return json.dumps(self.operations) def load_from_json(self, json_string: str): - """Loads the atom selection operations from a JSON string. + """Populate the operations sequence from the input string. + + Loads the atom selection operations from a JSON string. Adds the operations to the selection sequence. Parameters ---------- json_string : str A sequence of selection operations, encoded as a JSON string + """ json_setting = json.loads(json_string) for k0, v0 in json_setting.items(): diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py index 274a7770b..ca0ed3510 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py @@ -14,7 +14,8 @@ # along with this program. If not, see . # -from typing import Set, Sequence, Union +from collections.abc import Sequence +from typing import Union import numpy as np from scipy.spatial import KDTree @@ -29,8 +30,9 @@ def select_positions( position_minimum: Union[Sequence[float], None] = None, position_maximum: Union[Sequence[float], None] = None, **_kwargs: str, -) -> Set[int]: - """Selects atoms based on their positions at a specified frame number. +) -> set[int]: + """Select atoms based on their positions at a specified frame number. + Lower and upper limits of x, y and z coordinates can be given as input. Parameters @@ -46,8 +48,9 @@ def select_positions( Returns ------- - Set[int] - _description_ + set[int] + indicies of atoms with coordinates within limits + """ coordinates = trajectory.coordinates(frame_number) if position_minimum is None: @@ -71,8 +74,10 @@ def select_sphere( sphere_centre: Sequence[float], sphere_radius: float, **_kwargs: str, -) -> Set[int]: - """Selects atoms within a distance from a fixed point in space, +) -> set[int]: + """Select atoms within a sphere. + + Selects atoms at a given distance from a fixed point in space, based on coordinates at a specific frame number. Parameters @@ -88,8 +93,9 @@ def select_sphere( Returns ------- - Set[int] + set[int] set of indices of atoms inside the sphere + """ coordinates = trajectory.coordinates(frame_number) kdtree = KDTree(coordinates) diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py index 04825504e..23294ce48 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py @@ -14,17 +14,19 @@ # along with this program. If not, see . # -from json import JSONDecodeError from collections import Counter +from json import JSONDecodeError import numpy as np -from MDANSE.Framework.Configurators.IConfigurator import IConfigurator from MDANSE.Framework.AtomSelector.selector import ReusableSelection +from MDANSE.Framework.Configurators.IConfigurator import IConfigurator class AtomSelectionConfigurator(IConfigurator): - """This configurator allows the selection of a specific set of + """Selects atoms in trajectory based on the input string. + + This configurator allows the selection of a specific set of atoms on which the analysis will be performed. The defaults setting selects all atoms. @@ -32,6 +34,7 @@ class AtomSelectionConfigurator(IConfigurator): ---------- _default : str The defaults selection setting. + """ _default = "{}" @@ -43,6 +46,7 @@ def configure(self, value: str) -> None: ---------- value : str The selection setting in a json readable format. + """ self._original_input = value @@ -67,7 +71,7 @@ def configure(self, value: str) -> None: self.selector.load_from_json(value) indices = self.selector.select_in_trajectory(trajConfig["instance"]) - self["flatten_indices"] = sorted(list(indices)) + self["flatten_indices"] = sorted(indices) atoms = trajConfig["instance"].chemical_system.atom_list self["total_number_of_atoms"] = len(atoms) @@ -77,7 +81,7 @@ def configure(self, value: str) -> None: self["indices"] = [[idx] for idx in self["flatten_indices"]] self["elements"] = [[at] for at in selectedAtoms] - self["names"] = [at for at in selectedAtoms] + self["names"] = list(selectedAtoms) self["unique_names"] = sorted(set(self["names"])) self["masses"] = [ [trajConfig["instance"].get_atom_property(n, "atomic_weight")] @@ -86,29 +90,39 @@ def configure(self, value: str) -> None: if self["selection_length"] == 0: self.error_status = "The atom selection is empty." return - else: - self.error_status = "OK" + self.error_status = "OK" def get_natoms(self) -> dict[str, int]: - """ + """Count the selected atoms, per element. + Returns ------- dict A dictionary of the number of atom per element. + """ - nAtomsPerElement = Counter(self["names"]) - return nAtomsPerElement + return Counter(self["names"]) def get_total_natoms(self) -> int: - """ + """Count all the selected atoms. + Returns ------- int The total number of atoms selected. + """ return len(self["names"]) - def get_indices(self): + def get_indices(self) -> dict[str, list[int]]: + """Group atom indices per chemical element. + + Returns + ------- + dict[str, list[int]] + For each atom type, a list of indices of selected atoms + + """ indicesPerElement = {} for i, v in enumerate(self["names"]): if v in indicesPerElement: @@ -119,11 +133,13 @@ def get_indices(self): return indicesPerElement def get_information(self) -> str: - """ + """Create a text summary of the selection. + Returns ------- str - Some information on the atom selection. + Human-readable information on the atom selection. + """ if "selection_length" not in self: return "Not configured yet\n" @@ -133,12 +149,3 @@ def get_information(self) -> str: info.append(f"Selected elements:{self['unique_names']}") return "\n".join(info) + "\n" - - def get_selector(self) -> "ReusableSelection": - """ - Returns - ------- - ReusableSelection - the instance of the class which selects atoms in a trajectory - """ - return self.selector From 3b61207637d8adcbcc7dbc8066d58bcb3dd69700 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Wed, 5 Mar 2025 17:24:11 +0000 Subject: [PATCH 31/37] Format the GUI code according to linter --- .../MDANSE/Framework/AtomSelector/selector.py | 16 +- .../AtomSelector/spatial_selection.py | 6 +- .../AtomSelectionConfigurator.py | 2 - .../InputWidgets/AtomSelectionWidget.py | 161 +++++++++----- .../MDANSE_GUI/Widgets/SelectionWidgets.py | 199 ++++++++++++++++-- 5 files changed, 304 insertions(+), 80 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index 740fcb551..7ac449585 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -15,7 +15,7 @@ # import json -from typing import Any, Union +from typing import Any, Optional from MDANSE.Framework.AtomSelector.atom_selection import select_atoms from MDANSE.Framework.AtomSelector.general_selection import ( @@ -77,7 +77,7 @@ def reset(self): def set_selection( self, *, - number: Union[int, None] = None, + number: Optional[int] = None, function_parameters: dict[str, Any], ): """Append a new selection operation, or overwrite an existing one. @@ -144,10 +144,14 @@ def validate_selection_string( else: selection = temp_selection return bool( - (len(selection.difference(current_selection)) > 0 - and operation_type == "union") - or (len(current_selection.difference(selection)) > 0 - and operation_type != "union") + ( + len(selection.difference(current_selection)) > 0 + and operation_type == "union" + ) + or ( + len(current_selection.difference(selection)) > 0 + and operation_type != "union" + ) ) def select_in_trajectory(self, trajectory: Trajectory) -> set[int]: diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py index ca0ed3510..d192d6075 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py @@ -15,7 +15,7 @@ # from collections.abc import Sequence -from typing import Union +from typing import Optional import numpy as np from scipy.spatial import KDTree @@ -27,8 +27,8 @@ def select_positions( trajectory: Trajectory, *, frame_number: int = 0, - position_minimum: Union[Sequence[float], None] = None, - position_maximum: Union[Sequence[float], None] = None, + position_minimum: Optional[Sequence[float]] = None, + position_maximum: Optional[Sequence[float]] = None, **_kwargs: str, ) -> set[int]: """Select atoms based on their positions at a specified frame number. diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py index 23294ce48..065053c2e 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py @@ -17,8 +17,6 @@ from collections import Counter from json import JSONDecodeError -import numpy as np - from MDANSE.Framework.AtomSelector.selector import ReusableSelection from MDANSE.Framework.Configurators.IConfigurator import IConfigurator diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index e2ec1ae94..bd22bf5e3 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -14,59 +14,77 @@ # along with this program. If not, see . # -from typing import Set import json from enum import StrEnum -from qtpy.QtCore import Slot, Signal -from qtpy.QtGui import QStandardItemModel, QStandardItem +from MDANSE.Framework.AtomSelector.selector import ReusableSelection +from MDANSE.Framework.InputData.HDFTrajectoryInputData import HDFTrajectoryInputData +from qtpy.QtCore import Signal, Slot +from qtpy.QtGui import QStandardItem, QStandardItemModel from qtpy.QtWidgets import ( - QLineEdit, - QPushButton, + QAbstractItemView, QDialog, - QVBoxLayout, - QHBoxLayout, QGroupBox, + QHBoxLayout, QLabel, - QPlainTextEdit, - QWidget, + QLineEdit, QListView, - QAbstractItemView, + QPlainTextEdit, + QPushButton, QScrollArea, + QVBoxLayout, + QWidget, ) + from MDANSE_GUI.InputWidgets.WidgetBase import WidgetBase -from MDANSE_GUI.Tabs.Visualisers.View3D import View3D from MDANSE_GUI.MolecularViewer.MolecularViewer import MolecularViewerWithPicking -from MDANSE.Framework.InputData.HDFTrajectoryInputData import HDFTrajectoryInputData -from MDANSE.Framework.AtomSelector.selector import ReusableSelection +from MDANSE_GUI.Tabs.Visualisers.View3D import View3D from MDANSE_GUI.Widgets.SelectionWidgets import ( AllAtomSelection, AtomSelection, IndexSelection, + LabelSelection, MoleculeSelection, PatternSelection, - LabelSelection, PositionSelection, SphereSelection, ) class SelectionValidity(StrEnum): + """Strings for selection check results.""" + VALID_SELECTION = "Valid selection" USELESS_SELECTION = "Selection did not change. This operation is not needed." MALFORMED_SELECTION = "This is not a valid JSON string." class SelectionModel(QStandardItemModel): + """Stores the selection operations in the GUI view.""" + selection_changed = Signal() def __init__(self, trajectory): + """Assign the current trajectory to the model.""" super().__init__(None) self._trajectory = trajectory self._selection = ReusableSelection() self._current_selection = set() - def rebuild_selection(self, last_operation: str): + def rebuild_selection(self, last_operation: str) -> SelectionValidity: + """Update the current selection based on the text in the GUI. + + Parameters + ---------- + last_operation : str + Additional selection operation input by the user. + + Returns + ------- + SelectionValidity + result of the check on last_operation + + """ self._selection = ReusableSelection() self._current_selection = set() total_dict = {} @@ -80,21 +98,44 @@ def rebuild_selection(self, last_operation: str): if last_operation: try: valid = self._selection.validate_selection_string( - last_operation, self._trajectory, self._current_selection + last_operation, + self._trajectory, + self._current_selection, ) except json.JSONDecodeError: return SelectionValidity.MALFORMED_SELECTION if valid: self._selection.load_from_json(json_string) return SelectionValidity.VALID_SELECTION - else: - return SelectionValidity.USELESS_SELECTION + return SelectionValidity.USELESS_SELECTION + return None - def current_selection(self, last_operation: str = "") -> Set[int]: + def current_selection(self, last_operation: str = "") -> set[int]: + """Return the selected atom indices. + + Parameters + ---------- + last_operation : str, optional + Extra selection operation typed by the user, by default "" + + Returns + ------- + set[int] + indices of all the selected atoms + + """ self.rebuild_selection(last_operation) return self._selection.select_in_trajectory(self._trajectory) def current_steps(self) -> str: + """Return selection operations as a JSON string. + + Returns + ------- + str + one string with all the selection operations in sequence + + """ result = {} for row in range(self.rowCount()): index = self.index(row, 0) @@ -106,6 +147,7 @@ def current_steps(self) -> str: @Slot(str) def accept_from_widget(self, json_string: str): + """Add a selection operation sent from a selection widget.""" new_item = QStandardItem(json_string) new_item.setEditable(False) self.appendRow(new_item) @@ -119,6 +161,7 @@ class SelectionHelper(QDialog): ---------- _helper_title : str The title of the helper dialog window. + """ _helper_title = "Atom selection helper" @@ -131,7 +174,8 @@ def __init__( *args, **kwargs, ): - """ + """Create the selection dialog. + Parameters ---------- traj_data : tuple[str, HDFTrajectoryInputData] @@ -139,6 +183,7 @@ def __init__( field : QLineEdit The QLineEdit field that will need to be updated when applying the setting. + """ super().__init__(parent, *args, **kwargs) self.setWindowTitle(self._helper_title) @@ -176,23 +221,27 @@ def __init__( self.setLayout(helper_layout) self.all_selection = True - self.selected = set([]) + self.selected = set() self.reset() def closeEvent(self, a0): - """Hide the window instead of closing. Some issues occur in the + """Hide the window instead of closing. + + Some issues occur in the 3D viewer when it is closed and then reopened. """ a0.ignore() self.hide() def create_buttons(self) -> list[QPushButton]: - """ + """Add buttons to the dialog layout. + Returns ------- list[QPushButton] List of push buttons to add to the last layout from create_layouts. + """ apply = QPushButton("Use Setting") reset = QPushButton("Reset") @@ -203,11 +252,13 @@ def create_buttons(self) -> list[QPushButton]: return [apply, reset, close] def create_layouts(self) -> list[QVBoxLayout]: - """ + """Call functions creating other widgets. + Returns ------- list[QVBoxLayout] List of QVBoxLayout to add to the helper layout. + """ layout_3d = QVBoxLayout() layout_3d.addWidget(self.view_3d) @@ -223,24 +274,27 @@ def create_layouts(self) -> list[QVBoxLayout]: return [layout_3d, left, right] def right_widgets(self) -> list[QWidget]: - """ + """Create widgets visualising the selection results. + Returns ------- list[QWidget] List of QWidgets to add to the right layout from create_layouts. + """ return [self.selection_operations_view, self.selection_textbox] def left_widgets(self) -> list[QWidget]: - """ + """Create widgets for defining the selection. + Returns ------- list[QWidget] List of QWidgets to add to the left layout from create_layouts. - """ + """ select = QGroupBox("selection") select_layout = QVBoxLayout() scroll_area = QScrollArea() @@ -278,7 +332,7 @@ def left_widgets(self) -> list[QWidget]: self.selection_operations_view.setDragEnabled(True) self.selection_operations_view.setAcceptDrops(True) self.selection_operations_view.setDragDropMode( - QAbstractItemView.DragDropMode.InternalMove + QAbstractItemView.DragDropMode.InternalMove, ) self.selection_operations_view.setModel(self.selection_model) self.selection_model.selection_changed.connect(self.recalculate_selection) @@ -286,39 +340,36 @@ def left_widgets(self) -> list[QWidget]: @Slot() def recalculate_selection(self): + """Update atom indices after selection change.""" self.selected = self.selection_model.current_selection() self.view_3d._viewer.change_picked(self.selected) self.update_selection_textbox() - def update_from_3d_view(self, selection: set[int]) -> None: - """A selection/deselection was made in the 3d view, update the + def update_from_3d_view(self, _selection: set[int]) -> None: + """Update atom indices after an atom has been clicked. + + A selection/deselection was made in the 3d view, update the check_boxes, combo_boxes and textbox. Parameters ---------- selection : set[int] Selection indexes from the 3d view. - """ - self.update_selection_textbox() - @Slot(str) - def update_operation(self, new_json_string: str): - self.selection_line.setText(new_json_string) - self.view_3d._viewer.change_picked(self.selected) + """ self.update_selection_textbox() @Slot() def append_selection(self): + """Add a selection operation from the text input field.""" self.selection_line.setStyleSheet("") self.selection_line.setToolTip("") selection_text = self.selection_line.text() validation = self.selection_model.rebuild_selection(selection_text) - if validation == SelectionValidity.MALFORMED_SELECTION: - self.selection_line.setStyleSheet( - "QWidget#InputWidget { background-color:rgb(180,20,180); font-weight: bold }" - ) - self.selection_line.setToolTip(validation) - elif validation == SelectionValidity.USELESS_SELECTION: + if validation in ( + SelectionValidity.MALFORMED_SELECTION, + SelectionValidity.USELESS_SELECTION, + ): self.selection_line.setStyleSheet( "QWidget#InputWidget { background-color:rgb(180,20,180); font-weight: bold }" ) @@ -337,13 +388,11 @@ def update_selection_textbox(self) -> None: self.selection_textbox.setPlainText("".join(text)) def apply(self) -> None: - """Set the field of the AtomSelectionWidget to the currently - chosen setting in this widget. - """ + """Send the selection from the dialog to the main widget.""" self._field.setText(self.selection_model.current_steps()) def reset(self) -> None: - """Resets the helper to the default state.""" + """Reset the helper to the default state.""" self.selection_model.clear() self.recalculate_selection() @@ -356,6 +405,7 @@ class AtomSelectionWidget(WidgetBase): _tooltip_text = "Specify which atoms will be used in the analysis. The input is a JSON string, and can be created using the helper dialog." def __init__(self, *args, **kwargs): + """Create the main widget for atom selection.""" super().__init__(*args, **kwargs) self._value = self._default_value self._field = QLineEdit(self._default_value, self._base) @@ -377,9 +427,14 @@ def __init__(self, *args, **kwargs): self._field.setToolTip(self._tooltip_text) def create_helper( - self, traj_data: tuple[str, HDFTrajectoryInputData] + self, + traj_data: tuple[str, HDFTrajectoryInputData], ) -> SelectionHelper: - """ + """Create the selection dialog. + + It will be populated with selection widget which can be used + to create the complete atom selection string. + Parameters ---------- traj_data : tuple[str, HDFTrajectoryInputData] @@ -389,28 +444,30 @@ def create_helper( ------- SelectionHelper Create and return the selection helper QDialog. + """ return SelectionHelper(traj_data, self._field, self._base) @Slot() def helper_dialog(self) -> None: - """Opens the helper dialog.""" + """Open the helper dialog.""" if self.helper.isVisible(): self.helper.close() else: self.helper.show() def get_widget_value(self) -> str: - """ + """Return the current text in the input field. + Returns ------- str The JSON selector setting. + """ selection_string = self._field.text() if len(selection_string) < 1: self._empty = True return self._default_value - else: - self._empty = False + self._empty = False return selection_string diff --git a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py index a2bf234a9..e8dd3e721 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py @@ -14,30 +14,32 @@ # along with this program. If not, see . # -from typing import Dict, Any, TYPE_CHECKING, Tuple import json from enum import StrEnum +from typing import TYPE_CHECKING, Any import numpy as np +from MDANSE.MolecularDynamics.Trajectory import Trajectory from qtpy.QtCore import Signal, Slot -from qtpy.QtGui import QValidator, QDoubleValidator +from qtpy.QtGui import QDoubleValidator, QValidator from qtpy.QtWidgets import ( + QComboBox, QGroupBox, QHBoxLayout, - QPushButton, - QComboBox, QLabel, QLineEdit, + QPushButton, ) from MDANSE_GUI.InputWidgets.CheckableComboBox import CheckableComboBox -from MDANSE.MolecularDynamics.Trajectory import Trajectory if TYPE_CHECKING: from MDANSE_GUI.MolecularViewer.MolecularViewer import MolecularViewer class IndexSelectionMode(StrEnum): + """Valid atom selection modes for select_atoms.""" + LIST = "list" RANGE = "range" SLICE = "slice" @@ -45,6 +47,7 @@ class IndexSelectionMode(StrEnum): class XYZValidator(QValidator): """A custom validator for a QLineEdit. + It is intended to limit the input to a string of 3 comma-separated float numbers. @@ -54,8 +57,12 @@ class XYZValidator(QValidator): a preliminary step when typing in 3 numbers. """ - def validate(self, input_string: str, position: int) -> Tuple[int, str]: - """Implementation of the virtual method of QValidator. + PARAMETERS_NEEDED = 3 + + def validate(self, input_string: str, position: int) -> tuple[int, str]: + """Check the input string from a widget. + + Implementation of the virtual method of QValidator. It takes in the string from a QLineEdit and the cursor position, and an enum value of the validator state. Widgets will reject inputs which change the state to Invalid. @@ -75,6 +82,7 @@ def validate(self, input_string: str, position: int) -> Tuple[int, str]: Original input string. int Cursor position. + """ state = QValidator.State.Intermediate comma_count = input_string.count(",") @@ -82,14 +90,14 @@ def validate(self, input_string: str, position: int) -> Tuple[int, str]: try: values = [float(x) for x in input_string.split(",")] except (TypeError, ValueError): - if input_string.endswith(",") and comma_count < 3: + if input_string.endswith(",") and comma_count < self.PARAMETERS_NEEDED: state = QValidator.State.Intermediate else: state = QValidator.State.Invalid else: - if len(values) > 3: + if len(values) > self.PARAMETERS_NEEDED: state = QValidator.State.Invalid - elif len(values) == 3: + elif len(values) == self.PARAMETERS_NEEDED: state = QValidator.State.Acceptable else: state = QValidator.State.Intermediate @@ -97,9 +105,21 @@ def validate(self, input_string: str, position: int) -> Tuple[int, str]: class BasicSelectionWidget(QGroupBox): + """Base class for atom selection widgets.""" + new_selection = Signal(str) def __init__(self, parent=None, widget_label="Atom selection widget"): + """Create subwidgets common to atom selection. + + Parameters + ---------- + parent : QWidget, optional + parent in the Qt hierarchy, by default None + widget_label : str, optional + Text shown above the widget, by default "Atom selection widget" + + """ super().__init__(parent) layout = QHBoxLayout() self.setLayout(layout) @@ -107,17 +127,24 @@ def __init__(self, parent=None, widget_label="Atom selection widget"): self.add_specific_widgets() self.add_standard_widgets() - def parameter_dictionary(self) -> Dict[str, Any]: + def parameter_dictionary(self) -> dict[str, Any]: + """Collect and return selection function parameters.""" return {} def add_specific_widgets(self): - return None + """Add additional widgets to layout, depending on the selection function.""" + return def add_standard_widgets(self): + """Create widgets needed by all atom selection types. + + This creates a combo box for the set operation type, + and a button for making the selection. + """ self.mode_box = QComboBox(self) self.mode_box.setEditable(False) self.mode_box.addItems( - ["Add (union)", "Filter (intersection)", "Remove (difference)"] + ["Add (union)", "Filter (intersection)", "Remove (difference)"], ) self._mode_box_values = ["union", "intersection", "difference"] self.commit_button = QPushButton("Apply", self) @@ -127,19 +154,34 @@ def add_standard_widgets(self): self.commit_button.clicked.connect(self.create_selection) def get_mode(self) -> str: + """Get the current set operation type from the combo box.""" return self._mode_box_values[self.mode_box.currentIndex()] def create_selection(self): + """Collect the input values and emit them in a signal.""" funtion_parameters = self.parameter_dictionary() funtion_parameters["operation_type"] = self.get_mode() self.new_selection.emit(json.dumps(funtion_parameters)) class AllAtomSelection(BasicSelectionWidget): + """Widget for global atom selection, e.g. all atoms, no atoms.""" + def __init__(self, parent=None, widget_label="ALL ATOMS"): + """Pass inputs to the parent class init. + + Parameters + ---------- + parent : QWidget, optional + parent in the Qt hierarchy, by default None + widget_label : str, optional + Text over the widget, by default "ALL ATOMS" + + """ super().__init__(parent, widget_label) def add_specific_widgets(self): + """Add the INVERT button.""" layout = self.layout() inversion_button = QPushButton("INVERT selection", self) inversion_button.clicked.connect(self.invert_selection) @@ -147,16 +189,35 @@ def add_specific_widgets(self): layout.addWidget(QLabel("Add/remove ALL atoms")) def invert_selection(self): + """Emit the string for inverting the selection.""" self.new_selection.emit(json.dumps({"function_name": "invert_selection"})) def parameter_dictionary(self): + """Collect and return selection function parameters.""" return {"function_name": "select_all"} class AtomSelection(BasicSelectionWidget): + """GUI frontend for select_atoms.""" + def __init__( - self, parent=None, trajectory: Trajectory = None, widget_label="Select atoms" + self, + parent=None, + trajectory: Trajectory = None, + widget_label="Select atoms", ): + """Create the widgets for select_atoms. + + Parameters + ---------- + parent : QWidget, optional + parent from the Qt object hierarchy, by default None + trajectory : Trajectory, optional + The current trajectory object, by default None + widget_label : str, optional + Text shown over the widget, by default "Select atoms" + + """ self.atom_types = [] self.atom_names = [] if trajectory: @@ -172,6 +233,7 @@ def __init__( super().__init__(parent, widget_label) def add_specific_widgets(self): + """Create selection combo boxes.""" layout = self.layout() layout.addWidget(QLabel("Select atoms by atom")) self.selection_type_combo = QComboBox(self) @@ -186,6 +248,7 @@ def add_specific_widgets(self): @Slot(str) def switch_mode(self, new_mode: str): + """Change the contents of the second combo box.""" self.selection_field.clear() if new_mode == "type": self.selection_field.addItems(self.atom_types) @@ -195,6 +258,7 @@ def switch_mode(self, new_mode: str): self.selection_keyword = "atom_names" def parameter_dictionary(self): + """Collect and return selection function parameters.""" function_parameters = {"function_name": "select_atoms"} selection = self.selection_field.checked_values() function_parameters[self.selection_keyword] = selection @@ -202,11 +266,24 @@ def parameter_dictionary(self): class IndexSelection(BasicSelectionWidget): + """GUI frontend for select_atoms.""" + def __init__(self, parent=None, widget_label="Index selection"): + """Create all the widgets. + + Parameters + ---------- + parent : QWidget, optional + parent in the Qt object hierarchy, by default None + widget_label : str, optional + Text shown above the widget, by default "Index selection" + + """ super().__init__(parent, widget_label) self.selection_keyword = "index_list" def add_specific_widgets(self): + """Create the combo box and text input field.""" layout = self.layout() layout.addWidget(QLabel("Select atoms by index")) self.selection_type_combo = QComboBox(self) @@ -219,6 +296,7 @@ def add_specific_widgets(self): @Slot(str) def switch_mode(self, new_mode: str): + """Change the meaning of the text input field.""" self.selection_field.setText("") if new_mode == IndexSelectionMode.LIST: self.selection_field.setPlaceholderText("0,1,2") @@ -234,6 +312,7 @@ def switch_mode(self, new_mode: str): self.selection_separator = ":" def parameter_dictionary(self): + """Collect and return selection function parameters.""" function_parameters = {"function_name": "select_atoms"} selection = self.selection_field.text() function_parameters[self.selection_keyword] = [ @@ -243,18 +322,33 @@ def parameter_dictionary(self): class MoleculeSelection(BasicSelectionWidget): + """GUI frontend for select_molecule.""" + def __init__( self, parent=None, trajectory: Trajectory = None, widget_label="Select molecules", ): + """Create the widgets for select_atoms. + + Parameters + ---------- + parent : QWidget, optional + parent from the Qt object hierarchy, by default None + trajectory : Trajectory, optional + The current trajectory object, by default None + widget_label : str, optional + Text shown over the widget, by default "Select atoms" + + """ self.molecule_names = [] if trajectory: self.molecule_names = trajectory.chemical_system.unique_molecules() super().__init__(parent, widget_label) def add_specific_widgets(self): + """Create the combo box for molecule names.""" layout = self.layout() layout.addWidget(QLabel("Select molecules named: ")) self.selection_field = CheckableComboBox(self) @@ -262,6 +356,7 @@ def add_specific_widgets(self): self.selection_field.addItems(self.molecule_names) def parameter_dictionary(self): + """Collect and return selection function parameters.""" function_parameters = {"function_name": "select_molecules"} selection = self.selection_field.checked_values() function_parameters["molecule_names"] = selection @@ -269,18 +364,33 @@ def parameter_dictionary(self): class LabelSelection(BasicSelectionWidget): + """GUI frontend for select_label.""" + def __init__( self, parent=None, trajectory: Trajectory = None, widget_label="Select by label", ): + """Create the widgets for select_atoms. + + Parameters + ---------- + parent : QWidget, optional + parent from the Qt object hierarchy, by default None + trajectory : Trajectory, optional + The current trajectory object, by default None + widget_label : str, optional + Text shown over the widget, by default "Select atoms" + + """ self.labels = [] if trajectory: self.labels = list(trajectory.chemical_system._labels.keys()) super().__init__(parent, widget_label) def add_specific_widgets(self): + """Create the combo box for atom labels.""" layout = self.layout() layout.addWidget(QLabel("Select atoms with label: ")) self.selection_field = CheckableComboBox(self) @@ -288,6 +398,7 @@ def add_specific_widgets(self): self.selection_field.addItems(self.labels) def parameter_dictionary(self): + """Collect and return selection function parameters.""" function_parameters = {"function_name": "select_labels"} selection = self.selection_field.checked_values() function_parameters["atom_labels"] = selection @@ -295,11 +406,23 @@ def parameter_dictionary(self): class PatternSelection(BasicSelectionWidget): + """GUI frontend for select_pattern.""" + def __init__( self, parent=None, widget_label="SMARTS pattern matching", ): + """Create the widgets for select_atoms. + + Parameters + ---------- + parent : QWidget, optional + parent from the Qt object hierarchy, by default None + widget_label : str, optional + Text shown over the widget, by default "Select atoms" + + """ self.pattern_dictionary = { "primary amine": "[#7X3;H2;!$([#7][#6X3][!#6]);!$([#7][#6X2][!#6])](~[H])~[H]", "hydroxy": "[#8;H1,H2]~[H]", @@ -311,6 +434,7 @@ def __init__( super().__init__(parent, widget_label) def add_specific_widgets(self): + """Create the pattern text field.""" layout = self.layout() layout.addWidget(QLabel("Pick a group")) self.selection_field = QComboBox(self) @@ -324,10 +448,12 @@ def add_specific_widgets(self): @Slot(str) def update_string(self, key_string: str): + """Fill the input field with pre-defined text.""" if key_string in self.pattern_dictionary: self.input_field.setText(self.pattern_dictionary[key_string]) def parameter_dictionary(self): + """Collect and return selection function parameters.""" function_parameters = {"function_name": "select_pattern"} selection = self.input_field.text() function_parameters["rdkit_pattern"] = selection @@ -335,6 +461,8 @@ def parameter_dictionary(self): class PositionSelection(BasicSelectionWidget): + """GUI frontend for select_positions.""" + def __init__( self, parent=None, @@ -342,6 +470,20 @@ def __init__( molecular_viewer: "MolecularViewer" = None, widget_label="Select by position", ): + """Create the widgets for select_atoms. + + Parameters + ---------- + parent : QWidget, optional + parent from the Qt object hierarchy, by default None + trajectory : Trajectory, optional + The current trajectory object, by default None + molecular_viewer : MolecularViewer, optional + instance of the 3D viewer showing the current simulation frame + widget_label : str, optional + Text shown over the widget, by default "Select atoms" + + """ self._viewer = molecular_viewer self._lower_limit = np.zeros(3) self._upper_limit = np.linalg.norm(trajectory.unit_cell(0)._unit_cell, axis=1) @@ -350,15 +492,16 @@ def __init__( super().__init__(parent, widget_label) def add_specific_widgets(self): + """Create text input fields with validators.""" layout = self.layout() layout.addWidget(QLabel("Lower limits")) self._lower_limit_input = QLineEdit( - ",".join([str(round(x, 3)) for x in self._lower_limit]) + ",".join([str(round(x, 3)) for x in self._lower_limit]), ) layout.addWidget(self._lower_limit_input) layout.addWidget(QLabel("Upper limits")) self._upper_limit_input = QLineEdit( - ",".join([str(round(x, 3)) for x in self._upper_limit]) + ",".join([str(round(x, 3)) for x in self._upper_limit]), ) layout.addWidget(self._upper_limit_input) for field in [self._lower_limit_input, self._upper_limit_input]: @@ -367,6 +510,7 @@ def add_specific_widgets(self): @Slot() def check_inputs(self): + """Disable selection of invalid or incomplete input.""" enable = True try: self._current_lower_limit = [ @@ -386,6 +530,7 @@ def check_inputs(self): self.commit_button.setEnabled(enable) def parameter_dictionary(self): + """Collect and return selection function parameters.""" return { "function_name": "select_positions", "frame_number": self._viewer._current_frame, @@ -395,6 +540,8 @@ def parameter_dictionary(self): class SphereSelection(BasicSelectionWidget): + """GUI frontend for select_sphere.""" + def __init__( self, parent=None, @@ -402,16 +549,32 @@ def __init__( molecular_viewer: "MolecularViewer" = None, widget_label="Select in a sphere", ): + """Create the widgets for select_atoms. + + Parameters + ---------- + parent : QWidget, optional + parent from the Qt object hierarchy, by default None + trajectory : Trajectory, optional + The current trajectory object, by default None + molecular_viewer : MolecularViewer, optional + instance of the 3D viewer showing the current simulation frame + widget_label : str, optional + Text shown over the widget, by default "Select atoms" + + """ self._viewer = molecular_viewer self._current_sphere_centre = np.diag(trajectory.unit_cell(0)._unit_cell) * 0.5 self._current_sphere_radius = np.min(self._current_sphere_centre) super().__init__(parent, widget_label) def add_specific_widgets(self): + """Create the text input fields for sphere radius and centre.""" layout = self.layout() layout.addWidget(QLabel("Sphere centre")) self._sphere_centre_input = QLineEdit( - ",".join([str(round(x, 3)) for x in self._current_sphere_centre]), self + ",".join([str(round(x, 3)) for x in self._current_sphere_centre]), + self, ) layout.addWidget(self._sphere_centre_input) layout.addWidget(QLabel("Sphere radius (nm)")) @@ -424,6 +587,7 @@ def add_specific_widgets(self): @Slot() def check_inputs(self): + """Disable selection on invalid or incomplete input.""" enable = True try: self._current_sphere_centre = [ @@ -438,6 +602,7 @@ def check_inputs(self): self.commit_button.setEnabled(enable) def parameter_dictionary(self): + """Collect and return selection function parameters.""" return { "function_name": "select_sphere", "frame_number": self._viewer._current_frame, From 77f6c4c998ba504413aca82446c5883dea46d357 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Thu, 6 Mar 2025 10:20:13 +0000 Subject: [PATCH 32/37] Use keyword arguments in set_selection --- .../MDANSE/Framework/AtomSelector/selector.py | 2 +- .../UnitTests/test_reusable_selection.py | 34 +++++++++---------- 2 files changed, 18 insertions(+), 18 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index 7ac449585..b0140d47e 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -225,4 +225,4 @@ def load_from_json(self, json_string: str): json_setting = json.loads(json_string) for k0, v0 in json_setting.items(): if isinstance(v0, dict): - self.set_selection(k0, v0) + self.set_selection(number=k0, function_parameters=v0) diff --git a/MDANSE/Tests/UnitTests/test_reusable_selection.py b/MDANSE/Tests/UnitTests/test_reusable_selection.py index ecdf9bc4e..0e61fd70e 100644 --- a/MDANSE/Tests/UnitTests/test_reusable_selection.py +++ b/MDANSE/Tests/UnitTests/test_reusable_selection.py @@ -37,7 +37,7 @@ def trajectory(request): def test_select_all(trajectory): n_atoms = len(trajectory.chemical_system.atom_list) reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_all'}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'select_all'}) selection = reusable_selection.select_in_trajectory(trajectory.trajectory) assert len(selection) == n_atoms @@ -54,7 +54,7 @@ def test_empty_json_string_selects_all(trajectory): @pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj], indirect=True) def test_select_none(trajectory): reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_none'}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'select_none'}) selection = reusable_selection.select_in_trajectory(trajectory.trajectory) assert len(selection) == 0 @@ -62,16 +62,16 @@ def test_select_none(trajectory): @pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj], indirect=True) def test_inverted_all_is_none(trajectory): reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_all'}) - reusable_selection.set_selection(None, {'function_name': 'invert_selection'}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'select_all'}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'invert_selection'}) selection = reusable_selection.select_in_trajectory(trajectory.trajectory) assert len(selection) == 0 def test_json_saving_is_reversible(): reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_all'}) - reusable_selection.set_selection(None, {'function_name': 'invert_selection'}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'select_all'}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'invert_selection'}) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() another_selection.load_from_json(json_string) @@ -82,8 +82,8 @@ def test_json_saving_is_reversible(): @pytest.mark.parametrize("trajectory", [short_traj, mdmc_traj, com_traj], indirect=True) def test_selection_from_json_is_the_same_as_from_runtime(trajectory): reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_all'}) - reusable_selection.set_selection(None, {'function_name': 'invert_selection'}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'select_all'}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'invert_selection'}) selection = reusable_selection.select_in_trajectory(trajectory.trajectory) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() @@ -103,7 +103,7 @@ def test_selection_from_json_is_the_same_as_from_runtime(trajectory): )) def test_select_atoms_selects_by_element(trajectory, element, expected): reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {"function_name": "select_atoms", + reusable_selection.set_selection(number=None, function_parameters={"function_name": "select_atoms", "atom_types": element}) selection = reusable_selection.select_in_trajectory(trajectory) assert len(selection) == expected @@ -121,7 +121,7 @@ def test_select_atoms_selects_by_element(trajectory, element, expected): )) def test_select_atoms_selects_by_range(trajectory, index_range, expected): reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'index_range': index_range}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'select_atoms', 'index_range': index_range}) range_selection = reusable_selection.select_in_trajectory(trajectory.trajectory) assert len(range_selection) == expected @@ -133,7 +133,7 @@ def test_select_atoms_selects_by_range(trajectory, index_range, expected): )) def test_select_atoms_selects_by_slice(trajectory, index_slice, expected): reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'index_slice': index_slice}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'select_atoms', 'index_slice': index_slice}) range_selection = reusable_selection.select_in_trajectory(trajectory.trajectory) assert len(range_selection) == expected json_string = reusable_selection.convert_to_json() @@ -150,7 +150,7 @@ def test_select_atoms_selects_by_slice(trajectory, index_slice, expected): )) def test_select_molecules_selects_water(trajectory, molecule_names, expected): reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_molecules', 'molecule_names': molecule_names}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'select_molecules', 'molecule_names': molecule_names}) first_selection = reusable_selection.select_in_trajectory(trajectory.trajectory) assert len(first_selection) == expected json_string = reusable_selection.convert_to_json() @@ -163,8 +163,8 @@ def test_select_molecules_selects_water(trajectory, molecule_names, expected): @pytest.mark.parametrize("trajectory", [traj_2vb1], indirect=True) def test_select_molecules_inverted_selects_ions(trajectory): reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_molecules', 'molecule_names': ['C613 H959 N193 O185 S10', 'H2 O1']}) - reusable_selection.set_selection(None, {'function_name': 'invert_selection'}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'select_molecules', 'molecule_names': ['C613 H959 N193 O185 S10', 'H2 O1']}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'invert_selection'}) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() another_selection.load_from_json(json_string) @@ -175,7 +175,7 @@ def test_select_molecules_inverted_selects_ions(trajectory): @pytest.mark.parametrize("trajectory", [traj_2vb1], indirect=True) def test_select_pattern_selects_water(trajectory): reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_pattern', 'rdkit_pattern': "[#8X2;H2](~[H])~[H]"}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'select_pattern', 'rdkit_pattern': "[#8X2;H2](~[H])~[H]"}) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() another_selection.load_from_json(json_string) @@ -192,8 +192,8 @@ def test_selection_with_multiple_steps(trajectory): The selection is then saved to a JSON string, loaded from the string and applied to the trajectory. """ reusable_selection = ReusableSelection() - reusable_selection.set_selection(None, {'function_name': 'select_pattern', 'rdkit_pattern': "[#8X2;H2](~[H])~[H]"}) - reusable_selection.set_selection(None, {'function_name': 'select_atoms', 'atom_types': ['O'], 'operation_type': 'intersection'}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'select_pattern', 'rdkit_pattern': "[#8X2;H2](~[H])~[H]"}) + reusable_selection.set_selection(number=None, function_parameters={'function_name': 'select_atoms', 'atom_types': ['O'], 'operation_type': 'intersection'}) json_string = reusable_selection.convert_to_json() another_selection = ReusableSelection() another_selection.load_from_json(json_string) From 970c49bcc11f6ddb0338d825ca731f7872447e81 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Thu, 6 Mar 2025 11:53:24 +0000 Subject: [PATCH 33/37] Implement review suggestions --- .../Framework/AtomSelector/group_selection.py | 6 +- .../MDANSE/Framework/AtomSelector/selector.py | 86 +++++++++++-------- .../InputWidgets/AtomSelectionWidget.py | 2 +- 3 files changed, 56 insertions(+), 38 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py index 8a5937668..2bb0a6ad9 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/group_selection.py @@ -47,7 +47,8 @@ def select_labels( def select_pattern( trajectory: Trajectory, - rdkit_pattern: str = "", + *, + rdkit_pattern: str, **_kwargs: str, ) -> set[int]: """Select atoms according to the SMARTS string given as input. @@ -71,6 +72,5 @@ def select_pattern( """ selection = set() system = trajectory.chemical_system - if rdkit_pattern: - selection = system.get_substructure_matches(rdkit_pattern) + selection = system.get_substructure_matches(rdkit_pattern) return selection diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index b0140d47e..03e5a2c0c 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -93,12 +93,51 @@ def set_selection( if number is None: number = len(self.operations) else: - try: - number = int(number) - except TypeError: - number = len(self.operations) + number = int(number) self.operations[number] = function_parameters + def apply_single_selection( + self, + function_parameters: dict[str, Any], + trajectory: Trajectory, + selection: set[int], + ) -> set[int]: + """Modify the input selection based on input parameters. + + This method applied a single selection operation to + an already exising selection for a specific trajectory. + + Parameters + ---------- + function_parameters : dict[str, Any] + All the inputs needed to call an atom selection function + trajectory : Trajectory + Instance of the trajectory in which we are selecting atoms + selection : set[int] + indices of atoms that resulted from previous steps + + Returns + ------- + set[int] + indices of selected atoms from all operations so far + """ + function_name = function_parameters.get("function_name", "select_all") + if function_name == "invert_selection": + new_selection = self.all_idxs.difference(selection) + else: + operation_type = function_parameters.get("operation_type", "union") + function = function_lookup[function_name] + temp_selection = function(trajectory, **function_parameters) + if operation_type == "union": + new_selection = selection | temp_selection + elif operation_type == "intersection": + new_selection = selection & temp_selection + elif operation_type == "difference": + new_selection = selection - temp_selection + else: + new_selection = temp_selection + return new_selection + def validate_selection_string( self, json_string: str, @@ -122,27 +161,16 @@ def validate_selection_string( Returns ------- bool - True if the selection adds atoms, False otherwise + True if the operation changes selection, False otherwise """ function_parameters = json.loads(json_string) if not self.operations: return True - function_name = function_parameters.get("function_name", "select_all") - if function_name == "invert_selection": - selection = invert_selection(trajectory, current_selection) - return True operation_type = function_parameters.get("operation_type", "union") - function = function_lookup[function_name] - temp_selection = function(trajectory, **function_parameters) - if operation_type == "union": - selection = selection.union(temp_selection) - elif operation_type == "intersection": - selection = selection.intersection(temp_selection) - elif operation_type == "difference": - selection = selection.difference(temp_selection) - else: - selection = temp_selection + selection = self.apply_single_selection( + function_parameters, trajectory, current_selection + ) return bool( ( len(selection.difference(current_selection)) > 0 @@ -178,21 +206,9 @@ def select_in_trajectory(self, trajectory: Trajectory) -> set[int]: return self.all_idxs for number in sequence: function_parameters = self.operations[number] - function_name = function_parameters.get("function_name", "select_all") - if function_name == "invert_selection": - selection = self.all_idxs.difference(selection) - else: - operation_type = function_parameters.get("operation_type", "union") - function = function_lookup[function_name] - temp_selection = function(trajectory, **function_parameters) - if operation_type == "union": - selection = selection.union(temp_selection) - elif operation_type == "intersection": - selection = selection.intersection(temp_selection) - elif operation_type == "difference": - selection = selection.difference(temp_selection) - else: - selection = temp_selection + selection = self.apply_single_selection( + function_parameters, trajectory, selection + ) return selection def convert_to_json(self) -> str: @@ -226,3 +242,5 @@ def load_from_json(self, json_string: str): for k0, v0 in json_setting.items(): if isinstance(v0, dict): self.set_selection(number=k0, function_parameters=v0) + else: + raise TypeError(f"Selection {v0} is not a dictionary.") diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index bd22bf5e3..4e7af94d3 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -56,7 +56,7 @@ class SelectionValidity(StrEnum): VALID_SELECTION = "Valid selection" USELESS_SELECTION = "Selection did not change. This operation is not needed." - MALFORMED_SELECTION = "This is not a valid JSON string." + MALFORMED_SELECTION = "This is not a valid selection string." class SelectionModel(QStandardItemModel): From c9d8acc8f0fd24f870c60f09ee20540c39473919 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Thu, 6 Mar 2025 15:38:55 +0000 Subject: [PATCH 34/37] Add style corrections --- .../MDANSE/Framework/AtomSelector/selector.py | 21 +++++-------------- .../AtomSelector/spatial_selection.py | 18 +++++++++------- .../AtomSelectionConfigurator.py | 9 +++----- 3 files changed, 18 insertions(+), 30 deletions(-) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py index 03e5a2c0c..534a2d4ff 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/selector.py @@ -90,10 +90,7 @@ def set_selection( the dictionary of keyword arguments defining a selection operation """ - if number is None: - number = len(self.operations) - else: - number = int(number) + number = int(number) if number is not None else len(self.operations) self.operations[number] = function_parameters def apply_single_selection( @@ -171,15 +168,8 @@ def validate_selection_string( selection = self.apply_single_selection( function_parameters, trajectory, current_selection ) - return bool( - ( - len(selection.difference(current_selection)) > 0 - and operation_type == "union" - ) - or ( - len(current_selection.difference(selection)) > 0 - and operation_type != "union" - ) + return ((selection - current_selection) and operation_type == "union") or ( + (current_selection - selection) and operation_type != "union" ) def select_in_trajectory(self, trajectory: Trajectory) -> set[int]: @@ -240,7 +230,6 @@ def load_from_json(self, json_string: str): """ json_setting = json.loads(json_string) for k0, v0 in json_setting.items(): - if isinstance(v0, dict): - self.set_selection(number=k0, function_parameters=v0) - else: + if not isinstance(v0, dict): raise TypeError(f"Selection {v0} is not a dictionary.") + self.set_selection(number=k0, function_parameters=v0) diff --git a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py index d192d6075..00a3c0463 100644 --- a/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py +++ b/MDANSE/Src/MDANSE/Framework/AtomSelector/spatial_selection.py @@ -53,14 +53,16 @@ def select_positions( """ coordinates = trajectory.coordinates(frame_number) - if position_minimum is None: - lower_limits = np.array(3 * [-np.inf]) - else: - lower_limits = np.array(position_minimum) - if position_maximum is None: - upper_limits = np.array(3 * [np.inf]) - else: - upper_limits = np.array(position_maximum) + lower_limits = ( + np.array(position_minimum) + if position_minimum is not None + else np.array([-np.inf] * 3) + ) + upper_limits = ( + np.array(position_maximum) + if position_maximum is not None + else np.array([np.inf] * 3) + ) valid = np.where( ((coordinates > lower_limits) & (coordinates < upper_limits)).all(axis=1) ) diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py index 065053c2e..0602042a4 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py @@ -14,7 +14,7 @@ # along with this program. If not, see . # -from collections import Counter +from collections import Counter, defaultdict from json import JSONDecodeError from MDANSE.Framework.AtomSelector.selector import ReusableSelection @@ -121,12 +121,9 @@ def get_indices(self) -> dict[str, list[int]]: For each atom type, a list of indices of selected atoms """ - indicesPerElement = {} + indicesPerElement = defaultdict(list) for i, v in enumerate(self["names"]): - if v in indicesPerElement: - indicesPerElement[v].extend(self["indices"][i]) - else: - indicesPerElement[v] = self["indices"][i][:] + indicesPerElement[v].extend(self["indices"][i]) return indicesPerElement From d3dbca42a3e61d93a2aaa0086679c1f5f6084a83 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Thu, 6 Mar 2025 15:59:09 +0000 Subject: [PATCH 35/37] Validate SMARTS input --- .../MDANSE_GUI/InputWidgets/AtomSelectionWidget.py | 2 +- .../Src/MDANSE_GUI/Widgets/SelectionWidgets.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index 4e7af94d3..2002c376c 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -401,7 +401,7 @@ class AtomSelectionWidget(WidgetBase): """The atoms selection widget.""" _push_button_text = "Atom selection helper" - _default_value = '{"all": true}' + _default_value = "{}" _tooltip_text = "Specify which atoms will be used in the analysis. The input is a JSON string, and can be created using the helper dialog." def __init__(self, *args, **kwargs): diff --git a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py index e8dd3e721..09cfd1014 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/Widgets/SelectionWidgets.py @@ -30,6 +30,7 @@ QLineEdit, QPushButton, ) +from rdkit.Chem import MolFromSmarts from MDANSE_GUI.InputWidgets.CheckableComboBox import CheckableComboBox @@ -445,6 +446,17 @@ def add_specific_widgets(self): self.input_field.setPlaceholderText("can be edited") layout.addWidget(self.input_field) self.selection_field.currentTextChanged.connect(self.update_string) + self.input_field.textChanged.connect(self.check_inputs) + + @Slot() + def check_inputs(self): + """Disable selection of invalid or incomplete input.""" + enable = True + smarts_string = self.input_field.text() + temp_molecule = MolFromSmarts(smarts_string) + if temp_molecule is None: + enable = False + self.commit_button.setEnabled(enable) @Slot(str) def update_string(self, key_string: str): From 97b4a5af0260282bfe1afba8fdc90a3feac57eeb Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Thu, 6 Mar 2025 17:11:11 +0000 Subject: [PATCH 36/37] Select all atoms explicitly, disable dragging --- .../Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py index 2002c376c..2220c0c81 100644 --- a/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py +++ b/MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/AtomSelectionWidget.py @@ -329,11 +329,6 @@ def left_widgets(self) -> list[QWidget]: select.setLayout(select_layout) self.selection_operations_view = QListView(self) - self.selection_operations_view.setDragEnabled(True) - self.selection_operations_view.setAcceptDrops(True) - self.selection_operations_view.setDragDropMode( - QAbstractItemView.DragDropMode.InternalMove, - ) self.selection_operations_view.setModel(self.selection_model) self.selection_model.selection_changed.connect(self.recalculate_selection) return [scroll_area] @@ -394,6 +389,9 @@ def apply(self) -> None: def reset(self) -> None: """Reset the helper to the default state.""" self.selection_model.clear() + self.selection_model.accept_from_widget( + '{"function_name": "select_all", "operation_type": "union"}' + ) self.recalculate_selection() From b0221325d06399afbaea535f8c4f354f4e443377 Mon Sep 17 00:00:00 2001 From: Maciej Bartkowiak Date: Thu, 6 Mar 2025 17:17:10 +0000 Subject: [PATCH 37/37] Reset internal selection in the transmuter --- .../Framework/Configurators/AtomTransmutationConfigurator.py | 1 + 1 file changed, 1 insertion(+) diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py index 4bd55dd94..20c6f8d01 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/AtomTransmutationConfigurator.py @@ -86,6 +86,7 @@ def get_json_setting(self) -> str: def reset_setting(self) -> None: """Resets the transmutation setting.""" self._new_map = {} + self.selector.reset() class AtomTransmutationConfigurator(IConfigurator):