"""Utility functions for GASFIR package.
This module provides utility functions for calculations in the GASFIR package,
including window functions, zero crossing detection, momentum grid generation,
and other helper functions.
"""
import numpy as np
import numpy.typing as npt
from numba import njit, prange
from scipy.integrate import cumulative_simpson
from scipy.interpolate import CubicSpline
[docs]
class AtomicUnits:
"""Atomic units constants and unit conversion utilities.
All class attributes express *one atomic unit* of a given quantity in
practical/SI units. For example, ``AtomicUnits.nm = 0.05292`` means that
one atomic unit of length equals 0.05292 nm (the Bohr radius).
The static methods provide clearly named two-way conversions so callers
never have to reason about which direction to multiply or divide.
Attributes:
meter: 1 a.u. of length in metres (Bohr radius).
nm: 1 a.u. of length in nanometres.
angstrom: 1 a.u. of length in ångström.
second: 1 a.u. of time in seconds.
fs: 1 a.u. of time in femtoseconds.
Joule: 1 a.u. of energy in joules (Hartree).
eV: 1 a.u. of energy in electronvolts (Hartree).
Volts_per_meter: 1 a.u. of electric field in V/m.
Volts_per_Angstrom: 1 a.u. of electric field in V/Å.
speed_of_light: Speed of light in atomic units (= 1/α ≈ 137).
Coulomb: 1 a.u. of charge in coulombs (elementary charge).
PW_per_cm2_au: 1 PW/cm² expressed in a.u. of intensity.
Example:
>>> AtomicUnits.au_to_nm(1.0) # Bohr radius in nm
0.05292...
>>> AtomicUnits.nm_to_au(800) # 800 nm in a.u.
15117.6...
>>> AtomicUnits.au_to_eV(0.5) # hydrogen IP in eV
13.606...
>>> AtomicUnits.wavelength_nm_to_omega_au(800) # angular freq for 800 nm
0.05696...
"""
# --------------------------------------------------------------------- #
# Fundamental constants: 1 a.u. of <quantity> expressed in other units #
# (kept for backwards compatibility) #
# --------------------------------------------------------------------- #
meter: float = 5.2917720859e-11 # 1 a.u. length → m (Bohr radius)
nm: float = 5.2917721e-2 # 1 a.u. length → nm
angstrom: float = 5.2917721e-1 # 1 a.u. length → Å
second: float = 2.418884328e-17 # 1 a.u. time → s
fs: float = 2.418884328e-2 # 1 a.u. time → fs
Joule: float = 4.359743935e-18 # 1 a.u. energy → J (Hartree)
eV: float = 27.21138383 # 1 a.u. energy → eV (Hartree)
Volts_per_meter: float = 5.142206313e11 # 1 a.u. field → V/m
Volts_per_Angstrom: float = 51.42206313 # 1 a.u. field → V/Å
speed_of_light: float = 137.035999 # c in a.u. (= 1/α)
Coulomb: float = 1.60217646e-19 # 1 a.u. charge → C (elementary charge)
PW_per_cm2_au: float = 0.02849451308 # 1 PW/cm² in a.u. of intensity
# --------------------------------------------------------------------- #
# Length conversions #
# --------------------------------------------------------------------- #
@staticmethod
def au_to_nm(length_au: float) -> float:
"""Convert length from atomic units to nanometres."""
return length_au * AtomicUnits.nm
@staticmethod
def nm_to_au(length_nm: float) -> float:
"""Convert length from nanometres to atomic units."""
return length_nm / AtomicUnits.nm
@staticmethod
def au_to_m(length_au: float) -> float:
"""Convert length from atomic units to metres."""
return length_au * AtomicUnits.meter
@staticmethod
def m_to_au(length_m: float) -> float:
"""Convert length from metres to atomic units."""
return length_m / AtomicUnits.meter
@staticmethod
def au_to_angstrom(length_au: float) -> float:
"""Convert length from atomic units to ångström."""
return length_au * AtomicUnits.angstrom
@staticmethod
def angstrom_to_au(length_angstrom: float) -> float:
"""Convert length from ångström to atomic units."""
return length_angstrom / AtomicUnits.angstrom
# --------------------------------------------------------------------- #
# Time conversions #
# --------------------------------------------------------------------- #
@staticmethod
def au_to_fs(time_au: float) -> float:
"""Convert time from atomic units to femtoseconds."""
return time_au * AtomicUnits.fs
@staticmethod
def fs_to_au(time_fs: float) -> float:
"""Convert time from femtoseconds to atomic units."""
return time_fs / AtomicUnits.fs
@staticmethod
def au_to_s(time_au: float) -> float:
"""Convert time from atomic units to seconds."""
return time_au * AtomicUnits.second
@staticmethod
def s_to_au(time_s: float) -> float:
"""Convert time from seconds to atomic units."""
return time_s / AtomicUnits.second
# --------------------------------------------------------------------- #
# Energy conversions #
# --------------------------------------------------------------------- #
@staticmethod
def au_to_eV(energy_au: float) -> float:
"""Convert energy from atomic units (Hartree) to electronvolts."""
return energy_au * AtomicUnits.eV
@staticmethod
def eV_to_au(energy_eV: float) -> float:
"""Convert energy from electronvolts to atomic units (Hartree)."""
return energy_eV / AtomicUnits.eV
@staticmethod
def au_to_joule(energy_au: float) -> float:
"""Convert energy from atomic units (Hartree) to joules."""
return energy_au * AtomicUnits.Joule
@staticmethod
def joule_to_au(energy_J: float) -> float:
"""Convert energy from joules to atomic units (Hartree)."""
return energy_J / AtomicUnits.Joule
# --------------------------------------------------------------------- #
# Electric field conversions #
# --------------------------------------------------------------------- #
@staticmethod
def au_to_Vm(field_au: float) -> float:
"""Convert electric field from atomic units to V/m."""
return field_au * AtomicUnits.Volts_per_meter
@staticmethod
def Vm_to_au(field_Vm: float) -> float:
"""Convert electric field from V/m to atomic units."""
return field_Vm / AtomicUnits.Volts_per_meter
@staticmethod
def au_to_VA(field_au: float) -> float:
"""Convert electric field from atomic units to V/Å."""
return field_au * AtomicUnits.Volts_per_Angstrom
@staticmethod
def VA_to_au(field_VA: float) -> float:
"""Convert electric field from V/Å to atomic units."""
return field_VA / AtomicUnits.Volts_per_Angstrom
# --------------------------------------------------------------------- #
# Frequency / wavelength / photon-energy conversions #
# (all angular frequencies; in a.u. ω = E_photon since ħ = 1) #
# --------------------------------------------------------------------- #
@staticmethod
def wavelength_nm_to_omega_au(wavelength_nm: float) -> float:
"""Convert laser wavelength (nm) to angular frequency in atomic units.
Args:
wavelength_nm: Wavelength in nanometres.
Returns:
Angular frequency ω in atomic units (ħ = 1, so ω = E_photon).
"""
return (
2.0 * np.pi * AtomicUnits.speed_of_light / (wavelength_nm / AtomicUnits.nm)
)
@staticmethod
def omega_au_to_wavelength_nm(omega_au: float) -> float:
"""Convert angular frequency (a.u.) to laser wavelength in nm.
Args:
omega_au: Angular frequency in atomic units.
Returns:
Wavelength in nanometres.
"""
return 2.0 * np.pi * AtomicUnits.speed_of_light * AtomicUnits.nm / omega_au
@staticmethod
def photon_energy_eV_to_omega_au(photon_energy_eV: float) -> float:
"""Convert photon energy (eV) to angular frequency in atomic units.
In atomic units ħ = 1, so ω = E_photon (both in Hartree).
Args:
photon_energy_eV: Photon energy in electronvolts.
Returns:
Angular frequency in atomic units.
"""
return photon_energy_eV / AtomicUnits.eV
@staticmethod
def omega_au_to_photon_energy_eV(omega_au: float) -> float:
"""Convert angular frequency (a.u.) to photon energy in eV.
Args:
omega_au: Angular frequency in atomic units.
Returns:
Photon energy in electronvolts.
"""
return omega_au * AtomicUnits.eV
# --------------------------------------------------------------------- #
# Laser intensity ↔ field amplitude conversions #
# --------------------------------------------------------------------- #
@staticmethod
def intensity_Wcm2_to_field_au(intensity_Wcm2: float) -> float:
"""Convert laser peak intensity (W/cm²) to electric field amplitude in a.u.
Relation: E₀ = sqrt(I × PW_per_cm2_au × 10⁻¹⁵)
Args:
intensity_Wcm2: Peak laser intensity in W/cm².
Returns:
Peak electric field amplitude in atomic units.
"""
return np.sqrt(intensity_Wcm2 * AtomicUnits.PW_per_cm2_au * 1e-15)
@staticmethod
def field_au_to_intensity_Wcm2(field_au: float) -> float:
"""Convert electric field amplitude (a.u.) to laser peak intensity (W/cm²).
Args:
field_au: Peak electric field amplitude in atomic units.
Returns:
Peak laser intensity in W/cm².
"""
return field_au**2 / (AtomicUnits.PW_per_cm2_au * 1e-15)
@staticmethod
def intensity_Wcm2_to_au(intensity_Wcm2: float) -> float:
"""Convert laser intensity from W/cm² to atomic units of *intensity*.
.. note::
This returns the **intensity** in atomic units (I_au = I × factor),
**not** the electric field amplitude. For the field amplitude use
:meth:`intensity_Wcm2_to_field_au`, which returns
``sqrt(I × factor)`` and is ~18× larger for typical laser values.
Args:
intensity_Wcm2: Laser intensity in W/cm².
Returns:
Intensity in atomic units (not field amplitude).
"""
return intensity_Wcm2 * AtomicUnits.PW_per_cm2_au * 1e-15
@staticmethod
def au_to_intensity_Wcm2(intensity_au: float) -> float:
"""Convert laser intensity from atomic units to W/cm².
.. note::
The inverse of :meth:`intensity_Wcm2_to_au`. If you have a field
amplitude in a.u. use :meth:`field_au_to_intensity_Wcm2` instead.
Args:
intensity_au: Laser intensity in atomic units.
Returns:
Intensity in W/cm².
"""
return intensity_au / (AtomicUnits.PW_per_cm2_au * 1e-15)
# --------------------------------------------------------------------- #
# Keldysh parameter inverses #
# (γ = ω√(2Ip) / E₀ → E₀ = ω√(2Ip) / γ) #
# --------------------------------------------------------------------- #
@staticmethod
def intensity_for_keldysh(
gamma: float, wavelength_nm: float, Ip_au: float
) -> float:
"""Return the peak intensity (W/cm²) that yields a target Keldysh parameter.
Derived from γ = ω√(2Iₚ) / E₀:
.. math::
I = \\left(\\frac{\\omega\\sqrt{2I_p}}{\\gamma}\\right)^2
\\times \\frac{1}{\\mathrm{PW\\_per\\_cm2\\_au} \\times 10^{-15}}
Args:
gamma: Target Keldysh parameter (γ < 1 → tunnel, γ > 1 → multiphoton).
wavelength_nm: Laser wavelength in nm.
Ip_au: Ionization potential in atomic units.
Returns:
Peak laser intensity in W/cm².
Example:
>>> # 800 nm, hydrogen, border between tunnel and multiphoton
>>> AtomicUnits.intensity_for_keldysh(1.0, 800, 0.5)
1.41...e+14
"""
omega = AtomicUnits.wavelength_nm_to_omega_au(wavelength_nm)
E0 = omega * np.sqrt(2.0 * Ip_au) / gamma
return AtomicUnits.field_au_to_intensity_Wcm2(E0)
@staticmethod
def wavelength_for_keldysh(
gamma: float, intensity_Wcm2: float, Ip_au: float
) -> float:
"""Return the laser wavelength (nm) that yields a target Keldysh parameter.
Derived from γ = ω√(2Iₚ) / E₀:
.. math::
\\lambda = \\frac{2\\pi c}{\\omega}, \\quad
\\omega = \\frac{\\gamma E_0}{\\sqrt{2 I_p}}
Args:
gamma: Target Keldysh parameter.
intensity_Wcm2: Peak laser intensity in W/cm².
Ip_au: Ionization potential in atomic units.
Returns:
Laser wavelength in nm.
Example:
>>> # What wavelength gives γ=1 for hydrogen at 1e14 W/cm²?
>>> AtomicUnits.wavelength_for_keldysh(1.0, 1e14, 0.5)
749...
"""
E0 = AtomicUnits.intensity_Wcm2_to_field_au(intensity_Wcm2)
omega = gamma * E0 / np.sqrt(2.0 * Ip_au)
return AtomicUnits.omega_au_to_wavelength_nm(omega)
@njit()
def soft_window(
x_grid: npt.NDArray[np.float64], x_begin: float, x_end: float
) -> npt.NDArray[np.float64]:
"""Compute a soft window function.
Creates a window function that smoothly transitions from 1 to 0 between
x_begin and x_end. The window is constant outside this range.
Args:
x_grid: Array of x values
x_begin: Start point of transition
x_end: End point of transition, can be smaller than x_begin
Returns:
Array containing window values at each x point
"""
window = np.zeros(len(x_grid))
x_min = min((x_begin, x_end))
x_max = max((x_begin, x_end))
# Find transition region indices
u = np.nonzero(x_grid < x_min)[0]
i1 = min((u[-1] + 1, len(x_grid))) if len(u) > 0 else 0
u = np.nonzero(x_grid > x_max)[0]
i2 = u[0] if len(u) > 0 else len(x_grid)
# Compute window values
if x_begin <= x_end:
window[:i1] = 1
if i2 > i1:
window[i1:i2] = (
np.cos(np.pi / 2.0 * (x_grid[i1:i2] - x_min) / (x_max - x_min)) ** 2
)
else:
window[i2:] = 1
if i2 > i1:
window[i1:i2] = (
np.sin(np.pi / 2.0 * (x_grid[i1:i2] - x_min) / (x_max - x_min)) ** 2
)
return window
@njit(fastmath={"contract"}, cache=True, inline="always")
def trapz_jit(y, x):
s = 0.0
for i in range(1, len(x)):
s += (y[i] + y[i - 1]) * (x[i] - x[i - 1]) / 2.0
return s
@njit(parallel=True, fastmath=False, cache=True)
def find_zero_crossings(X, Y):
"""Find all the zero crossings: y(x) = 0
Parameters
----------
X : a 1D float array of x-values sorted in ascending order;
the array may not have identical elements;
Y : a float array of the same shape as X;
Returns
-------
out : an array of x-values where the linearly interpolated function y(x)
has zero values (an empty list if there are no zero crossings).
"""
Z = Y[:-1] * Y[1:]
out = []
for i in np.nonzero(Z <= 0)[0]:
if Z[i] == 0:
if Y[i] == 0:
out.append(X[i])
else:
# there is a zero crossing between X[i] and X[i+1]
out.append((X[i] * Y[i + 1] - X[i + 1] * Y[i]) / (Y[i + 1] - Y[i]))
return np.array(out)
@njit(parallel=True, fastmath=False, cache=True)
def find_extrema_positions(X, Y):
"""Find all the extrema of the given function
Parameters
----------
X : a 1D float array of x-values sorted in ascending order;
the array may not have identical elements;
Y : a float array of the same shape as X;
Returns
-------
out : an array of x-values where the linearly interpolated y'(x)
has zero values (an empty list if there are no such x-values).
"""
dY_dX = (Y[1:] - Y[:-1]) / (X[1:] - X[:-1])
return find_zero_crossings(0.5 * (X[1:] + X[:-1]), dY_dX)
@njit(parallel=True, fastmath=False, cache=False)
def integrate_oscillating_function_jit_old(X, f, phi, phase_step_threshold=1e-3):
r"""The function evaluates \int dx f(x) exp[i phi(x)] using an algorithm
suitable for integrating quickly oscillating functions.
Parameters
----------
X: a vector of sorted x-values;
f: either a vector or a 2D matrix where each column contains
the values of a function f(x);
phi: either a vector or a 2D matrix where each column contains
the values of a complex-valued phase phi(x);
phase_step_threshold (float): a small positive number; the formula that
approximates the integration of an oscillating function over a
small interval contains phi(x+dx)-phi(x) in the denominator;
this parameter prevents divisions by very small numbers.
Returns
-------
result: a row vector where elements correspond to the columns in
'f' and 'phi'.
"""
# check that the input data is OK
assert X.shape[0] == f.shape[0]
assert X.shape[0] == phi.shape[0]
phi = -1j * phi
# evaluate the integral(s)
dx = X[1:] - X[:-1]
result = np.zeros((f.shape[1]), dtype=np.complex128)
for i in range(f.shape[1]):
f1 = f[:-1, i]
f2 = f[1:, i]
phi1 = phi[1:, i]
phi2 = phi[:-1, i]
df = f2 - f1
dphi = phi1 - phi2
Z = np.where(
np.abs(dphi).real < phase_step_threshold,
(0.5 * (f1 + f2) + 0.125 * 1j * dphi * df)
* np.exp(0.5 * 1j * (phi1 + phi2)),
1
/ dphi**2
* (
np.exp(1j * phi1) * (df - 1j * f2 * dphi)
- (df - 1j * f1 * dphi) * np.exp(1j * phi2)
),
)
result[i] = np.sum(Z * dx)
return result
# Ensure your aggressive_safe_fastmath dict is available in this file
aggressive_safe_fastmath = {"contract", "reassoc", "nnan", "ninf", "nsz", "arcp"}
@njit(parallel=False, fastmath=aggressive_safe_fastmath, cache=False)
def integrate_oscillating_function_jit(X, f, phi, phase_step_threshold=1e-3):
r"""The function evaluates \int dx f(x) exp[i phi(x)] using an algorithm
suitable for integrating quickly oscillating functions.
"""
# check that the input data is OK
assert X.shape[0] == f.shape[0]
assert X.shape[0] == phi.shape[0]
N_X = X.shape[0]
N_cols = f.shape[1]
result = np.zeros(N_cols, dtype=np.complex128)
# 1. PARALLELIZE THE COLUMNS (Multithreading activated!)
for i in prange(N_cols):
col_sum = 0.0 + 0.0j
# 2. SCALAR INNER LOOP (Zero memory allocation)
for k in range(N_X - 1):
dx = X[k + 1] - X[k]
f1 = f[k, i]
f2 = f[k + 1, i]
# Apply the -1j scaling directly to the scalars
# Note: based on your original slicing, phi1 corresponds to k+1, phi2 to k
phi1 = -1j * phi[k + 1, i]
phi2 = -1j * phi[k, i]
df = f2 - f1
dphi = phi1 - phi2
# 3. FAST BRANCHING (Evaluates ONLY the math it needs to!)
if abs(dphi.real) < phase_step_threshold:
Z = (0.5 * (f1 + f2) + 0.125j * dphi * df) * np.exp(
0.5j * (phi1 + phi2)
)
else:
Z = (1.0 / (dphi * dphi)) * (
np.exp(1j * phi1) * (df - 1j * f2 * dphi)
- (df - 1j * f1 * dphi) * np.exp(1j * phi2)
)
col_sum += Z * dx
result[i] = col_sum
return result
@njit(fastmath={"contract"}, inline="always")
def safe_diff_linear(base_arr, cum_arr, tm, tp, dT, threshold=10):
"""Hybrid integrator for linear arrays to prevent catastrophic cancellation."""
if (tp - tm) < threshold:
val = 0.0
for k in range(tm, tp):
val += base_arr[k]
return val * dT
else:
return cum_arr[tp] - cum_arr[tm]
@njit(fastmath={"contract"}, inline="always")
def safe_diff_squared(base_arr, cum_arr, tm, tp, dT, threshold=10):
"""Hybrid integrator for squared arrays to prevent catastrophic cancellation."""
if (tp - tm) < threshold:
val = 0.0
for k in range(tm, tp):
val += base_arr[k] * base_arr[k]
return val * dT
else:
return cum_arr[tp] - cum_arr[tm]
### the following are a set of function for an experimental high order oscillating function integral scheme ###
## as of right now, it is not used in the code and remains unoptimized ##
@njit(fastmath=False)
def calculate_c(dx, f, phi):
a0, a1, a2, a3 = (f[0], f[1], f[2], f[3])
b0, b1 = (phi[0], phi[1])
c0 = (
1
/ (2 * dx**5)
* (
-2 * (6 * a3 + 3 * a2 * dx + (a1 + a3 * b1) * dx**2)
+ (
12 * a3
+ 6 * a2 * dx
+ 2 * (a1 - 5 * a3 * b1) * dx**2
- 6 * (2 * a3 * b0 + a2 * b1) * dx**3
- 2 * (3 * a2 * b0 + b1 * (a1 - 2 * a3 * b1)) * dx**4
+ 2 * b1 * (a0 + 6 * a3 * b0 + 2 * a2 * b1) * dx**5
+ (6 * a0 * b0 + 9 * a3 * b0**2 + 4 * b1 * (3 * a2 * b0 + a1 * b1))
* dx**6
+ (9 * a2 * b0**2 + 4 * b1 * (3 * a1 * b0 + a0 * b1)) * dx**7
+ 3 * b0 * (3 * a1 * b0 + 4 * a0 * b1) * dx**8
+ 9 * a0 * b0**2 * dx**9
)
* np.exp(dx**2 * (b1 + b0 * dx))
)
)
c1 = (
1
/ dx**4
* (
dx * (8 * a2 + 3 * a1 * dx)
+ 3 * a3 * (5 + b1 * dx**2)
- (
15 * a3
+ 8 * a2 * dx
+ 3 * (a1 - 4 * a3 * b1) * dx**2
- (15 * a3 * b0 + 8 * a2 * b1) * dx**3
+ (-9 * a2 * b0 + 4 * b1 * (-a1 + a3 * b1)) * dx**4
+ (-3 * a1 * b0 + 4 * b1 * (3 * a3 * b0 + a2 * b1)) * dx**5
+ (3 * a0 * b0 + 9 * a3 * b0**2 + 4 * b1 * (3 * a2 * b0 + a1 * b1))
* dx**6
+ (9 * a2 * b0**2 + 4 * b1 * (3 * a1 * b0 + a0 * b1)) * dx**7
+ 3 * b0 * (3 * a1 * b0 + 4 * a0 * b1) * dx**8
+ 9 * a0 * b0**2 * dx**9
)
* np.exp(dx**2 * (b1 + b0 * dx))
)
)
c2 = (
1
/ (2 * dx**3)
* (
-20 * a3
- 12 * a2 * dx
- 6 * (a1 + a3 * b1) * dx**2
+ (
6 * dx * (2 * a2 + a1 * dx)
+ a3
* (
20
+ 4 * b1**2 * dx**4
+ 9 * b0 * dx**3 * (-2 + b0 * dx**3)
+ 2 * b1 * dx**2 * (-7 + 6 * b0 * dx**3)
)
+ dx**3
* (
-6 * a1 * dx * (b1 + b0 * dx)
+ dx**2 * (a2 + a1 * dx) * (2 * b1 + 3 * b0 * dx) ** 2
- 2 * a2 * (5 * b1 + 6 * b0 * dx)
+ a0
* (
2
+ 4 * b1**2 * dx**4
+ 9 * b0**2 * dx**6
+ 2 * b1 * dx**2 * (-1 + 6 * b0 * dx**3)
)
)
)
* np.exp(dx**2 * (b1 + b0 * dx))
)
)
c3 = a1 + a3 * b1
c4 = a2
c5 = a3
return {
"c0": c0,
"c1": c1,
"c2": c2,
"c3": c3,
"c4": c4,
"c5": c5,
"b2": phi[2],
"b3": phi[3],
}
# @jit(fastmath = False)
def calculate_expression(dx, b2, c0, c1, c2, c3, c4, c5):
term1 = 120 * c0 + b2 * (
-24 * c1 + b2 * (6 * c2 + b2 * (-2 * c3 + b2 * (c4 - b2 * c5)))
)
term3 = 120 * c0 + b2 * (-24 * c1 + b2 * (6 * c2 - 2 * b2 * c3 + b2**2 * c4))
term4 = -60 * c0 + b2 * (12 * c1 + b2 * (-3 * c2 + b2 * c3))
term5 = 20 * c0 + b2 * (-4 * c1 + b2 * c2)
term6 = -5 * c0 + b2 * c1
d = b2 * dx
return (
-term1 * np.expm1(d)
+ np.exp(d)
* d
* np.polynomial.polynomial.polyval(
d, [c0, term6, term5, term4, term3][::-1], tensor=False
)
) / b2**6
# @jit(fastmath = False)
def calculate_expression_appx(dx, b2, c0, c1, c2, c3, c4, c5):
term1 = c5 # * dx
term2 = 1 / 2 * (c4 + b2 * c5) # * dx**2
term3 = 1 / 6 * (2 * c3 + b2 * (2 * c4 + b2 * c5)) # * dx**3
term4 = 1 / 8 * (2 * c2 + b2 * (2 * c3 + b2 * c4)) # * dx**4
term5 = 1 / 10 * (2 * c1 + b2 * (2 * c2 + b2 * c3)) # * dx**5
term6 = 1 / 12 * (2 * c0 + b2 * (2 * c1 + b2 * c2)) # * dx**6
term7 = 1 / 14 * b2 * (2 * c0 + b2 * c1) # * dx**7
term8 = 1 / 16 * b2**2 * c0 # * dx**8
result = np.polynomial.polynomial.polyval(
dx, [term8, term7, term6, term5, term4, term3, term2, term1][::-1], tensor=False
)
return result * dx
# @jit(fastmath = False)
def calculate_expression_main(dx, b2, b3, c0, c1, c2, c3, c4, c5):
s = np.abs(b2) <= 1e-10
# print(s.shape)
result = np.zeros(b2.shape)
if np.any(s):
result = calculate_expression_appx(dx, b2, c0, c1, c2, c3, c4, c5)
s = np.logical_not(s)
if np.any(s):
dx, b2, c0, c1, c2, c3, c4, c5 = (
np.tile(dx, b2.shape[1])[s],
b2[s],
c0[s],
c1[s],
c2[s],
c3[s],
c4[s],
c5[s],
)
# print(dx.shape, s)
result[s] = calculate_expression(dx, b2, c0, c1, c2, c3, c4, c5)
return result * np.exp(b3)
def integrate_oscillating_function_splined(X, f, phi, phase_step_threshold=1e-3):
r"""The function evaluates \int dx f(x) exp[i phi(x)] using an algorithm
suitable for integrating quickly oscillating functions.
Parameters
----------
X: a vector of sorted x-values;
f: either a vector or a matrix where each column contains
the values of a function f(x);
phi: either a vector or a matrix where each column contains
the values of a real-valued phase phi(x);
phase_step_threshold (float): a small positive number; the formula that
approximates the integration of an oscillating function over a
small interval contains phi(x+dx)-phi(x) in the denominator;
this parameter prevents divisions by very small numbers.
Returns
-------
result: a row vector where elements correspond to the columns in
'f' and 'phi'.
"""
# check that the input data is OK
assert X.shape[0] == f.shape[0]
assert X.shape[0] == phi.shape[0]
# assert(np.all(np.imag(phi)) == 0)
# evaluate the integral(s)
dx = X[1:] - X[:-1]
# dx = np.append(X[1:] - X[:-1],0.)
dx = dx[:, None] # np.tile(dx, (f.shape[1],1)).T
fR = CubicSpline(X, f.real, bc_type="natural").c
fI = CubicSpline(X, f.imag, bc_type="natural").c
f = np.vectorize(complex)(fR, fI)
phiR = CubicSpline(X, phi.real, bc_type="natural").c
phiI = CubicSpline(X, phi.imag, bc_type="natural").c
phi = np.vectorize(complex)(phiR, phiI)
# print(np.polyval(f, dx)*np.exp(np.polyval(phi, dx)))
# print(np.sum(np.polyval(f, dx)*np.exp(np.polyval(phi, dx)), axis=0))
# print(f.shape, phi.shape, dx.shape)
cDict = calculate_c(dx, f, phi)
expr = calculate_expression_main(dx, **cDict)
plt.plot(X[1:], np.cumsum(expr[:, 0], axis=0))
plt.show()
plt.close()
Z = np.sum(expr, axis=0)
return Z
Z = np.sum(
np.where(
abs(cDict["b2"]) <= 1e-3,
calculate_expression_appx(dx, **cDict),
calculate_expression(dx, **cDict),
),
axis=0,
)
return Z
@njit()
def meshgrid(x, y):
xx = np.empty(shape=(x.size, y.size), dtype=x.dtype)
yy = np.empty(shape=(x.size, y.size), dtype=y.dtype)
for j in range(y.size):
for k in range(x.size):
xx[k, j] = x[k] # change to x[k] if indexing xy
yy[k, j] = y[j] # change to y[j] if indexing xy
return xx, yy
# @njit()
def get_momentum_grid(div_p, div_theta, laser_field, Ip):
tot_energy = laser_field.get_total_energy(Ip)
a1_injection = np.diff(laser_field.get_time_interval())[0]
p_phys_max = np.sqrt(2 * 3.17 * tot_energy)
dcostheta = 1 / (2**2 * p_phys_max) / div_theta
density_theta = int(np.log2(2 / dcostheta - 1)) + 1
length_cos = max(15, int(2 ** (density_theta)) + 1)
Theta_grid = np.linspace(0, np.pi, length_cos)
intA = laser_field.get_cummulative_vector_potential(laser_field.get_tgrid())
Int_max = max(max(intA) - min(intA), 1)
dp = (
min(10 * np.pi / (a1_injection), np.pi / abs(4 * Int_max)) / div_p
) # np.pi/(2*a1_injection*4)
p_grid = np.unique(
np.concatenate(
(
np.arange(0, p_phys_max + dp, dp),
np.arange(p_phys_max + dp, 2 * p_phys_max, dp * 8),
)
)
)
window = soft_window(p_grid, p_phys_max, 2 * p_phys_max)
return p_grid, Theta_grid, window
def FourierTransform(t, Y, omega=None, is_periodic=False, a0=None):
r"""Apply the FFT in an easy-to-use way.
Parameters
----------
t (N_t): time discretization of Y(t); the array must be sorted in
ascending order;
Y (N_t, N_data): time-dependent data to Fourier transform;
omega (N_omega): circular frequencies; the array must be sorted in
ascending order;
is_periodic: if 'true', then Y(t) is assumed to be periodic, and the
time step is expected to be constant, so that Y(t[-1]+dt) = Y(t[0]);
if 'false', the function takes some precaution to avoid artifacts
induced by the fact that discrete Fourier transform implicitly
assumes periodic data;
a0 (scalar or an array of length N_data): if the input data represents
a pulse that is not centered in the middle of its time window, then
knowing the center of the pulse helps avoid interpolation artifacts;
this parameter specifies the central time of each pulse;
Returns
-------
If omega is None, the function returns a tuple of two elements: the
Fourier-transformed 'Y' and an array of circular frequencies;
otherwise, the function returns F[Y] at the specified frequencies:
$ F[Y](\omega) = \int_{-\infty}^\infty dt Y(t) \exp(i \omega t). $
"""
max_N_fft = 2**20
N_t = len(t)
shape_length = len(Y.shape)
if shape_length == 1:
Y = Y.reshape((N_t, 1))
N_data = Y.shape[1]
assert N_t > 1
dt = np.min(np.diff(t))
if is_periodic:
T = t[-1] - t[0] + np.mean(np.diff(t)) # period assuming a constant step
if omega is None:
dw = 2 * np.pi / (dt * N_t)
else:
if len(omega) == 1: # no need to do FFT
integrand = Y * np.exp(1j * omega * t.reshape((N_t, 1)))
result = np.trapezoid(integrand, t, axis=0)
if is_periodic:
result += 0.5 * dt * (integrand[0] + integrand[-1])
return result
dw = np.min(np.diff(omega))
# FFT discretization
if is_periodic:
N_fft = N_t
while 2 * np.pi / (dt * N_fft) > 1.1 * dw and N_fft < max_N_fft:
N_fft *= 2
else:
N_fft = 2 ** int(round(np.log(2 * np.pi / (dt * dw)) / np.log(2.0)))
while (N_fft - 1) * dt < t[-1] - t[0] and N_fft < max_N_fft:
N_fft *= 2
if not (omega is None):
while (
omega[-1] > np.pi / dt * (1 - 2.0 / N_fft) or -omega[0] > np.pi / dt
) and N_fft < max_N_fft:
dt /= 2.0
N_fft *= 2
dw = 2 * np.pi / (dt * N_fft)
if a0 is None:
tc = 0.5 * (t[0] + t[-1])
i0 = np.argmin(np.abs(t - tc)) # this element will be treated as t=0
a0 = t[i0] * np.ones(N_data)
t_grid = dt * np.fft.ifftshift(np.arange(N_fft) - N_fft // 2)
Z = np.zeros((N_fft, N_data), dtype=Y.dtype)
for j in range(N_data):
if np.isscalar(a0):
tt = t - a0
else:
tt = t - a0[j]
if is_periodic:
Z[:, j] = np.interp(t_grid, tt, Y[:, j], period=T)
else:
Z[:, j] = np.interp(t_grid, tt, Y[:, j], left=0.0, right=0.0)
if np.isscalar(a0):
a0 = a0 * np.ones((1, N_data))
else:
a0 = np.reshape(a0, (1, N_data))
# FFT
Z = np.fft.fftshift(np.fft.ifft(Z, axis=0), axes=0) * (N_fft * dt)
w_grid = dw * (np.arange(N_fft) - N_fft // 2)
if omega is None:
return (Z * np.exp(1j * a0 * w_grid.reshape((N_fft, 1))), w_grid)
## return (Z[N_fft//2:,:], w_grid[N_fft//2:])
# interpolate the result
result = np.zeros((len(omega), N_data), dtype=complex)
for j in range(N_data):
result[:, j] = np.interp(omega, w_grid, Z[:, j], left=0.0, right=0.0)
# correct for a0
result = result * np.exp(1j * a0 * omega.reshape((len(omega), 1)))
if shape_length == 1:
result = result.reshape(len(omega))
return result
def get_cummulative_integral_spline(
t: npt.NDArray[np.float64], Y: npt.NDArray[np.float64 | np.complex128]
) -> CubicSpline:
"""Calculate the cumulative integral of a function Y(t) over a time grid.
This function computes the cumulative integral of Y(t) over the time grid
defined by t and returns a scipy PPoly instance.
Args:
t: Input time grid (1D array)..
Y: Function values at t_in (1D array).
Returns:
Cumulative integral values at t_out.
"""
integral = cumulative_simpson(Y, x=t, initial=0.0, axis=0)
# Interpolate the cumulative integral to the output time grid
return CubicSpline(t, integral, bc_type="natural")
def align_yaxis(ax1, ax2):
"""Align zeros of the two axes, zooming them out by same ratio"""
axes = (ax1, ax2)
extrema = [ax.get_ylim() for ax in axes]
tops = [extr[1] / (extr[1] - extr[0]) for extr in extrema]
# Ensure that plots (intervals) are ordered bottom to top:
if tops[0] > tops[1]:
axes, extrema, tops = [list(reversed(l)) for l in (axes, extrema, tops)]
# How much would the plot overflow if we kept current zoom levels?
tot_span = tops[1] + 1 - tops[0]
b_new_t = extrema[0][0] + tot_span * (extrema[0][1] - extrema[0][0])
t_new_b = extrema[1][1] - tot_span * (extrema[1][1] - extrema[1][0])
axes[0].set_ylim(extrema[0][0], b_new_t)
axes[1].set_ylim(t_new_b, extrema[1][1])
def align_yaxis_np(ax1, ax2):
"""Align zeros of the two axes, zooming them out by same ratio"""
axes = np.array([ax1, ax2])
extrema = np.array([ax.get_ylim() for ax in axes])
tops = extrema[:, 1] / (extrema[:, 1] - extrema[:, 0])
# Ensure that plots (intervals) are ordered bottom to top:
if tops[0] > tops[1]:
axes, extrema, tops = [a[::-1] for a in (axes, extrema, tops)]
# How much would the plot overflow if we kept current zoom levels?
tot_span = tops[1] + 1 - tops[0]
extrema[0, 1] = extrema[0, 0] + tot_span * (extrema[0, 1] - extrema[0, 0])
extrema[1, 0] = extrema[1, 1] + tot_span * (extrema[1, 0] - extrema[1, 1])
[axes[i].set_ylim(*extrema[i]) for i in range(2)]
def dr_dt_3D(t, r, E_inj, E_drive, beta, coul_fact=1):
"""Integration of the governing vector differential equation.
Assuming that the electric field along y is zero, the y component is completely ignored
E = External Electric field vector at time t
d2r_dt2 = -Electric_field -1/(R^3+beta)*r with d2r_dt2 and r as vecotors.
Initial position and velocitz are given.
r[0:1] = position components
r[2:] = velocity components
beta: soft_core coulomb potential set to 0.2 such that the coulomb potential truncates at -5 au (GS H at -0.5 au)
"""
x = r[0, :]
y = r[1, :]
z = r[2, :]
v_x = r[3, :]
v_y = r[4, :]
v_z = r[5, :]
r_ab = np.sqrt(x**2 + z**2 + (beta) ** 2 + y**2)
# coul_fact=coul_fact*1/(1+np.exp(2*(r_ab-55)))
# k=40
dx = v_x
dy = v_y
dz = v_z
# if r_ab<60:
dvx = -E_drive(t) - x * coul_fact / (r_ab**3) # *logistic.cdf(-2*k*r_ab)
dvz = -E_inj(t) - z * coul_fact / (r_ab**3) # *logistic.cdf(-2*k*r_ab)
dvy = -1 * y * coul_fact / (r_ab**3)
# else:
# dvx=-E_drive(t)
# dvz=-E_inj(t)
return [dx, dy, dz, dvx, dvy, dvz]
def dr_dt_3D_recol_coul_free(t, r, t_b, E_inj, E_drive, beta, coul_fact=1):
"""Integration of the governing vector differential equation.
Assuming that the electric field along y is zero, the y component is completely ignored
E = External Electric field vector at time t
d2r_dt2 = -Electric_field -1/(R^3+beta)*r with d2r_dt2 and r as vecotors.
Initial position and velocitz are given.
r[0:1] = position components
r[2:] = velocity components
beta: soft_core coulomb potential set to 2 such that the coulomb potential at r=0 is -0.5 au (GS H at -0.5 au)
"""
x = r[0]
y = r[1]
z = r[2]
v_x = r[3]
v_y = r[4]
v_z = r[5]
r_ab = np.sqrt(x**2 + z**2 + (beta) ** 2 + y**2)
# coul_fact_off=1
# if t-t_b>30:
coul_fact_off = 1 / (1 + np.exp(2 * (t - t_b - 50))) + 1 / (
1 + np.exp(-2 * (t - t_b - 110))
)
# k=40
dx = v_x
dy = v_y
dz = v_z
# if r_ab<60:
dvx = -E_drive(t) - x * coul_fact / (
r_ab**3
) # *coul_fact_off#*logistic.cdf(-2*k*r_ab)
dvz = -E_inj(t) - z * coul_fact / (
r_ab**3
) # *coul_fact_off#*logistic.cdf(-2*k*r_ab)
dvy = -1 * y * coul_fact / (r_ab**3) # *coul_fact_off
# else:
# dvx=-E_drive(t)
# dvz=-E_inj(t)
print(f"t={t}, x={dx}, z={dz}, dvx={dvx}, dvz={dvz}")
return [dx, dy, dz, dvx, dvy, dvz]
def Z_cross(t, r, E_inj, E_drive, beta, coul_fact=1, t_b=0):
return r[1]
def jacobian(t, r, E_inj, E_drive, beta, coul_fact=1, t_b=0):
x = r[0]
y = r[1]
z = r[2]
v_x = r[3]
v_y = r[4]
v_z = r[5]
r_ab = np.sqrt(x**2 + z**2 + (beta) ** 2 + y**2)
# k=40
dx = v_x
dy = v_y
dz = v_z
f_0 = [0, 0, 0, 1, 0, 0]
f_1 = [0, 0, 0, 0, 1, 0]
f_2 = [0, 0, 0, 0, 0, 1]
f_3 = [
-1 * coul_fact * (z**2 + beta**2 + y**2 - 2 * x**2) / r_ab**5,
3 * coul_fact * x * y / r_ab**5,
3 * coul_fact * x * z / r_ab**5,
0,
0,
0,
]
f_4 = [
3 * coul_fact * y * x / r_ab**5,
-1 * coul_fact * (x**2 + beta**2 + z**2 - 2 * y**2) / r_ab**5,
3 * coul_fact * y * z / r_ab**5,
0,
0,
0,
]
f_5 = [
3 * coul_fact * x * z / r_ab**5,
3 * coul_fact * y * z / r_ab**5,
-1 * coul_fact * (x**2 + beta**2 + y**2 - 2 * z**2) / r_ab**5,
0,
0,
0,
]
return [f_0, f_1, f_2, f_3, f_4, f_5]
def Z_birth(Field_inj, Field_drive, beta, GS):
"""predict the initial startig point of the electron trajectory for a given injection pulse
Field : field strength (and direction) at time t_birth
if Field>= <GS>^2/4 i.e. field is so strong that the ground state is in continuum,
z(t_birth) is equal to the the distance at which the total potential has a local minimum
else
z(t_birth)= sign(Field)*(<GS>+sqrt(<GS>^2-4*abs(Field)))/(2*abs(Field))"""
Field = np.sqrt(Field_inj**2 + Field_drive**2)
root = np.sqrt(abs(Field))
cos = Field_inj / Field
sin = Field_drive / Field
if Field == 0:
return (0, 0, 0, 0) # e not free
elif abs(GS) < 2 * root:
print(f"ABI triggered, Field={Field}")
z_shift = -1 / root
v_shift = 0 # -1*np.sqrt(2*abs(GS+2*root))
return (z_shift * sin, z_shift * cos, v_shift * sin, v_shift * cos)
else:
z_shift = (
-1 * (abs(GS) + np.sqrt(GS**2 - 4 * abs(Field))) / (2 * Field)
) # beta correction (Field+abs(GS))*beta
v_shift = 0 # -1*np.sqrt(2*abs(GS+2*root))
return (z_shift * sin, z_shift * cos, v_shift * sin, v_shift * cos)
def ddf(x, sig):
val = np.zeros_like(x)
val[(-(1 / (2 * sig)) <= x) & (x <= (1 / (2 * sig)))] = 1
return val
def delT_determine(vz):
return 0.3625