"""
Functions for creating spatial and temporal grids.
1. ``GridXYZ`` class that calculates the spatial grid and its
Fourier counterpart.
2. Helper function ``setup_grids`` that automatically creates
grids either dynamically or from given grid sizes.
.. warning::
The dynamic grid creation is experimental and should be used with caution.
Especially, when testing on new physical problems, it is recommended to
visualize participating Maxwell fields and check the grid sizes.
.. note::
Parts of dynamical grid creation for gaussian beams were initially
implemented by Alexander Blinne (``get_kmax_from_bw``, ``get_xyz_size``,
``get_t_size``).
"""
from collections.abc import Iterable
from copy import deepcopy
import numexpr as ne
import numpy as np
import pyfftw
from scipy.constants import c, pi
from quvac import config
[docs]
class GridXYZ:
"""
Calculates spatial grid and its Fourier counterpart.
Parameters
----------
grid : tuple of np.array
A tuple containing three numpy arrays (x, y, z) representing the spatial grid
for field discretization.
Attributes
----------
grid : tuple of np.array
The spatial grid for field discretization.
grid_shape : tuple of int
The shape of the grid.
xyz : tuple of np.array
The meshgrid of the spatial grid.
dxyz : tuple of float
The grid spacing in each dimension.
dV : float
The volume element of the grid.
kgrid : tuple of np.array
The k-space grid.
kgrid_shifted : tuple of np.array
The shifted k-space grid.
dkxkykz : list of float
The spacing in k-space for each dimension.
dVk : float
The volume element in k-space.
kmeshgrid : tuple of np.array
The meshgrid of the k-space grid.
kabs : np.array
The absolute value of the k-vectors.
e1x, e1y, e1z : np.array
The x, y, and z components of the first polarization vector.
e2x, e2y, e2z : np.array
The x, y, and z components of the second polarization vector.
"""
def __init__(self, grid):
# Define spatial grid
self.grid = grid
self.grid_shape = tuple(ax.size for ax in grid)
self.xyz = self.x, self.y, self.z = np.meshgrid(
*grid, indexing="ij", sparse=True
)
self.dxyz = tuple(ax[1] - ax[0] for ax in grid)
self.dV = np.prod(self.dxyz)
self.kgrid = None
[docs]
def get_k_grid(self):
"""
Calculate the k-space grid and polarization vectors.
Returns
-------
None
"""
if self.kgrid is None:
self._calculate_k_grid()
def _calculate_k_grid(self):
for i, ax in enumerate("xyz"):
Nx, dx = self.grid_shape[i], self.dxyz[i]
k = (2 * pi * np.fft.fftfreq(Nx, dx)).astype(config.FDTYPE)
setattr(self, f"k{ax}", k)
self.kgrid = (self.kx, self.ky, self.kz)
self.kgrid_shifted = tuple(
np.fft.fftshift(k) for k in (self.kx, self.ky, self.kz)
)
self.dkxkykz = [ax[1] - ax[0] for ax in self.kgrid]
self.dVk = np.prod(self.dkxkykz)
self.kmeshgrid = np.meshgrid(
self.kx, self.ky, self.kz, indexing="ij", sparse=True
)
kx, ky, kz = self.kmeshgrid
self.kabs = kabs = ne.evaluate("sqrt(kx**2 + ky**2 + kz**2)") # noqa: F841
kperp = ne.evaluate("sqrt(kx**2 + ky**2)") # noqa: F841
# Polarization vectors
self.e1x = ne.evaluate(
"where((kx==0) & (ky==0), 2*(kz>0)-1, kx * kz / (kperp*kabs))"
)
self.e1y = ne.evaluate("where((kx==0) & (ky==0), 0, ky * kz / (kperp*kabs))")
self.e1z = ne.evaluate("where((kx==0) & (ky==0), 0, -kperp / kabs)")
self.e2x = ne.evaluate("where((kx==0) & (ky==0), 0, -ky / kperp)")
self.e2y = ne.evaluate("where((kx==0) & (ky==0), 1, kx / kperp)")
# self.e2y = ne.evaluate("where((kx==0) & (ky==0), 2*(kz>0)-1, kx / kperp)")
self.e2z = 0
[docs]
def get_ek(theta, phi):
"""
Calculate k-vector for given spherical angles theta and phi.
Parameters
----------
theta : float
The polar angle in radians.
phi : float
The azimuthal angle in radians.
Returns
-------
numpy.ndarray
A 3-element array representing the k-vector.
"""
ek = np.array(
[np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)]
)
return ek
[docs]
def get_pol_basis(theta, phi):
"""
Calculate polarization basis for given spherical angles theta and phi.
Parameters
----------
theta : float
The polar angle in radians.
phi : float
The azimuthal angle in radians.
Returns
-------
tuple of numpy.ndarray
A tuple containing two 3-element arrays representing the polarization
basis vectors e1 and e2.
"""
e1 = np.array(
[np.cos(theta) * np.cos(phi), np.cos(theta) * np.sin(phi), -np.sin(theta)],
dtype=config.FDTYPE
)
e2 = np.array([-np.sin(phi), np.cos(phi), 0], dtype=config.FDTYPE)
return e1, e2
[docs]
def get_polarization_vector(theta, phi, beta):
"""
Calculate polarization vector for given spherical angles theta and phi
and polarization angle beta.
Parameters
----------
theta : float
The polar angle in radians.
phi : float
The azimuthal angle in radians.
beta : float
The polarization angle in radians.
Returns
-------
numpy.ndarray
3-element array representing the polarization vector.
"""
e1, e2 = get_pol_basis(theta, phi)
ebeta = e1*np.cos(beta) + e2*np.sin(beta)
return ebeta
[docs]
def gaussian_bandwidth(field_params):
"""
Calculate the gaussian bandwidth.
Parameters
----------
field_params : dict
Dictionary containing the field parameters. Required keys are:
- tau : float
Pulse duration.
- w0 : float, optional
Beam waist. If 'w0' is not provided, 'w0x' and 'w0y' are used.
Returns
-------
tuple of float
A tuple containing the perpendicular and longitudinal bandwidths
kbw_perp, kbw_long).
Notes
-----
The bandwidths are calculated from the Fourier transform of transverse and
longitudinal Gaussian profiles:
- exp(-r**2/w0**2) -> exp(-w0**2*kperp**2/4) -> kperp_bw ~ 2*2/w0
- exp(-t**2/(tau/2)**2) -> exp(-tau**2*k**2/16) -> k_long ~ 2*4/(c*tau)
"""
tau = field_params["tau"]
# gaussian might have circular or elliptic cross section
if "w0" in field_params:
w0 = field_params["w0"]
elif "w0x" in field_params:
w0 = min(field_params["w0x"], field_params["w0y"])
kbw_perp = 4 / w0
kbw_long = 8 / (c * tau)
return kbw_perp, kbw_long
[docs]
def dipole_bandwidth(field_params):
"""
Calculate the dipole bandwidth.
Parameters
----------
field_params : dict
Dictionary containing the field parameters. Required keys are:
- lam : float
Wavelength of the pulse.
- tau : float
Pulse duration.
Returns
-------
tuple of float
A tuple containing the perpendicular and longitudinal bandwidths
(kbw_perp, kbw_long).
Notes
-----
The length scales are taken from:
I. Gonoskov et al. "Dipole pulse theory: Maximizing the field amplitude
from 4π focused laser pulses." PRA 86.5 (2012): 053836.
l_para = 0.58 * lam
l_perp = 0.4 * lam
"""
lam, tau = [field_params[k] for k in "lam tau".split()]
l_long = 0.58*lam
l_perp = 0.4*lam
tau_bw = 8 / (c * tau)
# Add up bandwidth coming from time and length scales
kbw_perp = 4 / l_perp + tau_bw
kbw_long = 4 / l_long + tau_bw
return kbw_perp, kbw_long
[docs]
def get_bw(field_params):
"""
Calculate the bandwidths for a given field type.
Parameters
----------
field_params : dict
Dictionary containing the field parameters.
Returns
-------
tuple of float
A tuple containing the perpendicular and longitudinal bandwidths
(kbw_perp, kbw_long).
Raises
------
NotImplementedError
If the field type is not supported.
"""
ftype = field_params["field_type"]
if "gaussian" in ftype:
assert "w0" in field_params or "w0x" in field_params, (
"Gaussian field parameters must have either 'w0' or 'w0x' key"
)
bws = gaussian_bandwidth(field_params)
elif "dipole" in ftype:
bws = dipole_bandwidth(field_params)
return bws
def _check_keys(field_params, required_keys):
err_msg = f"Field parameters must have {required_keys} keys"
assert all(key in field_params for key in required_keys), err_msg
[docs]
def get_kmax_from_bw(field_params):
"""
Calculate the maximum k-vector from bandwidth along each dimension
for a given field.
Parameters
----------
field_params : dict
Dictionary containing the field parameters. Required keys are:
- lam : float
Wavelength of the pulse.
- tau : float
Pulse duration.
- theta : float
Polar angle in degrees.
- phi : float
Azimuthal angle in degrees.
Returns
-------
numpy.ndarray
A 3-element array representing the maximum k-vector along each dimension.
Notes
-----
The maximum k-vector is calculated based on the bandwidths derived from the
field parameters.
"""
required_keys = "lam tau theta phi".split()
_check_keys(field_params, required_keys)
lam, _, theta, phi = [field_params[k] for k in required_keys]
k = 2 * np.pi / lam
theta *= pi / 180
phi = phi * pi / 180 if np.sin(theta) != 0.0 else 0.0
ek = get_ek(theta, phi)
e1, e2 = get_pol_basis(theta, phi)
kbw_perp, kbw_long = get_bw(field_params)
k0 = ek * k
kmax = np.abs(k0 + ek * kbw_long)
for beta in np.linspace(0, 2 * np.pi, 64, endpoint=False):
kp = k0 + ek * kbw_long + kbw_perp * (np.cos(beta) * e1 + np.sin(beta) * e2)
kmax = np.maximum(kmax, np.abs(kp))
return kmax
[docs]
def get_kmax_grid(field_params):
"""
Calculate the maximum k-vector along each dimension for a given field type.
Parameters
----------
field_params : dict
Dictionary containing the field parameters.
Returns
-------
numpy.ndarray
A 3-element array representing the maximum k-vector along each dimension.
Raises
------
NotImplementedError
If the field type is not supported.
"""
ftype = field_params["field_type"]
if ("gaussian" in ftype) or ("dipole" in ftype):
kmax = get_kmax_from_bw(field_params)
else:
raise NotImplementedError(f"{ftype} field type is not supported")
return kmax
[docs]
def get_xyz_size(fields, box_size, grid_res=1, equal_resolution=False):
"""
Calculate necessary spatial resolution.
Parameters
----------
fields : dict or list of dict
Parameters of participating fields.
box_size : sequence of float
Box size for 3 dimensions.
grid_res : float, optional
Controls the resolution (default is 1).
equal_resolution : bool, optional
Flag indicating if equal resolution in every dimension is needed
(default is False).
Returns
-------
list of int
List containing the number of grid points along each dimension.
"""
if isinstance(fields, dict):
fields = list(fields.values())
box_size = np.array(box_size)
kmax = np.zeros((3,))
for field_params in fields:
kmax = np.maximum(kmax, get_kmax_grid(field_params))
# if a square box is required
if equal_resolution:
kmax = np.max(kmax) * np.ones(3)
N = grid_res * np.ceil(box_size * 3 * kmax / pi)
N = [pyfftw.next_fast_len(int(n)) for n in N]
return N
[docs]
def get_t_size(t_start, t_end, lam, grid_res=1):
"""
Calculate necessary temporal resolution.
Parameters
----------
t_start : float
Start time.
t_end : float
End time.
lam : float
Wavelength.
grid_res : float, optional
Factor to scale the resolution (default is 1).
Returns
-------
float
Number of time steps.
"""
fmax = c / lam
return int(np.ceil((t_end - t_start) * fmax * 6 * grid_res))
[docs]
def get_max_duration(fields_params):
tau_max = 0
for field in fields_params:
tau = field.get("tau", 0)
# spectral chirp increases effective pulse duration
if field["field_type"] == "paraxial_gaussian_spectral_maxwell":
alpha = field.get("alpha_chirp", 0)
tau *= np.sqrt(1 + alpha**2)
tau_max = max([tau_max, tau])
return tau_max
[docs]
def get_box_size(fields_params, grid_params):
"""
Calculate the necessary box size for the spatial grid.
Parameters
----------
fields_params : list of dict
Parameters of participating fields. Each dictionary should contain:
- field_type : str
Type of the field.
- w0 : float, optional
Beam waist for focused fields.
- tau : float, optional
Pulse duration.
grid_params : dict
Grid parameters. Required keys are:
- transverse_factor : float
Factor to scale the transverse size.
- longitudinal_factor : float
Factor to scale the longitudinal size.
Returns
-------
tuple of float
A tuple containing the transverse and longitudinal sizes.
Notes
-----
The box size is calculated based on the maximum beam waist and pulse duration
among the fields, scaled by the transverse and longitudinal factors.
"""
perp_max = 0
for field in fields_params:
ftype = field["field_type"]
if "gauss" in ftype:
w0, w0x, w0y = [field.get(key, 0) for key in ("w0", "w0x", "w0y")]
length = length = np.max([w0, w0x, w0y])
# length = field.get("w0", 0)
elif "dipole" in ftype:
length = c * field.get("tau", 0) / 4
else:
length = 0
perp_max = np.maximum(length, perp_max)
# tau_max = max([field.get("tau", 0) for field in fields_params])
tau_max = get_max_duration(fields_params)
transverse_size = perp_max * grid_params["transverse_factor"]
longitudinal_size = tau_max * c * grid_params["longitudinal_factor"]
return float(transverse_size), float(longitudinal_size)
[docs]
def create_dynamic_grid(fields_params, grid_params):
"""
Dynamically create grids from given laser parameters.
One should be careful with this option.
Parameters
----------
fields_params : list of dict
Parameters of participating fields.
grid_params : dict
Grid parameters.
Returns
-------
dict
Updated grid parameters including calculated box size, number of spatial points,
and number of temporal points.
"""
grid_params_upd = deepcopy(grid_params)
# Create spatial box
collision_geometry = grid_params.get("collision_geometry", "z")
tau_max = max([field.get("tau", 0) for field in fields_params])
lam_min = min([field.get("lam", 1e10) for field in fields_params])
transverse_size, longitudinal_size = get_box_size(fields_params, grid_params)
box_xyz = [0, 0, 0]
for i, ax in enumerate("xyz"):
box_xyz[i] = longitudinal_size if ax in collision_geometry else transverse_size
grid_params_upd["box_xyz"] = box_xyz
# Number of spatial pts
box_size = np.array(box_xyz) / 2
Nxyz_ = get_xyz_size(fields_params, box_size)
res = grid_params.get("spatial_resolution", 1)
if isinstance(res, Iterable):
assert len(res) == 3, "Spatial resolution must be a list of 3 values or a "
"single value"
Nxyz = [int(Nx * res) for Nx, res in zip(Nxyz_, res, strict=True)]
elif type(res) in (int, float):
Nxyz = [int(Nx * res) for Nx in Nxyz_]
grid_params_upd["Nxyz"] = Nxyz
# Create temporal box
t0 = tau_max * grid_params["time_factor"]
Nt = get_t_size(-t0 / 2, t0 / 2, lam_min)
# Number of temporal pts
grid_params_upd["box_t"] = t0
grid_params_upd["Nt"] = Nt * grid_params.get("time_resolution", 1)
return grid_params_upd
[docs]
def setup_grids(fields_params, grid_params):
"""
Create spatial and temporal grids from given sizes or
dynamically from field parameters (e.g., duration and waist size).
Parameters
----------
fields_params : list of dict
Parameters of participating fields. Each dictionary should contain:
- field_type : str
Type of the field.
- lam : float
Wavelength of the pulse.
- w0 : float, optional
Beam waist for focused fields.
- tau : float, optional
Pulse duration.
grid_params : dict
Grid parameters. Required keys are:
- mode : str
Mode of grid creation ('dynamic' or 'static').
Keys for 'static' mode:
- box_xyz : tuple of float
Box size for the spatial grid.
- Nxyz : tuple of int
Number of grid points along each spatial dimension.
- Nt : int
Number of temporal points.
- box_t : float or tuple of float
Time duration or start and end times for the temporal grid.
Keys for 'dynamic' mode:
- collision_geometry : str
Specifies the collision geometry ('x', 'y', 'z').
- transverse_factor : float
Factor to scale the transverse size.
- longitudinal_factor : float
Factor to scale the longitudinal size.
- time_factor : float
Factor to scale the time duration.
- spatial_resolution : float or list of float, optional
Controls the spatial resolution.
- time_resolution : float, optional
Controls the temporal resolution.
- ignore_idx : list of int, optional
Indices of fields to ignore for dynamic grid creation.
Returns
-------
tuple
A tuple containing:
- grid_xyz : quvac.grid.GridXYZ
The spatial grid object.
- grid_t : numpy.ndarray
The temporal grid array.
Notes
-----
The spatial and temporal grids are created based on the mode specified in the grid
parameters. If 'dynamic' mode is selected, the grids are created dynamically based
on the field parameters.
"""
if grid_params["mode"] == "dynamic":
if isinstance(fields_params, dict):
fields_params = list(fields_params.values())
# filter out fields that should not contribute to dynamic grid creation
ignore_idx = grid_params.get("ignore_idx", None)
if ignore_idx is not None:
fields_params = [field for idx,field in enumerate(fields_params)
if idx not in ignore_idx]
grid_params = create_dynamic_grid(fields_params, grid_params)
x0, y0, z0 = grid_params["box_xyz"]
Nx, Ny, Nz = grid_params["Nxyz"]
x = np.linspace(-0.5 * x0, 0.5 * x0, Nx, endpoint=Nx % 2, dtype=config.FDTYPE)
y = np.linspace(-0.5 * y0, 0.5 * y0, Ny, endpoint=Ny % 2, dtype=config.FDTYPE)
z = np.linspace(-0.5 * z0, 0.5 * z0, Nz, endpoint=Nz % 2, dtype=config.FDTYPE)
grid_xyz = GridXYZ((x, y, z))
Nt = grid_params["Nt"]
box_t = grid_params["box_t"]
# Allow non-symmetric time-grids
if isinstance(box_t, float):
Nt += int(1 - Nt % 2)
t0 = box_t
grid_t = np.linspace(
-0.5 * t0, 0.5 * t0, Nt, endpoint=True, dtype=config.FDTYPE
)
else:
t_start, t_end = box_t
grid_t = np.linspace(t_start, t_end, Nt, endpoint=True, dtype=config.FDTYPE)
return grid_xyz, grid_t