"""Simulation grid and FFT conventions.
Conventions (read this once, everything else follows)
------------------------------------------------------
* Envelope A(x, y, t) with carrier exp[i(k0 z - w0 t)], |A|^2 = intensity in W/m^2.
* Temporal transform A~(Om) = sum_t A(t) exp(+i Om t) -> grid.fft.ft_t
A(t) = (1/Nt) sum_Om A~ exp(-i Om t) -> grid.fft.ift_t
=> a spectral phase +GDD/2 Om^2 is an up-chirp, exp(+i Om t0) delays by t0.
* Spatial transform: standard fft2 (plane wave exp(+i kx x)), propagation exp(+i kz z).
* ALL arrays are kept in FFT order (index 0 is x=0, t=0, Om=0). Operations used by
the solver are local, so the order never matters there; use grid.xs, grid.lams,
grid.shift_xy(), grid.shift_w() for plotting.
* The solver's home domain is (x, y, Om): complex arrays of shape (Nx, Ny, Nt).
"""
import importlib.util
import os
from typing import Any
import numpy as np
import scipy.fft
from scipy.constants import c as C0
from .backend import get_xp, resolve_device, xp_of
class _FFTWPlans:
"""Planned FFTW transforms (pyFFTW) for 3D complex fields; one plan per (shape, dtype, axes,
direction), made on first use (FFTW_MEASURE: a few seconds for a new grid shape, then reused;
FFTW keeps the result for the session, so other grids of the same shape plan instantly).
The plans of one shape share a pair of planning buffers (2 field arrays) and are pointed back
at them after each call, so no reference to the caller's arrays is kept."""
def __init__(self, threads, effort):
import pyfftw
self.pyfftw, self.threads, self.effort, self.plans, self.buffers = pyfftw, threads, effort, {}, {}
def __call__(self, a, axes, direction, scale):
key = (a.shape, a.dtype.str, axes, direction)
plan = self.plans.get(key)
bufs = self.buffers.get(key[:2])
if bufs is None:
bufs = self.buffers[key[:2]] = (self.pyfftw.empty_aligned(a.shape, a.dtype),
self.pyfftw.empty_aligned(a.shape, a.dtype))
if plan is None:
plan = self.plans[key] = self.pyfftw.FFTW(*bufs, axes=axes, direction=direction,
flags=(self.effort,), threads=self.threads)
if not (a.flags.c_contiguous and self.pyfftw.is_byte_aligned(a, plan.input_alignment)):
a = self.aligned(a) # the plan needs aligned C-ordered input
out = self.pyfftw.empty_aligned(a.shape, a.dtype) # new array, like scipy.fft
plan.update_arrays(a, out) # out of place: a is preserved
plan.execute()
plan.update_arrays(*bufs) # drop the references to a and out
if scale != 1:
out *= out.real.dtype.type(scale)
return out
def aligned(self, a):
"""Copy of a in an aligned, C-ordered array (FFTW's SIMD alignment)."""
b = self.pyfftw.empty_aligned(a.shape, a.dtype)
b[...] = a
return b
[docs]
def fftw_available():
"""True if pyFFTW is installed (fft_backend='fftw' usable)."""
return importlib.util.find_spec("pyfftw") is not None
[docs]
class FFT:
"""Thin wrapper so the backend can be switched (scipy default, numpy, fftw, or cupyx on the GPU).
The arrays passed in must live on the device of the backend (host for scipy/numpy/fftw).
backend='fftw' (pyFFTW): planned FFTW transforms for the 3D complex fields of the solvers
(planning effort FFT.fftw_effort, threads = workers); any other array (1D/2D diagnostics,
real input) goes to scipy.fft. Same results to rounding. Ignored on the GPU (cuFFT)."""
mod: Any # scipy.fft, numpy.fft or cupyx.scipy.fft: same interface, no common type
fftw_effort = "FFTW_MEASURE" # or "FFTW_ESTIMATE" (no planning time), "FFTW_PATIENT"
def __init__(self, backend="auto", workers=None, device="cpu"):
self._fftw = None
if device == "gpu":
import cupyx.scipy.fft as mod
self._kw, self.mod, self.backend = {}, mod, "cupy"
return
if backend in ("auto", "scipy", "fftw"): # pocketfft C++: ~8x faster than numpy.fft for 2D
mod = scipy.fft
self._kw = {"workers": workers if workers is not None else -1}
if backend == "fftw":
if not fftw_available():
raise RuntimeError("fft_backend='fftw' needs pyFFTW (uv pip install pyfftw, "
"or thunder4d[fftw])")
n = os.cpu_count() or 1 # scipy convention: workers=-1 = all cores, -2 = all but one
threads = n if workers is None else (int(workers) if workers > 0 else max(1, n + 1 + int(workers)))
self._fftw = _FFTWPlans(threads, self.fftw_effort)
else:
backend = "scipy"
elif backend == "numpy":
mod = np.fft
self._kw = {}
else:
raise ValueError("backend must be 'auto', 'numpy', 'scipy' or 'fftw'")
self.mod = mod
self.backend = backend
def _planned(self, a):
"""The FFTW plans if they handle this array (3D complex host array), else None (scipy)."""
if isinstance(a, np.ndarray) and a.ndim == 3 and a.dtype in (np.complex64, np.complex128):
return self._fftw
return None
[docs]
def fft2(self, a):
if (plans := self._planned(a)) is not None:
return plans(a, (0, 1), "FFTW_FORWARD", 1)
return self.mod.fft2(a, axes=(0, 1), **self._kw)
[docs]
def ifft2(self, a):
if (plans := self._planned(a)) is not None:
return plans(a, (0, 1), "FFTW_BACKWARD", 1 / (a.shape[0] * a.shape[1]))
return self.mod.ifft2(a, axes=(0, 1), **self._kw)
[docs]
def ft_t(self, a): # time -> Omega
if (plans := self._planned(a)) is not None:
return plans(a, (2,), "FFTW_BACKWARD", 1)
return self.mod.ifft(a, axis=-1, norm="forward", **self._kw)
[docs]
def ift_t(self, a): # Omega -> time
if (plans := self._planned(a)) is not None:
return plans(a, (2,), "FFTW_FORWARD", 1 / a.shape[2])
return self.mod.fft(a, axis=-1, norm="forward", **self._kw)
def _next_fast(n):
"""Smallest 2^a 3^b 5^c >= n (fast sizes for pocketfft and cuFFT; real=True excludes 7 and 11)."""
return scipy.fft.next_fast_len(int(np.ceil(n)), real=True)
[docs]
class Grid:
"""(x, y, t) grid centred on lambda0.
Parameters
----------
Nx, dx : transverse points and step [m] (Ny, dy default to the same)
Nt : number of temporal points
lambda0: carrier wavelength [m] (defines the moving frame and Omega = 0)
dt : time step [s]; or give lambda_min (and optionally lambda_max) [m]
and dt is chosen so that the spectral window covers them.
dtype : np.complex64 (default, half the memory) or np.complex128.
device : 'auto' (GPU if CuPy works, else CPU), 'cpu' or 'gpu'. The field and the propagators
live on that device; the axes (x, lam, ...) are always NumPy arrays.
fft_backend: 'auto' / 'scipy' (pocketfft), 'numpy', or 'fftw' (pyFFTW, planned transforms for the
solvers' 3D fields; see FFT). workers: FFT threads (default: all cores).
"""
def __init__(self, Nx, dx, Nt, lambda0, *, dt=None, lambda_min=None, lambda_max=None,
Ny=None, dy=None, dtype=np.complex64, fft_backend="auto", workers=None,
device="auto"):
self.Nx, self.Ny = int(Nx), int(Ny or Nx)
self.dx, self.dy = float(dx), float(dy or dx)
self.Nt = int(Nt)
self.lambda0 = float(lambda0)
self.omega0 = 2 * np.pi * C0 / lambda0
if dt is None:
if lambda_min is None:
raise ValueError("give dt or lambda_min")
w_hi = 2 * np.pi * C0 / lambda_min - self.omega0
w_lo = self.omega0 - 2 * np.pi * C0 / lambda_max if lambda_max else w_hi
# fftfreq has (Nt-1)//2 bins above Omega = 0 and Nt//2 below
dt = 2 * np.pi / self.Nt * min(((self.Nt - 1) // 2) / w_hi, (self.Nt // 2) / w_lo)
self.dt = float(dt)
self.x = np.fft.fftfreq(self.Nx, 1.0 / (self.Nx * self.dx))
self.y = np.fft.fftfreq(self.Ny, 1.0 / (self.Ny * self.dy))
self.kx = 2 * np.pi * np.fft.fftfreq(self.Nx, self.dx)
self.ky = 2 * np.pi * np.fft.fftfreq(self.Ny, self.dy)
self.t = np.fft.fftfreq(self.Nt, 1.0 / (self.Nt * self.dt))
self.Omega = 2 * np.pi * np.fft.fftfreq(self.Nt, self.dt)
self.omega = self.omega0 + self.Omega
self.valid = self.omega > 0
with np.errstate(divide="ignore"):
self.lam = np.where(self.valid, 2 * np.pi * C0 / np.where(self.valid, self.omega, 1.0), np.inf)
self.dtype = np.dtype(dtype)
self.rdtype = np.float32 if self.dtype == np.complex64 else np.float64
self.device = resolve_device(device)
self.xp = get_xp(self.device)
self.fft = FFT(fft_backend, workers, self.device)
self.hfft = self.fft if self.device == "cpu" else FFT(fft_backend, workers, "cpu") # host arrays
# ------------------------------------------------------------------ helpers
@property
def shape(self):
return (self.Nx, self.Ny, self.Nt)
@property
def X(self):
return self.x[:, None]
@property
def Y(self):
return self.y[None, :]
@property
def r2(self):
return self.x[:, None] ** 2 + self.y[None, :] ** 2
@property
def dA(self):
return self.dx * self.dy
def asarray(self, a, dtype=None):
"""Array on the grid's device (no copy if it is already there)."""
return self.xp.asarray(a, dtype=dtype)
[docs]
def energy(self, Aw):
"""Energy [J] of a field in the (x, y, Omega) domain (host or device array)."""
return float(xp_of(Aw).vdot(Aw, Aw).real) * self.dx * self.dy * self.dt / self.Nt
[docs]
def energy_t(self, At):
return float(xp_of(At).vdot(At, At).real) * self.dx * self.dy * self.dt
# sorted axes for display
xs = property(lambda s: np.fft.fftshift(s.x))
ys = property(lambda s: np.fft.fftshift(s.y))
ts = property(lambda s: np.fft.fftshift(s.t))
kxs = property(lambda s: np.fft.fftshift(s.kx))
lams = property(lambda s: np.fft.fftshift(s.lam))
Omegas = property(lambda s: np.fft.fftshift(s.Omega))
[docs]
@staticmethod
def shift_xy(a):
return np.fft.fftshift(a, axes=(0, 1))
[docs]
@staticmethod
def shift_w(a, axis=-1):
return np.fft.fftshift(a, axes=axis)
[docs]
def describe(self):
mem = np.prod(self.shape) * self.dtype.itemsize / 1e6
lv = self.lam[self.valid]
return (f"Grid {self.Nx}x{self.Ny}x{self.Nt} ({mem:.0f} MB per field, {self.device}/{self.fft.backend}) | "
f"box {self.Nx*self.dx*1e3:.2f} x {self.Ny*self.dy*1e3:.2f} mm, dx = {self.dx*1e6:.1f} um | "
f"T = {self.Nt*self.dt*1e15:.0f} fs, dt = {self.dt*1e15:.2f} fs | "
f"lambda window {lv.min()*1e9:.0f}-{lv.max()*1e9:.0f} nm")
[docs]
def suggest_grid(cell, lambda0, lambda_min, lambda_max, pulse_fwhm, *, box_factor=8.0,
points_per_waist=8.0, time_window_factor=8.0, **grid_kw):
"""Pick (Nx, dx, Nt, dt) for a cell.
* box = box_factor x largest mirror spot radius (increase to ~12 for top-hat / aberrated inputs,
their higher-order modes are larger than the fundamental),
* dx = min(smallest waist / points_per_waist, Nyquist of the mirror curvature at the box corner),
* Nt from the spectral window and time_window_factor x input FWHM.
"""
wm = cell.linear_mode(lambda_max)["w_mirror"]
w0 = cell.linear_mode(lambda_min)["w0"]
L = box_factor * wm
dx_waist = w0 / points_per_waist
dx_curv = 0.8 * lambda_min * cell.R / (np.sqrt(2) * L)
dx = min(dx_waist, dx_curv)
return _sized_grid(L, dx, lambda0, lambda_min, lambda_max, pulse_fwhm, time_window_factor, **grid_kw)
def _sized_grid(L, dx, lambda0, lambda_min, lambda_max, pulse_fwhm, time_window_factor, **grid_kw):
"""Grid of width L with a step <= dx, spectral window [lambda_min, lambda_max] and
time window >= time_window_factor x pulse_fwhm (shared by suggest_grid and suggest_lens_grid)."""
Nx = _next_fast(L / dx)
w_hi = 2 * np.pi * C0 / lambda_min - 2 * np.pi * C0 / lambda0
w_lo = 2 * np.pi * C0 / lambda0 - 2 * np.pi * C0 / lambda_max
dt = np.pi / max(w_hi, w_lo)
Nt = _next_fast(time_window_factor * pulse_fwhm / dt + 2) # Grid shortens dt by up to 2 bins
return Grid(Nx, L / Nx, Nt, lambda0, lambda_min=lambda_min, lambda_max=lambda_max, **grid_kw)