import hashlib
from abc import ABC, abstractmethod
from typing import Tuple
import numpy as np
import numpy.typing as npt
from numba import njit
from numpy._typing import NDArray
from scipy.integrate import cumulative_simpson
from scipy.interpolate import CubicSpline
from scipy.special import sici
from .cosN_envelope_analytical import cosN_int_A, cosN_int_A2, cosN_int_E2
from .utils import AtomicUnits, FourierTransform, get_cummulative_integral_spline
# import matplotlib.pyplot as plt
def _hash_grid(arr: np.ndarray):
h = hashlib.sha1(arr.shape[0].to_bytes(8, "little"))
h.update(arr.dtype.str.encode())
h.update(arr.tobytes())
return h.hexdigest()
[docs]
class Pulse(ABC):
def __init__(self):
# Structure: {method_name: {grid_hash: result_array}}
self._cache = {}
def _get_from_cache(self, method: str, t: np.ndarray):
key = _hash_grid(np.ascontiguousarray(t))
if method not in self._cache:
self._cache[method] = {}
if key in self._cache[method]:
return self._cache[method][key], key
return None, key
def _update_cache(
self, method: str, key: str, value: float | npt.NDArray[np.float64]
):
# print("updated cache")
self._cache[method][key] = value
[docs]
def get_electric_field(
self, t: float | npt.NDArray[np.float64]
) -> float | npt.NDArray[np.float64]:
t = np.ascontiguousarray(t)
val, key = self._get_from_cache("e_field", t)
if val is not None:
return val
val = self._e_field(t)
self._update_cache("e_field", key, val)
return val
[docs]
def get_vector_potential(
self, t: float | npt.NDArray[np.float64]
) -> float | npt.NDArray[np.float64]:
t = np.ascontiguousarray(t)
val, key = self._get_from_cache("vpot", t)
if val is not None:
return val
val = self._vector_potential(t)
self._update_cache("vpot", key, val)
return val
[docs]
def get_cummulative_vector_potential(
self, t: float | npt.NDArray[np.float64]
) -> float | npt.NDArray[np.float64]:
t = np.ascontiguousarray(t)
val, key = self._get_from_cache("intA", t)
if val is not None:
return val
val = self._cummulative_vector_potential(t)
self._update_cache("intA", key, val)
return val
[docs]
def get_cummulative_vector_potential_squared(
self, t: float | npt.NDArray[np.float64]
) -> float | npt.NDArray[np.float64]:
t = np.ascontiguousarray(t)
val, key = self._get_from_cache("intA2", t)
if val is not None:
return val
val = self._cummulative_vector_potential_squared(t)
self._update_cache("intA2", key, val)
return val
[docs]
def get_cummulative_electric_field_squared(
self, t: float | npt.NDArray[np.float64]
) -> float | npt.NDArray[np.float64]:
t = np.ascontiguousarray(t)
val, key = self._get_from_cache("intE2", t)
if val is not None:
return val
val = self._cummulative_electric_field_squared(t)
self._update_cache("intE2", key, val)
return val
[docs]
def get_tgrid(self, dt: float = 0.25) -> npt.NDArray[np.float64]:
"""Get time grid for the pulse.
Args:
dt: Time step in atomic units
Returns:
Time grid as a numpy array
"""
t_min, t_max = self.get_time_interval()
return np.arange(np.floor(t_min), np.ceil(t_max) + dt, dt)
[docs]
def get_fft(
self,
omega: None | npt.NDArray[np.float64] = None,
center_of_pulse: float = None,
) -> (
npt.NDArray[np.float64]
| Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]
):
"""Get the Fourier transform of the pulse.
Returns:
If ``omega`` is None, a tuple ``(frequencies, amplitudes)`` in
atomic units. If ``omega`` is an array, the amplitude at the
given frequency(ies) in atomic units.
"""
tgrid = self.get_tgrid(dt=1)
vpot = self.get_vector_potential(tgrid)
if omega is None:
# Return frequencies and amplitudes
ft_w, w = FourierTransform(t=tgrid, Y=vpot, a0=center_of_pulse)
return w, np.transpose(ft_w)[0]
elif isinstance(omega, np.ndarray):
return FourierTransform(t=tgrid, Y=vpot, a0=center_of_pulse, omega=omega)
[docs]
def get_w_max(self, center_of_pulse: float = None) -> float:
"""Get the maximum angular frequency of the pulse.
Returns:
float:
angular frequency in atomic units
This corresponds to the maximum photon energy.
"""
photon_energy = np.linspace(0, 30, 1000) # eV
w = AtomicUnits.photon_energy_eV_to_omega_au(photon_energy)
ft_w = self.get_fft(omega=w, center_of_pulse=center_of_pulse)
# plt.plot(w*AtomicUnits.eV, np.abs(ft_w))
# plt.savefig("fft.png")
return max(w[np.abs(ft_w) > max(abs(ft_w)) / 20]) if ft_w.size > 0 else 0.0
[docs]
def get_ponderomotive_energy(self) -> float:
"""Calculate ponderomotive energy.
Returns:
Ponderomotive energy in atomic units
"""
Up = 0.0
vpot = self.get_vector_potential(self.get_tgrid())
Up = np.max(vpot) ** 2 / 4
return float(Up)
[docs]
def get_keldysh_parameter(self, Ip: float) -> float:
"""Calculate the Keldysh parameter γ for this pulse and ionization potential.
The Keldysh parameter separates the tunnel (γ ≪ 1) and multiphoton
(γ ≫ 1) ionization regimes:
.. math::
\\gamma = \\sqrt{\\frac{I_p}{2 U_p}}
= \\frac{\\omega_0 \\sqrt{2 I_p}}{E_0}
where :math:`U_p = E_0^2 / (4\\omega_0^2)` is the ponderomotive energy.
Args:
Ip: Ionization potential in atomic units.
Returns:
Dimensionless Keldysh parameter.
γ < 1 → tunnel ionization regime.
γ > 1 → multiphoton ionization regime.
γ ≈ 1 → intermediate regime.
Example:
>>> laser = create_pulse(800, 1e14, 0, 30)
>>> laser.get_keldysh_parameter(0.5) # hydrogen, ~0.5 a.u. IP
1.0...
"""
Up = self.get_ponderomotive_energy()
return float(np.sqrt(Ip / (2.0 * Up)))
[docs]
def get_average_initial_KE_energy(self, Ip: float) -> float:
"""Calculate initial kinetic energy.
Args:
Ip: Ionization potential in atomic units
Returns:
Initial kinetic energy in atomic units
"""
w0 = self.get_w0()
return np.ceil(Ip / w0) * w0 - Ip
[docs]
def get_max_initial_KE_energy(self, Ip: float) -> float:
"""Calculate maximum initial kinetic energy.
Args:
Ip: Ionization potential in atomic units
Returns:
Maximum initial kinetic energy in atomic units
"""
w = self.get_w_max()
return np.ceil(Ip / w) * w - Ip
[docs]
def get_total_energy(self, Ip: float) -> float:
"""Calculate total energy.
Args:
Ip: Ionization potential in atomic units
Returns:
Total energy in atomic units
"""
return self.get_max_initial_KE_energy(Ip) + self.get_ponderomotive_energy()
[docs]
def get_central_wavelength(self) -> float:
"""Get the central wavelength of the pulse.
Returns:
Central wavelength in nm
"""
return AtomicUnits.omega_au_to_wavelength_nm(self.get_w0())
[docs]
def get_max_field_strength(self) -> float:
"""Get the maximum electric field strength of the pulse.
Returns:
Maximum electric field strength in atomic units
"""
tgrid = self.get_tgrid()
e_field = self.get_electric_field(tgrid)
return np.max(np.abs(e_field))
[docs]
def get_intensity(self) -> float:
"""Calculate the peak intensity of the pulse.
Returns:
Peak intensity in W/cm^2
"""
e_field = self.get_electric_field(self.get_tgrid())
return AtomicUnits.field_au_to_intensity_Wcm2(np.max(np.abs(e_field)))
# --- Abstract "guts" ---
@abstractmethod
def _e_field(
self, t: float | npt.NDArray[np.float64]
) -> float | npt.NDArray[np.float64]:
"""Calculate the electric field at time t."""
pass
@abstractmethod
def _vector_potential(
self, t: float | npt.NDArray[np.float64]
) -> float | npt.NDArray[np.float64]:
"""Calculate the vector potential at time t."""
pass
@abstractmethod
def _cummulative_vector_potential(
self, t: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
"""Calculate the cumulative intregral of vector potential up to time t, starting from -inf(pulse begin)"""
pass
@abstractmethod
def _cummulative_vector_potential_squared(
self, t: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
"""Calculate the cumulative integral of squared vector potential up to time t, starting from -inf(pulse begin)"""
pass
@abstractmethod
def _cummulative_electric_field_squared(
self, t: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
"""Calculate the cumulative integral of squared electric field up to time t, starting from -inf(pulse begin)"""
pass
[docs]
@abstractmethod
def get_w0(self) -> float:
"""Get the central wavelength of the pulse.
Returns:
Central wavelength in nm
"""
pass
[docs]
@abstractmethod
def get_time_interval(self) -> Tuple[float, float]:
"""Get time interval containing the pulse.
Returns:
Tuple of (start_time, end_time) in atomic units
"""
pass
# --- End of abstract methods ---
[docs]
class CosNPulse(Pulse):
def __init__(
self,
central_wavelength: float,
peak_intensity: float,
cep: float,
FWHM: float,
envelope_N: int = 8,
t0: float = 0.0,
) -> None:
"""Add a new pulse to the field.
Args:
central_wavelength: Central wavelength in nm
peak_intensity: Peak intensity in W/cm^2
CEP: Carrier-envelope phase in radians
FWHM: Full width at half maximum in au
envelope_N: Power of cosine envelope (default: 8)
t0: Time offset in au (default: 0.0)
"""
super().__init__()
self.w0 = AtomicUnits.wavelength_nm_to_omega_au(central_wavelength)
self.A0 = AtomicUnits.intensity_Wcm2_to_field_au(peak_intensity) / self.w0
self.tau = np.pi * FWHM / (4 * np.arccos(2 ** (-1 / (2 * envelope_N))))
self.t0 = t0
self.N = envelope_N
self.cep = cep
# @njit(fastmath=True, cache=True)
def _vector_potential(self, t):
t = t - self.t0
return np.where(
np.abs(t) <= self.tau,
-self.A0
* np.cos(0.5 * np.pi * t / self.tau) ** self.N
* np.sin(self.w0 * t - self.cep),
0.0,
)
# @njit(fastmath=True, cache=True)
def _e_field(self, t):
t = t - self.t0
x = self.w0 * t - self.cep
y = 0.5 * np.pi * t / self.tau
return np.where(
np.abs(t) <= self.tau,
self.A0
* (np.cos(y)) ** (self.N - 1)
* (
self.w0 * np.cos(x) * np.cos(y)
- (self.N * np.pi * np.sin(x) * np.sin(y)) / (2.0 * self.tau)
),
0.0,
)
def _cummulative_vector_potential(self, t):
return cosN_int_A(t - self.t0, self.A0, self.w0, self.tau, self.cep, self.N)
def _cummulative_vector_potential_squared(self, t):
return cosN_int_A2(t - self.t0, self.A0, self.w0, self.tau, self.cep, self.N)
def _cummulative_electric_field_squared(self, t):
return cosN_int_E2(t - self.t0, self.A0, self.w0, self.tau, self.cep, self.N)
[docs]
def get_w0(self) -> float:
return self.w0
[docs]
def get_time_interval(self) -> Tuple[float, float]:
tmin = -self.tau + self.t0
tmax = self.tau + self.t0
return (tmin, tmax)
[docs]
class EllipticalCosNPulse(CosNPulse):
def __init__(
self,
central_wavelength: float,
peak_intensity: float,
FWHM: float,
cep: float,
peak_intensity2: float = 0.0,
envelope_N: int = 8,
t0: float = 0.0,
) -> None:
"""Create a circularly polarized cos^N pulse.
this is a special method and not really a circular pulse class
since the vector potential and electric field are defined in y direction only.
the values are the magnitudes of the vector potential and electric field in circular polarization.
Args:
central_wavelength: Central wavelength in nm
peak_intensity: Peak intensity in W/cm^2
CEP: Carrier-envelope phase in radians
FWHM: Full width at half maximum in au
envelope_N: Power of cosine envelope (default: 8)
t0: Time offset in au (default: 0.0)
"""
super().__init__(
central_wavelength=central_wavelength,
peak_intensity=peak_intensity,
cep=cep,
FWHM=FWHM,
envelope_N=envelope_N,
t0=t0,
)
self.pulse_x = CosNPulse(
central_wavelength=central_wavelength,
peak_intensity=peak_intensity,
cep=cep,
FWHM=FWHM,
envelope_N=envelope_N,
t0=t0,
)
# Y Component (Phase shifted by pi/2 relative to X)
self.pulse_y = CosNPulse(
central_wavelength=central_wavelength,
peak_intensity=peak_intensity2,
# Phase difference of pi/2 is the definition of circular polarization:
cep=cep - np.pi / 2.0,
FWHM=FWHM,
envelope_N=envelope_N,
t0=t0,
)
self.A0_2 = AtomicUnits.intensity_Wcm2_to_field_au(peak_intensity2) / self.w0
[docs]
def get_vector_potential_tuple(self, t):
return (
self.pulse_x.get_vector_potential(t - self.t0),
self.pulse_y.get_vector_potential(t - self.t0),
)
[docs]
def get_electric_field_tuple(self, t):
return (
self.pulse_x.get_electric_field(t - self.t0),
self.pulse_y.get_electric_field(t - self.t0),
)
[docs]
def get_cummulative_vector_potential_tuple(self, t):
return (
self.pulse_x.get_cummulative_vector_potential(t - self.t0),
self.pulse_y.get_cummulative_vector_potential(t - self.t0),
)
def _cummulative_vector_potential_squared(self, t):
"""Calculate the cumulative integral of vector potential up to time t, starting from -inf (pulse begin)."""
timeGrid = self.get_tgrid(dt=0.125)
return cosN_int_A2(
t - self.t0, self.A0, self.w0, self.tau, self.cep, self.N
) + cosN_int_A2(
t - self.t0, self.A0_2, self.w0, self.tau, self.cep - np.pi / 2, self.N
)
def _cummulative_electric_field_squared(self, t):
return cosN_int_E2(
t - self.t0, self.A0, self.w0, self.tau, self.cep, self.N
) + cosN_int_E2(
t - self.t0, self.A0_2, self.w0, self.tau, self.cep - np.pi / 2, self.N
)
class pump_probe_CosNPulse(CosNPulse):
def __init__(
self,
central_wavelength1: float,
peak_intensity1: float,
FWHM1: float,
cep1: float,
central_wavelength2: float,
direction2: str = "collinear",
peak_intensity2: float = 0.0,
FWHM2: float = 0.0,
cep2: float = 0.0,
envelope_N: int = 8,
t0_1: float = 0.0,
t0_2: float = 0.0,
) -> None:
"""Create a circularly polarized cos^N pulse.
this is a special method and not really a circular pulse class
since the vector potential and electric field are defined in y direction only.
the values are the magnitudes of the vector potential and electric field in circular polarization.
Args:
central_wavelength1: Central wavelength in nm for pulse 1
peak_intensity1: Peak intensity in W/cm^2 for pulse 1
FWHM1: Full width at half maximum in au for pulse 1
cep1: Carrier-envelope phase in radians for pulse 1
central_wavelength2: Central wavelength in nm for pulse 2
direction2: 'collinear' or 'perpendicular' for pulse 2 propagation direction
peak_intensity2: Peak intensity in W/cm^2 for pulse 2
FWHM2: Full width at half maximum in au for pulse 2
cep2: Carrier-envelope phase in radians for pulse 2
envelope_N: Power of cosine envelope (default: 8)
t0: Time offset in au (default: 0.0)
"""
super().__init__(
central_wavelength=central_wavelength1,
peak_intensity=peak_intensity1,
cep=cep1,
FWHM=FWHM1,
envelope_N=envelope_N,
t0=t0_1,
)
self.pulse_x = CosNPulse(
central_wavelength=central_wavelength1,
peak_intensity=peak_intensity1,
cep=cep1,
FWHM=FWHM1,
envelope_N=envelope_N,
t0=t0_1,
)
# Y Component (Phase shifted by pi/2 relative to X)
self.pulse_y = CosNPulse(
central_wavelength=central_wavelength2,
peak_intensity=peak_intensity2,
# Phase difference of pi/2 is the definition of circular polarization:
cep=cep2,
FWHM=FWHM2,
envelope_N=envelope_N,
t0=t0_2,
)
self.A0_2 = AtomicUnits.intensity_Wcm2_to_field_au(peak_intensity2) / self.w0
def get_vector_potential_tuple(self, t):
return (
self.pulse_x.get_vector_potential(t),
self.pulse_y.get_vector_potential(t),
)
def get_electric_field_tuple(self, t):
return (self.pulse_x.get_electric_field(t), self.pulse_y.get_electric_field(t))
def get_cummulative_vector_potential_tuple(self, t):
return (
self.pulse_x.get_cummulative_vector_potential(t),
self.pulse_y.get_cummulative_vector_potential(t),
)
def _cummulative_vector_potential_squared(self, t):
"""Calculate the cumulative integral of vector potential up to time t, starting from -inf (pulse begin)."""
timeGrid = self.get_tgrid(dt=0.125)
return cosN_int_A2(
t - self.t0, self.A0, self.w0, self.tau, self.cep, self.N
) + cosN_int_A2(
t - self.t0, self.A0_2, self.w0, self.tau, self.cep + np.pi / 2, self.N
)
def _cummulative_electric_field_squared(self, t):
return cosN_int_E2(
t - self.t0, self.A0, self.w0, self.tau, self.cep, self.N
) + cosN_int_E2(
t - self.t0, self.A0_2, self.w0, self.tau, self.cep + np.pi / 2, self.N
)
class FlatFrequencyPulse(Pulse):
def __init__(
self,
peak_intensity: float,
bandwidth_au: float,
time_window_au: float = 500.0,
taper_fraction: float = 0.2,
cep: float = 0.0,
t0: float = 0.0,
) -> None:
"""Create a 'Sinc' pulse with a Tukey window applied to eliminate boundary artifacts.
Args:
peak_intensity: Peak intensity in W/cm^2.
bandwidth_au: The maximum photon energy (angular frequency) in atomic units.
time_window_au: Total duration the pulse is allowed to exist before being forced to 0.
taper_fraction: Fraction of the window used to taper to 0 (e.g., 0.2 means 10%
of the time is spent smoothing down to zero on each side).
cep: Carrier-envelope phase in radians
t0: Time offset (delay) in au.
"""
super().__init__()
self.peak_intensity = peak_intensity
self.bandwidth_au = bandwidth_au
self.time_window_au = time_window_au
self.taper_fraction = taper_fraction
self.t0 = t0
self.cep = cep
self.E0 = AtomicUnits.intensity_Wcm2_to_field_au(self.peak_intensity)
self.w_c = self.bandwidth_au
self.w0 = 1e-10
def _tukey_window(
self, t: float | npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
"""Calculates a smooth Tukey (tapered cosine) window centered at t0."""
# Shift time relative to pulse center
t_centered = t - self.t0
t_abs = np.abs(t_centered)
L = self.time_window_au
alpha = self.taper_fraction
# Calculate the boundary where the flat top ends and the taper begins
flat_limit = (L / 2) * (1 - alpha)
window = np.zeros_like(t_centered)
# 1. Flat region (W = 1)
flat_mask = t_abs <= flat_limit
window[flat_mask] = 1.0
# 2. Taper region (Cosine roll-off to 0)
taper_mask = (t_abs > flat_limit) & (t_abs <= L / 2)
window[taper_mask] = 0.5 * (
1 + np.cos(np.pi * (t_abs[taper_mask] - flat_limit) / (alpha * L / 2))
)
return window
def _e_field(self, t: float | npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
"""E(t) = Window(t) * E0 * sinc(w_c * (t - t0) + cep)
cep handeled by combining the even and odd components of the sinc function, which are cos(cep) and sin(cep) weighted respectively.
"""
t_arr = np.atleast_1d(t)
# x = w_c * (t - t0)
x = self.w_c * (t_arr - self.t0)
# 1. The Even Component (using np.sinc)
# We divide by pi because np.sinc(x) calculates sin(pi*x)/(pi*x)
E_even = np.sinc(x / np.pi)
# 2. The Odd Component (The Hilbert partner)
# This part still needs a manual handle for the 1/x singularity
E_odd = np.zeros_like(x)
mask = np.abs(x) > 1e-10
E_odd[mask] = (1.0 - np.cos(x[mask])) / x[mask]
# E_odd[~mask] is already 0.0, which is the correct limit
# 3. Combine with User Phase
E_raw = self.E0 * (E_even * np.cos(self.cep) + E_odd * np.sin(self.cep))
# 4. Window and Return
return E_raw * self._tukey_window(t_arr)
def _vector_potential(
self, t: float | npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
"""Numerical integration for vector potential A(t) = - integral(E(t)dt)"""
# Create a dense grid strictly within the non-zero window to save memory
t_start, t_end = self.get_time_interval()
timeGrid = np.arange(t_start, t_end + 0.125, 0.125)
# Note: Assuming standard A = -int(E)dt. If your solver expects A = int(E)dt, remove the minus sign.
spline = get_cummulative_integral_spline(t=timeGrid, Y=-self._e_field(timeGrid))
# The spline is only valid inside timeGrid. Outside, it should flatline at the max/min.
A = np.where(t < t_start, 0.0, np.where(t > t_end, spline(t_end), spline(t)))
return A.astype(np.float64)
def _cummulative_vector_potential(
self, t: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
timeGrid = self.get_tgrid(dt=0.125)
return get_cummulative_integral_spline(
t=timeGrid, Y=self._vector_potential(timeGrid)
)(t).astype(np.float64)
def _cummulative_vector_potential_squared(
self, t: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
timeGrid = self.get_tgrid(dt=0.125)
return get_cummulative_integral_spline(
t=timeGrid, Y=self._vector_potential(timeGrid) ** 2
)(t).astype(np.float64)
def _cummulative_electric_field_squared(
self, t: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
timeGrid = self.get_tgrid(dt=0.125)
return get_cummulative_integral_spline(
t=timeGrid, Y=self._e_field(timeGrid) ** 2
)(t).astype(np.float64)
def get_w0(self) -> float:
return self.w0
def get_time_interval(self) -> Tuple[float, float]:
# Return exactly the edges of the Tukey window
half_width = self.time_window_au / 2.0
return (self.t0 - half_width, self.t0 + half_width)
class flatTopPulse(Pulse):
def __init__(
self,
central_wavelength: float,
peak_intensity: float,
cep: float,
rise_fall_time: float,
flat_time: float,
t0: float = 0.0,
) -> None:
"""Create a flat-top pulse with linear rise and fall.
Args:
central_wavelength: Central wavelength in nm
peak_intensity: Peak intensity in W/cm^2
CEP: Carrier-envelope phase in radians
rise_fall_time: Time for cubic rise/fall in atomic units
flat_time: Duration of the flat top in atomic units
t0: Time offset (default: 0.0)
"""
super().__init__()
self.w0 = AtomicUnits.wavelength_nm_to_omega_au(central_wavelength)
self.A0 = AtomicUnits.intensity_Wcm2_to_field_au(peak_intensity) / self.w0
self.t0 = t0
self.cep = cep
self.rise_fall_time = rise_fall_time
self.flat_time = flat_time
def _vector_potential(self, t):
"""flatTop pulse uses 3rd degree polynomial for rise and fall of vector potential
mathematical form S(q) = 3q^2 - 2q^3$$ where q is normalized time in rise/fall region [0,1]
q=(b-x)/(b-a)
"""
t = t - self.t0
carrier = self.A0 * np.sin(self.w0 * t - self.cep)
b = self.flat_time / 2 + self.rise_fall_time
a = self.flat_time / 2
q = (b - abs(t)) / (b - a)
return 1.0 * np.where(
q < 0, 0.0, np.where(q > 1, carrier, (3 * q**2 - 2 * q**3) * carrier)
)
# eturn rise_fall_envelope * carrier
# return np.where(abs(t)<=a, carrier, rise_fall_envelope * carrier)
# return np.where(abs(t)<= self.flat_time/2, carrier,
def _e_field(self, t):
t = t - self.t0
carrier = self.A0 * np.sin(self.w0 * t - self.cep)
d_carrier = self.A0 * self.w0 * np.cos(self.w0 * t - self.cep)
b = self.flat_time / 2 + self.rise_fall_time
a = self.flat_time / 2
q = (b - abs(t)) / (b - a)
dq_dt = -1 / (b - a) * np.sign(t)
return -1.0 * np.where(
q < 0,
0.0,
np.where(
q > 1,
d_carrier,
(6 * q - 6 * q**2) * carrier * dq_dt
+ (3 * q**2 - 2 * q**3) * d_carrier,
),
)
def _cummulative_vector_potential(
self, t: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
"""Calculate the cumulative integral of vector potential up to time t, starting from -inf (pulse begin)."""
timeGrid = self.get_tgrid(dt=0.125)
return get_cummulative_integral_spline(
t=timeGrid, Y=self._vector_potential(timeGrid)
)(t).astype(np.float64)
def _cummulative_vector_potential_squared(
self, t: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
"""Calculate the cumulative integral of squared vector potential up to time t, starting from -inf (pulse begin)."""
timeGrid = self.get_tgrid(dt=0.125)
return get_cummulative_integral_spline(
t=timeGrid, Y=self._vector_potential(timeGrid) ** 2
)(t).astype(np.float64)
def _cummulative_electric_field_squared(
self, t: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
"""Calculate the cumulative integral of squared electric field up to time t, starting from -inf (pulse begin)."""
timeGrid = self.get_tgrid(dt=0.125)
return get_cummulative_integral_spline(
t=timeGrid, Y=self._e_field(timeGrid) ** 2
)(t).astype(np.float64)
def get_w0(self) -> float:
return self.w0
def get_time_interval(self) -> Tuple[float, float]:
tmin = -(self.flat_time / 2 + self.rise_fall_time) + self.t0
tmax = (self.flat_time / 2 + self.rise_fall_time) + self.t0
return (tmin, tmax)
class SumOfPulses(Pulse):
def __init__(self, pulses: list[Pulse]) -> None:
"""Initialize a sum of multiple pulses.
Args:
pulses: List of Pulse instances
"""
super().__init__()
self.pulses = pulses
def _vector_potential(self, t):
vepot_array = np.asarray(
[pulse.get_vector_potential(t) for pulse in self.pulses]
)
return np.sum(vepot_array, axis=0)
def _e_field(self, t):
e_field_array = np.asarray(
[pulse.get_electric_field(t) for pulse in self.pulses]
)
return np.sum(e_field_array, axis=0)
def _cummulative_vector_potential(self, t):
intA_array = np.asarray(
[pulse.get_cummulative_vector_potential(t) for pulse in self.pulses]
)
return np.sum(intA_array, axis=0)
def _cummulative_vector_potential_squared(self, t):
# since the integral of A^2 is not linear, numerically integrate vector potential squared
timeGrid = self.get_tgrid(dt=0.125)
return get_cummulative_integral_spline(
t=timeGrid, Y=self._vector_potential(timeGrid) ** 2
)(t).astype(np.float64)
def _cummulative_electric_field_squared(self, t):
# since the integral of E^2 is not linear, numerically integrate electric field squared
timeGrid = self.get_tgrid(dt=0.125)
return get_cummulative_integral_spline(
t=timeGrid, Y=self._e_field(timeGrid) ** 2
)(t).astype(np.float64)
def get_w0(self) -> float:
# get the minimum w0 among the pulses, which corresponds to the longest wavelength and thus the lowest photon energy
return min(pulse.get_w0() for pulse in self.pulses)
def get_time_interval(self) -> Tuple[float, float]:
tmin = 0
tmax = 0
tmin = min(pulse.get_time_interval()[0] for pulse in self.pulses)
tmax = max(pulse.get_time_interval()[1] for pulse in self.pulses)
return (tmin, tmax)
class _DelayedPulse(Pulse):
"""Internal wrapper that shifts a pulse in time by a fixed delay.
All public ``Pulse`` calls on the wrapper are forwarded to the
underlying pulse with the time argument offset by ``delay``. The
underlying pulse retains its own cache, so repeated evaluations on the
same wrapped pulse are still cheap.
Args:
pulse: Any :class:`Pulse` instance.
delay: Time shift in atomic units (positive = probe arrives later).
"""
def __init__(self, pulse: "Pulse", delay: float) -> None:
super().__init__()
self._pulse = pulse
self._delay = float(delay)
def _vector_potential(self, t):
return self._pulse.get_vector_potential(np.asarray(t) - self._delay)
def _e_field(self, t):
return self._pulse.get_electric_field(np.asarray(t) - self._delay)
def _cummulative_vector_potential(self, t):
return self._pulse.get_cummulative_vector_potential(np.asarray(t) - self._delay)
def _cummulative_vector_potential_squared(self, t):
return self._pulse.get_cummulative_vector_potential_squared(
np.asarray(t) - self._delay
)
def _cummulative_electric_field_squared(self, t):
return self._pulse.get_cummulative_electric_field_squared(
np.asarray(t) - self._delay
)
def get_w0(self) -> float:
return self._pulse.get_w0()
def get_time_interval(self) -> Tuple[float, float]:
t_min, t_max = self._pulse.get_time_interval()
return (t_min + self._delay, t_max + self._delay)
[docs]
class TIPTOE:
"""Two-pulse delay scanner for pump–probe TIPTOE spectroscopy.
Generates :class:`SumOfPulses` objects at arbitrary pump–probe delays and
provides a Nyquist-sampled delay grid for scanning experiments.
The *pump* pulse is held fixed; the *probe* is shifted in time by the
requested delay using a lightweight :class:`_DelayedPulse` wrapper that
adds no computational overhead beyond the underlying pulse's own caching.
Args:
pump: Fixed pump pulse (any :class:`Pulse` subclass).
probe: Probe pulse to be delayed. The zero-delay configuration
corresponds to the probe at its original position (``t0``).
Example::
from gasfir import create_pulse
from gasfir.pulse import TIPTOE
pump = create_pulse(800, 3e14, 0, 6)
probe = create_pulse(800, 1e13, 0, 2) # weak probe
scanner = TIPTOE(pump, probe)
# Scan over the Nyquist-sampled delay grid
delays = scanner.return_delay_array()
for tau in delays:
field = scanner.at_delay(tau)
prob = get_diabatic_ionization_probability(pulse=field, param_dict=params)
"""
def __init__(self, pump: Pulse, probe: Pulse) -> None:
self.pump = pump
self.probe = probe
# ------------------------------------------------------------------
# Core methods
# ------------------------------------------------------------------
[docs]
def at_delay(self, delay: float) -> "SumOfPulses":
"""Return a :class:`SumOfPulses` with the probe delayed by *delay* a.u.
Args:
delay: Pump–probe delay in atomic units. Positive values place
the probe *after* the pump; negative values place it
*before*.
Returns:
:class:`SumOfPulses` of the fixed pump and the shifted probe.
"""
return SumOfPulses([self.pump, _DelayedPulse(self.probe, delay)])
[docs]
def return_delay_array(
self,
delay_min: float | None = None,
delay_max: float | None = None,
oversample: int = 4,
) -> npt.NDArray[np.float64]:
"""Nyquist-sampled delay grid covering the full pump–probe overlap region.
The sampling theorem requires at least 2 delay points per oscillation
period of the highest-frequency field component. This method uses the
maximum of the pump and probe carrier frequencies to set the step size::
Δτ = T₀ / oversample where T₀ = 2π / max(ω_pump, ω_probe)
The default range extends 1.5× beyond the point where the two pulses
just touch (positive and negative), ensuring both fully-separated and
fully-overlapping configurations are included.
Args:
delay_min: Start of the delay range in atomic units. ``None``
(default) sets it to −1.5 × (pump_half_duration +
probe_half_duration).
delay_max: End of the delay range in atomic units. ``None``
(default) mirrors ``delay_min``.
oversample: Number of delay samples per optical cycle of the
highest-frequency component (default 4 — 2× Nyquist margin).
Use 8 or more for carrier-envelope resolved scans.
Returns:
1-D NumPy array of delay values in atomic units.
Example::
delays = scanner.return_delay_array() # auto range
delays = scanner.return_delay_array(-500, 500, oversample=8)
"""
# Dominant frequency of each pulse
w_pump = self.pump.get_w0()
w_probe = self.probe.get_w0()
w_max = max(w_pump, w_probe)
# Step size from Nyquist: T₀ / oversample
T0 = 2.0 * np.pi / w_max
dt = T0 / oversample
# Default range: 1.5× the touching distance in each direction
if delay_min is None or delay_max is None:
p_min, p_max = self.pump.get_time_interval()
r_min, r_max = self.probe.get_time_interval()
pump_half = (p_max - p_min) / 2.0
probe_half = (r_max - r_min) / 2.0
half_range = 1.5 * (pump_half + probe_half)
if delay_min is None:
delay_min = -half_range
if delay_max is None:
delay_max = +half_range
n_steps = int(np.ceil((delay_max - delay_min) / dt)) + 1
return np.linspace(delay_min, delay_max, n_steps)
class ConstructPulseFromData(Pulse):
def __init__(
self,
*,
data_t: npt.NDArray[np.float64],
data_E: npt.NDArray[np.float64 | np.complex128],
desired_peak_intensity: None | float = None,
cep_shift: float = 0.0,
t0: float = 0.0,
) -> None:
"""Construct a pulse from given data.
Args:
data_t: Time data in atomic units
data_E: Electric field data in atomic units
desired_peak_intensity: Desired peak intensity in W/cm^2.
cep_shift: Carrier-envelope phase shift in radians (default: 0.0)
t0: Time offset (default: 0.0)
Note:
The data should be in atomic units. If desired_peak_intensity is None, it will be extracted from the data.
t0 is the time offset such that get_electric_field(t) = E(t-t0).
cep_shift is applied such that get_electric_field(t) = Re[E(t-t0)*exp(i*cep_shift)].
"""
super().__init__()
self.data_t = np.ascontiguousarray(data_t)
self.data_E = np.ascontiguousarray(data_E)
self.field_spline = get_cummulative_integral_spline(
t=self.data_t, Y=self.data_E
)
self.desired_peak_intensity = desired_peak_intensity
self.cep_shift = cep_shift
self.t0 = t0
def _e_field(self, t):
"""Calculate the electric field at time t."""
result = self.field_spline(t - self.t0).astype(np.float64)
if self.desired_peak_intensity is not None:
E0 = AtomicUnits.intensity_Wcm2_to_field_au(self.desired_peak_intensity)
result = result * E0 / max(abs(result))
return np.real(result * np.exp(1j * self.cep_shift))
def _vector_potential(self, t):
"""Calculate the vector potential at time t."""
# time grid with 0.125 au step
timeGrid = self.get_tgrid(dt=0.125)
# calculate cumulative integral of the electric field
return get_cummulative_integral_spline(t=timeGrid, Y=self._e_field(timeGrid))(
t - self.t0
).astype(np.float64)
def _cummulative_vector_potential(
self, t: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
"""Calculate the cumulative integral of vector potential up to time t, starting from -inf (pulse begin)."""
timeGrid = self.get_tgrid(dt=0.125)
return get_cummulative_integral_spline(
t=timeGrid, Y=self._vector_potential(timeGrid)
)(t - self.t0).astype(np.float64)
def _cummulative_vector_potential_squared(
self, t: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
"""Calculate the cumulative integral of squared vector potential up to time t, starting from -inf (pulse begin)."""
timeGrid = self.get_tgrid(dt=0.125)
return get_cummulative_integral_spline(
t=timeGrid, Y=self._vector_potential(timeGrid) ** 2
)(t - self.t0).astype(np.float64)
def _cummulative_electric_field_squared(
self, t: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
"""Calculate the cumulative integral of squared electric field up to time t, starting from -inf (pulse begin)."""
timeGrid = self.get_tgrid(dt=0.125)
return get_cummulative_integral_spline(
t=timeGrid, Y=self._e_field(timeGrid) ** 2
)(t - self.t0).astype(np.float64)
[docs]
class DataPulse(Pulse):
"""Pulse constructed from a user-supplied numerical field on a time grid.
Accepts either the **electric field** E(t) or the **vector potential** A(t)
sampled on a uniform (or non-uniform) time grid in atomic units.
Physical requirement:
A(t) must vanish at both ends of the time window (no net momentum kick).
When E(t) is supplied, A is derived by integration
``A(t) = −∫ E(t') dt'``; the integration constant is chosen so that
A = 0 at t = t[0], and the field should be essentially zero at both
boundaries for this to be accurate.
All four cumulative-integral splines (∫A, ∫A², ∫E²) are precomputed
once at construction time from the supplied data grid.
Args:
t: 1-D time array in atomic units, must be strictly increasing.
E: Electric field array E(t) in atomic units. Mutually exclusive
with ``A``.
A: Vector potential array A(t) in atomic units. E is derived as
``E(t) = −dA/dt`` via the cubic-spline derivative. Mutually
exclusive with ``E``.
t0: Time offset (default 0). ``get_electric_field(t)`` evaluates at
``t − t0``.
Raises:
ValueError: If neither or both of ``E`` and ``A`` are given.
Example — from A(t)::
t = np.linspace(-200, 200, 3200)
A = sum_pulse.get_vector_potential(t)
dp = DataPulse(t, A=A)
prob = get_diabatic_ionization_probability(pulse=dp, param_dict=params)
Example — from E(t)::
E = sum_pulse.get_electric_field(t)
dp = DataPulse(t, E=E)
"""
def __init__(
self,
t: npt.NDArray[np.float64],
*,
E: npt.NDArray[np.float64] | None = None,
A: npt.NDArray[np.float64] | None = None,
t0: float = 0.0,
) -> None:
super().__init__()
if (E is None) == (A is None):
raise ValueError(
"Provide exactly one of E (electric field) or A (vector potential)."
)
self.t0 = t0
_t = np.ascontiguousarray(t, dtype=np.float64)
if A is not None:
# ── primary: A(t) ────────────────────────────────────────────
_A = np.ascontiguousarray(A, dtype=np.float64)
_A_spl = CubicSpline(_t, _A, bc_type="natural")
# E = −dA/dt (exact derivative of the spline)
_E = -_A_spl(_t, 1) # 1st derivative
else:
# ── primary: E(t) ────────────────────────────────────────────
_E = np.ascontiguousarray(E, dtype=np.float64)
# A = −∫ E dt (A = 0 at t[0])
_A = -cumulative_simpson(_E, x=_t, initial=0.0)
# ── store interpolation splines ───────────────────────────────────
self._A_spline = CubicSpline(_t, _A, bc_type="natural")
self._E_spline = CubicSpline(_t, _E, bc_type="natural")
# ── precompute cumulative-integral splines ────────────────────────
self._intA_spline = get_cummulative_integral_spline(_t, _A)
self._intA2_spline = get_cummulative_integral_spline(_t, _A**2)
self._intE2_spline = get_cummulative_integral_spline(_t, _E**2)
# ── time interval and dominant frequency ──────────────────────────
self._t_min = float(_t[0])
self._t_max = float(_t[-1])
self._w0 = self._dominant_frequency(_t, _A)
@staticmethod
def _dominant_frequency(
t: npt.NDArray[np.float64],
A: npt.NDArray[np.float64],
) -> float:
"""Estimate dominant angular frequency from the FFT of A(t)."""
dt = float(np.mean(np.diff(t)))
amps = np.abs(np.fft.rfft(A))
amps[0] = 0.0 # suppress DC component
freqs = np.fft.rfftfreq(len(A), d=dt) * 2.0 * np.pi # angular
w0 = float(freqs[np.argmax(amps)])
return w0 if w0 > 0 else float(np.pi / dt) # fallback: Nyquist/2
# ── abstract interface ────────────────────────────────────────────────
def _vector_potential(self, t: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
return self._A_spline(t - self.t0).astype(np.float64)
def _e_field(self, t: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
return self._E_spline(t - self.t0).astype(np.float64)
def _cummulative_vector_potential(
self, t: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
return self._intA_spline(t - self.t0).astype(np.float64)
def _cummulative_vector_potential_squared(
self, t: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
return self._intA2_spline(t - self.t0).astype(np.float64)
def _cummulative_electric_field_squared(
self, t: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
return self._intE2_spline(t - self.t0).astype(np.float64)
[docs]
def get_w0(self) -> float:
return self._w0
[docs]
def get_time_interval(self) -> Tuple[float, float]:
return (self._t_min + self.t0, self._t_max + self.t0)
[docs]
def OptCyc_au(fwhmOC: float, wavel: float) -> float:
"""Convert a duration in optical cycles to atomic units.
Args:
fwhmOC: Duration in optical cycles.
wavel: Central wavelength in nm.
Returns:
Duration in atomic units of time.
"""
return fwhmOC * 2.0 * np.pi / AtomicUnits.wavelength_nm_to_omega_au(wavel)
[docs]
def create_pulse(wavel, intens, cep, fwhmCyc, N=8, t0=0.0):
"""Create a cos8 pulse with specified parameters.
Args:
wavel: Central wavelength in nm
intens: Peak intensity in W/cm^2
cep: Carrier-envelope phase in radians
fwhmCyc: Full-width at half-maximum in optical cycles
N: Envelope parameter for CosNPulse
Returns:
Pulse: An instance of the CosNPulse class with the specified parameters.
"""
fwhmAU = OptCyc_au(fwhmCyc, wavel)
return CosNPulse(wavel, intens, cep, fwhmAU, N, t0)
def create_pulse_ellip(wavel, intens, cep, fwhmCyc, N=8, ellip_frac=1):
"""Create a cos8 pulse with specified parameters.
Args:
wavel: Central wavelength in nm
intens: Peak intensity in W/cm^2
cep: Carrier-envelope phase in radians
fwhmCyc: Full-width at half-maximum in optical cycles
N: Envelope parameter for CosNPulse
Returns:
Pulse: An instance of the CosNPulse class with the specified parameters.
"""
fwhmAU = OptCyc_au(fwhmCyc, wavel)
return EllipticalCosNPulse(
central_wavelength=wavel,
peak_intensity=intens,
peak_intensity2=intens / ellip_frac,
cep=cep,
FWHM=fwhmAU,
envelope_N=N,
)
[docs]
def ret_pulse_from_pandas_table(data, N=8, pulse_type="CosN", ellip_frac=100):
"""Create a list of Pulse objects from a pandas dataframe."""
pulses = []
if pulse_type == "CosN":
if "FWHM_OC" in data:
for dat in data.iterrows():
dat = dat[1]
pulses.append(
create_pulse(
wavel=dat.wavel,
intens=dat.intens,
cep=dat.cep,
fwhmCyc=dat.FWHM_OC,
N=N,
)
)
return pulses
elif "fwhmau" in data:
for dat in data.iterrows():
dat = dat[1]
pulse = CosNPulse(
central_wavelength=dat.wavel,
peak_intensity=dat.intens,
cep=dat.cep,
FWHM=dat.fwhmau,
envelope_N=N,
)
pulses.append(pulse)
return pulses
else:
raise ValueError(
"The dataframe does not contain the required columns. Try modifying dataFrame or ret_pulse_from_pandas_table function to correctly inherit pulse duration"
)
elif pulse_type == "flatTop":
if "FWHM_OC" in data:
for dat in data.iterrows():
dat = dat[1]
fall_rise_time = OptCyc_au(
1, dat.wavel
) # 1 optical cycle rise/fall time
flat_time = OptCyc_au(
dat.FWHM_OC - 1, dat.wavel
) # flat time excluding rise/fall
pulses.append(
flatTopPulse(
central_wavelength=dat.wavel,
peak_intensity=dat.intens,
cep=dat.cep,
rise_fall_time=fall_rise_time,
flat_time=flat_time,
)
)
return pulses
else:
raise ValueError(
"The dataframe does not contain the required columns. Try modifying dataFrame or ret_pulse_from_pandas_table function to correctly inherit pulse duration"
)
elif pulse_type == "CircularCosN":
if "FWHM_OC" in data:
for dat in data.iterrows():
dat = dat[1]
pulses.append(
EllipticalCosNPulse(
central_wavelength=dat.wavel,
peak_intensity=dat.intens,
peak_intensity2=dat.intens,
cep=dat.cep,
FWHM=OptCyc_au(dat.FWHM_OC, dat.wavel),
envelope_N=N,
)
)
return pulses
else:
raise ValueError(
"The dataframe does not contain the required columns. Try modifying dataFrame or ret_pulse_from_pandas_table function to correctly inherit pulse duration"
)
elif pulse_type == "EllipticalCosN":
if "FWHM_OC" in data:
for dat in data.iterrows():
dat = dat[1]
pulses.append(
EllipticalCosNPulse(
central_wavelength=dat.wavel,
peak_intensity=dat.intens,
peak_intensity2=dat.intens / ellip_frac,
cep=dat.cep,
FWHM=OptCyc_au(dat.FWHM_OC, dat.wavel),
envelope_N=N,
)
)
return pulses
else:
raise ValueError(
"The dataframe does not contain the required columns. Try modifying dataFrame or ret_pulse_from_pandas_table function to correctly inherit pulse duration"
)
else:
raise ValueError(
"Unsupported pulse type. Supported types are 'CosN', 'flatTop', and 'CircularCosN'."
)