Adding autocorrelation function to sampling.py. (#92)

* Adding autocorrelation function to sampling.py.

* pin python < 3.12

* Adding test_autocorr1D to test_sampling.py.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix packmol seed in failing gsd tests

* ignore pre-commit error at line level

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add parameters and return to doc strings

---------

Co-authored-by: chrisjonesBSU <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

cmeutils/sampling.py

@@ -103,3 +103,23 @@ def is_equilibrated(data, threshold_fraction=0.50, threshold_neff=50, nskip=1):
         return [True, t0, g, Neff]
     else:
         return [False, None, None, None]
+
+
+def autocorr1D(array):
+    """Takes in a linear np array, performs autocorrelation
+    function and returns normalized array with half the length
+    of the input.
+
+    Parameters
+    ----------
+    data : numpy.typing.Arraylike, required
+        1-D series of data to perform autocorrelation on.
+
+    Returns
+    -------
+    1D np.array
+
+    """
+    ft = np.fft.rfft(array - np.average(array))
+    acorr = np.fft.irfft(ft * np.conjugate(ft)) / (len(array) * np.var(array))
+    return acorr[0 : len(acorr) // 2]  # noqa: E203

cmeutils/tests/test_gsd.py

@@ -50,7 +50,7 @@ def test_ellipsoid_gsd(self, butane_gsd):
     def test_create_rigid_snapshot(self):
         benzene = mb.load("c1ccccc1", smiles=True)
         benzene.name = "Benzene"
-        box = mb.fill_box(benzene, 5, box=[1, 1, 1])
+        box = mb.fill_box(benzene, 5, box=[3, 3, 3], seed=42)
         box.label_rigid_bodies(discrete_bodies="Benzene")
 
         rigid_init_snap = create_rigid_snapshot(box)
@@ -60,7 +60,7 @@ def test_create_rigid_snapshot(self):
     def test_update_rigid_snapshot(self):
         benzene = mb.load("c1ccccc1", smiles=True)
         benzene.name = "Benzene"
-        box = mb.fill_box(benzene, 5, box=[1, 1, 1])
+        box = mb.fill_box(benzene, 5, box=[3, 3, 3], seed=42)
         box.label_rigid_bodies(discrete_bodies="Benzene")
 
         rigid_init_snap = create_rigid_snapshot(box)

cmeutils/tests/test_sampling.py

@@ -1,7 +1,8 @@
 import numpy as np
 import pytest
+import scipy
 
-from cmeutils.sampling import equil_sample, is_equilibrated
+from cmeutils.sampling import autocorr1D, equil_sample, is_equilibrated
 from cmeutils.tests.base_test import BaseTest
 
 
@@ -82,3 +83,28 @@ def test_trim_high_threshold(self, correlated_data_tau100_n10000):
             [equil_data, uncorr_indices, prod_start, Neff] = equil_sample(
                 data, threshold_fraction=0.75, threshold_neff=10000
             )
+
+    def test_autocorr1D(self):
+        randoms = np.random.randint(100, size=(100))
+        integers = np.array(np.arange(100))
+
+        autocorrelation_randoms = autocorr1D(randoms)
+        autocorrelation_integers = autocorr1D(integers)
+
+        def expfunc(x, a):
+            return np.exp(-x / a)
+
+        x_values = np.array(range(len(autocorrelation_integers)))
+
+        exp_coeff_randoms = scipy.optimize.curve_fit(
+            expfunc, x_values, autocorrelation_randoms
+        )[0][0]
+        exp_coeff_int = scipy.optimize.curve_fit(
+            expfunc, x_values, autocorrelation_integers
+        )[0][0]
+
+        assert (
+            round(autocorrelation_integers[0], 10)
+            and round(autocorrelation_randoms[0], 10) == 1
+        )
+        assert exp_coeff_int > 5 and exp_coeff_randoms < 1

environment-dev.yml

@@ -10,7 +10,7 @@ dependencies:
   - mbuild >=0.16.4
   - numpy
   - pip
-  - python >= 3.10
+  - python >= 3.10,<3.12
   - matplotlib
   - pre-commit
   - pymbar >= 4.0

environment.yml

@@ -10,7 +10,7 @@ dependencies:
   - mbuild >=0.16.4
   - numpy
   - pip
-  - python >= 3.10
+  - python >= 3.10,<3.12
   - matplotlib
   - pymbar >= 4.0
   - rowan