Rob/mdanse trajectory filter

MDANSE/Src/MDANSE/Framework/Configurators/TrajectoryFilterConfigurator.py

@@ -0,0 +1,71 @@
+#    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 <https://www.gnu.org/licenses/>.
+#
+import json
+from MDANSE.Framework.Configurators.IConfigurator import IConfigurator
+from MDANSE.Mathematics.Signal import (
+    DEFAULT_FILTER,
+    filter_default_attributes,
+    filter_description_string,
+)
+
+
+class TrajectoryFilterConfigurator(IConfigurator):
+    """This configurator allows the application of a filter to the trajectory of atoms in the simulation.
+
+    Attributes
+    ----------
+    _default : str
+        The defaults selection setting.
+    """
+
+    _default_filter = DEFAULT_FILTER
+
+    _settings = filter_default_attributes()
+
+    @classmethod
+    def get_default(cls) -> str:
+        """Return the default filter string.
+
+        Returns
+        -------
+
+            A string representation of the default filter settings dictionary
+        """
+        return cls._default
+
+    _default = filter_description_string()
+
+    def configure(self, value: str):
+        """Configure an input value.
+
+        Parameters
+        ----------
+        value : str
+            The selection setting in a json readable format.
+        """
+        self._settings = value
+
+        try:
+            dict_value = json.loads(value)
+
+            if not {"filter", "attributes"} <= dict_value.keys():
+                self.error_status = f"The dictionary \n{dict_value}\n does not contain the expected keys"
+
+        except (TypeError, ValueError):
+            self.error_status = f"Value \n{value}\n in {self} is not of correct format (expected JSON string)"
+
+        self.error_status = "OK"
+        self["value"] = self._settings

MDANSE/Src/MDANSE/Framework/Jobs/TrajectoryFilter.py

@@ -0,0 +1,279 @@
+#    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 <https://www.gnu.org/licenses/>.
+#
+import collections
+import copy
+import json
+
+import numpy as np
+
+import h5py
+from MDANSE.Framework.Formats import HDFFormat
+from MDANSE.Chemistry.ChemicalSystem import ChemicalSystem
+from MDANSE.Framework.Jobs.IJob import IJob
+from MDANSE.Mathematics.Signal import FILTER_MAP
+from MDANSE.MolecularDynamics.Configuration import (
+    RealConfiguration,
+    PeriodicRealConfiguration,
+)
+from MDANSE.MolecularDynamics.Trajectory import TrajectoryWriter
+from MDANSE.MLogging import LOG
+
+
+class TrajectoryFilter(IJob):
+    """
+    Design and apply a filter for the atomic trajectories.
+    """
+
+    label = "Trajectory Filter"
+
+    category = (
+        "Analysis",
+        "Trajectory",
+    )
+
+    ancestor = ["hdf_trajectory", "molecular_viewer"]
+
+    settings = collections.OrderedDict()
+    settings["trajectory"] = ("HDFTrajectoryConfigurator", {})
+    settings["frames"] = (
+        "CorrelationFramesConfigurator",
+        {"dependencies": {"trajectory": "trajectory"}},
+    )
+    settings["instrument_resolution"] = (
+        "InstrumentResolutionConfigurator",
+        {"dependencies": {"trajectory": "trajectory", "frames": "frames"}},
+    )
+    settings["projection"] = (
+        "ProjectionConfigurator",
+        {"label": "project coordinates"},
+    )
+    settings["trajectory_filter"] = (
+        "TrajectoryFilterConfigurator",
+        {"dependencies": {"trajectory": "trajectory"}},
+    )
+    settings["atom_selection"] = (
+        "AtomSelectionConfigurator",
+        {"dependencies": {"trajectory": "trajectory"}},
+    )
+    settings["weights"] = (
+        "WeightsConfigurator",
+        {
+            "default": "atomic_weight",
+            "dependencies": {
+                "trajectory": "trajectory",
+                "atom_selection": "atom_selection",
+            },
+        },
+    )
+    settings["output_files"] = (
+        "OutputTrajectoryConfigurator",
+        {"format": "MDTFormat"},
+    )
+    settings["running_mode"] = ("RunningModeConfigurator", {})
+
+    def initialize(self):
+        """
+        Initialize the input parameters and analysis self variables
+        """
+        super().initialize()
+
+        self.numberOfSteps = self.configuration["atom_selection"]["selection_length"]
+
+        self._atoms = self.configuration["trajectory"][
+            "instance"
+        ].chemical_system.atom_list
+
+        self._selected_atoms = [
+            self._atoms[i]
+            for i in np.array(self.configuration["atom_selection"]["indices"]).flatten()
+        ]
+
+        # This stores the trajectory (position array) of atoms by x, y, z component, to be filtered
+        self.atomic_trajectory_array = np.zeros(
+            (len(self._selected_atoms), 3, len(self.configuration["frames"]["value"]))
+        )
+
+    def run_step(self, index):
+        """
+        Runs a single step of the job.
+
+        :Parameters:
+            #. index (int): The index of the step.
+        """
+        LOG.debug(f"Running step: {index}")
+        trajectory = self.configuration["trajectory"]["instance"]
+
+        # get atom index
+        indexes = self.configuration["atom_selection"]["indices"][index]
+
+        series = trajectory.read_com_trajectory(
+            indexes,
+            first=self.configuration["frames"]["first"],
+            last=self.configuration["frames"]["last"] + 1,
+            step=self.configuration["frames"]["step"],
+        )
+
+        self.atomic_trajectory_array[index] = series.T
+
+        return index, None
+
+    def combine(self, index, x):
+        """
+        Combines returned results of run_step.
+        :Parameters:
+            #. index (int): The index of the step.
+            #. x (any): The returned result(s) of run_step
+        """
+        pass
+
+    def finalize(self):
+        """
+        Finalizes the calculations (e.g. averaging the total term, output files creations ...).
+        """
+
+        # Get filter class and instantiate filter object
+        filter_config = json.loads(self.configuration["trajectory_filter"]["value"])
+
+        filter_class, filter_attributes = (
+            FILTER_MAP[filter_config["filter"]],
+            filter_config["attributes"],
+        )
+
+        if {"n_steps", "time_step_ps"} not in set(filter_attributes.keys()):
+            filter_attributes.update(
+                {
+                    "n_steps": self.configuration["trajectory"]["length"],
+                    "time_step_ps": self.configuration["trajectory"]["md_time_step"],
+                }
+            )
+
+        filter = filter_class(**filter_attributes)
+
+        trajectories = copy.deepcopy(self.atomic_trajectory_array)
+
+        # Apply filter (only apply initial position offset to atoms if filter is not lowpass and setting has been applied)
+        filtered_coords = apply(
+            filter,
+            trajectories,
+            apply_offsets=filter_attributes.get("attenuation_type", "bandpass")
+            != "lowpass",
+        )
+
+        # Create new chemical system for output trajectory
+        name = self.configuration["output_files"]["file"].name.split(".").pop(0)
+        if not isinstance(name, str):
+            name = "filtered_traj_chemical_system"
+        output_chemical_system = ChemicalSystem(name)
+        output_chemical_system.initialise_atoms(self._selected_atoms)
+
+        # Create trajectory writer object
+        self._output_trajectory = TrajectoryWriter(
+            self.configuration["output_files"]["file"],
+            output_chemical_system,
+            filter_attributes["n_steps"],
+            None,
+            positions_dtype=self.configuration["output_files"]["dtype"],
+            compression=self.configuration["output_files"]["compression"],
+        )
+
+        # Write trajectory
+        write_filtered_trajectory(
+            parent_configuration=self.configuration,
+            nsteps=filter_attributes["n_steps"],
+            filtered_coordinates=filtered_coords,
+            output_trajectory=self._output_trajectory,
+        )
+
+        # The input trajectory is closed.
+        self.configuration["trajectory"]["instance"].close()
+
+        # The output trajectory is closed.
+        self._output_trajectory.close()
+
+        # Write the filter metadata to output
+        outputFile = h5py.File(self.configuration["output_files"]["file"], "r+")
+        outputFile.create_group("metadata").create_dataset(
+            "trajectory_filter",
+            (1,),
+            data=filter.__str__(),
+            dtype=h5py.string_dtype(),
+        )
+
+        outputFile.close()
+
+        super().finalize()
+
+
+def apply(filter, trajectories, apply_offsets: bool) -> np.ndarray:
+    """ """
+    output_trajectory_array = np.zeros(trajectories.shape)
+
+    for at, (x, y, z) in enumerate(trajectories):
+        # Store initial positions
+        offsets = np.array([x[0], y[0], z[0]])
+
+        for i, component in enumerate((x, y, z)):
+            output = filter.apply(component)
+
+            if apply_offsets:
+                output += offsets[i]
+            output_trajectory_array[at][i] = output
+
+    return output_trajectory_array
+
+
+def write_filtered_trajectory(
+    parent_configuration,
+    nsteps: int,
+    filtered_coordinates: np.ndarray,
+    output_trajectory: TrajectoryWriter,
+) -> None:
+    """ """
+    time = parent_configuration["frames"]["time"]
+    dt = time[1] - time[0]
+    for index in range(nsteps):
+        frame_coordinates = [
+            (x[index], y[index], z[index]) for (x, y, z) in filtered_coordinates
+        ]
+
+        # The filtered configuration coordinates at the current frame index
+        filtered_configuration_coordinates = np.array(frame_coordinates)
+
+        filtered_configuration = get_output_configuration(
+            parent=parent_configuration["trajectory"]["instance"].configuration(
+                parent_configuration["frames"]["value"][0]
+            ),
+            output_chemical_system=output_trajectory.chemical_system,
+            output_coordinates=filtered_configuration_coordinates,
+        )
+
+        output_trajectory.chemical_system.configuration = filtered_configuration
+
+        output_trajectory.dump_configuration(
+            filtered_configuration,
+            dt * index,
+            units={"time": "ps", "unit_cell": "nm", "coordinates": "nm"},
+        )
+
+
+def get_output_configuration(parent, output_chemical_system, output_coordinates):
+    """ """
+    if parent.is_periodic:
+        return PeriodicRealConfiguration(
+            output_chemical_system, output_coordinates, parent.unit_cell
+        )
+
+    return RealConfiguration(output_chemical_system, output_coordinates)