Save Q vectors in the output file

MDANSE/Src/MDANSE/Framework/Configurators/QVectorsConfigurator.py

@@ -58,12 +58,12 @@ def configure(self, value):
         trajConfig = self._configurable[self._dependencies["trajectory"]]
         if isinstance(value, tuple):
             try:
-                generator, parameters = value
+                generator_name, parameters = value
             except ValueError:
                 self.error_status = f"Invalid q vectors settings {value}"
                 return
             generator = IQVectors.create(
-                generator, trajConfig["instance"].chemical_system
+                generator_name, trajConfig["instance"].chemical_system
             )
             try:
                 generator.setup(parameters)
@@ -99,6 +99,7 @@ def configure(self, value):
         self["shells"] = list(self["q_vectors"].keys())
         self["n_shells"] = len(self["q_vectors"])
         self["value"] = self["q_vectors"]
+        self["generator"] = generator
         self.error_status = "OK"
 
     def preview_output_axis(self):

MDANSE/Src/MDANSE/Framework/Jobs/CurrentCorrelationFunction.py

@@ -369,6 +369,10 @@ def finalize(self):
         """Normalize, Fourier transform and write the results out to
         the MDA files.
         """
+        self.configuration["q_vectors"]["generator"].write_vectors_to_file(
+            self._outputData
+        )
+
         nAtomsPerElement = self.configuration["atom_selection"].get_natoms()
         for pair in self._elementsPairs:
             at1, at2 = pair

MDANSE/Src/MDANSE/Framework/Jobs/DynamicCoherentStructureFactor.py

@@ -24,6 +24,8 @@
 from MDANSE.Framework.Jobs.IJob import IJob
 from MDANSE.Mathematics.Arithmetic import weight
 from MDANSE.Mathematics.Signal import get_spectrum
+from MDANSE.Framework.QVectors.IQVectors import IQVectors
+from MDANSE.MolecularDynamics.UnitCell import UnitCell
 
 
 class DynamicCoherentStructureFactorError(Error):
@@ -181,6 +183,18 @@ def initialize(self):
             main_result=True,
         )
 
+        self._average_unit_cell = UnitCell(
+            np.mean(
+                [
+                    self.configuration["trajectory"]["instance"]
+                    .unit_cell(frame)
+                    ._unit_cell
+                    for frame in self.configuration["frames"]["value"]
+                ],
+                axis=0,
+            )
+        )
+
     def run_step(self, index):
         """
         Runs a single step of the job.\n
@@ -213,7 +227,16 @@ def run_step(self, index):
 
             # loop over the trajectory time steps
             for i, frame in enumerate(self.configuration["frames"]["value"]):
-                qVectors = self.configuration["q_vectors"]["value"][shell]["q_vectors"]
+                # unit_cell = traj.unit_cell(frame)
+                unit_cell = self._average_unit_cell
+                try:
+                    hkls = self.configuration["q_vectors"]["value"][shell]["hkls"]
+                except KeyError:
+                    qVectors = self.configuration["q_vectors"]["value"][shell][
+                        "q_vectors"
+                    ]
+                else:
+                    qVectors = IQVectors.hkl_to_qvectors(hkls, unit_cell)
 
                 coords = traj.configuration(frame)["coordinates"]
 
@@ -247,6 +270,9 @@ def finalize(self):
         """
         Finalizes the calculations (e.g. averaging the total term, output files creations ...)
         """
+        self.configuration["q_vectors"]["generator"].write_vectors_to_file(
+            self._outputData
+        )
 
         nAtomsPerElement = self.configuration["atom_selection"].get_natoms()
         for pair in self._elementsPairs:

MDANSE/Src/MDANSE/Framework/Jobs/DynamicIncoherentStructureFactor.py

@@ -253,6 +253,9 @@ def finalize(self):
         """
         Finalizes the calculations (e.g. averaging the total term, output files creations ...)
         """
+        self.configuration["q_vectors"]["generator"].write_vectors_to_file(
+            self._outputData
+        )
 
         nAtomsPerElement = self.configuration["atom_selection"].get_natoms()
         for element, number in list(nAtomsPerElement.items()):

MDANSE/Src/MDANSE/Framework/Jobs/ElasticIncoherentStructureFactor.py

@@ -193,6 +193,10 @@ def finalize(self):
         Finalizes the calculations (e.g. averaging the total term, output files creations ...)
         """
 
+        self.configuration["q_vectors"]["generator"].write_vectors_to_file(
+            self._outputData
+        )
+
         nAtomsPerElement = self.configuration["atom_selection"].get_natoms()
         for element, number in list(nAtomsPerElement.items()):
             self._outputData["eisf_%s" % element][:] /= number

MDANSE/Src/MDANSE/Framework/QVectors/ApproximateDispersionQVectors.py

@@ -69,7 +69,7 @@ def _generate(self):
             np.array(qStart)[:, np.newaxis] + np.outer(n, np.arange(0, nSteps)) * qStep
         )
 
-        hkls = np.rint(np.dot(self._directUnitCell, vects))
+        hkls = self.qvectors_to_hkl(vects, self._unit_cell)
 
         dists = np.sqrt(np.sum(vects**2, axis=0))
         dists = list(zip(range(len(dists)), dists))

MDANSE/Src/MDANSE/Framework/QVectors/CircularLatticeQVectors.py

@@ -61,7 +61,7 @@ def _generate(self):
             ]
         )
 
-        qVects = np.dot(self._inverseUnitCell, hkls)
+        qVects = self.hkl_to_qvectors(hkls, self._unit_cell)
 
         qMax = (
             self._configuration["shells"]["last"]
@@ -109,11 +109,8 @@ def _generate(self):
                 self._configuration["q_vectors"][q]["q_vectors"] = vects[:, hits]
                 self._configuration["q_vectors"][q]["n_q_vectors"] = n
                 self._configuration["q_vectors"][q]["q"] = q
-                self._configuration["q_vectors"][q]["hkls"] = np.rint(
-                    np.dot(
-                        self._directUnitCell,
-                        self._configuration["q_vectors"][q]["q_vectors"],
-                    )
+                self._configuration["q_vectors"][q]["hkls"] = self.qvectors_to_hkl(
+                    vects[:, hits], self._unit_cell
                 )
 
             if self._status is not None:

MDANSE/Src/MDANSE/Framework/QVectors/DispersionLatticeQVectors.py

@@ -48,7 +48,7 @@ def _generate(self):
         )
 
         # The k matrix (3,n_hkls)
-        vects = np.dot(self._inverseUnitCell, hkls)
+        vects = self.hkl_to_qvectors(hkls, self._unit_cell)
 
         dists = np.sqrt(np.sum(vects**2, axis=0))
 

MDANSE/Src/MDANSE/Framework/QVectors/GridQVectors.py

@@ -59,7 +59,7 @@ def _generate(self):
         hkls = hkls.reshape(3, nh * nk * nl)
 
         # The k matrix (3,n_hkls)
-        vects = np.dot(self._inverseUnitCell, hkls)
+        vects = self.hkl_to_qvectors(hkls, self._unit_cell)
 
         dists = np.sqrt(np.sum(vects**2, axis=0))
 

MDANSE/Src/MDANSE/Framework/QVectors/IQVectors.py

@@ -14,12 +14,19 @@
 #    along with this program.  If not, see <https://www.gnu.org/licenses/>.
 #
 import abc
+from typing import TYPE_CHECKING
+
+import numpy as np
 
 from MDANSE.Core.Error import Error
 from MDANSE.Framework.Configurable import Configurable
 from MDANSE.Core.SubclassFactory import SubclassFactory
 from MDANSE.MLogging import LOG
 
+if TYPE_CHECKING:
+    from MDANSE.MolecularDynamics.UnitCell import UnitCell
+    from MDANSE.Framework.OutputVariables.IOutputVariable import OutputData
+
 
 class QVectorsError(Error):
     pass
@@ -54,3 +61,90 @@ def generate(self) -> bool:
 
     def setStatus(self, status):
         self._status = status
+
+    @classmethod
+    def qvectors_to_hkl(
+        self, vector_array: np.array, unit_cell: "UnitCell"
+    ) -> np.ndarray:
+        """Using a unit cell definition, recalculates an array
+        of q vectors to an equivalent array of HKL Miller indices.
+
+        Parameters
+        ----------
+        vector_array : np.array
+            a (3,N) array of scattering vectors
+        unit_cell : UnitCell
+            an instance of UnitCell class describing the simulation box
+
+        Returns
+        -------
+        np.ndarray
+            A (3,N) array of HKL values (Miller indices)
+        """
+        return np.dot(unit_cell.direct, vector_array) / (2 * np.pi)
+
+    @classmethod
+    def hkl_to_qvectors(self, hkls: np.array, unit_cell: "UnitCell") -> np.ndarray:
+        """Converts an array of HKL values to scattering vectors
+        based on the definition of a unit cell.
+
+        Parameters
+        ----------
+        hkls : np.array
+            A (3,N) array of HKL values (Miller indices)
+        unit_cell : UnitCell
+            An instance of UnitCell class describing the simulation box shape
+
+        Returns
+        -------
+        np.ndarray
+            a (3, N) array of Q vectors (scattering vectors)
+        """
+        return 2 * np.pi * np.dot(unit_cell.inverse, hkls)
+
+    def write_vectors_to_file(self, output_data: "OutputData"):
+        """Writes a summary of the generated vectors to the output
+        file using an OutputData class instance.
+
+        Parameters
+        ----------
+        output_data : OutputData
+            An object managing the writeout to one or many output files
+        """
+        q_values = [float(x) for x in self._configuration["q_vectors"].keys()]
+        output_data.add(
+            "vector_generator_q",
+            "LineOutputVariable",
+            q_values,
+            units="1/nm",
+        )
+        qvector_lengths = [
+            self._configuration["q_vectors"][q]["q_vectors"].shape[1] for q in q_values
+        ]
+        qarray_maxlength = np.max(qvector_lengths)
+        output_data.add(
+            "vector_generator_qvector_array",
+            "VolumeOutputVariable",
+            (len(q_values), 3, qarray_maxlength),
+            units="1/nm",
+        )
+        output_data["vector_generator_qvector_array"][:] = 0.0
+        for nq, q in enumerate(q_values):
+            output_data["vector_generator_qvector_array"][
+                nq, :, : qvector_lengths[nq]
+            ] = self._configuration["q_vectors"][q]["q_vectors"]
+        try:
+            hkl_arrays = [self._configuration["q_vectors"][q]["hkls"] for q in q_values]
+        except KeyError:
+            return
+        output_data.add(
+            "vector_generator_hkl_array",
+            "VolumeOutputVariable",
+            (len(q_values), 3, qarray_maxlength),
+            units="au",
+        )
+        output_data["vector_generator_hkl_array"][:] = 0.0
+        for nq, q in enumerate(q_values):
+            output_data["vector_generator_hkl_array"][nq, :, : qvector_lengths[nq]] = (
+                hkl_arrays[nq]
+            )

MDANSE/Src/MDANSE/Framework/QVectors/LatticeQVectors.py

@@ -33,10 +33,4 @@ def __init__(self, chemical_system, status=None):
                 "The universe must be periodic for building lattice-based Q vectors"
             )
 
-        self._inverseUnitCell = (
-            2.0 * np.pi * self._chemical_system.configuration.unit_cell.inverse
-        )
-
-        self._directUnitCell = (
-            2.0 * np.pi * self._chemical_system.configuration.unit_cell.direct
-        )
+        self._unit_cell = self._chemical_system.configuration.unit_cell

MDANSE/Src/MDANSE/Framework/QVectors/LinearLatticeQVectors.py

@@ -51,7 +51,9 @@ def _generate(self):
             random.seed(self._configuration["seed"]["value"])
 
         # The Q vector corresponding to the input hkl.
-        qVect = np.dot(self._inverseUnitCell, self._configuration["axis"]["vector"])
+        qVect = self.hkl_to_qvectors(
+            self._configuration["axis"]["vector"], self._unit_cell
+        )
 
         qMax = (
             self._configuration["shells"]["last"]
@@ -95,11 +97,8 @@ def _generate(self):
                 self._configuration["q_vectors"][q]["q_vectors"] = vects[:, hits]
                 self._configuration["q_vectors"][q]["n_q_vectors"] = n
                 self._configuration["q_vectors"][q]["q"] = q
-                self._configuration["q_vectors"][q]["hkls"] = np.rint(
-                    np.dot(
-                        self._directUnitCell,
-                        self._configuration["q_vectors"][q]["q_vectors"],
-                    )
+                self._configuration["q_vectors"][q]["hkls"] = self.qvectors_to_hkl(
+                    vects[:, hits], self._unit_cell
                 )
 
             if self._status is not None:

MDANSE/Src/MDANSE/Framework/QVectors/MillerIndicesQVectors.py

@@ -70,7 +70,7 @@ def _generate(self):
         hkls = hkls.reshape(3, int(round(hkls.size / 3)))
 
         # The k matrix (3,n_hkls)
-        vects = np.dot(self._inverseUnitCell, hkls)
+        vects = self.hkl_to_qvectors(hkls, self._unit_cell)
 
         dists2 = np.sum(vects**2, axis=0)
 
@@ -96,11 +96,8 @@ def _generate(self):
                 self._configuration["q_vectors"][q]["q_vectors"] = vects[:, hits]
                 self._configuration["q_vectors"][q]["n_q_vectors"] = nHits
                 self._configuration["q_vectors"][q]["q"] = q
-                self._configuration["q_vectors"][q]["hkls"] = np.rint(
-                    np.dot(
-                        self._directUnitCell,
-                        self._configuration["q_vectors"][q]["q_vectors"],
-                    )
+                self._configuration["q_vectors"][q]["hkls"] = self.qvectors_to_hkl(
+                    vects[:, hits], self._unit_cell
                 )
 
             if self._status is not None:

MDANSE/Src/MDANSE/Framework/QVectors/SphericalLatticeQVectors.py

@@ -43,27 +43,27 @@ def _generate(self):
         if self._configuration["seed"]["value"] != 0:
             np.random.seed(self._configuration["seed"]["value"])
             random.seed(self._configuration["seed"]["value"])
-
         qMax = (
             self._configuration["shells"]["last"]
             + 0.5 * self._configuration["width"]["value"]
         )
 
-        hklMax = (
-            np.ceil([qMax / np.sqrt(np.sum(v**2)) for v in self._inverseUnitCell.T]) + 1
-        )
+        hklMax = np.ceil(self.qvectors_to_hkl(qMax * np.eye(3), self._unit_cell)) + 1
 
-        vects = np.mgrid[
-            -hklMax[0] : hklMax[0] + 1,
-            -hklMax[1] : hklMax[1] + 1,
-            -hklMax[2] : hklMax[2] + 1,
+        hkl_vects = np.mgrid[
+            -hklMax[0, 0] : hklMax[0, 0] + 1,
+            -hklMax[1, 1] : hklMax[1, 1] + 1,
+            -hklMax[2, 2] : hklMax[2, 2] + 1,
         ]
 
-        vects = vects.reshape(
-            3, int(2 * hklMax[0] + 1) * int(2 * hklMax[1] + 1) * int(2 * hklMax[2] + 1)
+        hkl_vects = hkl_vects.reshape(
+            3,
+            int(2 * hklMax[0, 0] + 1)
+            * int(2 * hklMax[1, 1] + 1)
+            * int(2 * hklMax[2, 2] + 1),
         )
 
-        vects = np.dot(self._inverseUnitCell, vects)
+        vects = self.hkl_to_qvectors(hkl_vects, self._unit_cell)
 
         dists2 = np.sum(vects**2, axis=0)
 
@@ -96,11 +96,8 @@ def _generate(self):
                 self._configuration["q_vectors"][q]["q_vectors"] = vects[:, hits]
                 self._configuration["q_vectors"][q]["n_q_vectors"] = n
                 self._configuration["q_vectors"][q]["q"] = q
-                self._configuration["q_vectors"][q]["hkls"] = np.rint(
-                    np.dot(
-                        self._directUnitCell,
-                        self._configuration["q_vectors"][q]["q_vectors"],
-                    )
+                self._configuration["q_vectors"][q]["hkls"] = self.qvectors_to_hkl(
+                    vects[:, hits], self._unit_cell
                 )
 
             if self._status is not None: