Source code for lambdapic.callback.laser

from typing import Optional
import numpy as np

from ..core.boundary.cpml import PMLXmin
from ..core.boundary.utils import get_pml
from ..core.utils.logger import logger
from ..core.fields import Fields
from ..core.patch.patch import Patch, Patch2D, Patch3D
from numba import njit, prange, typed
from scipy.constants import c, e, epsilon_0, m_e, mu_0, pi
from scipy.special import genlaguerre, factorial
from numpy.typing import NDArray
 
from ..simulation import Simulation, Simulation3D, Simulation2D
from ..core.utils.jit_spinner import jit_spinner

@jit_spinner
@njit(parallel=True, cache=True)
def _update_laser_bfields_2d(
    laserpos,
    ex, ey, ez,
    bx, by, bz,
    jx, jy, jz, 
    dx, dy, dt,
    iy_start, iy_end,
    ey_source: NDArray[np.float64], ez_source: NDArray[np.float64], 
):
    for iy in prange(iy_start, iy_end):
        bx[laserpos-1, iy] = bx[0, iy]
    for iy in prange(iy_start, iy_end):
        bz[laserpos-1, iy] = 1 / ((c*dt / dx + 1)*c) * (
            + 4 * ey_source[iy]
            + 2 * (ey[0, iy] + c * 0.5*(bz[0, iy] + bz[-1, iy]))
            - 2 * ey[laserpos, iy]
            + dt/epsilon_0 * jy[laserpos, iy]
            + (c*dt / dx - 1)*c * bz[laserpos, iy]
        )
        by[laserpos-1, iy] = 1 / ((c*dt / dx + 1)*c) * (
            - 4 * ez_source[iy]
            - 2 * (ez[0, iy] - c * 0.5*(by[0, iy] + by[-1, iy]))
            + 2 * ez[laserpos, iy]
            - (dt*c**2) * (bx[laserpos, iy] - bx[laserpos, iy-1])/dy
            - dt/epsilon_0 * jz[laserpos, iy]
            + (c*dt / dx - 1)*c * by[laserpos, iy]
        )

@jit_spinner
@njit(parallel=True, cache=True)
def _update_laser_bfields_3d(
    laserpos,
    ex, ey, ez,
    bx, by, bz,
    jx, jy, jz, 
    dx, dy, dz, dt,
    iy_start, iy_end, iz_start, iz_end,
    ey_source: NDArray[np.float64], ez_source: NDArray[np.float64], 
):
    for iy in prange(iy_start, iy_end):
        bx[laserpos-1, iy, :] = bx[0, iy, :]
    for iy in prange(iy_start, iy_end):
        for iz in range(iz_start, iz_end):
            bz[laserpos-1, iy, iz] = 1 / ((c*dt / dx + 1)*c) * (
                + 4 * ey_source[iy, iz]
                + 2 * (ey[0, iy, iz] + c * 0.5*(bz[0, iy, iz] + bz[-1, iy, iz]))
                - 2 * ey[laserpos, iy, iz]
                - (dt*c**2) * (bx[laserpos, iy, iz] - bx[laserpos, iy, iz-1])/dz
                + dt/epsilon_0 * jy[laserpos, iy, iz]
                + (c*dt / dx - 1)*c * bz[laserpos, iy, iz]
            )
            by[laserpos-1, iy, iz] = 1 / ((c*dt / dx + 1)*c) * (
                - 4 * ez_source[iy, iz]
                - 2 * (ez[0, iy, iz] - c * 0.5*(by[0, iy, iz] + by[-1, iy, iz]))
                + 2 * ez[laserpos, iy, iz]
                - (dt*c**2) * (bx[laserpos, iy, iz] - bx[laserpos, iy-1, iz])/dy
                - dt/epsilon_0 * jz[laserpos, iy, iz]
                + (c*dt / dx - 1)*c * by[laserpos, iy, iz]
            )

        
class Laser:
    def __init__(self) -> None:
        self.stage = "_laser"
        self.disabled = False
        self.side = "xmin"
        self.tstop = np.inf
        self.y0: float|None = None
        self.z0: float|None = None

    def _get_r(self, sim, patch: Patch) -> NDArray[np.float64]:
        """Calculate the radial distance from the center of the laser beam."""
        raise NotImplementedError
    
    def _get_phi(self, sim, patch: Patch) -> NDArray[np.float64]:
        """Calculate the azimuthal angle from the center of the laser beam."""
        raise NotImplementedError
        
    def _get_boundary_coordinates(self, sim: Simulation, patch: Patch):
        """Get the coordinates of the boundary. Centered to y0 and z0"""
        raise NotImplementedError
    
    def _update_bfields(self, laserpos: int, patch: Patch, ey_source: NDArray[np.float64], ez_source: NDArray[np.float64], dt: float):
        raise NotImplementedError
    
    def _calculate_bound_fields(self, sim: Simulation, patch: Patch) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
        """Calculate ey_source and ez_source for this laser."""
        raise NotImplementedError
    
    def __call__(self, sim: Simulation):
        if self.disabled:
            return

        time = sim.time
        # Stop injecting after twice the pulse duration for smooth falloff
        if c*time >= self.tstop:
            self.disabled = True
            return

        if self.side == "xmin":
            ipatch_x = 0
            laserpos = sim.cpml_thickness + 2
            patches = list(filter(lambda p: p.ipatch_x == ipatch_x, sim.patches))
            n_pmlxmin = sum(isinstance(pml, PMLXmin) for p in patches for pml in p.pml_boundary)
            if n_pmlxmin < len(patches):
                logger.info("Disabling laser for lacking of PML. Maybe due to MovingWindow starts.")
                self.disabled = True
                return
        else:
            raise ValueError("Invalid side: only 'xmin' is supported.")

        # Inject the laser from the left boundary
        for p in sim.patches:
            # Only inject from the leftmost patch
            if p.ipatch_x == ipatch_x:
                ey_source, ez_source = self._calculate_bound_fields(sim, p)
                if ey_source is not None:
                    self._update_bfields(laserpos, p, ey_source, ez_source, dt=sim.dt)

    def __add__(self, other):
        """Add two lasers together."""
        if self.side != other.side:
            raise TypeError(f"Cannot add lasers from different sides: {self.side} and {other.side}")
        if not isinstance(other, Laser):
            raise TypeError(f"Cannot add Laser with {type(other)}")
        if isinstance(self, Laser2D) and isinstance(other, Laser2D):
            return _CombinedLaser2D(self, other)
        elif isinstance(self, Laser3D) and isinstance(other, Laser3D):
            return _CombinedLaser3D(self, other)
        else:
            raise TypeError("Cannot add 2D and 3D laser")


class Laser2D(Laser):
    def _get_r(self, sim: Simulation, patch: Patch2D) -> NDArray[np.float64]:
        f = patch.fields
        y0 = self.y0 or sim.Ly/2
        r = abs(f.yaxis[0, :] - sim.dy/2 - y0)
        return r
    
    def _get_phi(self, sim: Simulation2D, patch: Patch2D) -> NDArray[np.float64]:
        f = patch.fields
        y0 = self.y0 or sim.Ly/2
        phi = np.arctan2(0.0, f.yaxis[0, :] - sim.dy/2 - y0)
        return phi
        
    def _get_boundary_coordinates(self, sim: Simulation2D, patch: Patch2D):
        f = patch.fields
        y0 = self.y0 or sim.Ly/2
        y = f.yaxis[0, :] - sim.dy/2 - y0
        z = 0.0
        r = abs(y)
        return y, z, r

    def _update_bfields(self, laserpos: int, patch: Patch2D, ey_source: NDArray[np.float64], ez_source: NDArray[np.float64], dt: float):
        f = patch.fields
        iy_start, iy_end = 0, f.ny
        if (pml := get_pml(patch.pml_boundary, "ymin")) is not None:
            iy_start = pml.thickness
        if (pml := get_pml(patch.pml_boundary, "ymax")) is not None:
            iy_end = f.ny - pml.thickness
        _update_laser_bfields_2d(
            laserpos,
            f.ex, f.ey, f.ez,
            f.bx, f.by, f.bz,
            f.jx, f.jy, f.jz,
            f.dx, f.dy, dt,
            iy_start, iy_end,
            ey_source,
            ez_source
        )

class Laser3D(Laser):
    def _get_r(self, sim: Simulation3D, patch: Patch3D) -> NDArray[np.float64]:
        f = patch.fields
        y0 = self.y0 or sim.Ly/2
        z0 = self.z0 or sim.Lz/2
        r = ((f.yaxis[0, :, :] - sim.dy/2 - y0)**2 + 
             (f.zaxis[0, :, :] - sim.dz/2 - z0)**2)**0.5
        return r
    
    def _get_phi(self, sim: Simulation3D, patch: Patch3D) -> NDArray[np.float64]:
        f = patch.fields
        y0 = self.y0 or sim.Ly/2
        z0 = self.z0 or sim.Lz/2
        phi = np.arctan2(f.zaxis[0, :, :] - sim.dz/2 - z0,
                         f.yaxis[0, :, :] - sim.dy/2 - y0)
        return phi
        
    def _get_boundary_coordinates(self, sim: Simulation3D, patch: Patch3D):
        f = patch.fields
        y0 = self.y0 or sim.Ly/2
        z0 = self.z0 or sim.Lz/2
        
        y = f.yaxis[0, :, :] - sim.dy/2 - y0
        z = f.zaxis[0, :, :] - sim.dz/2 - z0
        r = np.sqrt(y**2 + z**2)
        return y, z, r
    
    def _update_bfields(self, laserpos: int, patch: Patch3D, ey_source: NDArray[np.float64], ez_source: NDArray[np.float64], dt: float):
        f = patch.fields
        iy_start, iy_end = 0, f.ny
        iz_start, iz_end = 0, f.nz
        if (pml := get_pml(patch.pml_boundary, "ymin")) is not None:
            iy_start = pml.thickness
        if (pml := get_pml(patch.pml_boundary, "ymax")) is not None:
            iy_end = f.ny - pml.thickness
        if (pml := get_pml(patch.pml_boundary, "zmin")) is not None:
            iz_start = pml.thickness
        if (pml := get_pml(patch.pml_boundary, "zmax")) is not None:
            iz_end = f.nz - pml.thickness
        _update_laser_bfields_3d(
            laserpos,
            f.ex, f.ey, f.ez,
            f.bx, f.by, f.bz,
            f.jx, f.jy, f.jz,
            f.dx, f.dy, f.dz, dt,
            iy_start, iy_end, iz_start, iz_end,
            ey_source,
            ez_source
        )

class _CombinedLaser(Laser):
    """A laser that combines two laser sources by adding their fields."""

    def __init__(self, laser1: Laser, laser2: Laser):
        super().__init__()
        self.laser1 = laser1
        self.laser2 = laser2

        self.side = self.laser1.side

        self.tstop = max(laser1.tstop, laser2.tstop)

    def _calculate_bound_fields(self, sim: Simulation, patch: Patch) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
        # Calculate source fields from both lasers and sum them
        ey1, ez1 = self.laser1._calculate_bound_fields(sim, patch)
        ey2, ez2 = self.laser2._calculate_bound_fields(sim, patch)

        if ey1 is None and ey2 is None:
            return None, None
        elif ey1 is None:
            return ey2, ez2
        elif ey2 is None:
            return ey1, ez1
        else:
            return ey1 + ey2, ez1 + ez2
        

class _CombinedLaser2D(Laser2D,_CombinedLaser): 
    ...
class _CombinedLaser3D(Laser3D,_CombinedLaser): 
    ...
[docs] class SimpleLaser(Laser): """ A simple laser pulse implementation with basic spatial and temporal profiles. This class provides a straightforward way to inject a laser pulse into the simulation from the left boundary with a Gaussian transverse profile and a smooth temporal envelope. Use SimpleLaser2D or SimpleLaser3D for 2D and 3D simulations, respectively. Note: This is a simplified laser implementation suitable for basic simulations. For more accurate physics including proper beam evolution, wavefront curvature, and Gouy phase, use the GaussianLaser class instead. """ def __init__(self, a0: float, w0: float, ctau: float, y0: float|None=None, z0: float|None=None, angle_y: float=0, angle_z: float=0, tstop: float|None=None, pol_angle: float = 0.0, cep: float = 0.0, l0: float=0.8e-6, side="xmin"): """ Parameters: sim: Simulation object that this laser will be injected into a0: Normalized vector potential amplitude w0: Laser waist size ctau: Pulse duration (c*tau) y0: y position of the laser center (default: Ly/2) z0: z position of the laser center (default: Lz/2). No effect for 2D laser. angle_y: incident angle with boundary normal in y direction (default: 0) angle_z: NOT IMPLEMENTED. Must be 0. tstop: Time at which the laser pulse should stop (default: 2*ctau) pol_angle: Polarization angle in radians (default: 0.0 for y-polarization) cep: Carrier envelope phase (default: 0.0) l0: Laser wavelength (default: 800nm) Raises: ValueError: If parameters are invalid (negative or zero values) """ super().__init__() # Parameter validation if any(p <= 0 for p in [a0, l0, w0, ctau]): raise ValueError("All parameters (a0, l0, w0, ctau) must be positive") if side not in ["xmin"]: raise NotImplementedError("Invalid side: only 'xmin' is supported.") if abs(angle_y) >= pi/2: raise ValueError("Angle_y must be in range (-pi/2, pi/2)") if angle_z != 0: raise NotImplementedError("Angle_z is not implemented") self.a0 = a0 self.l0 = l0 self.omega0 = 2 * pi * c / l0 self.w0 = w0 self.ctau = ctau self.y0 = y0 self.z0 = z0 self.angle_y = angle_y self.angle_z = angle_z self.tstop = 2*ctau if tstop is None else c*tstop self.E0 = a0 * m_e * c * self.omega0 / e self.pol_angle = pol_angle self.cep = cep self.side = side # Wavevector components for incident angles self.k0 = self.omega0 / c self.ky = self.k0 * np.sin(self.angle_y) self.kz = 0 def _calculate_bound_fields(self, sim: Simulation, patch: Patch) -> tuple[NDArray[np.float64], NDArray[np.float64]]: """Calculate ey_source and ez_source for this laser.""" time = sim.time if c*time >= self.tstop: return None, None if self.side == "xmin": y, z, r = self._get_boundary_coordinates(sim, patch) r_rot = np.sqrt((y/np.cos(self.angle_y))**2 + z**2) transverse_phase = -(self.ky * y + self.kz * z) else: raise NotImplementedError("Invalid side: only 'xmin' is supported.") # Calculate temporal profile (sin² envelope for smooth turn-on/off) t_rot = c*time - y * np.sin(self.angle_y) tprof = np.sin(t_rot/(2*self.ctau)*pi)**2 * (t_rot < 2*self.ctau) # Calculate base field amplitude with: # - Gaussian transverse profile: exp(-r²/w0²) # - Temporal oscillation: sin(ω₀t + transverse_phase) # - Smooth temporal envelope: tprof efield = self.E0 * np.exp(-r_rot**2/self.w0**2) * np.sin(self.omega0 * time + self.cep + transverse_phase) * tprof ey_source = efield * np.cos(self.pol_angle) * np.cos(self.angle_y) ez_source = efield * np.sin(self.pol_angle) * np.cos(self.angle_z) return ey_source, ez_source
[docs] class SimpleLaser2D(Laser2D, SimpleLaser): ...
[docs] class SimpleLaser3D(Laser3D, SimpleLaser): ...
[docs] class GaussianLaser(Laser): r""" Implementation of a proper Gaussian laser beam with full physics including: - Gaussian temporal and spatial profiles - Proper beam waist evolution (:math:`w(z) = w_0\sqrt{1 + (z/z_R)^2}`) - Gouy phase (:math:`tan^{-1}(z/z_R)`) - Wavefront curvature (:math:`R(z) = z(1 + (z_R/z)^2)`) - Correct phase evolution including propagation and curvature terms - Laguerre-Gaussian (LG) beam modes via ``l`` and ``p`` parameters Use GaussianLaser2D or GaussianLaser3D for 2D and 3D simulations, respectively. Note: This implementation provides more accurate physics than SimpleLaser, including proper beam evolution and phase effects. Use this for realistic simulations where these effects matter. """ def __init__(self, a0: float, l0: float, w0: float, ctau: float, x0: float|None=None, y0: float|None=None, z0: float|None=None, tstop: float|None=None, pol_angle: float = 0.0, cep: float = 0.0, focus_position: float = 0.0, side: str = "xmin", l: int=0, p: int=0, # Laguerre Gaussian laser index ): """ Parameters: a0: Normalized vector potential amplitude l0: Laser wavelength w0: Waist size at focus ctau: Pulse duration (c*tau) x0: Pulse center position from boundary (default: 3*ctau) y0: y position of the laser center (default: 0) z0: z position of the laser center (default: 0). No effect for 2D laser. tstop: Time to stop injection (default: 6*ctau) pol_angle: Polarization angle in radians (default: 0.0 for y-polarization) cep: Carrier envelope phase (default: 0.0) focus_position: Position of laser focus relative to boundary (default: 0.0) side: Injection boundary ('xmin' or 'xmax') (default: 'xmin') l: Azimuthal index of Laguerre-Gaussian laser (default: 0) p: Number of radial nodes of Laguerre-Gaussian laser (default: 0) """ super().__init__() # Parameter validation if any(par <= 0 for par in [a0, l0, w0, ctau]): raise ValueError("All parameters (a0, l0, w0, ctau) must be positive") if side not in ["xmin"]: raise ValueError("Invalid side: only 'xmin' is implemented.") if not isinstance(p, int) or p < 0: raise ValueError("Number of radial nodes p must be a non-negative integer") if not isinstance(l, int): raise ValueError("Azimuthal index l must be an integer") self.a0 = a0 self.l0 = l0 self.omega0 = 2 * pi * c / l0 self.k0 = self.omega0 / c self.w0 = w0 self.ctau = ctau self.x0 = 3*ctau if x0 is None else x0 self.y0 = y0 self.z0 = z0 self.tstop = 6*ctau if tstop is None else c*tstop self.E0 = a0 * m_e * c * self.omega0 / e self.pol_angle = pol_angle self.cep = cep self.focus_position = focus_position self.side = side # Derived parameters self.zR = pi * w0**2 / l0 # Rayleigh length # LG parameters self._is_lg = False self.l = l self.p = p if l != 0 or p > 0: self._is_lg = True self.lg_norm = np.sqrt(2 * factorial(p) / (pi * factorial(p + abs(l)))) self.lg_norm /= np.sqrt(2/pi) # let GS (p=0, l=0) norm=1 self.laguerre = genlaguerre(self.p, abs(self.l)) def _gaussian_beam_params(self, z): """Calculate Gaussian beam parameters at position z""" # Normalized distance from focus z = z - self.focus_position # Beam radius w = self.w0 * np.sqrt(1 + (z/self.zR)**2) # Radius of curvature R = z * (1 + (self.zR/z)**2) if abs(z) > 1e-10 else np.inf # Gouy phase psi = np.arctan(z/self.zR) return w, R, psi def _calculate_bound_fields(self, sim: Simulation, patch: Patch) -> tuple[NDArray[np.float64], NDArray[np.float64]]: """Calculate ey_source and ez_source for this laser.""" time = sim.time if c*time >= self.tstop: return None, None # Temporal envelope (Gaussian) tprof = np.exp(-(c*time - self.x0)**2 / self.ctau**2) # Calculate boundary parameters x_rel = sim.cpml_thickness * sim.dx boundary_w, boundary_R, boundary_psi = self._gaussian_beam_params(x_rel) # r is 2D or 3D depending on the simulation dimension r = self._get_r(sim, patch) # lg mode if self._is_lg: phi = self._get_phi(sim, patch) amp_lg = self.lg_norm * (np.sqrt(2)*r/boundary_w)**abs(self.l) \ * self.laguerre((np.sqrt(2)*r/boundary_w)**2) phase_lg = self.l * phi else: amp_lg = 1.0 phase_lg = 0.0 # Calculate amplitude and phase amp = self.E0 * (self.w0/boundary_w) * np.exp(-r**2/boundary_w**2) * amp_lg phase_curv = self.k0 * r**2/(2*boundary_R) phase = (self.omega0 * time + self.cep - # Oscillation self.k0 * x_rel - # Propagation phase_curv - # Curvature (2*self.p+abs(self.l)+1) * boundary_psi # Gouy phase - phase_lg) # Vortex # Set fields based on polarization efield = amp * np.sin(phase) * tprof ey_source = efield * np.cos(self.pol_angle) ez_source = efield * np.sin(self.pol_angle) return ey_source, ez_source
[docs] class GaussianLaser2D(Laser2D, GaussianLaser): ...
[docs] class GaussianLaser3D(Laser3D, GaussianLaser): ...