chebyshev scalar and shift tests

classical/chebyshev.py

@@ -7,11 +7,14 @@ from propagators import PolynomialPropagator
 from liouville import LiouvillePowersQP
 
 def cheb_coef(bf, nterms, alpha):
+    """ chebyshev expansion coefficients for the function exp(z*t) """
     coef = bf(range(nterms), np.full(nterms, alpha))
     coef[1:] = coef[1:]*2.
     return coef
 
 def besf(fpprec, signed, prec=None):
+    """ select the proper bessel function according to requested precision
+        and selected axis: if signed if True [-1, +1] otherwise [-i, +i] """
     if fpprec == 'fixed':
         bessel = iv if signed else jn
     elif fpprec == 'multi':
@@ -24,6 +27,31 @@ def besf(fpprec, signed, prec=None):
         bessel = None
     return bessel
 
+def chebyshev_scalar(z, t, delta=None, lambda_min=None, nterms=None,
+                     signed=False, fpprec='fixed', prec=None):
+    """ Chebyshev polynomial expansion of the scalar function exp(z*t) """
+
+    assert isinstance(z, np.ndarray)
+    delta = max(abs(z)) if delta is None else delta 
+    alpha = delta*t/2.
+    lambda_min = -delta/2. if lambda_min is None else lambda_min
+    nterms = int(np.ceil(abs(alpha)))+1 if nterms is None else nterms
+    assert nterms > 1
+    scale = 2./delta
+    shift = -(2.*lambda_min/delta+1.)
+    sign = -1 if signed else 1
+    z = z*scale+shift
+    exp = np.exp if fpprec == 'fixed' else lambda x: sp.exp(x).evalf(prec)
+    scaling = exp((lambda_min+delta/2.)*t)
+    bessel_func = besf(fpprec, signed, prec=prec)
+    ch_coef = cheb_coef(bessel_func, nterms, alpha)*scaling
+    ch_poly = np.empty((nterms, len(z)), dtype=z.dtype)
+    ch_poly[0, :] = 1.
+    ch_poly[1, :] = z
+    for order in range(2, nterms):
+        ch_poly[order, :] = 2.*z*ch_poly[order-1, :]+sign*ch_poly[order-2, :]
+    return np.dot(ch_coef, ch_poly)
+
 
 class Chebyshev(PolynomialPropagator):
     """ Chebyshev propagator for classical particle dynamics """

classical/tests.py

@@ -463,6 +463,18 @@ class SymplecticPropagatorTest(unittest.TestCase):
 class ChebyshevAssessmentTest(unittest.TestCase):
     """ Compare Chebyshev to velocity Verlet and analytical solution """
 
+    def chebyshev_scalar_test(self):
+        """ chebyshev expansion of the scalar function exp(z*t) """
+        from chebyshev import chebyshev_scalar
+        from itertools import product
+
+        sample_length = 10
+        lambdas = (np.random.rand((sample_length))-0.5)*6*np.random.rand()
+        for sign, fact in product((False, True), (1., 1.j)):
+            apr = chebyshev_scalar(lambdas*fact, 1.1, nterms=16, signed=sign)
+            ref = np.exp(lambdas*1.1*fact)
+            self.assertTrue(np.allclose(apr, ref))
+
     def chebyshev_symplectic_1p_test(self):
         """ with one-particle system: anharmonic oscillator """
         params = {
@@ -545,6 +557,19 @@ class ChebyshevAssessmentTest(unittest.TestCase):
         self.assertTrue(np.allclose(pt, ptr, atol=1e-3))
         self.assertLess(np.max(np.fabs(prop.er)), 1e-3)
 
+    def chebyshev_shift_test(self):
+        """ chebyshev propagator with scale and shift """
+        syst = Harmonic(q0=1., p0=0.)
+        params = {'delta': 3., 'lambda_min': -2., 'nterms': 16,
+                  'signed': False, 'tstart': 0., 'tend': 10.0, 'tstep': 0.5}
+        prop = Chebyshev(syst=syst, **params)
+        prop.propagate()
+        prop.analyse()
+        t, q, p = prop.get_trajectory()
+        qr, pr = syst.solution(times=t)
+        self.assertTrue(np.allclose(q, qr)) 
+        self.assertLess(np.max(np.fabs(prop.er)), 1e-6)
+
     def chebyshev_symplectic_2p_lj_test(self):
         """ two particles in Lennard-Jones potential """