revised and fixed the issue with the stepping symmetric propagator --…

… big success (#27)
hightower8083 · Sep 30, 2024 · 8d31d4e · 8d31d4e
1 parent fff449d
commit 8d31d4e
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 4 deletions.
diff --git a/axiprop/simulation/lib.py b/axiprop/simulation/lib.py
@@ -1,3 +1,4 @@
+import numpy as np
 from scipy.special import jn
 
 from ..lib import PropagatorResampling
@@ -12,7 +13,7 @@ class PropagatorResamplingStepping(
     based on `PropagatorResampling`
     """
 
-    def init_TST_resampled(self):
+    def init_TST_stepping(self):
         """
         Setup DHT transform
         """
@@ -32,10 +33,10 @@ def TST_stepping(self):
         Forward QDHT transform.
         """
         if not hasattr(self, 'TST_resampled_matmul'):
-            self.init_TST_resampled()
+            self.init_TST_stepping()
 
         self.u_ht = self.TST_resampled_matmul(
-            self.TM_resampled, self.u_loc, self.u_ht)
+            self.TM_resampled, self.u_iht, self.u_ht)
 
 
 class PropagatorSymmetricStepping(
@@ -45,3 +46,36 @@ class PropagatorSymmetricStepping(
     based on `PropagatorSymmetric`
     """
     TST_stepping = PropagatorSymmetric.TST
+
+    def init_TST_stepping(self):
+        Rmax = self.Rmax
+        Nr = self.Nr
+        Nr_new = self.Nr_new
+        dtype = self.dtype
+        mode = self.mode
+        alpha = self.alpha
+        alpha_np1 = self.alpha_np1
+
+        self._j_stepping = self.bcknd.to_device( np.abs(jn(mode+1, alpha)) / Rmax )[:Nr_new]
+        denominator = alpha_np1 * np.abs(jn(mode+1, alpha[:,None]) \
+                                       * jn(mode+1, alpha[None,:]))
+
+        TM = 2 * jn(mode, alpha[:,None] * alpha[None,:] / alpha_np1)\
+                     / denominator
+
+        self.TM_stepping = self.bcknd.to_device(TM[:,:Nr_new], dtype)
+
+        self.TST_stepping_matmul = self.bcknd.make_matmul(self.TM_stepping, self.u_iht, self.u_ht)
+
+    def TST_stepping(self):
+        """
+        """
+        if not hasattr(self, 'TST_stepping_matmul'):
+            self.init_TST_stepping()
+
+        self.u_iht /= self._j_stepping
+
+        self.u_ht = self.TST_stepping_matmul(
+            self.TM_stepping, self.u_iht, self.u_ht)
+
+
diff --git a/axiprop/simulation/steppers.py b/axiprop/simulation/steppers.py
@@ -156,10 +156,11 @@ def perform_transfer_TST(self, u, image=None, stepping=True):
             image = self.bcknd.zeros( out_shape, self.dtype )
 
         for ikz in range(self.Nkz):
-            self.u_loc = self.bcknd.to_device(u[ikz,:].copy())
             if stepping:
+                self.u_iht = self.bcknd.to_device(u[ikz,:].copy())
                 self.TST_stepping()
             else:
+                self.u_loc = self.bcknd.to_device(u[ikz,:].copy())
                 self.TST()
             image[ikz] = self.u_ht.copy()