From 8ad8d477f987f3ee0f70b602d37af411ba5cb1f2 Mon Sep 17 00:00:00 2001
From: Paul Saxe <paul@thesaxes.net>
Date: Mon, 15 Jul 2024 14:32:57 -0400
Subject: [PATCH] Fixed bugs in Helfand Moments approach.

---
 HISTORY.rst                     |  3 +++
 diffusivity_step/analysis.py    |  2 +-
 diffusivity_step/diffusivity.py | 15 ++++++++++++---
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/HISTORY.rst b/HISTORY.rst
index 84aa607..c13333d 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -1,6 +1,9 @@
 =======
 History
 =======
+2024.7.15 -- Bugfix: Significant error in Helfand Moment approach
+    * Now fixed and seems to be working.
+      
 2024.7.4 -- Improved fitting of curves
     * Removed weighting of the fit by the stdev since it is too biased to the beginning
     * Added control over the portion of the data to fit in order to avoid the initial
diff --git a/diffusivity_step/analysis.py b/diffusivity_step/analysis.py
index 4a68ddf..f0de51d 100644
--- a/diffusivity_step/analysis.py
+++ b/diffusivity_step/analysis.py
@@ -94,7 +94,7 @@ def create_helfand_moments(v, species, m=None):
 
     n, nmols, _ = v.shape
     if m is None:
-        m = min(n // 20, 10000)
+        m = min(n // 20, 1000)
 
     Ms = []
     M_errs = []
diff --git a/diffusivity_step/diffusivity.py b/diffusivity_step/diffusivity.py
index 0b077b6..304c0e4 100644
--- a/diffusivity_step/diffusivity.py
+++ b/diffusivity_step/diffusivity.py
@@ -8,6 +8,7 @@
 from pathlib import Path
 import pkg_resources
 import sys
+import time
 import traceback
 
 import numpy as np
@@ -387,6 +388,8 @@ def analyze(
             for spec, smiles in enumerate(self.species.keys()):
                 # Fit the slopes
                 fit = []
+                # Convert units and remember the factor of 1/3 in the Einstein equation
+                factor = Q_("Å^2/ps").m_as("m^2/s") / 3
                 for i in range(nalpha):
                     if weighted:
                         slope, err, xs, ys = fit_msd(
@@ -403,8 +406,8 @@ def analyze(
                             start=P["helfand_fit_start"],
                             end=P["helfand_fit_end"],
                         )
-                    d_coeff = slope
-                    err = err
+                    d_coeff = slope * factor
+                    err = err * factor
                     if self._scale is None:
                         # Set a scale factor to make the numbers managable
                         self._scale = 10 ** floor(log10(d_coeff))
@@ -863,7 +866,10 @@ def process_run(self, run, run_dir):
 
             metadata, result = read_vector_trajectory(paths[0])
             self._msd_dt = Q_(metadata["dt"], metadata["tunits"])
+            tic = time.perf_counter_ns()
             msd, err = compute_msd(result, species)
+            toc = time.perf_counter_ns()
+            self.logger.info(f"compute_msd: {(toc-tic)/1e+9:.3f}")
             for i in range(n_species):
                 self._msds[i].append(msd[i])
                 self._msd_errs[i].append(err[i])
@@ -902,9 +908,12 @@ def process_run(self, run, run_dir):
 
             # Convert units and remember the factor of 2 in the Helfand moments
             v_sq = Q_("Å^2/fs^2")
-            constants = (v_sq * self._velocity_dt.units).m_as("m^2/s") / 2
+            constants = (v_sq * self._velocity_dt.to("fs") ** 2).m_as("Å^2") / 2
 
+            tic = time.perf_counter_ns()
             M, err = create_helfand_moments(result, species, m=m)
+            toc = time.perf_counter_ns()
+            self.logger.info(f"create_helfand_moments: {(toc-tic)/1e+9:.3f}")
             for i in range(n_species):
                 M[i] *= constants
                 err[i] *= constants