two-step chebyshev propagator, an example and three one-particle tests

examples/test_cheb_mo_1p_1d_twostep.py

+""" One particle in Morse potential - Chebyshev and two-step Chebyshev """
+import numpy as np
+import matplotlib.pyplot as plt
+from rotagaporp.chebyshev import Chebyshev, TwoStepChebyshev
+from rotagaporp.symplectic import Symplectic
+from rotagaporp.systems import Morse
+
+params = {'tstart': 0., 'tend': 500., 'tstep': 0.05, 'nterms': 4, 'tqdm': True}
+syspar0 = {'q0': 3., 'p0': 0.}
+system = Morse(**syspar0)
+
+# find the initial values for the two step propagator
+ini_params = params.copy()
+ini_params['tend'] = params['tstart'] + params['tstep']
+ini_params['tstep'] = 0.001
+ini_params['nterms'] = 5
+ini_prop = Chebyshev(syst=system, **ini_params)
+ini_prop.propagate()
+ini_prop.analyse()
+print('energy drift of initializer', np.max(np.fabs(ini_prop.er)))
+
+assert np.isclose(ini_prop.time[-1], params['tstart'] + params['tstep'])
+assert len(ini_prop.sq) == len(ini_prop.time)
+syspar = {'q0': ini_prop.sq[-1], 'p0': ini_prop.sp[-1]}
+print(syspar)
+system = Morse(**syspar)
+
+# velocity Verlet
+params_verlet = params.copy()
+del params_verlet['nterms']
+vv_prop = Symplectic(syst=system, **params_verlet)
+vv_prop.propagate()
+vv_prop.analyse()
+edrift_vv = vv_prop.er
+
+# one-step propagator
+ch_prop = Chebyshev(syst=system, **params)
+ch_prop.propagate()
+time_ch, q_ch, p_ch = ch_prop.get_trajectory()
+ch_prop.analyse()
+edrift_ch = ch_prop.er
+
+# two-step propagator
+params['q_1'] = syspar0['q0']
+params['p_1'] = syspar0['p0']
+params['parity'] = 0
+# params['signed'] = False
+ts_prop = TwoStepChebyshev(syst=system, **params)
+ts_prop.propagate()
+time_ts, q_ts, p_ts = ts_prop.get_trajectory()
+ts_prop.analyse()
+edrift_ts = ts_prop.er
+
+assert len(time_ts) == len(time_ch)
+assert np.allclose(time_ts, time_ch)
+q_an, p_an = system.solution(times=time_ts)
+
+fig = plt.figure()
+plot1 = fig.add_subplot(311)
+plot1.plot(time_ch, np.maximum.accumulate(edrift_ch), label='Chebyshev')
+plot1.plot(time_ts, np.maximum.accumulate(edrift_ts), label='Two-step Chebyshev')
+plot1.plot(time_ts, np.maximum.accumulate(edrift_vv), label='Velocity Verlet')
+plot1.set_ylabel('maximum energy drift')
+plot1.set_yscale('log')
+plot1.set_xscale('log')
+plot1.legend()
+
+plot2 = fig.add_subplot(312)
+plot2.plot(time_ts, q_ch, label='Chebyshev')
+plot2.plot(time_ts, q_ts, label='Two-step Chebyshev')
+plot2.plot(time_ts, q_an, label='Analytic solution')
+plot2.set_xlabel(r'$t$')
+plot2.set_ylabel(r'$q(t)$')
+plot2.set_xscale('log')
+plot2.legend()
+
+plot3 = fig.add_subplot(313)
+plot3.plot(time_ts, p_ch, label='Chebyshev')
+plot3.plot(time_ts, p_ts, label='Two-step Chebyshev')
+plot3.plot(time_ts, p_an, label='Analytic solution')
+plot3.set_xlabel(r'$t$')
+plot3.set_ylabel(r'$p(t)$')
+plot3.set_xscale('log')
+plot3.legend()
+
+plt.show()

rotagaporp/chebyshev.py

@@ -4,6 +4,7 @@ import numpy as np
 import sympy as sp
 from scipy.special import jn, iv
 from propagators import PolynomialPropagator
+from propagators import TwoStepPolynomialPropagator
 
 
 def cheb_coef(bf, nterms, alpha):
@@ -87,3 +88,35 @@ class Chebyshev(PolynomialPropagator):
         self.set_polynomial(system=self.syst, scale=self.scale,
                             shift=self.shift, recursion='chebyshev',
                             sign=sign)
+
+
+class TwoStepChebyshev(TwoStepPolynomialPropagator):
+    """ Two-step Chebyshev propagator for classical particle dynamics """
+    name = 'two-step chebyshev'
+    def __init__(self, delta=1., nterms=None, signed=False, **kwargs):
+        super().__init__(**kwargs)
+
+        self.delta = delta
+        alpha = self.delta*self.tstep/2.
+        self.nterms = int(np.ceil(alpha)) if nterms is None else nterms
+        self.efficiency = alpha/self.nterms
+        bf = besself(self.fpprec, signed, prec=self.prec)
+        sign = -1 if signed else 1
+        assert self.substep is None, 'substep not implemented'
+        self.coeffs = 2.*cheb_coef(bf, self.nterms, alpha)[self.parity::2]
+        self.set_polynomial(system=self.syst, scale=2./self.delta,
+                            recursion='chebyshev', sign=sign)
+
+    def step_nopbc(self, time, q, p):
+        """ perform a time step with initial coordinate and momentum """
+        # the get_val_float() method has to be adapted to evaluate only
+        # either the odd or the even powers
+        qvecaux, pvecaux = self.iln.get_val_float(q, p)
+        qvec = qvecaux[self.parity::2]
+        pvec = pvecaux[self.parity::2]
+        #
+        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
+        return (qt, pt)

rotagaporp/propagators.py

@@ -148,3 +148,17 @@ class PolynomialPropagator(Propagator):
         import cloudpickle as pkl
         with open(filename, 'wb') as fh:
             pkl.dump(self, fh)
+
+
+class TwoStepPolynomialPropagator(PolynomialPropagator):
+    """ Base class for two-step polynomial classical propagators """
+
+    name = 'two-step polynomial'
+
+    def __init__(self, q_1, p_1, parity=0, **kwargs):
+        assert parity in (0, 1)
+        super().__init__(**kwargs)
+        self.parity = parity
+        self.prefact = 2*parity-1
+        self.q_1 = q_1
+        self.p_1 = p_1

tests/sp_tests.py

@@ -5,7 +5,7 @@ import sympy as sp
 from rotagaporp.liouville import q, p
 from rotagaporp.liouville import LiouvillePowers, LiouvillePowersQP, LiouvillePowersNR
 from rotagaporp.symplectic import Symplectic
-from rotagaporp.chebyshev import Chebyshev
+from rotagaporp.chebyshev import Chebyshev, TwoStepChebyshev
 from rotagaporp.newtonian import Newtonian
 from rotagaporp.newtonian import leja_points_ga, leja_points_jb
 from rotagaporp.newtonian import chebyshev_points
@@ -24,6 +24,8 @@ from rotagaporp.combinatorics import cnk_chebyshev_T3
 from rotagaporp.combinatorics import cnk_chebyshev_U1, cnk_chebyshev_U2
 from rotagaporp.combinatorics import cnk_chebyshev_U3
 
+atol32 = np.finfo(np.float32).resolution
+
 class LiouvillePowersTest(unittest.TestCase):
     """ Test of the one-particle Liouville powers class """
 
@@ -283,7 +285,6 @@ class ChebyshevPropagatorTest(unittest.TestCase):
         (ti, qt, pt) = prop.get_trajectory()
         (qr, pr) = syst.solution(self.syspar['q0'], self.syspar['p0'], ti)
         self.assertTrue(np.allclose(qt, qr, rtol=1e-7))
-        atol32 = np.finfo(np.float32).resolution
         self.assertTrue(np.allclose(pt, pr, atol=atol32))
 
     def chebyshev_propagator_signed_test(self):
@@ -322,6 +323,75 @@ class ChebyshevPropagatorTest(unittest.TestCase):
         self.assertTrue(np.allclose(pt, pr, rtol=1e-7))
 
 
+class TwoStepChebyshevPropagatorTest(unittest.TestCase):
+    """ Tests for the two-step Chebyshev polynomial propagator """
+
+    params = {'tstart': 0.0, 'tend': 5.0, 'tstep': 0.2, 'nterms': 8}
+    syspar = {'q0': 1.1, 'p0': 0.0}
+
+    def chebyshev_propagator_ballistic_and_harmonic_test(self):
+        """ test with ballistic motion and harmonic oscillator """
+        for potential in (Harmonic, Ballistic):
+            sys = potential(**self.syspar)
+            t_1 = self.params['tstart'] - self.params['tstep']
+            q_1, p_1 = sys.solution(self.syspar['q0'], self.syspar['p0'], t_1)
+            prop = TwoStepChebyshev(syst=sys, q_1=q_1, p_1=p_1, **self.params)
+            prop.propagate()
+            prop.analyse()
+            self.assertLess(np.max(np.fabs(prop.er)), 1e-8)
+            (ti, qt, pt) = prop.get_trajectory()
+            (qr, pr) = sys.solution(self.syspar['q0'], self.syspar['p0'], ti)
+            self.assertTrue(np.allclose(qt, qr, rtol=1e-8))
+            self.assertTrue(np.allclose(pt, pr, rtol=1e-8))
+
+    def chebyshev_propagator_morse_test(self):
+        """ test with Morse oscillator """
+        syst = Morse(**self.syspar)
+        init_params = self.params.copy()
+        init_params['tend'] = self.params['tstart'] + self.params['tstep']
+        init_times = [init_params['tstart'], init_params['tend']]
+        qinit, pinit = syst.solution(times=init_times)
+        syst = Morse(q0=qinit[1], p0=pinit[1])
+        params = self.params.copy()
+        params['q_1'] = qinit[0]
+        params['p_1'] = pinit[0]
+        for parity in (0, 1):
+            for signed in (True, False):
+                prop = TwoStepChebyshev(syst=syst, parity=parity,
+                                        signed=signed, **params)
+                prop.propagate()
+                prop.analyse()
+                self.assertLess(np.max(np.fabs(prop.er)), 1e-8)
+                ti, qt, pt = prop.get_trajectory()
+                qr, pr = syst.solution(times=ti)
+                self.assertTrue(np.allclose(qt, qr, rtol=1e-8))
+                self.assertTrue(np.allclose(pt, pr, atol=atol32))
+
+    def chebyshev_propagator_lj_test(self):
+        """ test with a particle in Lennard-Jones potential """
+        syspar = {'p0': -1.0, 'q0': 2}
+        params = {'tstart': 0.0, 'tend': 1.8, 'tstep': 0.02, 'nterms': 6}
+        syst = LennardJones(**syspar)
+        ref_prop = Chebyshev(syst=syst, **params)
+        ref_prop.propagate()
+        time_ref, q_ref, p_ref = ref_prop.get_trajectory()
+        syst = LennardJones(q0=q_ref[1], p0=p_ref[1])
+        params['q_1'] = q_ref[0]
+        params['p_1'] = p_ref[0]
+        params['tstart'] = params['tstart'] + params['tstep']
+        for parity in (0, 1):
+            for signed in (False, True):
+                prop = TwoStepChebyshev(syst=syst, parity=parity,
+                                        signed=signed, **params)
+                prop.propagate()
+                prop.analyse()
+                self.assertLess(np.max(np.fabs(prop.er)), 1e-3)
+                time, qt, pt = prop.get_trajectory()
+                assert len(time) == len(time_ref[1:])
+                self.assertTrue(np.allclose(qt, q_ref[1:], rtol=1e-3))
+                self.assertTrue(np.allclose(pt, p_ref[1:], rtol=1e-2))
+
+
 class PropagatorRestartTest(unittest.TestCase):
     """ Tests with restarting polynomial propagators """
 
@@ -500,7 +570,6 @@ class NewtonianPropagatorTest(unittest.TestCase):
         (ti, qt, pt) = prop.get_trajectory()
         qr, pr = syst.solution(self.syspar['q0'], self.syspar['p0'], ti)
         self.assertTrue(np.allclose(qt, qr, rtol=1e-7))
-        atol32 = np.finfo(np.float32).resolution
         self.assertTrue(np.allclose(pt, pr, atol=atol32))
 
 
@@ -547,7 +616,6 @@ class SymplecticPropagatorTest(unittest.TestCase):
         (ti, qt, pt) = prop.get_trajectory()
         qr, pr = syst.solution(self.syspar['q0'], self.syspar['p0'], ti)
         self.assertTrue(np.allclose(qt, qr, rtol=1e-2))
-        atol32 = np.finfo(np.float32).resolution
         self.assertTrue(np.allclose(pt, pr, atol=atol32))
 
 
@@ -661,7 +729,7 @@ class ChebyshevAssessmentTest(unittest.TestCase):
         prop.analyse()
         t, q, p = prop.get_trajectory()
         qr, pr = syst.solution(times=t)
-        self.assertTrue(np.allclose(q, qr)) 
+        self.assertTrue(np.allclose(q, qr))
         self.assertLess(np.max(np.fabs(prop.er)), 1e-6)
 
     def chebyshev_symplectic_2p_lj_test(self):