Merge pull request #677 from ISISNeutronMuon/chi/gdisf-add-msd-to-output

Save the MSD results with equal weights to GDISF output

MDANSE/Src/MDANSE/Framework/Configurators/WeightsConfigurator.py

@@ -13,6 +13,7 @@
 #    You should have received a copy of the GNU General Public License
 #    along with this program.  If not, see <https://www.gnu.org/licenses/>.
 #
+from typing import Optional
 from collections import defaultdict
 import itertools
 
@@ -109,7 +110,23 @@ def configure(self, value):
         self["property"] = value
         self.error_status = "OK"
 
-    def get_weights(self):
+    def get_weights(self, prop: Optional[str] = None):
+        """Generate a dictionary of weights.
+
+        Parameters
+        ----------
+        prop : str or None, optional
+            The property to generate the weights from, if None then the
+            property set in this configurator will be used.
+
+        Returns
+        -------
+        dict[str, float]
+            The dictionary of the weights.
+        """
+        if not prop:
+            prop = self["property"]
+
         atom_selection_configurator = self._configurable[
             self._dependencies["atom_selection"]
         ]
@@ -123,7 +140,7 @@ def get_weights(self):
             atom_selection_configurator["selection_length"],
         ):
             weights[name] += sum(
-                self._trajectory.get_atom_property(element, self["property"])
+                self._trajectory.get_atom_property(element, prop)
                 for element in elements
             )
 

MDANSE/Src/MDANSE/Framework/Jobs/GaussianDynamicIncoherentStructureFactor.py

@@ -166,6 +166,15 @@ def initialize(self):
                 main_result=True,
                 partial_result=True,
             )
+            self._outputData.add(
+                f"msd_{element}",
+                "LineOutputVariable",
+                (self._nFrames,),
+                axis="time",
+                units="nm2",
+                main_result=True,
+                partial_result=True,
+            )
 
         self._outputData.add(
             "f(q,t)_total",
@@ -182,20 +191,32 @@ def initialize(self):
             units="nm2/ps",
             main_result=True,
         )
+        self._outputData.add(
+            "msd_total",
+            "LineOutputVariable",
+            (self._nFrames,),
+            axis="time",
+            units="nm2",
+            main_result=True,
+        )
 
         self._atoms = self.configuration["trajectory"][
             "instance"
         ].chemical_system.atom_list
 
-    def run_step(self, index):
-        """
-        Runs a single step of the job.\n
+    def run_step(self, index: int):
+        """Calculates the GDISF and MSD of an atom.
 
-        :Parameters:
-            #. index (int): The index of the step.
-        :Returns:
-            #. index (int): The index of the step.
-            #. atomicSF (np.array): The atomic structure factor
+        Parameters
+        ----------
+        index : int
+            The index of the atom that the calculation will be run over.
+
+        Returns
+        -------
+        tuple[int, tuple[np.ndarray, np.ndarray]]
+            A tuple which contains the job index and a tuple of the
+            GDISF and MSD of an atom.
         """
 
         # get atom index
@@ -218,29 +239,29 @@ def run_step(self, index):
 
         for i, q2 in enumerate(self._kSquare):
             gaussian = np.exp(-msd * q2 / 6.0)
-
             atomicSF[i, :] += gaussian
 
-        return index, atomicSF
+        return index, (atomicSF, msd)
 
-    def combine(self, index, x):
-        """
-        Combines returned results of run_step.\n
-        :Parameters:
-            #. index (int): The index of the step.\n
-            #. x (any): The returned result(s) of run_step
-        """
+    def combine(self, index: int, x: tuple[np.ndarray, np.ndarray]):
+        """Add the results to the output files.
 
-        # The symbol of the atom.
+        Parameters
+        ----------
+        index : int
+            The atom index that the calculation was run over.
+        x : tuple[np.ndarray, np.ndarray]
+            A tuple of the GDISF and MSD of an atom.
+        """
         element = self.configuration["atom_selection"]["names"][index]
-
-        self._outputData[f"f(q,t)_{element}"] += x
+        atomicSF, msd = x
+        self._outputData[f"f(q,t)_{element}"] += atomicSF
+        self._outputData[f"msd_{element}"] += msd
 
     def finalize(self):
         """
         Finalizes the calculations (e.g. averaging the total term, output files creations ...)
         """
-
         nAtomsPerElement = self.configuration["atom_selection"].get_natoms()
         for element, number in nAtomsPerElement.items():
             self._outputData[f"f(q,t)_{element}"][:] /= number
@@ -250,6 +271,8 @@ def finalize(self):
                 self.configuration["instrument_resolution"]["time_step"],
                 axis=1,
             )
+            self._outputData[f"msd_{element}"][:] /= number
+
         weights = self.configuration["weights"].get_weights()
         weight_dict = get_weights(weights, nAtomsPerElement, 1)
         assign_weights(self._outputData, weight_dict, "f(q,t)_%s")
@@ -259,13 +282,23 @@ def finalize(self):
             weight_dict,
             "f(q,t)_%s",
         )
-
         self._outputData["s(q,f)_total"][:] = weighted_sum(
             self._outputData,
             weight_dict,
             "s(q,f)_%s",
         )
 
+        # since GDISF ~ exp(-msd * q2 / 6.0) the MSD isn't weighted in
+        # the exp lets save the MSD with equal weights
+        weights = self.configuration["weights"].get_weights("equal")
+        weight_dict = get_weights(weights, nAtomsPerElement, 1)
+        assign_weights(self._outputData, weight_dict, "msd_%s")
+        self._outputData["msd_total"][:] = weighted_sum(
+            self._outputData,
+            weight_dict,
+            "msd_%s",
+        )
+
         self._outputData.write(
             self.configuration["output_files"]["root"],
             self.configuration["output_files"]["formats"],

MDANSE/Tests/UnitTests/Analysis/test_scattering.py

@@ -291,12 +291,12 @@ def test_gdisf(traj_info):
     result_file = os.path.join(result_dir, f"gdisf_{traj_info[0]}.mda")
     with h5py.File(temp_name + ".mda") as actual, h5py.File(result_file) as desired:
         keys = [
-            i for i in desired.keys() if any([j in i for j in ["f(q,t)", "s(q,f)"]])
+            i for i in desired.keys() if any([j in i for j in ["f(q,t)", "s(q,f)", "msd"]])
         ]
         for key in keys:
             np.testing.assert_array_almost_equal(
                 actual[f"/{key}"] * actual[f"/{key}"].attrs["scaling_factor"],
-                desired[f"/{key}"],
+                desired[f"/{key}"] * desired[f"/{key}"].attrs["scaling_factor"],
             )
 
     os.remove(temp_name + ".mda")