Add lags to vafs

Author: oerc0122

janus_core/helpers/post_process.py

@@ -190,6 +190,7 @@ def compute_vaf(
     fft: bool = False,
     index: SliceLike = (0, None, 1),
     filter_atoms: MaybeSequence[MaybeSequence[int]] = ((),),
+    time_step: float = 1.,
 ) -> NDArray[float64]:
     """
     Compute the velocity autocorrelation function (VAF) of `data`.
@@ -212,6 +213,9 @@ def compute_vaf(
     filter_atoms : MaybeSequence[MaybeSequence[int]]
         Compute the VAF averaged over subsets of the system.
         Default is all atoms.
+    time_step : float
+        Time step for scaling lags to align with input data.
+        Default is 1 (i.e. no scaling).
 
     Returns
     -------
@@ -238,9 +242,7 @@ def compute_vaf(
     data = data[slice(*index)]
 
     if use_velocities:
-        momenta = np.asarray(
-            [datum.get_momenta() / datum.get_masses() for datum in data]
-        )
+        momenta = np.asarray([datum.get_velocities() for datum in data])
     else:
         momenta = np.asarray([datum.get_momenta() for datum in data])
 
@@ -270,17 +272,21 @@ def compute_vaf(
 
     vafs /= n_steps - np.arange(n_steps)
 
+    lags = np.arange(n_steps) * time_step
+
     if fft:
         vafs = np.fft.fft(vafs, axis=0)
+        lags = 1. / lags
 
-    vafs = [
+    vafs = (lags, [
         np.average([vafs[used_atoms[i]] for i in atoms], axis=0)
         for atoms in filter_atoms
-    ]
+    ])
 
     if filenames:
-        for filename, vaf in zip(filenames, vafs):
+        for vaf, filename in zip(vafs[1], filenames):
             with open(filename, "w", encoding="utf-8") as out_file:
-                print(*vaf, file=out_file, sep="\n")
+                for lag, dat in zip(lags, vaf):
+                    print(lag, dat, file=out_file)
 
     return vafs