Merge pull request #339 from ISISNeutronMuon/chi/atom-mapping

Atom mapping update

MDANSE/Src/MDANSE/Framework/AtomMapping/__init__.py

@@ -0,0 +1,5 @@
+from .atom_mapping import guess_element
+from .atom_mapping import get_element_from_mapping
+from .atom_mapping import fill_remaining_labels
+from .atom_mapping import check_mapping_valid
+from .atom_mapping import AtomLabel

MDANSE/Src/MDANSE/Framework/AtomMapping/atom_mapping.py

@@ -0,0 +1,170 @@
+from typing import Union
+import re
+import numpy as np
+
+from MDANSE.Chemistry import ATOMS_DATABASE
+
+
+class AtomLabel:
+
+    def __init__(self, atm_label, **kwargs):
+        self.atm_label = atm_label
+        self.grp_label = f""
+        if kwargs:
+            for k, v in kwargs.items():
+                self.grp_label += f"{k}={v};"
+            self.grp_label = self.grp_label[:-1]
+        self.mass = kwargs.get("mass", None)
+        if self.mass is not None:
+            self.mass = float(self.mass)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, AtomLabel):
+            AssertionError(f"{other} should be an instance of AtomLabel.")
+        if self.grp_label == other.grp_label and self.atm_label == other.atm_label:
+            return True
+
+
+def guess_element(atm_label: str, mass: Union[float, int, None] = None) -> str:
+    """From an input atom label find a match to an element in the atom
+    database.
+
+    Parameters
+    ----------
+    atm_label : str
+        The atom label.
+    mass : Union[float, int, None]
+        The atomic weight in atomic mass units.
+
+    Returns
+    -------
+    str
+        The symbol of the guessed element.
+
+    Raises
+    ------
+    AttributeError
+        Error if unable to match to an element.
+    """
+    if mass is not None and mass == 0.0:
+        return "Du"
+
+    guesses = []
+    guess_0 = re.findall("([A-Za-z][a-z]?)", atm_label)
+    if len(guess_0) != 0:
+        guess = guess_0[0].capitalize()
+        guesses.append(guess)
+        if len(guess) == 2:
+            guesses.append(guess[0])
+
+    # using the guess match to the atom and then match to the mass
+    # if available
+    best_match = None
+    best_diff = np.inf
+    for guess in guesses:
+        if guess in ATOMS_DATABASE:
+            if mass is None:
+                return guess
+            num = ATOMS_DATABASE[guess]["proton"]
+            atms = ATOMS_DATABASE.match_numeric_property("proton", num)
+            for atm in atms:
+                atm_mass = ATOMS_DATABASE[atm]["atomic_weight"]
+                diff = abs(mass - atm_mass)
+                if diff < best_diff:
+                    best_match = atm
+                    best_diff = diff
+    if best_match is not None:
+        return best_match
+
+    # try to match based on mass if available and guesses failed
+    best_diff = np.inf
+    if mass is not None:
+        for atm, properties in ATOMS_DATABASE._data.items():
+            atm_mass = properties.get("atomic_weight", None)
+            if atm_mass is None:
+                continue
+            diff = abs(mass - atm_mass)
+            if diff < best_diff:
+                best_match = atm
+                best_diff = diff
+        return best_match
+
+    raise AttributeError(f"Unable to guess: {atm_label}")
+
+
+def get_element_from_mapping(
+    mapping: dict[str, dict[str, str]], label: str, **kwargs
+) -> str:
+    """Determine the symbol of the element from the atom label and
+    the information from the kwargs.
+
+    Parameters
+    ----------
+    mapping : dict[str, dict[str, str]]
+        A dict which maps group and atom labels to an element from the
+        atom database.
+    label : str
+        The atom label.
+
+    Returns
+    -------
+    str
+        The symbol of the element from the MDANSE atom database.
+    """
+    label = AtomLabel(label, **kwargs)
+    grp_label = label.grp_label
+    atm_label = label.atm_label
+    if grp_label in mapping and atm_label in mapping[grp_label]:
+        element = mapping[grp_label][atm_label]
+    elif "" in mapping and atm_label in mapping[""]:
+        element = mapping[""][atm_label]
+    else:
+        element = guess_element(atm_label, label.mass)
+    return element
+
+
+def fill_remaining_labels(
+    mapping: dict[str, dict[str, str]], labels: list[AtomLabel]
+) -> None:
+    """Given a list of labels fill the remaining labels in the mapping
+    dictionary.
+
+    Parameters
+    ----------
+    mapping : dict[str, dict[str, str]]
+        The atom mapping dictionary.
+    labels : list[AtomLabel]
+        A list of atom labels.
+    """
+    for label in labels:
+        grp_label = label.grp_label
+        atm_label = label.atm_label
+        if grp_label not in mapping:
+            mapping[grp_label] = {}
+        if atm_label not in mapping[grp_label]:
+            mapping[grp_label][atm_label] = guess_element(atm_label, label.mass)
+
+
+def check_mapping_valid(mapping: dict[str, dict[str, str]], labels: list[AtomLabel]):
+    """Given a list of labels check that the mapping is valid.
+
+    Parameters
+    ----------
+    mapping : dict[str, dict[str, str]]
+        The atom mapping dictionary.
+    labels : list[AtomLabel]
+        A list of atom labels.
+
+    Returns
+    -------
+    bool
+        True if the mapping is valid.
+    """
+    for label in labels:
+        grp_label = label.grp_label
+        atm_label = label.atm_label
+        if grp_label not in mapping or atm_label not in mapping[grp_label]:
+            return False
+        if mapping[grp_label][atm_label] not in ATOMS_DATABASE:
+            return False
+    return True

MDANSE/Src/MDANSE/Framework/AtomSelector/group_selectors.py

@@ -85,7 +85,9 @@ def select_methly(
         return system.get_substructure_matches(pattern)
 
 
-def select_phosphate(system: ChemicalSystem, check_exists: bool = False) -> set[int]:
+def select_phosphate(
+    system: ChemicalSystem, check_exists: bool = False
+) -> Union[set[int], bool]:
     """Selects the P and O atoms of all phosphate groups.
 
     Parameters

MDANSE/Src/MDANSE/Framework/Configurators/ASEFileConfigurator.py

@@ -0,0 +1,51 @@
+# **************************************************************************
+#
+# MDANSE: Molecular Dynamics Analysis for Neutron Scattering Experiments
+#
+# @file      Src/Framework/Configurators/InputFileConfigurator.py
+# @brief     Implements module/class/test InputFileConfigurator
+#
+# @homepage  https://www.isis.stfc.ac.uk/Pages/MDANSEproject.aspx
+# @license   GNU General Public License v3 or higher (see LICENSE)
+# @copyright Institut Laue Langevin 2013-now
+# @copyright ISIS Neutron and Muon Source, STFC, UKRI 2021-now
+# @authors   Scientific Computing Group at ILL (see AUTHORS)
+#
+# **************************************************************************
+from ase.io import iread, read
+from ase.io.trajectory import Trajectory as ASETrajectory
+
+from MDANSE.Framework.AtomMapping import AtomLabel
+from .FileWithAtomDataConfigurator import FileWithAtomDataConfigurator
+
+
+class ASEFileConfigurator(FileWithAtomDataConfigurator):
+    """
+    This Configurator allows to set an input file.
+    """
+
+    def parse(self):
+
+        try:
+            self._input = ASETrajectory(self["filename"])
+        except:
+            self._input = iread(self["filename"], index="[:]")
+            first_frame = read(self["filename"], index=0)
+        else:
+            first_frame = self._input[0]
+
+        self["element_list"] = first_frame.get_chemical_symbols()
+
+    def get_atom_labels(self) -> list[AtomLabel]:
+        """
+        Returns
+        -------
+        list[AtomLabel]
+            An ordered list of atom labels.
+        """
+        labels = []
+        for atm_label in self["element_list"]:
+            label = AtomLabel(atm_label)
+            if label not in labels:
+                labels.append(label)
+        return labels

MDANSE/Src/MDANSE/Framework/Configurators/AtomMappingConfigurator.py

@@ -0,0 +1,82 @@
+# **************************************************************************
+#
+# MDANSE: Molecular Dynamics Analysis for Neutron Scattering Experiments
+#
+# @file      MDANSE/Framework/Configurators/AtomMappingConfigurator.py
+# @brief     Implements module/class/test AtomMappingConfigurator
+#
+# @homepage  https://www.isis.stfc.ac.uk/Pages/MDANSEproject.aspx
+# @license   GNU General Public License v3 or higher (see LICENSE)
+# @copyright Institut Laue Langevin 2013-now
+# @copyright ISIS Neutron and Muon Source, STFC, UKRI 2021-now
+# @authors   Scientific Computing Group at ILL (see AUTHORS)
+#
+# **************************************************************************
+import json
+
+from MDANSE.Framework.AtomMapping import fill_remaining_labels, check_mapping_valid
+from MDANSE.Framework.Configurators.IConfigurator import IConfigurator
+
+
+class AtomMappingConfigurator(IConfigurator):
+    """The atom mapping configurator.
+
+    Attributes
+    ----------
+    _default : dict
+        The default atom map setting JSON string.
+    """
+
+    _default = "{}"
+
+    def configure(self, value) -> None:
+        """
+        Parameters
+        ----------
+        value : str
+            The atom map setting JSON string.
+        """
+        if value is None:
+            value = self._default
+
+        if not isinstance(value, str):
+            self.error_status = "Invalid input value."
+            return
+
+        try:
+            value = json.loads(value)
+        except json.decoder.JSONDecodeError:
+            self.error_status = "Unable to load JSON string."
+            return
+
+        file_configurator = self._configurable[self._dependencies["input_file"]]
+        if not file_configurator._valid:
+            self.error_status = "Input file not selected."
+            return
+
+        labels = file_configurator.get_atom_labels()
+        try:
+            fill_remaining_labels(value, labels)
+        except AttributeError:
+            self.error_status = "Unable to map all atoms."
+            return
+
+        if not check_mapping_valid(value, labels):
+            self.error_status = "Atom mapping is not valid."
+            return
+
+        self.error_status = "OK"
+        self["value"] = value
+
+    def get_information(self) -> str:
+        """Returns some information on the atom mapping configurator.
+
+        Returns
+        -------
+        str
+            The atom map JSON string.
+        """
+        if "value" not in self:
+            return "Not configured yet\n"
+
+        return str(self["value"])

MDANSE/Src/MDANSE/Framework/Configurators/AtomSelectionConfigurator.py

@@ -59,9 +59,13 @@ def configure(self, value):
             self.error_status = "Invalid input value."
             return
 
+        selector = Selector(trajConfig["instance"].chemical_system)
+        if not selector.check_valid_json_settings(value):
+            self.error_status = "Invalid JSON string."
+            return
+
         self["value"] = value
 
-        selector = Selector(trajConfig["instance"].chemical_system)
         selector.update_from_json(value)
         indexes = selector.get_idxs()