Update molecule mapping

mdsuite/graph_modules/molecular_graph.py

@@ -51,6 +51,7 @@ def __init__(
         from_smiles: bool = False,
         from_configuration: bool = True,
         smiles_string: str = None,
+        reference_dict: dict = None,
         species: list = None,
     ):
         """
@@ -65,7 +66,17 @@ def __init__(
         from_configuration : bool
                 Build graphs from a configuration.
         smiles_string : str
-                SMILES string to read
+                SMILES string to read and used in the construction of a reference
+                graph with molecule information.
+        reference_dict : dict
+                Alternatively to the SMILES string a reference dict can be given to
+                construct a reference molecule. This dict should be of the form:
+
+                .. code-block: python
+
+                   {'PF6': {'P': 1, 'F': 6}}}
+
+                for a PF6 molecule.
         species : list
                 List of species to build a graph with.
         """
@@ -74,6 +85,7 @@ def __init__(
         self.from_configuration = from_configuration
 
         self.smiles_string = smiles_string
+        self.reference_dict = reference_dict
         self.species = species
 
         self.database = Database(self.experiment.database_path / "database.hdf5")
@@ -101,18 +113,11 @@ def get_neighbour_list(positions: tf.Tensor, cell: list = None) -> tf.Tensor:
             If periodic boundary conditions are used, please supply the cell
             dimensions, e.g. [13.97, 13.97, 13.97]. If the cell is provided
             minimum image convention will be applied!
-        batch_size: int
-            Has to be evenly divisible by the the number of configurations.
 
         Returns
         -------
-        generator object which results all distances for the current batch of
-        time steps
-
-        To get the real r_ij matrix for one time_step you can use the following:
-            r_ij_mat = np.zeros((n_atoms, n_atoms, 3))
-            r_ij_mat[np.triu_indices(n_atoms, k = 1)] = get_neighbour_list(``*args``)
-            r_ij_mat -= r_ij_mat.transpose(1, 0, 2)
+        neighbour_list : tf.Tensor
+                Neighbour list for a single configuration.
 
         """
         r_ij_matrix = tf.reshape(positions, (1, len(positions), 3)) - tf.reshape(
@@ -163,8 +168,9 @@ def build_configuration_graph(self, cutoff: float, adjacency: bool = True):
         Parameters
         ----------
         cutoff : float
+                Cutoff over which to look for molecules.
         adjacency : bool
-                If true, the adjacent matrix is returned.
+                If true, the adjacency matrix is returned.
         Returns
         -------
 
@@ -252,9 +258,9 @@ def check_a_in_b(a, b):
         else:
             if len(molecules) != n_molecules:
                 raise ValueError(
-                    "Expected number of molecules does not "
-                    "match the amount computed, please adjust"
-                    "parameters."
+                    f"Expected number of molecules ({n_molecules}) does not "
+                    f"match the amount computed ({len(molecules)}), please adjust "
+                    f"parameters."
                 )
             else:
                 return molecules

mdsuite/transformations/map_molecules.py

@@ -34,6 +34,11 @@
 from mdsuite.transformations.transformations import Transformations
 from mdsuite.utils.meta_functions import join_path
 
+import logging
+
+log = logging.getLogger(__name__)
+
+
 switcher_transformations = {
     "Translational_Dipole_Moment": "TranslationalDipoleMoment",
     "Ionic_Current": "IonicCurrent",
@@ -53,8 +58,6 @@ class MolecularMap(Transformations):
     scale_function : dict
             A dictionary referencing the memory/time scaling function of the
             transformation.
-    experiment : object
-            Experiment object to work within.
     molecules : dict
             Molecule dictionary to use as reference. e.g.
 
@@ -67,19 +70,8 @@ class MolecularMap(Transformations):
     """
 
     def __init__(self):
-        """Constructor for the MolecularMap class.
-
-        Parameters
-        ----------
-        molecules : dict
-                Molecule dictionary to use as reference. e.g, the input for
-                emim-PF6 ionic liquid would be.
-
-                .. code-block::
-
-                   {'emim': {'smiles': 'CCN1C=C[N+](+C1)C', 'amount': 20},
-                   'PF6': {'smiles': 'F[P-](F)(F)(F)(F)F', 'amount': 20}}
-
+        """
+        Constructor for the MolecularMap class.
         """
         super().__init__()
         self.molecules = None
@@ -88,9 +80,16 @@ def __init__(self):
         self.dependency = "Unwrapped_Positions"
         self.scale_function = {"quadratic": {"outer_scale_factor": 5}}
 
-    def _prepare_database_entry(self, species, number_of_molecules: int) -> dict:
+    def _prepare_database_entry(self, species: str, number_of_molecules: int) -> dict:
         """
         Call some housekeeping methods and prepare for the transformations.
+
+        Parameters
+        ----------
+        species
+                Name of the species to be added
+        number_of_molecules : int
+                Number of molecules to be added to the database.
         Returns
         -------
         data_structure : dict
@@ -240,20 +239,38 @@ def _update_species_type_dict(
 
     def _build_reference_graphs(self):
         """
-        Build the reference graphs from the SMILES strings.
+        Build the reference graphs from the reference input.
+        .
         Returns
         -------
-        Nothing.
+        Updates the reference molecules attribute of the class. It should be populated
+        after this call and look something like:
+
+        {'bmim': {'species': ['C', 'H', 'N'], 'mass': 200, 'graph': tf.Tensor}}
+
+        If a dict is used as a reference, the graph will not exist and stronger
+        isomorphism checks will not be possible.
         """
         for item in self.molecules:
             self.reference_molecules[item] = {}
-            mol, species = MolecularGraph(
-                self.experiment,
-                from_smiles=True,
-                smiles_string=self.molecules[item]["smiles"],
-            ).build_smiles_graph()
+            if 'smiles' in self.molecules[item]:
+                mol, species = MolecularGraph(
+                    self.experiment,
+                    from_smiles=True,
+                    smiles_string=self.molecules[item]["smiles"],
+                ).build_smiles_graph()
+                self.reference_molecules[item]["graph"] = mol
+            elif 'reference' in self.molecules[item]:
+                species = self.molecules[item]['reference']
+            else:
+                error_msg = "The minimum amount of data was not given to the mapping." \
+                            "Either provide a reference key with information about" \
+                            "Which species and the number of them are in the molecule," \
+                            "or provide a SMILES string that can be used to compute " \
+                            "this information."
+                log.info(error_msg)
+                raise ValueError("Provide either a reference key or smiles string")
             self.reference_molecules[item]["species"] = list(species)
-            self.reference_molecules[item]["graph"] = mol
             self.reference_molecules[item]["mass"] = self._get_molecular_mass(species)
 
     def _get_molecular_mass(self, species_dict: dict) -> float:
@@ -399,6 +416,7 @@ def _map_molecules(self):
                 for t, molecule in enumerate(
                     self.adjacency_graphs[molecule_name]["molecules"]
                 ):
+                    # indices of the atoms inside the specific molecule.
                     indices = list(
                         self.adjacency_graphs[molecule_name]["molecules"][
                             molecule
@@ -425,7 +443,7 @@ def _build_indices_dict(self, indices: List[int], species: List[str]) -> dict:
         indices : list[int]
                 Indices of atoms belonging in the molecule
         species : list[str]
-                Elements in the molecules, in the order that they are loaded.
+                Species in the molecules, in the order that they are loaded.
         Returns
         -------
         group_dict : dict
@@ -458,6 +476,28 @@ def _build_indices_dict(self, indices: List[int], species: List[str]) -> dict:
     def run_transformation(self, molecules: dict):
         """
         Perform the transformation.
+
+        Parameters
+        ----------
+        molecules : dict
+                Molecule dictionary to use as reference. The reference component is the
+                most critical part. One can either use a smiles string or a reference
+                dict as demonstrated below.
+
+                e.g, the input for emim-PF6 ionic liquid would be:
+
+                .. code-block::
+
+                   {'emim': {'smiles': 'CCN1C=C[N+](+C1)C', 'amount': 20},
+                   'PF6': {'smiles': 'F[P-](F)(F)(F)(F)F', 'amount': 20}}
+
+                or:
+
+                .. code-block::
+
+                   {'emim': {'reference': {'C': , 'N': 'H': }}, 'amount': 20},
+                   'PF6': {'reference': {'P': 1, 'F': 6}, 'amount': 20}}
+
         Returns
         -------
         Update the experiment database.

mdsuite/transformations/transformation_dict.py

@@ -25,7 +25,7 @@
 -------
 """
 
-# TODO remvoe this!
+# TODO remove when molecule mapping is adjusted.
 # map_molecules is the only transformation still working with it
 
 from mdsuite.transformations.integrated_heat_current import IntegratedHeatCurrent

mdsuite/visualizer/znvis_visualizer.py

@@ -65,7 +65,7 @@ def __init__(
         """
         self.counter = 0
         # Particle information
-        self.database = Database(name=join_path(database_path, "database.hdf5"))
+        self.database = Database(join_path(database_path, "database.hdf5"))
         self.frame_rate = frame_rate
         self.species = species
         if unwrapped:
@@ -138,5 +138,5 @@ def run_visualization(self):
         Opens the ZnVis app and runs the visualization.
         """
         particle_list = self._prepare_species()
-        visualizer = znvis.Visualizer(particles=particle_list, frame_rate=50)
+        visualizer = znvis.Visualizer(particles=particle_list, frame_rate=17)
         visualizer.run_visualization()