Energy conservation bug

axiprop/physproc/ionization_inline.py

@@ -70,8 +70,8 @@ def get_plasma_ADK( E_laser, A_laser, dt, n_gas, pack_ADK, Uion,
                     all_new_events += new_events
 
                     if ionization_current:
-                        IonizationEnergy[ir] += n_gas_loc \
-                            * new_events * Uion[ion_state]
+                        #IonizationEnergy[ir] += n_gas_loc \
+                        #    * new_events * Uion[ion_state]
                         J_ion = n_gas_loc * new_events \
                             * Uion[ion_state] / E_ir_t[it] / dt
                         J_ir_t[it] += J_ion
@@ -83,6 +83,8 @@ def get_plasma_ADK( E_laser, A_laser, dt, n_gas, pack_ADK, Uion,
         n_e[ir] = n_e_slice[-1]
         J_plasma[:, ir] = J_ir_t.copy()
 
+        IonizationEnergy[ir] = 0.5 * np.real((J_ir_t * np.conj(E_ir_t))).sum() * dt
+
     return J_plasma, n_e, IonizationEnergy
 
 
@@ -150,3 +152,92 @@ def get_OFI_heating(
         Xi[ir, :] += ion_fracts_loc
 
     return n_e, T_e, Xi
+
+@jit(nopython=True, cache=True, parallel=True)
+def get_plasma_ADK_ref( E_laser, A_laser, t_axis, omega0,
+                        n_gas, pack_ADK, Uion,
+                        Z_init, Zmax, ionization_current ):
+
+    num_ions = pack_ADK[0].size + 1
+    Nt, Nr = E_laser.shape
+
+    dt = t_axis[1] - t_axis[0]
+    phase_env = np.exp(-1j * omega0 * t_axis)
+
+    J_loc_t_re = np.zeros((Nt, Nr), dtype=np.double)
+    n_e = np.zeros(Nr)
+    T_e = np.zeros(Nr)
+    Xi = np.zeros( (Nr, num_ions) )
+
+    for ir in prange(Nr):
+
+        n_gas_loc = n_gas[ir]
+
+        E_ir_t_re = np.real( E_laser[:,ir] * phase_env )
+        A_ir_t_re = np.real( A_laser[:,ir] * phase_env )
+
+        J_ir_t_re = np.zeros_like(E_ir_t_re)
+        n_e_slice = np.zeros(Nt)
+
+        P_ir_t_re = e * A_ir_t_re
+        v_ir_t_re = P_ir_t_re / m_e / \
+            np.sqrt( 1 + P_ir_t_re**2 / (m_e*c)**2 )
+
+        if Z_init>0:
+            J_ir_t_re += -e * Z_init * n_gas_loc * v_ir_t_re
+
+        ion_fracts_loc = np.zeros(num_ions)
+        ion_fracts_loc[Z_init] = 1.
+
+        all_new_events = 0.0
+        electron_temp_dens = 0.0
+
+        for it in range(Nt):
+
+            if np.abs(ion_fracts_loc[Zmax]-1)<1e-16:
+                break
+
+            probs = get_ADK_probability( np.abs( E_ir_t_re[it] ), dt, *pack_ADK )
+
+            W_ir = m_e * c**2 * (
+                np.sqrt( 1 + P_ir_t_re[it]**2 / (m_e*c)**2 ) - 1.0
+            )
+
+            for ion_state in range(num_ions-1):
+
+                if probs[ion_state]<1e-8:
+                    continue
+
+                if ion_fracts_loc[ion_state]<1e-16:
+                    continue
+
+                new_events = ion_fracts_loc[ion_state] * probs[ion_state]
+                if new_events>0:
+                    ion_fracts_loc[ion_state+1] += new_events
+                    ion_fracts_loc[ion_state] -= new_events
+                    all_new_events += new_events
+
+                    # ideal electrons gas is assumed so `E = 3/2 k_B T`
+                    electron_temp_dens += 2. / 3. * W_ir * new_events
+
+                    if ionization_current:
+                        J_ion = n_gas_loc * new_events \
+                            * Uion[ion_state] / E_ir_t_re[it] / dt
+                        J_ir_t_re[it] += J_ion
+
+                    n_e_slice[it:] += new_events * n_gas_loc
+
+        J_ir_t_re = -e * n_e_slice * v_ir_t_re
+        n_e[ir] = n_e_slice[-1]
+
+        J_loc_t_re[:, ir] = J_ir_t_re.copy()
+
+        if all_new_events>0:
+            T_e[ir] = electron_temp_dens / all_new_events
+        else:
+            T_e[ir] = 0.0
+
+        Xi[ir, :] += ion_fracts_loc
+
+    return J_loc_t_re, n_e, T_e, Xi
+

axiprop/physproc/plasma.py

@@ -3,13 +3,16 @@
 from scipy.constants import e, m_e, c, pi, epsilon_0
 from scipy.constants import mu_0, fine_structure
 from scipy.special import gamma as gamma_func
+from scipy.signal import hilbert
 from mendeleev import element as table_element
 
 from ..containers import ScalarFieldEnvelope
 from ..utils import refine1d, refine1d_TR
 from .ionization_inline import get_plasma_ADK
+from .ionization_inline import get_plasma_ADK_ref
 from .ionization_inline import get_OFI_heating
 
+
 r_e = e**2 / m_e / c**2 / 4 / pi /epsilon_0
 mc_m2 = 1. / ( m_e * c )**2
 omega_a = fine_structure**3 * c / r_e
@@ -122,12 +125,18 @@ def __init__( self, n_gas, dens_func, sim, my_element,
         super().__init__(n_gas, dens_func, sim)
         self.n_gas = n_gas
 
+        self.T_e = 0.0
         self.Z_init = Z_init
         self.Zmax = Zmax
         self.dt = sim.t_axis[1] - sim.t_axis[0]
         self.make_ADK(my_element, ionization_mode)
         self.ionization_current = ionization_current
 
+        omega_shift = sim.prop.kz[:, None] * c - sim.omega0
+        self.lowpass_filt = np.cos(
+            0.5 * np.pi * omega_shift / omega_shift.max()
+        )
+
     def make_ADK(self, my_element, ionization_mode):
         """
         Prepare the useful data for ADK probabplity calculation for a given
@@ -197,6 +206,8 @@ def get_RHS(self, E_ts, dz=0.0 ):
         Jp_obj.t_loc += sim.dt_shift
         Jp_ft = Jp_obj.import_field(Jp_loc_t).Field_ft
 
+        Jp_ft *= self.lowpass_filt
+
         Jp_ts = prop.perform_transfer_TST( Jp_ft )
 
         if dz != 0.0:
@@ -272,3 +283,78 @@ def get_data(self, E_obj):
             E_loc_t, P_loc_t, t_axis, sim.omega0, n_gas_loc,
             (self.adk_power, self.adk_prefactor, self.adk_exp_prefactor),
             self.Z_init, self.Zmax)
+
+
+class PlasmaIonizationRefine(PlasmaIonization):
+
+    def __init__( self, n_gas, dens_func, sim, my_element,
+                  Z_init=0, Zmax=-1, ionization_current=True,
+                  ionization_mode='DC', refine_ord=8, **kw_args):
+
+        super().__init__(n_gas, dens_func, sim,
+            my_element=my_element, Z_init=Z_init, Zmax=Zmax,
+            ionization_current=ionization_current,
+            ionization_mode=ionization_mode)
+
+        self.refine_ord = refine_ord
+
+    def get_RHS(self, E_ts, dz=0.0 ):
+        sim = self.sim
+        prop = self.sim.prop
+        omega = sim.prop.kz[:, None] * c
+
+        n_gas = self.n_gas * self.dens_func( sim.z_loc + dz, prop.r_new )
+
+        if dz != 0.0:
+            sim.t_axis += dz / c
+            E_loc = prop.step_and_iTST_transfer(E_ts, dz)
+        else:
+            E_loc = prop.perform_iTST_transfer(E_ts)
+
+        A_loc = -1j * prop.omega_inv * E_loc
+        A_loc *= (omega>0.0)
+
+        E_loc_obj = ScalarFieldEnvelope(*sim.EnvArgs)
+        E_loc_obj.t += sim.dt_shift
+        E_loc_obj.t_loc += sim.dt_shift
+        E_loc_t = E_loc_obj.import_field_ft(E_loc).Field
+
+        A_loc_obj = ScalarFieldEnvelope(*sim.EnvArgs)
+        A_loc_obj.t += sim.dt_shift
+        A_loc_obj.t_loc += sim.dt_shift
+        A_loc_t = A_loc_obj.import_field_ft(A_loc).Field
+
+        if self.refine_ord>1:
+            E_loc_t = refine1d_TR(E_loc_t, self.refine_ord)
+            A_loc_t = refine1d_TR(A_loc_t, self.refine_ord)
+            t_axis  = refine1d(sim.t_axis, self.refine_ord)
+        else:
+            t_axis = sim.t_axis.copy()
+
+        Jp_loc_t_re, self.n_e, self.T_e, self.Xi = get_plasma_ADK_ref(
+            E_loc_t, A_loc_t, t_axis, sim.omega0, n_gas,
+            (self.adk_power, self.adk_prefactor, self.adk_exp_prefactor),
+            self.Uion, self.Z_init, self.Zmax, self.ionization_current
+        )
+
+        Jp_loc_t = np.exp(1j * sim.omega0 * sim.t_axis)[:, None] \
+            * np.conj(
+                hilbert(Jp_loc_t_re, axis=0)[::self.refine_ord]
+              )
+
+        Jp_obj = ScalarFieldEnvelope(*sim.EnvArgs)
+        Jp_obj.t += sim.dt_shift
+        Jp_obj.t_loc += sim.dt_shift
+        Jp_ft = Jp_obj.import_field(Jp_loc_t).Field_ft
+
+        Jp_ft *= self.lowpass_filt
+
+        Jp_ts = prop.perform_transfer_TST( Jp_ft )
+
+        if dz != 0.0:
+            sim.t_axis -= dz / c
+            Jp_ts = prop.step_simple(Jp_ts, -dz)
+
+        Jp_ts *= self.coef_RHS
+
+        return Jp_ts

axiprop/physproc/simulation.py

@@ -86,8 +86,8 @@ def __init__(self, prop, t_axis, k0, z_0,
         self.verbose = verbose
 
     def run(self, E0, Lz, dz0, N_diags,
-            method='RK4', adjust_dz=True, dz_min=2e-7, err_max=1e-2,
-            growth_rate=None):
+            method='RK4', adjust_dz=True, dz_min=2e-7,
+            err_max=1e-2, growth_rate=None):
 
         physprocs = self.physprocs
         for diag_str in ['n_e', 'T_e', 'Xi']:
@@ -109,7 +109,7 @@ def run(self, E0, Lz, dz0, N_diags,
         if self.verbose:
             self._pbar_init(Lz)
 
-        while (self.z_loc < self.z_0 + Lz) :
+        while (self.z_loc <= self.z_0 + Lz) :
             # record diagnostics data
             if do_diag_next:
                 self._record_diags(En_ts, physprocs)
@@ -254,7 +254,7 @@ def _record_diags(self, E_fb, physprocs):
             self.diags['E_ft'].append(E_obj.Field_ft)
 
         if 'E_ft_onax' in self.diags.keys():
-            self.diags['E_ft'].append(E_obj.Field_ft[:, 0])
+            self.diags['E_ft_onax'].append(E_obj.Field_ft[:, 0])
 
         if 'E_t_env' in self.diags.keys():
             self.diags['E_t_env'].append(E_obj.Field)

axiprop/utils.py

@@ -27,8 +27,9 @@ def func_wrp(*args, **kw_args):
 
 def refine1d(A, refine_ord):
     refine_ord = int(refine_ord)
-    x = np.arange(A.size, dtype=np.double)
-    x_new = np.linspace(x.min(), x.max(), x.size*refine_ord)
+    Nx = A.size
+    x = np.arange(Nx, dtype=np.double)
+    x_new = np.linspace(x.min(), x.max(), refine_ord * (Nx-1) + 1 )
 
     if A.dtype == np.double:
         interp_fu = interp1d(x, A, assume_sorted=True)
@@ -49,9 +50,10 @@ def refine1d(A, refine_ord):
 
 def refine1d_TR(A, refine_ord):
     refine_ord = int(refine_ord)
+    Nx = A.shape[0]
 
-    t = np.arange(A.shape[0], dtype=np.double)
-    t_new = np.linspace(t.min(), t.max(), t.size*refine_ord)
+    t = np.arange(Nx, dtype=np.double)
+    t_new = np.linspace(t.min(), t.max(), refine_ord * (Nx-1) + 1)
 
     A_new = np.zeros((t_new.size, A.shape[1]), dtype=A.dtype)