added theano and lambdify, added expression graph to examples

examples/graph_cheb_mo_1p.py

+""" Expression graph of Chebyshev propagator """
+import sympy as sp
+from rotagaporp.chebyshev import Chebyshev
+from rotagaporp.systems import Morse
+
+syspar = {'q0': 3.0, 'p0': 0.0}
+system = Morse(**syspar)
+params = {'tstart': 0.0, 'tend': 100.0, 'tstep': 0.5, 'tqdm': True}
+cheb = Chebyshev(nterms=4, syst=system, **params)
+print(sp.dotprint(cheb.iln.ilnqcart[-1]))
+# print(sum(i.count_ops() for i in cheb.iln.ilnqcart))
+# print(sum(i.count_ops() for i in cheb.iln.ilnpcart))

examples/graph_symp_mo_1p.py

+""" Expression graph of high-order symplectic integrator """
+import sympy as sp
+from rotagaporp.symplectic import HighOrderSymplectic
+from rotagaporp.systems import Morse
+
+syspar = {'q0': 3.0, 'p0': 0.0}
+system = Morse(**syspar)
+params = {'tstart': 0.0, 'tend': 100.0, 'tstep': 0.5, 'tqdm': True}
+symp = HighOrderSymplectic(order=4, syst=system, **params)
+print(sp.dotprint(symp.qt))
+# print(symp.qt.count_ops())
+# print(symp.pt.count_ops())

examples/test_symp_mo_1p_1d.py

@@ -6,14 +6,22 @@ from rotagaporp.systems import Morse
 
 syspar = {'q0': 3.0, 'p0': 0.0}
 system = Morse(**syspar)
+numeric = {'tool': 'lambdify', 'modules': 'math'}
+# numeric = {'tool': 'theano'}
+# numeric = None
 
-params = {'tstart': 0.0, 'tend': 10.0, 'tstep': 0.1, 'tqdm': True}
-prop = HighOrderSymplectic(order=4, syst=system, **params)
-# prop = Symplectic(syst=system, **params)
+params = {'tstart': 0.0, 'tend': 100.0, 'tstep': 0.2, 'tqdm': True}
+prop = HighOrderSymplectic(order=4, syst=system, numeric=numeric, **params)
 prop.propagate()
 prop.analyse()
 (ti, qt, pt) = prop.get_trajectory()
-print('energy drift of symplectic propagator', np.max(np.fabs(prop.er)))
+print('energy drift of ho symplectic propagator', np.max(np.fabs(prop.er)))
+
+prop_vv = Symplectic(syst=system, **params)
+prop_vv.propagate()
+prop_vv.analyse()
+(ti_vv, qt_vv, pt_vv) = prop.get_trajectory()
+print('energy drift of velocity verlet propagator', np.max(np.fabs(prop.er)))
 
 q_an, p_an = system.solution(times=ti)
 en0 = system.energy(system.q0, system.p0)
@@ -23,21 +31,25 @@ print('energy drift of analytic solution', np.max(np.fabs(err_an)))
 
 fig = plt.figure()
 plot = fig.add_subplot(311)
-# plt.plot(ti, qt, label='q symplectic')
-# plt.plot(ti, q_an, label='q analytic')
-plt.plot(ti, abs(qt-q_an), label='q diff')
+plt.plot(ti, qt, label='q ho symplectic')
+plt.plot(ti_vv, qt_vv, label='q velocity verlet')
+plt.plot(ti, q_an, label='q analytic')
+plt.plot(ti, abs(qt-q_an), label='q diff ho symplectic')
+plt.plot(ti, abs(qt_vv-q_an), label='q diff velocity verlet')
 plot.set_ylabel(r'$q(t)$')
 plt.legend()
 
 plot = fig.add_subplot(312)
-# plt.plot(ti, pt, label='p symplectic')
-# plt.plot(ti, p_an, label='p analytic')
+plt.plot(ti, pt, label='p ho symplectic')
+plt.plot(ti_vv, pt_vv, label='p velocity Verlet')
+plt.plot(ti, p_an, label='p analytic')
 plt.plot(ti, abs(pt-p_an), label='p diff')
 plot.set_ylabel(r'$p(t)$')
 plt.legend()
 
 plot = fig.add_subplot(313)
-plt.plot(ti, np.maximum.accumulate(abs(prop.er)), label='energy drift')
+plt.plot(ti, np.maximum.accumulate(abs(prop.er)), label='energy drift ho symplectic')
+plt.plot(ti, np.maximum.accumulate(abs(prop_vv.er)), label='energy drift velocity Verlet')
 plot.set_ylabel(r'$\Delta E$')
 plot.set_xlabel('time')
 plt.legend()

rotagaporp/symplectic.py

@@ -36,14 +36,43 @@ class Symplectic(Propagator):
 
 class HighOrderSymplectic(Propagator):
     """ High-order symplectic propagator for classical particle dynamics """
-    name = 'high order symplectic'
+    name = 'high-order symplectic'
+    numeric = {'tool': 'subs', 'modules': None}
+
+    def __init__(self, order=None, numeric=None, fpprec='fixed', prec=None,
+                 **kwargs):
+
+        self.fpprec = fpprec
+        self.prec = prec
+        from generic import validate_numeric
+        if numeric is not None:
+            self.numeric = numeric
+        assert self.numeric['tool'] in ['subs', 'theano', 'lambdify']
+        validate_numeric(self.fpprec, self.numeric)
+
+        if self.numeric['tool'] == 'theano':
+            self.step_nopbc = self.step_theano
+        elif self.numeric['tool'] == 'lambdify':
+            self.step_nopbc = self.step_lambdify
+        else:
+            self.step_nopbc = self.step_subs
 
-    def __init__(self, order=None, **kwargs):
         super().__init__(**kwargs)
+
         assert order > 0, 'order must be positive'
         assert order%2 == 0, 'order must be even integer'
         self.order = order
         self.set_expression_1()
+        if self.numeric['tool'] == 'theano':
+            from sympy.printing.theanocode import theano_function as tf
+            self.z = (self.syst.q, self.syst.p)
+            self.qttf = tf(self.z, [self.qt], on_unused_input='ignore')
+            self.pttf = tf(self.z, [self.pt], on_unused_input='ignore')
+        if self.numeric['tool'] == 'lambdify':
+            from generic import lambdify_long as lambdify
+            self.z = (self.syst.q, self.syst.p)
+            self.qttf = lambdify(self.z, self.qt, self.numeric['modules'])
+            self.pttf = lambdify(self.z, self.pt, self.numeric['modules'])
 
     def set_expression_1(self):
         """ First version, H. Yoshida, Phys. Lett. A 150, 262 (1990) """
@@ -57,7 +86,7 @@ class HighOrderSymplectic(Propagator):
             self.pt -= self.tstep*ci*A.subs(self.syst.q, self.qt)
 
     def get_coefficients(self):
-        """ Decomposition coefficients prod_i exp(ci*tau*A)*exp(di*tau*B)
+        """ Decomposition coefficients in prod_i exp(ci*tau*A)*exp(di*tau*B)
             H. Yoshida, Phys. Lett. A 150, 262 (1990) """
         if self.order > 2:
             A = sp.Symbol('A', commutative=False)
@@ -94,8 +123,8 @@ class HighOrderSymplectic(Propagator):
         self.pt = self.pt.xreplace({tau: self.tstep})
 
     def get_decomposition(self, A, B, tau):
-        """ Symmetric decomposition of the propagator exp(tau*(A+B))
-            H. Yoshida, Phys. Lett. A 150, 262 (1990) """
+        """ Symmetric decomposition of the propagator exp(tau*(A+B)) where
+            [A, B] != 0, source H. Yoshida, Phys. Lett. A 150, 262 (1990) """
         S = sp.exp(tau*A/2)*sp.exp(tau*B)*sp.exp(tau*A/2)
         for n in range(1, self.order//2):
             z0 = -2**(1/(2*n+1))/(2-2**(1/(2*n+1)))
@@ -103,7 +132,17 @@ class HighOrderSymplectic(Propagator):
             S = S.subs(tau, z1*tau)*S.subs(tau, z0*tau)*S.subs(tau, z1*tau)
         return S
 
-    def step_nopbc(self, time, q0, p0):
+    def step_subs(self, time, q0, p0):
         """ Time stepping """
         repl = {self.syst.q: q0, self.syst.p: p0}
         return self.qt.xreplace(repl), self.pt.xreplace(repl)
+
+    def step_theano(self, time, q0, p0):
+        """ Time stepping """
+        repl = (q0, p0)
+        return self.qttf(*repl), self.pttf(*repl)
+
+    def step_lambdify(self, time, q0, p0):
+        """ Time stepping """
+        repl = (q0, p0)
+        return self.qttf(repl), self.pttf(repl)