schr.utils.fft module

FFT utilities for quantum simulations.

schr.utils.fft.fft_derivative(f: Array, dx: float, axis: int = -1, order: int = 1) → Array

Compute spectral derivative using FFT.

Spectral accuracy with periodic boundaries. Derivatives in Fourier space:

• 1st order: multiply by \(ik\)

• 2nd order: multiply by \(-k^2\)

Parameters:

• f – Function values on uniform grid.

• dx – Grid spacing (a.u.).

• axis – Axis for derivative (default: -1 = last axis).

• order – Derivative order (1 or 2).

Returns:

\(\partial^n f/\partial x^n\).

Raises:

ValueError – If order not in {1, 2}.

schr.utils.fft.fftfreq_grid(n: int, d: float, dtype: ~numpy.dtype = <class 'jax.numpy.float32'>) → Array

Create frequency grid for FFT operations.

Parameters:

• n – Number of grid points.

• d – Grid spacing (a.u.).

• dtype – JAX dtype (default: float32).

Returns:

Frequency array (rad/a.u.).

schr.utils.fft.momentum_operator(shape: tuple[int, ...], dx: float, hbar: float = 1.0, dtype: ~numpy.dtype = <class 'jax.numpy.float32'>) → tuple[Array, ...]

Create momentum operators in Fourier space.

For kinetic energy: \(\hat{T} = \frac{\hbar^2k^2}{2m}\).

Parameters:

• shape – Grid shape (nx,) for 1D, (ny, nx) for 2D, (nz, ny, nx) for 3D.

• dx – Grid spacing (a.u., uniform).

• hbar – Reduced Planck constant (a.u., default: 1.0).

• dtype – JAX dtype (default: float32).

Returns:

• 1D: (kx,) of shape (nx,)

• 2D: (kx, ky) of shapes (ny, nx) each

• 3D: (kx, ky, kz) of shapes (nz, ny, nx) each

Return type:

Tuple of momentum operators (one per dimension)

Raises:

ValueError – If ndim > 3.