From 6fcd5a83d3d161b86c25a4aedc21a06486b1c307 Mon Sep 17 00:00:00 2001 From: Archis Joglekar Date: Thu, 24 Oct 2024 21:03:16 -0700 Subject: [PATCH] LPSE2D Bugs (#84) * Factor of 2 in ky because of ny * intensity bug * updated config --- adept/_lpse2d/core/epw.py | 48 +---------- adept/_lpse2d/core/integrator.py | 134 ------------------------------- adept/_lpse2d/helpers.py | 22 ++--- configs/envelope-2d/tpd.yaml | 33 ++++---- 4 files changed, 25 insertions(+), 212 deletions(-) delete mode 100644 adept/_lpse2d/core/integrator.py diff --git a/adept/_lpse2d/core/epw.py b/adept/_lpse2d/core/epw.py index 7971d20..bb035a6 100644 --- a/adept/_lpse2d/core/epw.py +++ b/adept/_lpse2d/core/epw.py @@ -97,7 +97,6 @@ def tpd(self, t: float, y: Array, args: Dict) -> Array: tpd1 = E0[..., 1] * jnp.conj(ey) tpd1 = jnp.fft.ifft2(jnp.fft.fft2(tpd1) * self.low_pass_filter) - # tpd1 = E0_Ey divE_true = jnp.fft.ifft2(self.k_sq * jnp.fft.fft2(phi)) E0_divE_k = jnp.fft.fft2(E0[..., 1] * jnp.conj(divE_true)) @@ -110,49 +109,6 @@ def tpd(self, t: float, y: Array, args: Dict) -> Array: return dphi - def calc_tpd1(self, t: float, y: Array, args: Dict) -> Array: - """ - Calculates the first term of the two plasmon decay - - Args: - t (float): time - y (Array): phi(x, y) - args (Dict): dictionary containing E0 - - Returns: - Array: dphi(x, y) - """ - E0 = args["E0"] - phi = y - - _, ey = self.calc_fields_from_phi(phi) - - tpd1 = E0[..., 1] * jnp.conj(ey) - return self.tpd_const * tpd1 - - def calc_tpd2(self, t: float, y: Array, args: Dict) -> Array: - """ - Calculates the second term of the two plasmon decay - - Args: - t (float): time - y (Array): phi(x, y) - args (Dict): dictionary containing E0 - - Returns: - Array: dphi(x, y) - - """ - phi = y - E0 = args["E0"] - - divE_true = jnp.fft.ifft2(self.k_sq * jnp.fft.fft2(phi)) - E0_divE_k = jnp.fft.fft2(E0[..., 1] * jnp.conj(divE_true)) - - tpd2 = 1j * self.ky[None, :] * self.one_over_ksq * E0_divE_k - tpd2 = jnp.fft.ifft2(tpd2) - return self.tpd_const * tpd2 - def get_noise(self): 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)) @@ -175,8 +131,8 @@ def __call__(self, t: float, y: Dict[str, Array], args: Dict) -> Array: # density gradient if self.cfg["terms"]["epw"]["density_gradient"]: ex, ey = self.calc_fields_from_phi(phi) - ex *= jnp.exp(-1j * self.wp0 / 2.0 * (1 - background_density / self.envelope_density) * self.dt) - ey *= jnp.exp(-1j * self.wp0 / 2.0 * (1 - background_density / self.envelope_density) * self.dt) + ex *= jnp.exp(-1j * self.wp0 / 2.0 * (background_density / self.envelope_density - 1) * self.dt) + ey *= jnp.exp(-1j * self.wp0 / 2.0 * (background_density / self.envelope_density - 1) * self.dt) phi = self.calc_phi_from_fields(ex, ey) if self.cfg["terms"]["epw"]["source"]["noise"]: diff --git a/adept/_lpse2d/core/integrator.py b/adept/_lpse2d/core/integrator.py deleted file mode 100644 index d002bcc..0000000 --- a/adept/_lpse2d/core/integrator.py +++ /dev/null @@ -1,134 +0,0 @@ -from typing import Dict - -from jax import numpy as jnp -import numpy as np - -from adept._base_ import get_envelope -from adept._lpse2d.core import epw, laser - - -def get_laser_t_coeff(t: float, args: Dict) -> float: - """ - Calculate the laser time coefficient. - - This function computes the time-dependent coefficient for a laser pulse - based on the provided parameters. - - Args: - t (float): The time at which to evaluate the coefficient. - args (dict): A dictionary containing the parameters for the laser pulse. - It should have the following structure: - { - "drivers": { - "E0": { - "tc": float, # Center time of the pulse - "tw": float, # Full width at half maximum (FWHM) of the pulse - "tr": float # Rise time of the pulse - } - } - } - - Returns: - float: The time-dependent coefficient for the laser pulse. - """ - t_L = args["drivers"]["E0"]["tc"] - args["drivers"]["E0"]["tw"] * 0.5 - t_R = args["drivers"]["E0"]["tc"] + args["drivers"]["E0"]["tw"] * 0.5 - t_wL = args["drivers"]["E0"]["tr"] - t_wR = args["drivers"]["E0"]["tr"] - t_coeff = get_envelope(t_wL, t_wR, t_L, t_R, t) - return t_coeff - - -class SplitStep: - """ - This class contains the function that updates the state - - 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: - """ - - def __init__(self, cfg): - super().__init__() - 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 = 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): - new_y = {} - for k in y.keys(): - if k in self.complex_state_vars: - new_y[k] = y[k].view(jnp.complex128) - else: - new_y[k] = y[k].view(jnp.float64) - return new_y - - def pack_y(self, y, new_y): - for k in y.keys(): - y[k] = y[k].view(jnp.float64) - new_y[k] = new_y[k].view(jnp.float64) - - return y, new_y - - def light_split_step(self, t, y, args): - t_coeff = get_laser_t_coeff(t, args) - y["E0"] = self.boundary_envelope[..., None] * t_coeff * self.light.laser_update(t, y, args["E0"]) - - # if self.cfg["terms"]["light"]["update"]: - # y["E0"] = y["E0"] + self.dt * jnp.real(k1_E0) - - # y["E0"] = t_coeff * self.light.laser_update(t + 0.5 * self.dt, y, args["E0"]) - # if self.cfg["terms"]["light"]["update"]: - # y["E0"] = y["E0"] + 1j * self.dt * jnp.imag(k1_E0) - - return y - - def landau_damping(self, epw, vte_sq): - gammaLandauEpw = ( - np.sqrt(np.pi / 8) - * self.wp0**4 - * self.one_over_ksq**1.5 - / (vte_sq**1.5) - * jnp.exp(-self.wp0**2.0 * self.one_over_ksq / (2 * vte_sq)) - ) - - return jnp.fft.ifft2(jnp.fft.fft2(epw) * jnp.exp(-(gammaLandauEpw + self.nu_coll) * self.dt)) - - def __call__(self, t: float, y: Dict, args: Dict) -> Dict: - # unpack y into complex128 - 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) - - # landau and collisional damping - 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"] = new_y["epw"] * self.boundary_envelope - - # pack y into float64 - y, new_y = self.pack_y(y, new_y) - - return new_y diff --git a/adept/_lpse2d/helpers.py b/adept/_lpse2d/helpers.py index fb0c64a..1c3b04c 100644 --- a/adept/_lpse2d/helpers.py +++ b/adept/_lpse2d/helpers.py @@ -34,9 +34,7 @@ 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 / 2 - ) ## NB - the factor of 2 enables matching to the Short scaling + I0 = _Q(cfg["units"]["laser intensity"]).to("W/cm^2").value envelopeDensity = cfg["units"]["envelope density"] # Scaled constants @@ -185,26 +183,16 @@ def get_derived_quantities(cfg: Dict) -> Dict: cfg_grid["ymin"] = _Q(cfg_grid["ymin"]).to("um").value cfg_grid["dx"] = _Q(cfg_grid["dx"]).to("um").value - cfg_grid["nx"] = int(cfg_grid["xmax"] / cfg_grid["dx"]) + 1 - # cfg_grid["dx"] = cfg_grid["xmax"] / cfg_grid["nx"] - cfg_grid["dy"] = cfg_grid["dx"] # cfg_grid["ymax"] / cfg_grid["ny"] - cfg_grid["ny"] = int(cfg_grid["ymax"] / cfg_grid["dy"]) + 1 + cfg_grid["nx"] = int((cfg_grid["xmax"] - cfg_grid["xmin"]) / cfg_grid["dx"]) + 1 - # midpt = (cfg_grid["xmax"] + cfg_grid["xmin"]) / 2 + cfg_grid["dy"] = cfg_grid["dx"] + cfg_grid["ny"] = int((cfg_grid["ymax"] - cfg_grid["ymin"]) / cfg_grid["dy"]) + 1 - # max_density = cfg["density"]["val at center"] + (cfg["grid"]["xmax"] - midpt) / L - - # norm_n_max = np.abs(nmax / cfg["units"]["envelope density"] - 1) - # cfg_grid["dt"] = 0.05 * cfg_grid["dx"] cfg_grid["dt"] = _Q(cfg_grid["dt"]).to("ps").value cfg_grid["tmax"] = _Q(cfg_grid["tmax"]).to("ps").value cfg_grid["nt"] = int(cfg_grid["tmax"] / cfg_grid["dt"] + 1) cfg_grid["tmax"] = cfg_grid["dt"] * cfg_grid["nt"] - # if cfg_grid["nt"] > 1e6: - # cfg_grid["max_steps"] = int(1e6) - # print(r"Only running $10^6$ steps") - # else: cfg_grid["max_steps"] = cfg_grid["nt"] + 4 # change driver parameters to the right units @@ -308,7 +296,7 @@ def get_solver_quantities(cfg: Dict) -> Dict: ) cfg_grid["zero_mask"] = ( - np.where(cfg_grid["kx"] == 0, 0, 1)[:, None] * np.ones_like(cfg_grid["ky"])[None, :] + np.where(np.sqrt(cfg_grid["kx"][:, None] ** 2 + cfg_grid["ky"][None, :] ** 2) == 0, 0, 1) if cfg["terms"]["zero_mask"] else 1 ) diff --git a/configs/envelope-2d/tpd.yaml b/configs/envelope-2d/tpd.yaml index dc213ec..70dc2d2 100644 --- a/configs/envelope-2d/tpd.yaml +++ b/configs/envelope-2d/tpd.yaml @@ -1,7 +1,7 @@ density: basis: linear - gradient scale length: 200.0um - max: 0.3 + gradient scale length: 62.5um + max: 0.28 min: 0.2 noise: max: 1.0e-09 @@ -14,39 +14,41 @@ drivers: delta_omega_max: 0.015 num_colors: 1 envelope: - tw: 40ps + tw: 50ps tr: 0.25ps - tc: 20.5ps + tc: 20.05ps xr: 0.2um xw: 1000um xc: 50um yr: 0.2um yw: 1000um yc: 50um + params: + {} grid: boundary_abs_coeff: 1.0e4 - boundary_width: 1um - low_pass_filter: 0.3 - dt: 0.002ps - dx: 20nm - tmax: 25.0ps + boundary_width: 1.5um + low_pass_filter: 0.66 + dt: 0.0005ps + dx: 40nm + tmax: 1.5ps tmin: 0.0ns ymax: 3um ymin: -3um mlflow: experiment: tpd - run: 1.5e14 + run: test save: fields: t: - dt: 1ps - tmax: 25ps + dt: 0.1ps + tmax: 1.5ps tmin: 0ps x: - dx: 80nm + dx: 50nm y: - dy: 128nm + dy: 50nm solver: envelope-2d terms: epw: @@ -62,11 +64,12 @@ terms: noise: true tpd: true zero_mask: true + units: atomic number: 40 envelope density: 0.25 ionization state: 6 - laser intensity: 2.0e+14W/cm^2 + laser intensity: 1.2e+15W/cm^2 laser_wavelength: 351nm reference electron temperature: 2000.0eV reference ion temperature: 1000eV \ No newline at end of file