added stepping with pbc and 9 tests for many-particle classes

rotagaporp/propagators.py

@@ -119,16 +119,16 @@ class PolynomialPropagator(Propagator):
             params = {'fpprec': self.fpprec, **kwargs, **self.liou_params}
             self.iln = self.liouville(self.nterms, **params)
 
-    def step_nopbc(self, time, q, p):
+    def step_nopbc(self, time, qval, pval):
         """ perform a time step with initial coordinate and momentum """
-        qvec, pvec = self.iln.get_val_float(q, p)
+        qvec, pvec = self.iln.get_val_float(qval, pval)
         qt = np.dot(self.coeffs, qvec)
         pt = np.dot(self.coeffs, pvec)
         return (qt, pt)
 
-    def step_pbc(self, time, q, p):
+    def step_pbc(self, time, qval, pval):
         """ perform a time step with periodic boundary conditions """
-        qvec, pvec = self.iln.get_val_float(q, p)
+        qvec, pvec = self.iln.get_val_float(qval, pval)
         qt = np.dot(self.coeffs, qvec)
         pt = np.dot(self.coeffs, pvec)
         return (self.syst.wrap_q(qt), pt)
@@ -162,6 +162,9 @@ class TwoStepPolynomialPropagator(PolynomialPropagator):
     iln = None
 
     def __init__(self, q_1, p_1, parity=0, **kwargs):
+        """ parity: even: 0, odd: 1
+            q_1: the position(s) at time=tstart-tstep (float or np.array)
+            p_1: the momentum(a) at time=tstart-tstep (float or np.array) """
         assert parity in (0, 1)
         super().__init__(**kwargs)
         self.parity = parity
@@ -169,11 +172,20 @@ class TwoStepPolynomialPropagator(PolynomialPropagator):
         self.q_1 = q_1
         self.p_1 = p_1
 
-    def step_nopbc(self, time, q, p):
+    def step_nopbc(self, time, qval, pval):
         """ perform a time step with initial coordinate and momentum """
-        qvec, pvec = self.iln.get_val_float(q, p)
+        qvec, pvec = self.iln.get_val_float(qval, pval)
         qt = np.dot(self.coeffs, qvec) + self.prefact*self.q_1
         pt = np.dot(self.coeffs, pvec) + self.prefact*self.p_1
-        self.q_1 = q
-        self.p_1 = p
+        self.q_1 = qval
+        self.p_1 = pval
         return (qt, pt)
+
+    def step_pbc(self, time, qval, pval):
+        """ perform a time step with periodic boundary conditions """
+        qvec, pvec = self.iln.get_val_float(qval, pval)
+        qt = np.dot(self.coeffs, qvec) + self.prefact*self.q_1
+        pt = np.dot(self.coeffs, pvec) + self.prefact*self.p_1
+        self.q_1 = qval
+        self.p_1 = pval
+        return (self.syst.wrap_q(qt), pt)

tests/ts_tests.py

+""" Compare one-step Chebyshev and two-step Chebyshev for two particles """
+import unittest
+import itertools
+import numpy as np
+import sympy as sp
+from rotagaporp.systems import MPPS1D, MPPS3D
+from rotagaporp.chebyshev import Chebyshev, TwoStepChebyshev
+from rotagaporp.mpliouville import MPLP
+from rotagaporp.mpliouville import MPLPQP, MPLPVQ, MPLPDR, MPLPDV
+
+# NOTABENE: MPLPAQ and MPLPNR currently cannot work with one-parity Chebyshev
+
+class TestMPLP(unittest.TestCase):
+    """ Compare odd and even polynomials generated with and without parity """
+
+    def test_chebyshev_parity_2p_1d(self):
+        """ compare Chebyshev polynomials of same parity with and without
+            parity parameter for different sign, valid for any H and f """
+
+        # nterms = 7 # 1690 seconds
+        nterms = 5 # 16 seconds
+        Q = sp.Array(sp.symbols('q1 q2', real=True))
+        P = sp.Array(sp.symbols('p1 p2', real=True))
+        pot = sp.Function('V', real=True)(*Q)
+        kin = sp.Function('T', real=True)(*P)
+        ham = (pot, kin)
+        for sign in (-1, 1):
+            ilnalls = MPLP(nterms, ham, Q, P, Q, recursion='chebyshev',
+                           sign=sign, parity=None)
+            ilneven = MPLP(nterms, ham, Q, P, Q, recursion='chebyshev',
+                           sign=sign, parity=0)
+            for ind in range(nterms//2):
+                alls = [lf.factor() for lf in ilnalls.Lf[2*ind]]
+                even = [lf.factor() for lf in ilneven.Lf[ind]]
+                self.assertListEqual(alls, even)
+
+            ilnodds = MPLP(nterms+1, ham, Q, P, Q, recursion='chebyshev',
+                           sign=sign, parity=1)
+            for ind in range(nterms//2):
+                alls = [lf.factor() for lf in ilnalls.Lf[2*ind+1]]
+                odds = [lf.factor() for lf in ilnodds.Lf[ind]]
+                self.assertListEqual(alls, odds)
+
+
+class TestTwoStepChebyshevPropagatorMP(unittest.TestCase):
+    """ Test the two-step Chebyshev propagator for many-particle systems """
+
+    def compare(self, nterms, params, syst, rch, pch, atol):
+        """ compare trajectories with one- and two-step Chebyshev """
+        for parity, signed in itertools.product((0, 1), (False, True)):
+            params['nterms'] = nterms + parity
+            params['parity'] = parity
+            params['signed'] = signed
+            ts = TwoStepChebyshev(syst=syst, **params)
+            ts.propagate()
+            ts.analyse()
+            (qts, pts) = ts.get_trajectory()[1:]
+            rts = np.array([syst.get_distances(q)[0] for q in qts])
+            self.assertTrue(np.allclose(rch, rts, atol=atol))
+            self.assertTrue(np.allclose(pch, pts, atol=atol))
+            self.assertTrue(np.allclose(ts.er, 0, atol=atol))
+
+    def test_00_2p_1d(self):
+        """ Two non-interacting particles in one dimension """
+
+        syspar = {'symbols': ['l', 'r'], 'q0': [[1.02], [2.98]],
+                  'p0': [[1.], [-1.]], 'masses': [1., 1.],
+                  'numeric': {'tool': 'subs'}}
+        ffparams = [{'class': 'FFZero', 'pair': ((0,), (1,))}]
+        syst = MPPS1D(ffparams, **syspar)
+        params = {'tstart': 0.01, 'tend': 8.0, 'tstep': 0.02, 'nterms': 5,
+                  'liouville': MPLPQP, 'signed': False, 'pbc': False}
+
+        ch = Chebyshev(syst=syst, **params)
+        ch.propagate()
+        (qch, pch) = ch.get_trajectory()[1:]
+        rch = np.array([syst.get_distances(q)[0] for q in qch])
+
+        params['q_1'] = np.array([1., 3.])
+        params['p_1'] = np.array([1., -1.])
+        nterms = 3
+        self.compare(nterms, params, syst, rch, pch, 1e-4)
+
+    def test_00_2p_1d_pbc(self):
+        """ Two non-interacting particles in one dimension """
+
+        syspar = {'symbols': ['l', 'r'], 'q0': [[1.05], [2.95]],
+                  'p0': [[1.], [-1.]], 'masses': [1., 1.], 'pbc': True,
+                  'cell': [5.], 'numeric': {'tool': 'subs'}}
+        ffparams = [{'class': 'FFZero', 'pair': ((0,), (1,))}]
+        syst = MPPS1D(ffparams, **syspar)
+        params = {'tstart': 0.01, 'tend': 8.0, 'tstep': 0.05, 'nterms': 5,
+                  'liouville': MPLPDV, 'signed': False, 'pbc': True}
+
+        ch = Chebyshev(syst=syst, **params)
+        ch.propagate()
+        (qch, pch) = ch.get_trajectory()[1:]
+        rch = np.array([syst.get_distances(q)[0] for q in qch])
+
+        params['q_1'] = np.array([1., 3.])
+        params['p_1'] = np.array([1., -1.])
+        nterms = 5
+        self.compare(nterms, params, syst, rch, pch, 1e-5)
+
+    def test_00_2p_3d(self):
+        """ Two non-interacting particles in three periodic dimensions """
+
+        ffpars = [{'class': 'FFZero', 'pair': ((0, 1, 2), (3, 4, 5))}]
+        syspar = {'symbols': ['0x:z', '1x:z'], 'masses': [1., 1.],
+                  'q0': [[0.04, 0., 0.], [2.96, 0., 0.]],
+                  'p0': [[1., 0., 0.], [-1., 0., 0.]], 'pbc': False,
+                  'numeric': {'tool': 'subs'}}
+        syst = MPPS3D(ffpars, **syspar)
+        params = {'tstart': 0.0, 'tend': 8.0, 'tstep': 0.04, 'nterms': 4,
+                  'liouville': MPLPVQ, 'pbc': False}
+
+        ch = Chebyshev(syst=syst, **params)
+        ch.propagate()
+        (qch, pch) = ch.get_trajectory()[1:]
+        rch = np.array([syst.get_distances(q)[0] for q in qch])
+
+        params['q_1'] = np.array([0., 0., 0., 3., 0., 0.])
+        params['p_1'] = np.array([1., 0., 0., -1., 0., 0.])
+        nterms = 3
+        self.compare(nterms, params, syst, rch, pch, 1e-4)
+
+    def test_00_2p_3d_pbc(self):
+        """ Two non-interacting particles in three periodic dimensions """
+
+        ffpars = [{'class': 'FFZero', 'pair': ((0, 1, 2), (3, 4, 5))}]
+        syspar = {'symbols': ['0x:z', '1x:z'], 'masses': [1., 1.],
+                  'q0': [[0.04, 0., 0.], [2.96, 0., 0.]],
+                  'p0': [[1., 0., 0.], [-1., 0., 0.]],
+                  'pbc': (True, True, True), 'cell': [5, 2, 2],
+                  'numeric': {'tool': 'lambdify', 'modules': 'math'}}
+        syst = MPPS3D(ffpars, **syspar)
+        params = {'tstart': 0.0, 'tend': 8.0, 'tstep': 0.04, 'nterms': 4,
+                  'liouville': MPLPDV, 'pbc': True}
+
+        ch = Chebyshev(syst=syst, **params)
+        ch.propagate()
+        (qch, pch) = ch.get_trajectory()[1:]
+        rch = np.array([syst.get_distances(q)[0] for q in qch])
+
+        params['q_1'] = np.array([0., 0., 0., 3., 0., 0.])
+        params['p_1'] = np.array([1., 0., 0., -1., 0., 0.])
+        nterms = 3
+        self.compare(nterms, params, syst, rch, pch, 1e-4)
+
+    def test_lj_2p_1d(self):
+        """ Two particles with LJ interaction in one dimension """
+
+        positions = [[1.], [3.]]
+        momenta = [[0.], [-1.]]
+        params = {'tstart': 0.0, 'tend': 2.0, 'tstep': 0.01, 'nterms': 4,
+                  'liouville': MPLPDR, 'pbc': False}
+        ffpars = [{'class': 'FFLennardJones', 'pair': ((0,), (1,))}]
+        syspar = {'symbols': ['l', 'r'], 'q0': positions, 'p0': momenta,
+                  'masses': [1., 1.]}
+        syst = MPPS1D(ffpars, **syspar)
+        ini_params = params.copy()
+        ini_params['tend'] = params['tstart'] + params['tstep']
+        ini_params['tstep'] = 0.001
+        ch_ini = Chebyshev(syst=syst, **ini_params)
+        ch_ini.propagate()
+        t_ini, qch_ini, pch_ini = ch_ini.get_trajectory()
+        assert np.isclose(t_ini[-1], params['tstart']+params['tstep'])
+
+        syspar['q0'] = qch_ini[-1].reshape(-1, 1)
+        syspar['p0'] = pch_ini[-1].reshape(-1, 1)
+        params['tstart'] = params['tstart'] + params['tstep']
+        syst = MPPS1D(ffpars, **syspar)
+        ch = Chebyshev(syst=syst, **params)
+        ch.propagate()
+        qch, pch = ch.get_trajectory()[1:]
+        rch = np.array([syst.get_distances(q)[0] for q in qch])
+
+        params['q_1'] = np.array(positions).flatten()
+        params['p_1'] = np.array(momenta).flatten()
+        nterms = 3
+        self.compare(nterms, params, syst, rch, pch, 1e-1)
+
+    def test_lj_2p_1d_pbc(self):
+        """ Two particles with LJ interaction in one periodic dimension """
+
+        positions = [[1.], [3.]]
+        momenta = [[0.], [-1.]]
+        params = {'tstart': 0.0, 'tend': 2.0, 'tstep': 0.01, 'nterms': 4,
+                  'liouville': MPLPDV, 'pbc': True}
+        ffpars = [{'class': 'FFLennardJones', 'pair': ((0,), (1,))}]
+        syspar = {'symbols': ['l', 'r'], 'q0': positions, 'p0': momenta,
+                  'masses': [1., 1.], 'pbc': [True], 'cell': [5]}
+        syst = MPPS1D(ffpars, **syspar)
+        ini_params = params.copy()
+        ini_params['tend'] = params['tstart'] + params['tstep']
+        ini_params['tstep'] = 0.001
+        ch_ini = Chebyshev(syst=syst, **ini_params)
+        ch_ini.propagate()
+        t_ini, qch_ini, pch_ini = ch_ini.get_trajectory()
+        assert np.isclose(t_ini[-1], params['tstart']+params['tstep'])
+
+        syspar['q0'] = qch_ini[-1].reshape(-1, 1)
+        syspar['p0'] = pch_ini[-1].reshape(-1, 1)
+        params['tstart'] = params['tstart'] + params['tstep']
+        syst = MPPS1D(ffpars, **syspar)
+        ch = Chebyshev(syst=syst, **params)
+        ch.propagate()
+        qch, pch = ch.get_trajectory()[1:]
+        rch = np.array([syst.get_distances(q)[0] for q in qch])
+
+        params['q_1'] = np.array(positions).flatten()
+        params['p_1'] = np.array(momenta).flatten()
+        nterms = 3
+        self.compare(nterms, params, syst, rch, pch, 1e-1)
+
+    def test_lj_2p_3d(self):
+        """ Two particles with LJ interaction in three dimensions """
+
+        numeric = {'tool': 'theano'}
+        positions = np.array([[1., 2., 3.], [3., 2., 3.]])
+        momenta = np.array([[0., 0., 0.], [-1., 0., 0.]])
+        syspar = {'symbols': ['0x:z', '1x:z'], 'q0': positions, 'p0': momenta,
+                  'masses': [1., 1.], 'pbc': False, 'numeric': numeric}
+        ffpars = [{'class': 'FFLennardJones', 'pair': ((0, 1, 2), (3, 4, 5))}]
+        syst = MPPS3D(ffpars, **syspar)
+        params = {'tstart': 0.0, 'tend': 2.0, 'tstep': 0.01, 'nterms': 4,
+                  'liouville': MPLPDR, 'pbc': False}
+
+        ini_params = params.copy()
+        ini_params['tend'] = params['tstart'] + params['tstep']
+        ini_params['tstep'] = 0.001
+        ch_ini = Chebyshev(syst=syst, **ini_params)
+        ch_ini.propagate()
+        t_ini, qch_ini, pch_ini = ch_ini.get_trajectory_3d()
+        assert np.isclose(t_ini[-1], ini_params['tend'])
+        syspar['q0'] = qch_ini[-1]
+        syspar['p0'] = pch_ini[-1]
+        syst = MPPS3D(ffpars, **syspar)
+
+        params['tstart'] = params['tstart'] + params['tstep']
+        ch = Chebyshev(syst=syst, **params)
+        ch.propagate()
+        (qch, pch) = ch.get_trajectory()[1:]
+        rch = np.array([syst.get_distances(q)[0] for q in qch])
+
+        params['q_1'] = positions.flatten()
+        params['p_1'] = momenta.flatten()
+        nterms = 3
+        self.compare(nterms, params, syst, rch, pch, 1e-1)
+
+    def test_lj_2p_3d_pbc(self):
+        """ Two particles with LJ interaction in three dimensions """
+
+        numeric = {'tool': 'lambdify', 'modules': 'math'}
+        positions = np.array([[1., 2., 3.], [3., 2., 3.]])
+        momenta = np.array([[0., 0., 0.], [-1., 0., 0.]])
+        syspar = {'symbols': ['0x:z', '1x:z'], 'q0': positions, 'p0': momenta,
+                  'masses': [1., 1.], 'pbc': (True, True, True),
+                  'cell': [10, 10, 10], 'numeric': numeric}
+        ffpars = [{'class': 'FFLennardJones', 'pair': ((0, 1, 2), (3, 4, 5))}]
+        syst = MPPS3D(ffpars, **syspar)
+        params = {'tstart': 0.0, 'tend': 2.0, 'tstep': 0.01, 'nterms': 4,
+                  'liouville': MPLPDV, 'pbc': True}
+
+        ini_params = params.copy()
+        ini_params['tend'] = params['tstart'] + params['tstep']
+        ini_params['tstep'] = 0.001
+        ch_ini = Chebyshev(syst=syst, **ini_params)
+        ch_ini.propagate()
+        t_ini, qch_ini, pch_ini = ch_ini.get_trajectory_3d()
+        assert np.isclose(t_ini[-1], ini_params['tend'])
+        syspar['q0'] = qch_ini[-1]
+        syspar['p0'] = pch_ini[-1]
+        syst = MPPS3D(ffpars, **syspar)
+
+        params['tstart'] = params['tstart'] + params['tstep']
+        ch = Chebyshev(syst=syst, **params)
+        ch.propagate()
+        (qch, pch) = ch.get_trajectory()[1:]
+        rch = np.array([syst.get_distances(q)[0] for q in qch])
+
+        params['q_1'] = positions.flatten()
+        params['p_1'] = momenta.flatten()
+        nterms = 3
+        self.compare(nterms, params, syst, rch, pch, 1e-1)