schr.qm.wavefunction module¶
Wavefunction utilities for quantum mechanics.
- schr.qm.wavefunction.gaussian_wavepacket(x: ~jax.Array | tuple[~jax.Array, ...], x0: float | tuple[float, ...], p0: float | tuple[float, ...], sigma: float, hbar: float = 1.0, dtype: ~numpy.dtype = <class 'jax.numpy.complex64'>) Array[source]¶
Create Gaussian wave packet.
\[\psi(\mathbf{r}) = (2\pi\sigma^2)^{-d/4} \exp\left[-\frac{(\mathbf{r}-\mathbf{r}_0)^2}{4\sigma^2}\right] \exp\left[\frac{i\mathbf{p}_0\cdot(\mathbf{r}-\mathbf{r}_0)}{\hbar}\right]\]where \(d\) is the dimensionality.
- Parameters:
x – Spatial grid. 1D: array (nx,). 2D: tuple (X, Y) of shape (ny, nx). 3D: tuple (X, Y, Z).
x0 – Center position (a.u.). Scalar for 1D, tuple for 2D/3D.
p0 – Initial momentum (a.u.). Scalar for 1D, tuple for 2D/3D.
sigma – Width parameter (a.u.).
hbar – Reduced Planck constant (default: 1.0 a.u.).
dtype – JAX dtype (default: complex64).
- Returns:
Normalized Gaussian wave packet.
- Raises:
ValueError – If dimensions are inconsistent.
- schr.qm.wavefunction.normalize(psi: Array, dx: float) Array[source]¶
Normalize wavefunction to \(\int|\psi|^2 dV = 1\).
- Parameters:
psi – Wavefunction.
dx – Grid spacing (a.u., uniform in all dimensions).
- Returns:
Normalized wavefunction.
- schr.qm.wavefunction.plane_wave(x: ~jax.Array | tuple[~jax.Array, ...], k: float | tuple[float, ...], dtype: ~numpy.dtype = <class 'jax.numpy.complex64'>) Array[source]¶
Create plane wave \(\exp(i\mathbf{k}\cdot\mathbf{r})\).
- Parameters:
x – Spatial grid (1D array or tuple of arrays).
k – Wave vector (a.u.). Scalar for 1D, tuple for 2D/3D.
dtype – JAX dtype (default: complex64).
- Returns:
Plane wave (not normalized).
- Raises:
ValueError – If dimensions are inconsistent.
- schr.qm.wavefunction.probability_current_1d(psi: Array, dx: float, mass: float = 1.0, hbar: float = 1.0) Array[source]¶
Compute probability current density in 1D.
\[j = \frac{\hbar}{m} \text{Im}\left[\psi^* \frac{\partial\psi}{\partial x}\right]\]- Parameters:
psi – Wavefunction (1D).
dx – Grid spacing (a.u.).
mass – Particle mass (default: 1.0 a.u.).
hbar – Reduced Planck constant (default: 1.0 a.u.).
- Returns:
Probability current density.