feat(vlasov1d): point source EM driver + driver normalization work for SRS

adept/_vlasov1d/datamodel.py

@@ -1,3 +1,5 @@
+from typing import Literal
+
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 from adept.functions import EnvelopeConfig, SpaceTimeEnvelopeConfig
@@ -107,6 +109,7 @@ def check_w_or_k(self) -> "AKWDriverConfig":
 class EMDriverConfig(BaseModel):
     params: IntensityWavelengthDriverConfig | AKWDriverConfig
     envelope: SpaceTimeEnvelopeConfig
+    source_type: Literal["extended", "point"] = "extended"
 
 
 class EMDriverSetConfig(BaseModel):

adept/_vlasov1d/grid.py

@@ -56,7 +56,7 @@ def __init__(
         # Override dt for EM wave stability if needed
         if should_override_dt_for_em_waves:
             c_light = 1.0 / beta
-            self.dt = float(0.95 * self.dx / c_light)
+            self.dt = min(dt_requested, float(0.95 * self.dx / c_light))
         else:
             self.dt = dt_requested
 

adept/_vlasov1d/helpers.py

@@ -210,14 +210,16 @@ def post_process(result: Solution, cfg: dict, td: str, args: dict):
                     field_name = nm.split("-", 1)[1] if "-" in nm else nm
 
                     # Spacetime plot
-                    fld.plot()
-                    plt.savefig(os.path.join(species_dir, f"spacetime_{field_name}.png"), bbox_inches="tight")
+                    fld.plot(figsize=(12, 8))
+                    plt.savefig(os.path.join(species_dir, f"spacetime_{field_name}.png"), bbox_inches="tight", dpi=150)
                     plt.close()
 
                     # Log plot
-                    np.log10(np.abs(fld)).plot()
+                    np.log10(np.abs(fld)).plot(figsize=(12, 8))
                     plt.savefig(
-                        os.path.join(species_dir, "logplots", f"spacetime_log_{field_name}.png"), bbox_inches="tight"
+                        os.path.join(species_dir, "logplots", f"spacetime_log_{field_name}.png"),
+                        bbox_inches="tight",
+                        dpi=150,
                     )
                     plt.close()
 

adept/_vlasov1d/simulation.py

@@ -35,6 +35,7 @@ class EMDriver(eqx.Module):
     w0: float
     dw0: float
     envelope: SpaceTimeEnvelopeFunction
+    is_point_source: bool = False
 
     @staticmethod
     def from_config(cfg: EMDriverConfig, norm: PlasmaNormalization | None = None) -> "EMDriver":
@@ -52,7 +53,8 @@ def from_config(cfg: EMDriverConfig, norm: PlasmaNormalization | None = None) ->
                 else:
                     k0, w0 = params.k0, params.w0
 
-                return EMDriver(params.a0, k0, w0, params.dw0, envelope)
+                is_point = cfg.source_type == "point"
+                return EMDriver(params.a0, k0, w0, params.dw0, envelope, is_point_source=is_point)
 
             case IntensityWavelengthDriverConfig(intensity=intensity, wavelength=wavelength, leftgoing=leftgoing):
                 intensity = UREG.Quantity(intensity).to("W/m^2")
@@ -79,7 +81,8 @@ def from_config(cfg: EMDriverConfig, norm: PlasmaNormalization | None = None) ->
 
                 dw0 = 0.0  # ???
 
-                return EMDriver(a0, k0, w0, dw0, envelope)
+                is_point = cfg.source_type == "point"
+                return EMDriver(a0, k0, w0, dw0, envelope, is_point_source=is_point)
 
 
 class EMDriverSet(eqx.Module):

adept/_vlasov1d/solvers/pushers/field.py

@@ -8,29 +8,83 @@
 from adept._vlasov1d.simulation import EMDriver
 
 
-class Driver:
+class LongitudinalElectricFieldDriver:
+    """Computes normalized E_tilde = ω * a0 * sin(kx - ωt) from EM drivers."""
+
     def __init__(self, xax, drivers: list[EMDriver]):
         self.xax = xax
         self.drivers = drivers
 
-    def get_this_pulse(self, this_pulse: EMDriver, current_time: jnp.float64):
-        kk = this_pulse.k0
-        ww = this_pulse.w0
-        dw = this_pulse.dw0
-
-        factor = this_pulse.envelope(self.xax, current_time)
-
-        return factor * jnp.abs(kk) * this_pulse.a0 * jnp.sin(kk * self.xax - (ww + dw) * current_time)
+    def _single_driver_field(self, driver: EMDriver, current_time):
+        kk = driver.k0
+        ww = driver.w0
+        dw = driver.dw0
+        factor = driver.envelope(self.xax, current_time)
+        return factor * (ww + dw) * driver.a0 * jnp.sin(kk * self.xax - (ww + dw) * current_time)
 
     def __call__(self, t, args):
         total_de = jnp.zeros_like(self.xax)
-
         for pulse in self.drivers:
-            total_de += self.get_this_pulse(pulse, t)
-
+            total_de += self._single_driver_field(pulse, t)
         return total_de
 
 
+class TransverseCurrentSourceDriver:
+    """Computes source term for the transverse wave equation.
+
+    For extended sources (default): S = -ω² * a0 * envelope(x,t) * sin(kx - ωt)
+    For point sources: S = (F0/dx) * time_envelope(t) * δ_i0 * sin(ωt)
+
+    For point sources, the amplitude scaling uses vacuum dispersion. For a 1D wave
+    equation with point source S = F0 * delta(x - x0) * sin(wt), the Green's function
+    gives outgoing waves with amplitude F0 / (2*k*c^2). To get amplitude a0, we need
+    F0 = 2*k*c^2*a0. Using vacuum k = w/c, this gives F0 = 2*w*c*a0. In plasma
+    (k_plasma < k_vac), the actual amplitude will be slightly larger than a0 by a
+    factor k_vac / k_plasma.
+
+    Note: point sources radiate equally in both directions. Place the source near a
+    boundary with absorbing BCs to get a unidirectional wave.
+    """
+
+    def __init__(self, xax, drivers: list[EMDriver], c: float = 0.0):
+        self.xax = xax
+        self.drivers = drivers
+        dx = float(xax[1] - xax[0])
+
+        self.point_source_masks = []
+        self.point_source_scales = []
+        for driver in drivers:
+            if driver.is_point_source:
+                center = driver.envelope.space_envelope.center
+                i0 = jnp.argmin(jnp.abs(xax - center))
+                mask = jnp.zeros_like(xax).at[i0].set(1.0)
+                w_total = driver.w0 + driver.dw0
+                F0 = 2.0 * w_total * c * driver.a0
+                self.point_source_masks.append(mask)
+                self.point_source_scales.append(F0 / dx)
+            else:
+                self.point_source_masks.append(None)
+                self.point_source_scales.append(None)
+
+    def _single_driver_source(self, driver: EMDriver, mask, scale, current_time):
+        ww = driver.w0
+        dw = driver.dw0
+        w_total = ww + dw
+        if driver.is_point_source:
+            time_env = driver.envelope.time_envelope(current_time)
+            return scale * time_env * mask * jnp.sin(w_total * current_time)
+        else:
+            kk = driver.k0
+            factor = driver.envelope(self.xax, current_time)
+            return -factor * w_total**2 * driver.a0 * jnp.sin(kk * self.xax - w_total * current_time)
+
+    def __call__(self, t, args):
+        total = jnp.zeros_like(self.xax)
+        for driver, mask, scale in zip(self.drivers, self.point_source_masks, self.point_source_scales, strict=True):
+            total += self._single_driver_source(driver, mask, scale, t)
+        return total
+
+
 class WaveSolver:
     def __init__(self, c: jnp.float64, dx: jnp.float64, dt: jnp.float64):
         super().__init__()

adept/_vlasov1d/solvers/pushers/vlasov.py

@@ -60,23 +60,34 @@ def __init__(self, species_grids, species_params, parallel=False):
         if parallel:
             self.mesh = Mesh(np.array(jax.devices()), ("device",))
 
-    def push(self, f_dict, e, dt):
+    def push(self, f_dict, e, pond, dt):
         result = {}
         for species_name, f in f_dict.items():
             kv_real = self.species_grids[species_name]["kvr"]
-            qm = self.species_params[species_name]["charge_to_mass"]
+            q = self.species_params[species_name]["charge"]
+            m = self.species_params[species_name]["mass"]
+            # force = q*E + (q²/m)*pond where pond = -(1/2)*grad(a²)
+            # accel = force / m
+            force = q * e + (q**2 / m) * pond
+            accel = force / m
             result[species_name] = jnp.real(
-                jnp.fft.irfft(jnp.exp(-1j * kv_real[None, :] * dt * qm * e[:, None]) * jnp.fft.rfft(f, axis=1), axis=1)
+                jnp.fft.irfft(
+                    jnp.exp(-1j * kv_real[None, :] * dt * accel[:, None]) * jnp.fft.rfft(f, axis=1),
+                    axis=1,
+                )
             )
         return result
 
-    def __call__(self, f_dict, e, dt):
+    def __call__(self, f_dict, e, pond, dt):
         if self.parallel:
             return shard_map(
-                self.push, mesh=self.mesh, in_specs=(P("device", None), P("device"), P()), out_specs=P("device", None)
-            )(f_dict, e, dt)
+                self.push,
+                mesh=self.mesh,
+                in_specs=(P("device", None), P("device"), P("device"), P()),
+                out_specs=P("device", None),
+            )(f_dict, e, pond, dt)
         else:
-            return self.push(f_dict, e, dt)
+            return self.push(f_dict, e, pond, dt)
 
 
 class VelocityCubicSpline:
@@ -88,24 +99,30 @@ def __init__(self, species_grids, species_params, parallel=False):
         if self.parallel:
             self.mesh = Mesh(np.array(jax.devices()), ("device",))
 
-    def push(self, f_dict, e, dt):
+    def push(self, f_dict, e, pond, dt):
         result = {}
         for species_name, f in f_dict.items():
             v = self.species_grids[species_name]["v"]
-            qm = self.species_params[species_name]["charge_to_mass"]
+            q = self.species_params[species_name]["charge"]
+            m = self.species_params[species_name]["mass"]
             nx = f.shape[0]
             v_repeated = jnp.repeat(v[None, :], repeats=nx, axis=0)
-            vq = v_repeated - qm * e[:, None] * dt
+            force = q * e + (q**2 / m) * pond
+            accel = force / m
+            vq = v_repeated - accel[:, None] * dt
             result[species_name] = self.interp(xq=vq, x=v_repeated, f=f)
         return result
 
-    def __call__(self, f_dict, e, dt):
+    def __call__(self, f_dict, e, pond, dt):
         if self.parallel:
             return shard_map(
-                self.push, mesh=self.mesh, in_specs=(P("device", None), P("device"), P()), out_specs=P("device", None)
-            )(f_dict, e, dt)
+                self.push,
+                mesh=self.mesh,
+                in_specs=(P("device", None), P("device"), P("device"), P()),
+                out_specs=P("device", None),
+            )(f_dict, e, pond, dt)
         else:
-            return self.push(f_dict, e, dt)
+            return self.push(f_dict, e, pond, dt)
 
 
 class HouLiFilter:

adept/_vlasov1d/solvers/vector_field.py

@@ -85,7 +85,7 @@ def __call__(self, f_dict: dict, a: Array, dex_array: Array, prev_ex: Array) ->
         else:
             f_for_field = f_after_v
         pond, e = self.field_solve(f_dict=f_for_field, a=a, prev_ex=prev_ex, dt=self.dt)
-        f_dict = self.edfdv(f_after_v, e=pond + e + dex_array[0], dt=self.dt)
+        f_dict = self.edfdv(f_after_v, e=e + dex_array[0], pond=pond, dt=self.dt)
 
         return e, f_dict
 
@@ -149,39 +149,33 @@ def __call__(self, f_dict: dict, a: Array, dex_array: Array, prev_ex: Array) ->
         Returns:
             Tuple of (electric_field[nx], updated_f_dict)
         """
-        ponderomotive_force, self_consistent_ex = self.field_solve(f_dict=f_dict, a=a, prev_ex=None, dt=None)
-        force = ponderomotive_force + dex_array[0] + self_consistent_ex
-        f_dict = self.edfdv(f_dict, e=force, dt=self.D1 * self.dt)
+        pond, self_consistent_ex = self.field_solve(f_dict=f_dict, a=a, prev_ex=None, dt=None)
+        f_dict = self.edfdv(f_dict, e=dex_array[0] + self_consistent_ex, pond=pond, dt=self.D1 * self.dt)
 
         f_dict = self.vdfdx(f_dict, dt=self.a1 * self.dt)
-        ponderomotive_force, self_consistent_ex = self.field_solve(f_dict=f_dict, a=a, prev_ex=None, dt=None)
-        force = ponderomotive_force + dex_array[1] + self_consistent_ex
+        pond, self_consistent_ex = self.field_solve(f_dict=f_dict, a=a, prev_ex=None, dt=None)
 
-        f_dict = self.edfdv(f_dict, e=force, dt=self.D2 * self.dt)
+        f_dict = self.edfdv(f_dict, e=dex_array[1] + self_consistent_ex, pond=pond, dt=self.D2 * self.dt)
 
         f_dict = self.vdfdx(f_dict, dt=self.a2 * self.dt)
-        ponderomotive_force, self_consistent_ex = self.field_solve(f_dict=f_dict, a=a, prev_ex=None, dt=None)
-        force = ponderomotive_force + dex_array[2] + self_consistent_ex
+        pond, self_consistent_ex = self.field_solve(f_dict=f_dict, a=a, prev_ex=None, dt=None)
 
-        f_dict = self.edfdv(f_dict, e=force, dt=self.D3 * self.dt)
+        f_dict = self.edfdv(f_dict, e=dex_array[2] + self_consistent_ex, pond=pond, dt=self.D3 * self.dt)
 
         f_dict = self.vdfdx(f_dict, dt=self.a3 * self.dt)
-        ponderomotive_force, self_consistent_ex = self.field_solve(f_dict=f_dict, a=a, prev_ex=None, dt=None)
-        force = ponderomotive_force + dex_array[3] + self_consistent_ex
+        pond, self_consistent_ex = self.field_solve(f_dict=f_dict, a=a, prev_ex=None, dt=None)
 
-        f_dict = self.edfdv(f_dict, e=force, dt=self.D3 * self.dt)
+        f_dict = self.edfdv(f_dict, e=dex_array[3] + self_consistent_ex, pond=pond, dt=self.D3 * self.dt)
 
         f_dict = self.vdfdx(f_dict, dt=self.a2 * self.dt)
-        ponderomotive_force, self_consistent_ex = self.field_solve(f_dict=f_dict, a=a, prev_ex=None, dt=None)
-        force = ponderomotive_force + dex_array[4] + self_consistent_ex
+        pond, self_consistent_ex = self.field_solve(f_dict=f_dict, a=a, prev_ex=None, dt=None)
 
-        f_dict = self.edfdv(f_dict, e=force, dt=self.D2 * self.dt)
+        f_dict = self.edfdv(f_dict, e=dex_array[4] + self_consistent_ex, pond=pond, dt=self.D2 * self.dt)
 
         f_dict = self.vdfdx(f_dict, dt=self.a1 * self.dt)
-        ponderomotive_force, self_consistent_ex = self.field_solve(f_dict=f_dict, a=a, prev_ex=None, dt=None)
-        force = ponderomotive_force + dex_array[5] + self_consistent_ex
+        pond, self_consistent_ex = self.field_solve(f_dict=f_dict, a=a, prev_ex=None, dt=None)
 
-        f_dict = self.edfdv(f_dict, e=force, dt=self.D1 * self.dt)
+        f_dict = self.edfdv(f_dict, e=dex_array[5] + self_consistent_ex, pond=pond, dt=self.D1 * self.dt)
 
         return self_consistent_ex, f_dict
 
@@ -279,11 +273,12 @@ def __init__(
         self.vpfp = VlasovPoissonFokkerPlanck(cfg, grid)
         # beta = 1/c_norm is still read from cfg["grid"] for now
         beta = cfg["grid"]["beta"]
-        self.wave_solver = field.WaveSolver(c=1.0 / beta, dx=grid.dx, dt=grid.dt)
+        c = 1.0 / beta
+        self.wave_solver = field.WaveSolver(c=c, dx=grid.dx, dt=grid.dt)
 
         self.dt = grid.dt
-        self.ey_driver = field.Driver(grid.x_a, drivers=drivers.ey)
-        self.ex_driver = field.Driver(grid.x, drivers=drivers.ex)
+        self.ey_driver = field.TransverseCurrentSourceDriver(grid.x_a, drivers=drivers.ey, c=c)
+        self.ex_driver = field.LongitudinalElectricFieldDriver(grid.x, drivers=drivers.ex)
 
     def compute_electron_charge_density(self, f_dict):
         """Compute charge density from the electron distribution function."""