# Copyright (c) 2024 Manoram Agarwal
#
# -*- coding:utf-8 -*-
# @Script: kernels.py
# @Author: Manoram Agarwal
# @Email: manoram.agarwal@mpq.mpg.de
# @Create At: 2024-07-19 15:00:20
# @Last Modified By: Manoram Agarwal
# @Last Modified At: 2025-01-24 12:29:04
# @Description: This python file contains the function that provides three different methods to compute the kernel, rates and probabilities
import warnings
from typing import Dict, Tuple, Union
import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt
from numba import njit, prange, typed, types
from scipy.integrate import simpson
from .pulse import EllipticalCosNPulse
from .utils import (
find_extrema_positions,
find_zero_crossings,
get_momentum_grid,
)
from .utils import integrate_oscillating_function_jit as IOF
from .utils import (
meshgrid,
)
"""Kernel functions for computing ionization rates and probabilities.
This module provides functions to compute ionization rates and probabilities using
different kernel methods. The kernels are used to model the ionization process in
atoms and molecules under strong laser fields.
The module implements three different methods:
1. GASFIR (General Approximator for Strong Field Ionization Rates)
2. Exact SFA (Strong Field Approximation)
3. QS (Quasi-Static)
Parameters can be obtained either by using pre-defined values for common gases
through get_parameters() or by specifying custom values.
Typical usage example:
from gasfir import create_pulse, get_parameters
laser = create_pulse(800, 1e14, 0, 30) # 800nm, 1e14 W/cm2, CEP=0, 30 optical cycles
# Using pre-defined parameters
params = get_parameters("Hydrogen_SFA")
rates = get_diabatic_ionization_rate(t_grid, laser, params)
# Using custom parameters
custom_params = {'E_g': 0.5, 'a1': 1.0, ...}
rates = get_diabatic_ionization_rate(t_grid, laser, custom_params)
"""
########################
@njit(fastmath=False, cache=True)
def Kernel_f_term(
EFp: float | npt.NDArray[np.float64],
EFm: float | npt.NDArray[np.float64],
EFp_r: float | npt.NDArray[np.float64],
EFm_r: float | npt.NDArray[np.float64],
term0: complex | npt.NDArray[np.complex128],
term1: float | npt.NDArray[np.float64],
term2: float | npt.NDArray[np.float64],
term1_rotated: float | npt.NDArray[np.float64],
term2_rotated: float | npt.NDArray[np.float64],
Ti: float | npt.NDArray[np.float64],
params: Dict[str, float],
z1: float | npt.NDArray[np.float64],
z2: float | npt.NDArray[np.float64],
x1: float | npt.NDArray[np.float64],
x2: float | npt.NDArray[np.float64],
) -> complex | npt.NDArray[np.complex128]:
"""Compute the pre-exponential factor of the kernel.
Args:
EFp: Electric field at t+Ti
EFm: Electric field at t-Ti
term0: Complex term = 1j/(T+1j*a1)
term1: xi1=(A(t+Ti)-A(t-Ti))**2/4
term2: xi2=(A(t+Ti)/2 + A(t-Ti)/2-A_bar)**2
Ti: Time before and after moment of ionization
params: Dictionary containing the kernel parameters:
a0: Normalization factor
a1: Time constant related to smearing of Gaussian
internal_par_xi1: Multiplier of xi1**2 in numerator
a5: Multiplier of xi2**2 in numerator
a4: Power of denominator (Coulomb effects)
internal_param_pxi1: Power of xi1 term in numerator
internal_param_pxi2: Power of xi2 term in numerator
Returns:
Complex value representing the pre-exponential factor
"""
a0 = params["a0"]
a5 = params.get("a5", 0.0) # Default to 0 if not provided
a4 = params["a4"]
internal_param_xi1 = params.get("internal_param_xi1", 1)
internal_param_pxi1 = params.get("internal_param_pxi1", 1)
internal_param_pxi2 = params.get("internal_param_pxi2", 1)
multiplier = (1j * np.pi) ** 1.5 * a0 * term0 ** (1.5 + a4)
ZZ = (
EFp
* EFm
* (
2 * term0
- 4 * internal_param_xi1 * term1**internal_param_pxi1
- 4 * a5 * (term2 * term0**2 * Ti**2) ** internal_param_pxi2
)
)
XX = (
EFp_r
* EFm_r
* (
2 * term0
- 4 * internal_param_xi1 * term1_rotated**internal_param_pxi1
- 4 * a5 * (term2_rotated * term0**2 * Ti**2) ** internal_param_pxi2
)
)
term0 = term0 * 1j
ZX = -EFp * EFm_r * (z1 - term0 * z2 * Ti) * (x1 + term0 * x2 * Ti)
XZ = -EFp_r * EFm * (z1 + term0 * z2 * Ti) * (x1 - term0 * x2 * Ti)
return (ZX + XZ + ZZ + XX) * multiplier
@njit(fastmath=False, cache=True)
def Kernel_phase_term(
term0: complex | npt.NDArray[np.complex128],
term1: float | npt.NDArray[np.float64],
term2: float | npt.NDArray[np.float64],
DelAbar: float | npt.NDArray[np.float64],
DelA2bar: float | npt.NDArray[np.float64],
E2diff: float | npt.NDArray[np.float64],
Ti: float | npt.NDArray[np.float64],
params: Dict[str, float],
) -> complex | npt.NDArray[np.complex128]:
"""Compute the complex argument of the exponential function of the kernel.
Args:
term0: Complex term = 1j/(T+1j*a1)
term1: xi1=(A(t+Ti)-A(t-Ti))**2/4
term2: xi2=(A(t+Ti)/2 + A(t-Ti)/2-A_bar)**2
DelAbar: (A(t+Ti)-A(t-Ti))/2
DelA2bar: (A(t+Ti)**2-A(t-Ti)**2)/4
E2diff: E^2(t+Ti)-E^2(t-Ti)
Ti: Time before and after moment of ionization
params: Dictionary containing the kernel parameters:
E_g: Band gap or ionization potential
αPol: Polarization for Stark shift
a1: Time constant for Gaussian smearing
a2: Exponential decay factor for term1
a3: Exponential decay factor for term2
Returns:
Complex value representing the phase term
"""
E_g = params["E_g"]
αPol = params.get("αPol", 0.0) # Default to 0 if not provided
a1 = params["a1"]
a2 = params["a2"]
a3 = params.get("a3", 0.0) # Default to 0 if not provided
# for QS function, nothing needs to be updated as long as
# 1j*(Ti*(2*E_g+DelA2bar-DelAbar**2) + 3*np.pi/4 + αPol * E2diff/2 ) - a1*(a2*term1) is UNTOUCHED
# term2 is zero in QS case so changes there are irrelevant
return (
1j * (Ti * (2 * E_g + DelA2bar - DelAbar**2) + αPol * E2diff / 2)
- a1 * (a2 * term1 + a3 * term2 * Ti / (Ti + 1j * a1))
- 3j * np.pi / 4
)
@njit(parallel=True, fastmath=False, cache=True)
def Kernel_jit_helper(
tar,
Tar,
params,
EF,
EF_rotated,
EF2,
VP,
VP_rotated,
intA,
intA_rotated,
intA2,
dT,
N,
n,
nmin,
Ti_ar,
):
"""return the kernel_GASFIR(t_grid,T_grid) for a given laser field computed with provided parameters using a jit optimized implementation
Args:
tar (np.ndarray): the grid of moment of ionization
Tar (np.ndarray): array time T before and after time t that affects that moment of ionization
EF (np.ndarray): Electric field
EF2 (np.ndarray): Cummulative Electric field squared
VP (np.ndarray): Vector potential
A (np.ndarray): cummulative of the vector potential
A2 (np.ndarray): cummulative squared of the vector potential
dT (float64): step size of dense arrays used for EF, EF2 etc
n (int64): number of steps of dT needed to reach next t[i+1]
Ti_ar (np.ndarray): indices of time array T for which the values of kernel will actually be stored.
params: dictionary containing the parameters for the kernel
Returns:
f0 (np.ndarray, shape=(T_grid.size, t_grid.size)): 2d grid to store pre-exponential
phase0 (np.ndarray, shape=(T_grid.size, t_grid.size)):2d grid storing complex agument of the exponential function
"""
f0 = np.zeros((Tar.size, tar.size), dtype=np.cdouble)
phase0 = np.zeros((Tar.size, tar.size), dtype=np.cdouble)
for i in prange(Tar.size):
Ti = Ti_ar[i]
for j in range(tar.size):
tj = N + nmin + j * n
tp = tj + Ti
tm = tj - Ti
if tp >= 0 and tp < EF.size and tm >= 0 and tm < EF.size:
T = Ti * dT
a1 = params["a1"]
if T != 0.0:
E2diff = EF2[tp] - EF2[tm]
# term1 = ((VP[tp] - VP[tm])**2+(VP_rotated[tp] - VP_rotated[tm])**2)/4
z1 = VP[tp] - VP[tm]
term1 = (VP[tp] - VP[tm]) ** 2 / 4
x1 = VP_rotated[tp] - VP_rotated[tm]
term1_rotated = (VP_rotated[tp] - VP_rotated[tm]) ** 2 / 4
# xi2=0
# term2 = (VP[tp]/2 + VP[tm]/2-(intA[tp] - intA[tm])/2/T)**2+(VP_rotated[tp]/2 + VP_rotated[tm]/2-(intA_rotated[tp] - intA_rotated[tm])/2/T)**2
z2 = VP[tp] + VP[tm] - (intA[tp] - intA[tm]) / T
term2 = (
VP[tp] / 2 + VP[tm] / 2 - (intA[tp] - intA[tm]) / 2 / T
) ** 2
x2 = (
VP_rotated[tp]
+ VP_rotated[tm]
- (intA_rotated[tp] - intA_rotated[tm]) / T
)
term2_rotated = (
VP_rotated[tp] / 2
+ VP_rotated[tm] / 2
- (intA_rotated[tp] - intA_rotated[tm]) / 2 / T
) ** 2
DelAbar = (
np.sqrt(
(intA[tp] - intA[tm]) ** 2
+ (intA_rotated[tp] - intA_rotated[tm]) ** 2
)
/ 2
/ T
)
DelA2bar = (intA2[tp] - intA2[tm]) / 2 / T # +VPt**2-2*VPt*DelAbar
# DelAbar = DelAbar - VPt
term0 = 1j / (T + 1j * a1)
else:
E2diff = 0
term1 = 0
DelAbar = 0
term2 = 0
term1_rotated = 0
term2_rotated = 0
DelA2bar = 0
term0 = 1 / a1
z1, z2, x1, x2 = 0.0, 0.0, 0.0, 0.0
f0[i, j] = Kernel_f_term(
EF[tp],
EF[tm],
EF_rotated[tp],
EF_rotated[tm],
term0,
term1,
term2,
term1_rotated,
term2_rotated,
T,
params,
z1,
z2,
x1,
x2,
)
phase0[i, j] = Kernel_phase_term(
term0,
term1 + term1_rotated,
term2 + term2_rotated,
DelAbar,
DelA2bar,
E2diff,
T,
params,
)
return f0, phase0
def Kernel_jit(
t_grid: npt.NDArray[np.float64],
T_grid: npt.NDArray[np.float64],
pulse: EllipticalCosNPulse,
param_dict: Dict[str, float],
kernel_type: str = "GASFIR",
) -> Tuple[npt.NDArray[np.complex128], npt.NDArray[np.complex128]]:
"""Compute the kernel for ionization rate calculation.
Args:
t_grid: Time grid for the calculation
T_grid: Integration time grid
pulse: EllipticalCosNPulse object defining the laser field
param_dict: Dictionary containing the kernel parameters
kernel_type: Type of kernel to use ("GASFIR" or "exact_SFA")
Returns:
Tuple of (kernel_real, kernel_imag) arrays
"""
t = t_grid
T = T_grid
dt = min(np.round(np.diff(t), 4)) # rounding just to remove any numerical artifact
dT = min(np.diff(T)) # /2
t_min, t_max = pulse.get_time_interval()
a1_injection = int(max(abs(t_min), abs(t_max))) + 1
# print(dt, dT)
assert np.isclose(
dt % dT, 0, atol=1e-8
), f"dt should be approximately a multiple of dT, got dt={dt}, dT={dT}, dt%dT={dt%dT}"
n = int(dt // dT)
N = int(a1_injection // dT) + 1
# nmax=int(t[-1]//dT)
nmin = int(t[0] // dT)
tAr = np.arange(-N, N + 1, 1.0) * dT
# VP=pulse.get_vector_potential(tAr)
# VP_rotated=pulse.get_vector_potential(tAr+cep_shift)
VP, VP_rotated = pulse.get_vector_potential_tuple(tAr)
# EF=pulse.get_electric_field(tAr)
# EF_rotated=pulse.get_electric_field(tAr+cep_shift)
EF, EF_rotated = pulse.get_electric_field_tuple(tAr)
# intA=pulse.get_cummulative_vector_potential(tAr) # np.cumsum(VP*dT)
# intA_rotated=pulse.get_cummulative_vector_potential(tAr+cep_shift) # np.cumsum(VP_rotated*dT)
intA, intA_rotated = pulse.get_cummulative_vector_potential_tuple(tAr)
intA2 = pulse.get_cummulative_vector_potential_squared(tAr) # np.cumsum(VP**2*dT)
EF2 = pulse.get_cummulative_electric_field_squared(tAr) # np.cumsum(EF**2*dT)
Ti_ar = (T // dT).astype(np.int64)
# f0 = np.zeros((T.size, t.size), dtype=np.cdouble)
# phase0 = np.zeros((T.size, t.size), dtype=np.cdouble)
params = typed.Dict.empty(key_type=types.unicode_type, value_type=types.complex128)
for key, value in param_dict.items():
params[key] = value
if kernel_type == "GASFIR":
f0, phase0 = Kernel_jit_helper(
t,
T,
params,
EF,
EF_rotated,
EF2,
VP,
VP_rotated,
intA,
intA_rotated,
intA2,
dT,
N,
n,
nmin,
Ti_ar,
)
else:
raise NotImplementedError
return f0, phase0
[docs]
def get_diabatic_ionization_rate(
t_grid: npt.NDArray[np.float64],
pulse: EllipticalCosNPulse,
param_dict: Dict[str, float],
dT: float = 0.25,
kernel_type: str = "GASFIR",
ret_tgrid: bool = False,
) -> Union[
npt.NDArray[np.float64], Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]
]:
"""Compute ionization rates for a defined pulse.
Args:
t_grid: Time grid for ionization moments
pulse: EllipticalCosNPulse object defining the laser field
param_dict: Dictionary containing kernel parameters
dT: Time step for integration (default: 0.25)
kernel_type: Type of kernel to use ("GASFIR" or "exact_SFA")
ret_tgrid: Whether to return the time grid (default: False)
Returns:
If ``ret_tgrid`` is False, an array of ionization rates at each time
point. If ``ret_tgrid`` is True, a tuple ``(time_grid, rates)``.
Notes:
The function handles various input formats for t_grid:
- None: Uses laser field's default time grid
- pandas Series: Converts to numpy array
- list: Converts to numpy array
- numpy array: Used as is
"""
t_min, t_max = pulse.get_time_interval()
a1_injection = int(max(abs(t_min), abs(t_max))) + 1
T_grid = np.ascontiguousarray(np.arange(0, a1_injection + dT, dT, dtype=np.float64))
if type(t_grid) is str:
t_grid = np.array(json.loads(t_grid), dtype=np.float64)
elif t_grid is None:
t_grid = pulse.get_tgrid(dt=1.0) # get_tGrid_dt(dt=1)
ret_tgrid = True
else:
t_grid = np.array(t_grid, dtype=np.float64)
if not np.all(np.diff(t_grid) % dT == 0):
user_grid = t_grid
t_grid = np.arange(user_grid[0], user_grid[-1] + dT * 4, dT * 4)
t_grid = np.unique(t_grid)
f, phase = Kernel_jit(
t_grid, T_grid, pulse, param_dict, kernel_type=kernel_type
)
rate = 2 * np.real(IOF(T_grid, f, phase))
return np.interp(user_grid, t_grid, rate)
else:
f, phase = Kernel_jit(
t_grid, T_grid, pulse, param_dict, kernel_type=kernel_type
)
rate = 2 * np.real(IOF(T_grid, f, phase))
# rate=2*np.real(np.trapz(f*np.exp(phase), x=T_grid, axis=0))
# the factor two is to account for the fact that the kernel
# is symmetric in T and we integrate from 0 to inf
if ret_tgrid:
return t_grid, np.array(rate)
else:
return np.array(rate)
# @njit(parallel=True,fastmath = False, cache=True)
[docs]
def get_quasi_static_rate_for_field(
field: float | npt.NDArray[np.float64], param_dict: Dict[str, float]
) -> npt.NDArray[np.float64]:
"""return the ionization rate for a define pulse computed with provided parameters
Args:
field (float64/np.ndarray): the grid of electric field strengths
param_dict: dictionary defining the medium's parameters
Returns:
np.ndarray, shape=(field.size): the ionization rates for given array of electric field strengths
"""
params = typed.Dict.empty(key_type=types.unicode_type, value_type=types.float64)
for key, value in param_dict.items():
params[key] = value
E_g = params["E_g"]
αPol = params.get("αPol", 0.0) # Default to 0 if not provided
a1 = params["a1"]
a2 = params["a2"]
# Ensure field is always a numpy array
if np.isscalar(field):
field = np.array([field], dtype=np.float64)
scalar_input = True
else:
field = np.asarray(field, dtype=np.float64)
scalar_input = False
tmp = np.zeros(field.size, dtype=np.complex128)
field = np.abs(field)
cond = field > 0
field = field[cond]
T_saddle = 1j * (-a2 * a1 + np.sqrt(αPol + a2**2 * a1**2 + 2 * E_g / field**2))
E2diff = 2 * field**2 * T_saddle
term1 = (-2 * field * T_saddle) ** 2 / 4
DelAbar = 0
term2 = 0
DelA2bar = (field * T_saddle) ** 2 / 3
term0 = 1j / (T_saddle + 1j * a1)
# saddle_contribution=np.sqrt(np.pi/(a2*a1-1j*T_saddle))/field
saddle_contribution = np.sqrt(
np.pi / field / np.sqrt(2 * E_g + field**2 * (a2**2 * a1**2 + αPol))
)
tmp[cond] = (
np.exp(
Kernel_phase_term(
term0, term1, term2, DelAbar, DelA2bar, E2diff, T_saddle, params
)
)
* Kernel_f_term(field, field, term0, term1, term2, T_saddle, params)
* saddle_contribution
)
rate = np.real(tmp)
assert np.all(abs(np.imag(rate)) <= 1e-10)
return rate
# @njit(parallel=True,fastmath = False, cache=True)
[docs]
def get_rate_quasi_static_limit(
t_grid: npt.NDArray[np.float64],
pulse: EllipticalCosNPulse,
param_dict: Dict[str, float],
) -> npt.NDArray[np.float64]:
"""Calculate the quasi-static ionization rate.
Args:
t_grid: Time grid for the calculation
pulse: EllipticalCosNPulse object defining the laser field
param_dict: Dictionary containing the parameters
Returns:
Array of ionization rates at each time point
"""
field = np.abs(pulse.get_electric_field(t_grid))
return get_quasi_static_rate_for_field(field, param_dict)
# @njit(parallel=True,fastmath = False, cache=True)
[docs]
def get_probability_quasi_static_limit(
pulse: EllipticalCosNPulse, param_dict: Dict[str, float], dt: float = 2.0
) -> float:
"""Calculate the quasi-static ionization probability.
Args:
pulse: EllipticalCosNPulse object defining the laser field
param_dict: Dictionary containing the parameters
dt: Time step for integration
Returns:
Ionization probability as a float
"""
t_min, t_max = pulse.get_time_interval()
t_grid = np.arange(t_min, t_max + dt, dt)
tmp = get_rate_quasi_static_limit(t_grid, pulse, param_dict)
return float(simpson(tmp, x=t_grid, axis=-1))
[docs]
def get_diabatic_ionization_probability(
pulse: EllipticalCosNPulse,
param_dict: Dict[str, float],
dt: float = 2.0,
dT: float = 0.25,
filterTreshold: float = 0.0,
kernel_type: str = "GASFIR",
ret_Rate: bool = False,
) -> float | Tuple[float, npt.NDArray[np.float64], npt.NDArray[np.float64]]:
"""Compute ionization probability for a defined pulse.
Args:
pulse: EllipticalCosNPulse object defining the pulse
param_dict: Dictionary containing kernel parameters
dt: Time step for rate calculation (default: 2.0)
dT: Time step for kernel integration (default: 0.25)
filterTreshold: Threshold for filtering rates (default: 0.0)
kernel_type: Type of kernel to use ("GASFIR" or "exact_SFA")
ret_Rate: Whether to return rates array (default: False)
Returns:
If ``ret_Rate`` is False, the total ionization probability (float).
If ``ret_Rate`` is True, a tuple ``(probability, t_grid, rates)``.
"""
t_min, t_max = pulse.get_time_interval()
a1_injection = int(max(abs(t_min), abs(t_max))) + 1
if np.any(np.array(pulse.get_central_wavelength()) < 140):
dt = 0.5
dT = 0.25
t_grid = np.ascontiguousarray(np.arange(int(t_min), int(t_max) + dt, dt))
if filterTreshold > 0:
### filter out the t_grid such that |E(t_grid)|>=1% of max(E(t_grid)) ###
ElecField = lambda t: EllipticalCosNPulse.get_electric_field(t)
extr = find_extrema_positions(t_grid, ElecField(t_grid))
Fextr = ElecField(extr)
extr = extr[np.abs(Fextr) >= max(np.abs(Fextr)) * filterTreshold]
if extr[0] > 0 and extr[-1] < 0:
extr = extr[::-1]
### smartly take the t_grid uptill the correspondin zero crossing rather than abruptly ending a a sub-cycle peak
zeroCr = find_zero_crossings(t_grid, ElecField(t_grid))
if zeroCr[0] > 0 and zeroCr[-1] < 0:
zeroCr = zeroCr[::-1]
t_grid = np.arange(
np.floor(zeroCr[zeroCr < extr[0]][-1]),
np.ceil(zeroCr[zeroCr > extr[-1]][0]) + dt,
dt,
dtype=np.float64,
)
rate_result = get_diabatic_ionization_rate(
t_grid, pulse, param_dict, dT, kernel_type=kernel_type
)
if isinstance(rate_result, tuple):
t_grid_out, rate = rate_result
else:
t_grid_out, rate = t_grid, rate_result
ionization_probability = 1 - np.exp(
-np.double(simpson(rate, x=t_grid_out, axis=-1))
)
# if ionization_probability < 0:
# old_prediction= ionization_probability
# warnings.warn(f"Ionization probability, {ionization_probability} is negative at {pulse.get_central_wavelength()} nm, I={pulse.get_peak_intensity()}. {param_dict}")
if ret_Rate:
return ionization_probability, t_grid_out, rate
else:
return ionization_probability