TPD threshold test (#51)

* TPD threshold test * TPD threshold test * factor of 2 in intensity reproduces short result
ergodicio · Jul 22, 2024 · 2f6e93d · 2f6e93d
1 parent 7e4ac72
commit 2f6e93d
Show file tree

Hide file tree

Showing 10 changed files with 186 additions and 369 deletions.
diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
@@ -35,4 +35,4 @@ jobs:
 
     - name: Test with pytest
       run: |
-        pytest tests/
+        CPU_ONLY=True pytest tests/
diff --git a/adept/lpse2d/core/epw.py b/adept/lpse2d/core/epw.py
@@ -11,183 +11,6 @@
 from adept.lpse2d.core.trapper import ParticleTrapper
 
 
-class SpectralPotential_old(eqx.Module):
-    wp0: float
-    ld_rates: jax.Array
-    n0: float
-    w0: float
-    nuei: float
-    dx: float
-    dy: float
-    trapper: eqx.Module
-    driver: eqx.Module
-    kx: jax.Array
-    ky: jax.Array
-    kax_sq: jax.Array
-    one_over_ksq: jax.Array
-    dt: float
-    cfg: Dict
-    wr_corr: jax.Array
-    num_substeps: int
-    dt_substeps: float
-
-    def __init__(self, cfg):
-        self.cfg = cfg
-        self.dt = cfg["grid"]["dt"]
-        self.wp0 = cfg["plasma"]["wp0"]
-        self.n0 = np.sqrt(self.wp0)
-        self.w0 = cfg["drivers"]["E0"]["w0"]
-        self.nuei = 0.0  # -cfg["units"]["derived"]["nuei_norm"]
-        self.kx = cfg["grid"]["kx"]
-        self.ky = cfg["grid"]["ky"]
-        self.dx = cfg["grid"]["dx"]
-        self.dy = cfg["grid"]["dy"]
-        self.kax_sq = self.kx[:, None] ** 2.0 + self.ky[None, :] ** 2.0
-        self.one_over_ksq = cfg["grid"]["one_over_ksq"]
-
-        self.driver = Driver(cfg)
-        self.trapper = ParticleTrapper(cfg)
-
-        self.dt_substeps = 0.005
-        self.num_substeps = int(self.dt / self.dt_substeps) + 4
-
-        table_wrs, table_wis, table_klds = electrostatic.get_complex_frequency_table(
-            1024, cfg["terms"]["epw"]["kinetic real part"]
-        )
-        all_ks = jnp.sqrt(self.kax_sq).flatten()
-        self.ld_rates = jnp.interp(all_ks, table_klds, table_wis, left=0.0, right=table_wis[-1]).reshape(
-            self.kax_sq.shape
-        )
-        wrs = jnp.interp(all_ks, table_klds, table_wrs, left=1.0, right=table_wrs[-1]).reshape(self.kax_sq.shape)
-        if self.cfg["terms"]["epw"]["kinetic real part"]:
-            self.wr_corr = (wrs - 1.0) * self.one_over_ksq / 1.5
-        else:
-            self.wr_corr = 1.0
-
-    def get_eh_x(self, phi: jax.Array) -> jax.Array:
-        ehx = -jnp.fft.ifft2(1j * self.kx[:, None] * phi)
-        ehy = -jnp.fft.ifft2(1j * self.ky[None, :] * phi)
-
-        return jnp.concatenate([ehx[..., None], ehy[..., None]], axis=-1) * self.kx.size * self.ky.size / 4
-
-    def get_phi_from_eh(self, eh: jax.Array) -> jax.Array:
-        ekx = jnp.fft.fft2(eh[..., 0])
-        eky = jnp.fft.fft2(eh[..., 1])
-        phi = (
-            (1j * self.kx[:, None] * ekx + 1j * self.ky[None, :] * eky)
-            * self.one_over_ksq
-            * 4
-            / self.kx.size
-            / self.ky.size
-        )
-        return phi
-
-    def calc_linear_step(self, temperature: jax.Array) -> Tuple[jax.Array, jax.Array]:
-        damping_term = self.nuei / 2.0 + self.ld_rates
-        osc_term = self.wr_corr * 1j * 1.5 * temperature * self.kax_sq / self.wp0
-        return osc_term, damping_term
-
-    def calc_density_pert_step(self, nb: jax.Array, dn: jax.Array) -> jax.Array:
-        return -1j / 2.0 * self.wp0 * (1.0 - nb / self.n0 - dn / self.n0)
-
-    def _calc_div_(self, arr):
-        arrk = jnp.fft.fft2(arr, axes=[0, 1])
-        divk = self.kx[:, None] * arrk[..., 0] + self.ky[None, :] * arrk[..., 1]
-        return jnp.fft.ifft2(divk, axes=[0, 1])
-
-    def calc_tpd_source_step(self, phi: jax.Array, e0: jax.Array, nb: jax.Array, t: float) -> jax.Array:
-        """
-        Returns i*e/(4 m_e w_0) *exp(i(2w_p0 - w_0)t) *
-        (
-        F(n_b / n_0 * E_h^* dot E_0 ) -
-        (1 - w_0/w_{p0})*ik/k^2 dot (F(n_b / n_0) E_0 div E_h^*)
-        )
-
-
-        :param phi:
-        :param e0:
-        :param nb:
-        :param t:
-        :return:
-        """
-        eh = self.get_eh_x(phi)
-        eh_star = jnp.conj(eh)
-        div_ehstar = self._calc_div_(eh_star)
-        coeff = -1j / 4.0 / self.w0 * jnp.exp(1j * (2.0 * self.wp0 - self.w0) * t)
-        term1 = jnp.fft.fft2(nb / self.n0 * jnp.sum(eh_star * e0, axis=-1))  # 2d array
-
-        term2_temp = jnp.fft.fft2(nb[..., None] / self.n0 * (e0 * div_ehstar[..., None]), axes=[0, 1])  # 3d array
-        term2 = (
-            (1.0 - self.wp0 / self.w0)
-            * 1j
-            * self.one_over_ksq
-            * (self.kx[:, None] * term2_temp[..., 0] + self.ky[None, :] * term2_temp[..., 1])
-        )
-        return coeff * (term1 + term2)
-
-    def update_potential(self, t, y):
-        # do the equation first --- this is equation 54 so far
-        if self.cfg["terms"]["epw"]["linear"]:
-            osc_term, damping_term = self.calc_linear_step(y["temperature"])
-            new_phi = y["phi"] * jnp.exp(osc_term * self.dt)
-            new_phi = new_phi + self.dt * jnp.fft.fft2(jnp.fft.ifft2(damping_term * new_phi) / (1 + y["delta"] ** 2.0))
-
-            if self.cfg["terms"]["epw"]["density_gradient"]:
-                eh_x = self.get_eh_x(new_phi)
-                density_step = self.calc_density_pert_step(y["nb"], y["dn"])[..., None]
-                eh_x = eh_x * jnp.exp(density_step * self.dt)
-                new_phi = self.get_phi_from_eh(eh_x)
-
-            if self.cfg["terms"]["epw"]["source"]["tpd"]:
-                new_phi += self.dt * self.calc_tpd_source_step(y["phi"], y["e0"], y["nb"], t)
-
-        else:
-            raise NotImplementedError("The linear term is necessary to run the code")
-
-        return new_phi
-
-    def update_delta(self, t, y, args):
-        # return y["delta"] + self.dt * self.trapper(t, y["delta"], args)
-        return diffrax.diffeqsolve(
-            terms=diffrax.ODETerm(self.trapper),
-            solver=diffrax.Tsit5(),
-            t0=t,
-            t1=t + self.dt,
-            max_steps=self.num_substeps,
-            dt0=self.dt_substeps,
-            y0=y["delta"],
-            args=args,
-        ).ys[0]
-
-    def __call__(self, t, y, args):
-        # push the equation of motion for the potential
-        y["phi"] = self.update_potential(t, y)
-
-        if ("E2" in self.cfg["drivers"].keys()) or self.cfg["terms"]["epw"]["trapping"]["active"]:
-            eh = self.get_eh_x(y["phi"])
-
-        # then push the trapper if desired
-        if self.cfg["terms"]["epw"]["trapping"]["active"]:
-            y["delta"] = self.update_delta(t, y, args={"eh": eh, "nu_g": args["nu_g"]})
-
-        # then push the driver if any
-        if "E2" in self.cfg["drivers"].keys():
-            eh += jnp.concatenate(
-                [temp := self.driver(args["drivers"]["E2"], t)[..., None], jnp.zeros_like(temp)], axis=-1
-            )
-            y["phi"] = self.get_phi_from_eh(eh)
-
-        if (
-            self.cfg["terms"]["epw"]["boundary"]["x"] == "absorbing"
-            or self.cfg["terms"]["epw"]["boundary"]["y"] == "absorbing"
-        ):
-            y["phi"] = self.get_phi_from_eh(
-                self.get_eh_x(y["phi"]) * self.cfg["grid"]["absorbing_boundaries"][..., None]
-            )
-
-        return y
-
-
 class SpectralPotential:
     def __init__(self, cfg) -> None:
 
@@ -204,10 +27,8 @@ def __init__(self, cfg) -> None:
         self.dt = cfg["grid"]["dt"]
         self.cfg = cfg
         self.amp_key, self.phase_key = jax.random.split(jax.random.PRNGKey(np.random.randint(2**20)), 2)
-        self.low_pass_filter = cfg["grid"][
-            "low_pass_filter"
-        ]  # np.where(np.sqrt(self.k_sq) < 2.0 / 3.0 * np.amax(self.kx), 1, 0)
-        zero_mask = 1.0  # cfg["grid"]["zero_mask"]
+        self.low_pass_filter = cfg["grid"]["low_pass_filter"]
+        zero_mask = cfg["grid"]["zero_mask"]
         self.low_pass_filter = self.low_pass_filter * zero_mask
         self.nx = cfg["grid"]["nx"]
         self.ny = cfg["grid"]["ny"]
@@ -232,6 +53,7 @@ def calc_fields_from_phi(self, phi: Array) -> Tuple[Array, Array]:
         Returns:
             A Tuple containing ex(x, y) and ey(x, y)
         """
+
         phi_k = jnp.fft.fft2(phi)
         phi_k *= self.low_pass_filter
 
@@ -251,6 +73,7 @@ def calc_phi_from_fields(self, ex: Array, ey: Array) -> Array:
             Array: phi(x, y)
 
         """
+
         ex_k = jnp.fft.fft2(ex)
         ey_k = jnp.fft.fft2(ey)
         divE_k = 1j * (self.kx[:, None] * ex_k + self.ky[None, :] * ey_k)
@@ -336,7 +159,7 @@ def calc_tpd2(self, t: float, y: Array, args: Dict) -> Array:
         return self.tpd_const * tpd2
 
     def get_noise(self):
-        random_amps = 1000.0  # jax.random.uniform(self.amp_key, (self.nx, self.ny))
+        random_amps = 1.0  # jax.random.uniform(self.amp_key, (self.nx, self.ny))
         random_phases = 2 * np.pi * jax.random.uniform(self.phase_key, (self.nx, self.ny))
         return jnp.fft.ifft2(random_amps * jnp.exp(1j * random_phases) * self.low_pass_filter)
 

diff --git a/adept/lpse2d/helpers.py b/adept/lpse2d/helpers.py
@@ -39,7 +39,9 @@ def write_units(cfg: Dict) -> Dict:
     Z = cfg["units"]["ionization state"]
     A = cfg["units"]["atomic number"]
     lam0 = _Q(cfg["units"]["laser wavelength"]).to("um").value
-    I0 = _Q(cfg["units"]["laser intensity"]).to("W/cm^2").value
+    I0 = (
+        _Q(cfg["units"]["laser intensity"]).to("W/cm^2").value / 2
+    )  ## NB - the factor of 2 enables matching to the Short scaling
     envelopeDensity = cfg["units"]["envelope density"]
 
     # Scaled constants
@@ -311,7 +313,9 @@ def get_solver_quantities(cfg: Dict) -> Dict:
     )
 
     cfg_grid["zero_mask"] = (
-        np.where(cfg_grid["kx"][:, None] * cfg_grid["ky"][None, :] == 0, 0, 1) if cfg["terms"]["zero_mask"] else 1
+        np.where(cfg_grid["kx"] == 0, 0, 1)[:, None] * np.ones_like(cfg_grid["ky"])[None, :]
+        if cfg["terms"]["zero_mask"]
+        else 1
     )
     # sqrt(kx**2 + ky**2) < 2/3 kmax
     cfg_grid["low_pass_filter"] = np.where(

diff --git a/adept/lpse2d/vector_field.py b/adept/lpse2d/vector_field.py
@@ -1,4 +1,4 @@
-from typing import Dict, List
+from typing import Dict
 
 from jax import numpy as jnp, Array
 import numpy as np
@@ -13,8 +13,6 @@ class SplitStep:
 
     All the pushers are chosen and initialized here and a single time-step is defined here.
 
-    Note that EPW can be either the charge density (div E) or the potential
-
     :param cfg:
     :return:
     """
@@ -24,18 +22,17 @@ def __init__(self, cfg):
         self.cfg = cfg
         self.dt = cfg["grid"]["dt"]
         self.wp0 = cfg["units"]["derived"]["wp0"]
-        # self.epw = epw.FDChargeDensity(cfg)
         self.epw = epw.SpectralPotential(cfg)
         self.light = laser.Light(cfg)
         self.complex_state_vars = ["E0", "epw"]
         self.boundary_envelope = cfg["grid"]["absorbing_boundaries"]
         self.one_over_ksq = cfg["grid"]["one_over_ksq"]
-        self.zero_mask = 1.0  # cfg["grid"]["zero_mask"]
+        self.zero_mask = cfg["grid"]["zero_mask"]
         self.low_pass_filter = cfg["grid"]["low_pass_filter"]
 
         self.nu_coll = cfg["units"]["derived"]["nu_coll"]
 
-    def unpack_y(self, y: Dict[str, Array]) -> Dict[str, Array]:
+    def _unpack_y_(self, y: Dict[str, Array]) -> Dict[str, Array]:
         new_y = {}
         for k in y.keys():
             if k in self.complex_state_vars:
@@ -44,7 +41,7 @@ def unpack_y(self, y: Dict[str, Array]) -> Dict[str, Array]:
                 new_y[k] = y[k].view(jnp.float64)
         return new_y
 
-    def pack_y(self, y: Dict[str, Array], new_y: Dict[str, Array]) -> tuple[Dict[str, Array], Dict[str, Array]]:
+    def _pack_y_(self, y: Dict[str, Array], new_y: Dict[str, Array]) -> tuple[Dict[str, Array], Dict[str, Array]]:
         for k in y.keys():
             y[k] = y[k].view(jnp.float64)
             new_y[k] = new_y[k].view(jnp.float64)
@@ -85,29 +82,27 @@ def landau_damping(self, epw: Array, vte_sq: float):
 
     def __call__(self, t, y, args):
         # unpack y into complex128
-        new_y = self.unpack_y(y)
+        new_y = self._unpack_y_(y)
 
         # split step
         new_y = self.light_split_step(t, new_y, args["drivers"])
-        new_y["epw"] = self.epw(t, new_y, args)
 
         if "E2" in args["drivers"]:
             new_y["epw"] += self.dt * self.epw.driver(args["drivers"]["E2"], t)
+        new_y["epw"] = self.epw(t, new_y, args)
 
         # landau and collisional damping
         if self.cfg["terms"]["epw"]["damping"]["landau"]:
             new_y["epw"] = self.landau_damping(epw=new_y["epw"], vte_sq=y["vte_sq"])
 
-        new_y["epw"] = jnp.fft.ifft2(self.zero_mask * jnp.fft.fft2(new_y["epw"]))
-
         # boundary damping
         ex, ey = self.epw.calc_fields_from_phi(new_y["epw"])
         ex = ex * self.boundary_envelope
         ey = ey * self.boundary_envelope
         new_y["epw"] = self.epw.calc_phi_from_fields(ex, ey)
-        # new_y["epw"] = jnp.fft.ifft2(self.zero_mask * self.low_pass_filter * jnp.fft.fft2(new_y["epw"]))
+        new_y["epw"] = jnp.fft.ifft2(self.zero_mask * self.low_pass_filter * jnp.fft.fft2(new_y["epw"]))
 
         # pack y into float64
-        y, new_y = self.pack_y(y, new_y)
+        y, new_y = self._pack_y_(y, new_y)
 
         return new_y
diff --git a/configs/envelope-2d/tpd.yaml b/configs/envelope-2d/tpd.yaml
@@ -14,9 +14,9 @@ drivers:
     delta_omega_max: 0.015
     num_colors: 1
     envelope:
-      tw: 20ns
-      tr: 0.5ns
-      tc: 11fs
+      tw: 20ps
+      tr: 0.25ps
+      tc: 10.5ps
       xr: 0.2um
       xw: 1000um
       xc: 50um
@@ -27,27 +27,27 @@ drivers:
 grid:
   boundary_abs_coeff: 1.0e4
   boundary_width: 1um
-  low_pass_filter: 0.7
+  low_pass_filter: 0.3
   dt: 0.002ps
   dx: 20nm
-  tmax: 2.0ps
+  tmax: 4.0ps
   tmin: 0.0ns
-  ymax: 4um
-  ymin: -4um
+  ymax: 3um
+  ymin: -3um
 machine:
   calculator: gpu
 mlflow:
-  experiment: test-lpse
-  run: test
+  experiment: tpd
+  run: 1.5e14
 save:
   t:
     dt: 100fs
-    tmax: 2ps
+    tmax: 4ps
     tmin: 0ps
   x:
-    dx: 48nm
+    dx: 80nm
   y:
-    dy: 160nm
+    dy: 128nm
 solver: envelope-2d
 terms:
   epw:
@@ -67,7 +67,7 @@ units:
   atomic number: 40
   envelope density: 0.25
   ionization state: 6
-  laser intensity: 3.5e+14W/cm^2
+  laser intensity: 2.0e+14W/cm^2
   laser wavelength: 351nm
   reference electron temperature: 2000.0eV
   reference ion temperature: 1000eV