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[Optional[int]]] = ((None),),
+    time_step: float = 1.0,
 ) -> NDArray[float64]:
     """
     Compute the velocity autocorrelation function (VAF) of `data`.
@@ -212,6 +213,9 @@ def compute_vaf(
     filter_atoms : MaybeSequence[MaybeSequence[Optional[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
     -------
@@ -270,17 +274,24 @@ 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)
-
-    vafs = [
-        np.average([vafs[used_atoms[i]] for i in atoms], axis=0)
-        for atoms in filter_atoms
-    ]
+        lags = 1. / lags
+
+    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