refactoring SymmetricChebyshev -> SymmetricPolynomialPropagator

rotagaporp/chebyshev.py

@@ -98,8 +98,6 @@ class SymmetricChebyshev(SymmetricPolynomialPropagator):
     maxiter = 4
 
     def __init__(self, signed=False, **kwargs):
-        """ """
-
         super().__init__(**kwargs)
 
         bf = besself(self.fpprec, signed, prec=self.prec)
@@ -108,31 +106,6 @@ class SymmetricChebyshev(SymmetricPolynomialPropagator):
         assert self.substep is None, 'substep not implemented'
         self.coeffs = cheb_coef(bf, self.nterms, self.alpha)
         self.reverse_coeffs = cheb_coef(bf, self.nterms, -self.alpha)
-        if self.fpprec == 'multi':
-            self.coeffs = sp.Matrix(self.coeffs)
-            self.reverse_coeffs = sp.Matrix(self.reverse_coeffs)
         self.set_polynomial(system=self.syst, scale=2./self.delta, sign=sign,
                             recursion='chebyshev')
-
-        # the rest of the function has to be moved to the superclass
-        gq = self.syst.q + self.reverse_coeffs.dot(self.iln.qobj.Lf)
-        gp = self.syst.p + self.reverse_coeffs.dot(self.iln.pobj.Lf)
-        self.z = sp.Matrix([self.syst.q, self.syst.p])
-        jacobian = sp.Matrix([gq, gp]).jacobian(self.z)
-        if self.numeric['tool'] == 'lambdify':
-            modules = self.numeric['modules']
-            if self.numeric['modules'] == 'mpmath':
-                gvar = sp.Matrix([gq, gp])
-            else:
-                gvar = sp.Array([gq, gp])
-            self.jacobian = sp.lambdify(self.z, jacobian, modules)
-            self.gvar = sp.lambdify(self.z, gvar, modules)
-        elif self.numeric['tool'] == 'theano':
-            from sympy.printing.theanocode import theano_function as tf
-            gvar = sp.Matrix([gq, gp])
-            self.jacobian = tf(self.z, jacobian, on_unused_input='ignore')
-            self.gvar = tf(self.z, gvar, on_unused_input='ignore')
-            self.jshape = jacobian.shape
-        else:
-            self.jacobian = jacobian
-            self.gvar = sp.Matrix([gq, gp])
+        self.set_gvar_jacobian()

rotagaporp/propagators.py

@@ -230,6 +230,35 @@ class SymmetricPolynomialPropagator(PolynomialPropagator):
         if self.fpprec == 'fixed':
             self.prec = np.finfo(float).precision
 
+    def set_gvar_jacobian(self):
+        """ initialize the G_var vector function and the Jacobian matrix """
+
+        if self.fpprec == 'multi':
+            self.coeffs = sp.Matrix(self.coeffs)
+            self.reverse_coeffs = sp.Matrix(self.reverse_coeffs)
+
+        gq = self.syst.q + self.reverse_coeffs.dot(self.iln.qobj.Lf)
+        gp = self.syst.p + self.reverse_coeffs.dot(self.iln.pobj.Lf)
+        self.z = sp.Matrix((self.syst.q, self.syst.p))
+        jacobian = sp.Matrix([gq, gp]).jacobian(self.z)
+        if self.numeric['tool'] == 'lambdify':
+            modules = self.numeric['modules']
+            if self.numeric['modules'] == 'mpmath':
+                gvar = sp.Matrix([gq, gp])
+            else:
+                gvar = sp.Array([gq, gp])
+            self.jacobian = sp.lambdify(self.z, jacobian, modules)
+            self.gvar = sp.lambdify(self.z, gvar, modules)
+        elif self.numeric['tool'] == 'theano':
+            from sympy.printing.theanocode import theano_function as tf
+            gvar = sp.Matrix([gq, gp])
+            self.jacobian = tf(self.z, jacobian, on_unused_input='ignore')
+            self.gvar = tf(self.z, gvar, on_unused_input='ignore')
+            self.jshape = jacobian.shape
+        else:
+            self.jacobian = jacobian
+            self.gvar = sp.Matrix([gq, gp])
+
     def step_sympy(self, time, qval, pval):
         """ this version works with sympy objects """
         zval = sp.Matrix([qval, pval])
@@ -295,7 +324,7 @@ class SymmetricPolynomialPropagator(PolynomialPropagator):
             jacval = np.array(self.jacobian(*zval1))
             gvar = np.array(self.gvar(*zval1))
             if self.numeric['tool'] == 'theano':
-               jacval.shape = self.jshape
+                jacval.shape = self.jshape
             correction = np.matmul(np.linalg.inv(jacval), gvar-gcon)
             zval1 = zval1 - correction
             residue = np.linalg.norm(correction)