Update waterdynamics.py

AUTHORS.md

@@ -54,3 +54,4 @@ The rules for this file:
 
 **2024**
 - Fiona Naughton <@fiona-naughton>
+- Rodolphe Pollet <@rodpollet>

CHANGELOG.md

@@ -20,8 +20,10 @@ The rules for this file:
 
 ### Authors
 - fiona-naughton
+- rodpollet
 
 ### Added
+- first-order Legendre polynomial for water orientation relaxation analysis
 
 ### Fixed
 

waterdynamics/tests/test_waterdynamics.py

@@ -43,15 +43,29 @@ def universe():
 
 def test_WaterOrientationalRelaxation(universe):
     wor = waterdynamics.WaterOrientationalRelaxation(
-        universe, SELECTION1, 0, 5, 2)
+        universe, SELECTION1, 0, 5, 2, order=2)
     wor.run()
     assert_almost_equal(wor.timeseries[1][2], 0.35887,
                         decimal=5)
 
 
 def test_WaterOrientationalRelaxation_zeroMolecules(universe):
     wor_zero = waterdynamics.WaterOrientationalRelaxation(
-        universe, SELECTION2, 0, 5, 2)
+        universe, SELECTION2, 0, 5, 2, order=2)
+    wor_zero.run()
+    assert_almost_equal(wor_zero.timeseries[1], (0.0, 0.0, 0.0))
+
+def test_WaterOrientationalRelaxation(universe):
+    wor = waterdynamics.WaterOrientationalRelaxation(
+        universe, SELECTION1, 0, 5, 2, order=1)
+    wor.run()
+    assert_almost_equal(wor.timeseries[1][2], 0.71486,
+                        decimal=5)
+
+
+def test_WaterOrientationalRelaxation_zeroMolecules(universe):
+    wor_zero = waterdynamics.WaterOrientationalRelaxation(
+        universe, SELECTION2, 0, 5, 2, order=1)
     wor_zero.run()
     assert_almost_equal(wor_zero.timeseries[1], (0.0, 0.0, 0.0))
 
@@ -273,4 +287,4 @@ def test_SurvivalProbability_stepEqualDtMax(universe):
         sp = waterdynamics.SurvivalProbability(universe, "")
         sp.run(tau_max=4, step=5, stop=10, verbose=True)
         # all frames from 0, with 9 inclusive
-        assert_equal(select_atoms_mock.call_count, 10)
+        assert_equal(select_atoms_mock.call_count, 10)

waterdynamics/waterdynamics.py

@@ -62,7 +62,8 @@ class WaterOrientationalRelaxation(object):
         C_{\hat u}(\tau)=\langle \mathit{P}_2[\mathbf{\hat{u}}(t_0)\cdot\mathbf{\hat{u}}(t_0+\tau)]\rangle
 
     where :math:`P_2=(3x^2-1)/2` is the second-order Legendre polynomial and :math:`\hat{u}` is
-    a unit vector along HH, OH or dipole vector.
+    a unit vector along HH, OH or dipole vector. Another option is to select the first-order Legendre
+    polynomial, :math:`P_1=x`.
 
 
     Parameters
@@ -77,15 +78,21 @@ class WaterOrientationalRelaxation(object):
       frame where analysis ends
     dtmax : int
       Maximum dt size, `dtmax` < `tf` or it will crash.
+    order : 1 or 2 (default)
+      first- or second-order Legendre polynomial
     """
 
-    def __init__(self, universe, select, t0, tf, dtmax, nproc=1):
+    def __init__(self, universe, select, t0, tf, dtmax, nproc=1, order=2):
         self.universe = universe
         self.selection = select
         self.t0 = t0
         self.tf = tf
         self.dtmax = dtmax
         self.nproc = nproc
+        if order != 1 and order != 2:
+            raise ValueError(f"order = {order} but only first- or second-order Legendre polynomial is allowed.")
+        else:
+            self.order = order
         self.timeseries = None
 
     def _repeatedIndex(self, selection, dt, totalFrames):
@@ -161,9 +168,14 @@ def _getOneDeltaPoint(self, universe, repInd, i, t0, dt):
                               dipVectorp[1] / normdipVectorp,
                               dipVectorp[2] / normdipVectorp]
 
-            valOH += self.lg2(np.dot(unitOHVector0, unitOHVectorp))
-            valHH += self.lg2(np.dot(unitHHVector0, unitHHVectorp))
-            valdip += self.lg2(np.dot(unitdipVector0, unitdipVectorp))
+            if self.order == 1:
+                valOH += self.lg1(np.dot(unitOHVector0, unitOHVectorp))
+                valHH += self.lg1(np.dot(unitHHVector0, unitHHVectorp))
+                valdip += self.lg1(np.dot(unitdipVector0, unitdipVectorp))
+            else:
+                valOH += self.lg2(np.dot(unitOHVector0, unitOHVectorp))
+                valHH += self.lg2(np.dot(unitHHVector0, unitHHVectorp))
+                valdip += self.lg2(np.dot(unitdipVector0, unitdipVectorp))
             n += 1
         return  (valOH/n, valHH/n, valdip/n) if n > 0 else (0, 0, 0)
 
@@ -213,6 +225,11 @@ def _selection_serial(self, universe, selection_str):
             selection.append(universe.select_atoms(selection_str))
         return selection
 
+    @staticmethod
+    def lg1(x):
+        """First Legendre polynomial"""
+        return x
+
     @staticmethod
     def lg2(x):
         """Second Legendre polynomial"""