replaced spectral laplacian by finite-difference one

Author: hightower8083

axiprop/simulation/inline_methods.py

@@ -2,17 +2,35 @@
 from numba import jit, prange
 from scipy.constants import e, m_e, c
 
-@jit(nopython=True, cache=True, parallel=True)
 def wake_kernel_integrate(kernel_t, delta_n, k_p, xi_ax):
     dxi = xi_ax[1] - xi_ax[0]
-    for it in prange(1, xi_ax.size):
+    for it in range(1, xi_ax.size):
         xi = xi_ax[it]
         delta_n[it] = dxi * np.sum(
             kernel_t[:it] * np.sin( k_p * (xi - xi_ax[:it, None]) ),
             axis=0
         )
     return delta_n
 
+def laplacian_fd(A, xi_ax, r_ax):
+    r_inv = 1.0 / r_ax
+    dxi_inv = 1.0 / ( xi_ax[1] - xi_ax[0] )
+    dxi2_inv = dxi_inv * dxi_inv
+    dr_inv = 1.0 / ( r_ax[1] - r_ax[0] )
+    dr2_inv = dr_inv * dr_inv
+
+    nabla2_A = np.zeros_like(A)
+
+    nabla2_A[1:-1, 1:-1] = \
+        dxi2_inv * (A[2:, 1:-1] - 2* A[1:-1, 1:-1] + A[:-2, 1:-1]) \
+      + dr2_inv * (A[1:-1, 2:] - 2* A[1:-1, 1:-1] + A[1:-1, :-2,]) \
+      + r_inv[None, 1:-1] * 0.5 * dr_inv * (A[1:-1, 2:] - A[1:-1, :-2] )
+
+    nabla2_A[:,0] = nabla2_A[:,1]
+
+    return nabla2_A
+
+
 @jit(nopython=True, cache=True)
 def get_ADK_probability(E_fld, dt, adk_power, \
                         adk_prefactor, adk_exp_prefactor):

axiprop/simulation/plasma.py

@@ -12,6 +12,7 @@
 from .inline_methods import get_plasma_ADK_ref
 from .inline_methods import get_OFI_heating
 from .inline_methods import wake_kernel_integrate
+from .inline_methods import laplacian_fd
 
 r_e = e**2 / m_e / c**2 / 4 / pi /epsilon_0
 mc_m2 = 1. / ( m_e * c )**2
@@ -123,46 +124,25 @@ def get_density(self, E_ft ):
         sim = self.sim
         prop = self.sim.prop
         n_pe = self.n_pe
+        xi_ax = c * ScalarFieldEnvelope(*sim.EnvArgs).t
+        r_ax = prop.r_new
         k_p = np.sqrt( 4 * np.pi * r_e * n_pe )
 
-        # normalized vector potential
+        # get normalized vector potential
         A_ft = -1j * prop.omega_inv * e_mc * E_ft
         A_ft *= (prop.kz[:, None]>0.0)
 
-        # from frequency to time
+        # bring from frequency to time domain
         A_obj = ScalarFieldEnvelope(*sim.EnvArgs)
         A_obj.t += sim.dt_shift
         A_obj.t_loc += sim.dt_shift
         A_t = A_obj.import_field_ft(A_ft).Field
 
         # cycle-averaged square of normalized vector potential
-        A2_t = 0.5 * np.astype( np.abs(A_t)**2, A_t.dtype )
-
-        # from time to frequency
-        A2_obj = ScalarFieldEnvelope(*sim.EnvArgs)
-        A2_obj.t += sim.dt_shift
-        A2_obj.t_loc += sim.dt_shift
-        A2_ft = A2_obj.import_field(A2_t).Field_ft
-
-        # transverse spectral transform
-        A2_ts = prop.perform_transfer_TST( A2_ft )
-
-        # wakefield kernel (spectral calculation of Laplacian)
-        k_env = ScalarFieldEnvelope(*sim.EnvArgs).k_freq_base
-        laplacian_ts = -( k_env[:, None]**2 + prop.kr[None,:]**2 )
-        laplacian_ts = prop.bcknd.to_device(laplacian_ts)
-
-        kernel_ts = 0.5 / k_p * laplacian_ts *  A2_ts
-        kernel_ft = prop.perform_iTST_transfer(kernel_ts)
-
-        kernel_obj = ScalarFieldEnvelope(*sim.EnvArgs)
-        kernel_obj.t += sim.dt_shift
-        kernel_obj.t_loc += sim.dt_shift
-        kernel_t = kernel_obj.import_field_ft(kernel_ft).Field
-        kernel_t = np.abs(kernel_t)
+        A2_t = 0.5 * np.abs(A_t)**2
+        kernel_t = 0.5 / k_p * laplacian_fd(A2_t, xi_ax, r_ax)
 
         # do the integral
-        xi_ax = c * kernel_obj.t.copy()
         delta_n = np.zeros_like(kernel_t)
         delta_n = wake_kernel_integrate(kernel_t, delta_n, k_p, xi_ax)