schr.qm package¶

Module contents¶

Quantum mechanics module for schr package.

This module provides implementations of quantum mechanical systems, including Hamiltonians, wavefunctions, and Schrödinger equation solvers.

class schr.qm.CrankNicolson(hamiltonian: ~schr.core.base.Hamiltonian, hbar: float = 1.0, max_iter: int = 100, tol: float = 1e-08, dtype: ~numpy.dtype = <class 'jax.numpy.complex64'>)[source]¶

Bases: Solver

Crank-Nicolson implicit solver.

Solves \((I + i\hat{H}dt/(2\hbar))|\psi(t+dt)\rangle = (I - i\hat{H}dt/(2\hbar))|\psi(t)\rangle\)

Unconditionally stable, preserves unitarity. Accuracy: \(O(dt^2)\).

step(psi: Array, t: float, dt: float) → Array[source]¶

Perform single Crank-Nicolson time step.

Uses fixed-point iteration for implicit equation.

Parameters:

psi – Wavefunction at time \(t\).
t – Current time (a.u.).
dt – Time step (a.u.).

Returns:

Wavefunction at time \(t + dt\).

class schr.qm.FreeParticle(grid_shape: tuple, dx: float, mass: float = 1.0, hbar: float = 1.0, dtype: ~numpy.dtype = <class 'jax.numpy.complex64'>)[source]¶

Bases: ParticleInPotential

Free particle Hamiltonian (\(V = 0\)).

\[\hat{H} = -\frac{\hbar^2}{2m}\nabla^2\]

class schr.qm.HarmonicOscillator(omega: float, grid_shape: tuple, dx: float, grid_coords: ~jax.Array, mass: float = 1.0, hbar: float = 1.0, dtype: ~numpy.dtype = <class 'jax.numpy.complex64'>)[source]¶

Bases: ParticleInPotential

Quantum harmonic oscillator Hamiltonian.

\[\hat{H} = -\frac{\hbar^2}{2m}\nabla^2 + \frac{1}{2}m\omega^2r^2\]

class schr.qm.ParticleInPotential(potential: ~collections.abc.Callable[[~jax.Array, float], ~jax.Array], grid_shape: tuple, dx: float, mass: float = 1.0, hbar: float = 1.0, dtype: ~numpy.dtype = <class 'jax.numpy.complex64'>)[source]¶

Bases: Hamiltonian

Hamiltonian for particle in arbitrary potential.

\[\hat{H} = -\frac{\hbar^2}{2m}\nabla^2 + V(\mathbf{r}, t)\]

potential¶: Callable V(r, t) returning potential energy.

mass¶: Particle mass (a.u., default: 1.0 = electron mass).

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

dx¶: Grid spacing (a.u.).

kinetic_operator¶: Momentum space kinetic energy \(\hbar^2k^2/(2m)\).

apply(psi: Array, t: float) → Array[source]¶

Apply \(\hat{H}|\psi\rangle = (\hat{T} + \hat{V})|\psi\rangle\) using split-operator method.

Parameters:

psi – Wavefunction.
t – Time (a.u.).

Returns:

\(\hat{H}|\psi\rangle\).

energy(psi: Array, t: float) → float[source]¶

Compute energy expectation \(\langle\psi|\hat{H}|\psi\rangle\).

Parameters:

psi – Normalized wavefunction.
t – Time (a.u.).

Returns:

Energy expectation value (a.u.).

kinetic_energy(psi: Array) → float[source]¶

Compute kinetic energy expectation \(\langle\psi|\hat{T}|\psi\rangle\).

Parameters:: psi – Normalized wavefunction.
Returns:: Kinetic energy (a.u.).

potential_energy(psi: Array, t: float) → float[source]¶

Compute potential energy expectation \(\langle\psi|\hat{V}|\psi\rangle\).

Parameters:

psi – Normalized wavefunction.
t – Time (a.u.).

Returns:

Potential energy (a.u.).

class schr.qm.RungeKutta4(hamiltonian: ~schr.core.base.Hamiltonian, hbar: float = 1.0, dtype: ~numpy.dtype = <class 'jax.numpy.complex64'>)[source]¶

Bases: Solver

Fourth-order Runge-Kutta solver.

General-purpose ODE solver for \(i\hbar\partial_t|\psi\rangle = \hat{H}|\psi\rangle\). Accuracy: \(O(dt^4)\).

Warning

RK4 has limited stability for the Schrödinger equation. For best results:

Use small timesteps: \(dt < 0.05\) (in atomic units)
For grid-based problems: \(dt \lesssim \Delta x^2\)
For strong potentials, reduce \(dt\) further
Consider SplitStepFourier for better stability and speed

Notes

Unlike implicit methods (Crank-Nicolson) or operator splitting (SSFM), standard RK4 is not unconditionally stable for the Schrödinger equation. Large timesteps or strong potentials can cause exponential error growth leading to NaN values.

step(psi: Array, t: float, dt: float) → Array[source]¶

Perform single RK4 time step.

Parameters:

psi – Wavefunction at time \(t\).
t – Current time (a.u.).
dt – Time step (a.u.).

Returns:

Wavefunction at time \(t + dt\).

Note

This method uses the standard RK4 formula for the Schrödinger equation \(i\hbar\partial_t|\psi\rangle = \hat{H}|\psi\rangle\):

\[\begin{split}k_1 &= -\frac{i}{\hbar}\hat{H}|\psi(t)\rangle \\ k_2 &= -\frac{i}{\hbar}\hat{H}\left(|\psi(t)\rangle + \frac{dt}{2}k_1\right) \\ k_3 &= -\frac{i}{\hbar}\hat{H}\left(|\psi(t)\rangle + \frac{dt}{2}k_2\right) \\ k_4 &= -\frac{i}{\hbar}\hat{H}\left(|\psi(t)\rangle + dt \cdot k_3\right) \\ |\psi(t+dt)\rangle &= |\psi(t)\rangle + \frac{dt}{6}(k_1 + 2k_2 + 2k_3 + k_4)\end{split}\]

Stability: Not unconditionally stable. Becomes unstable for \(dt \gtrsim 0.05\) a.u. or with strong potentials.

class schr.qm.SplitStepFourier(hamiltonian: ~schr.core.base.Hamiltonian, hbar: float = 1.0, dtype: ~numpy.dtype = <class 'jax.numpy.complex64'>)[source]¶

Bases: Solver

Split-Step Fourier Method (SSFM) solver.

Pseudo-spectral method alternating between position and momentum space:

\[\psi(t+dt) \approx e^{-iV dt/(2\hbar)} e^{-iT dt/\hbar} e^{-iV dt/(2\hbar)} \psi(t)\]

Accuracy: \(O(dt^3)\), preserves unitarity.

hamiltonian¶: Hamiltonian operator.

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

kinetic_operator¶: Precomputed \(T(k) = \hbar^2k^2/(2m)\).

step(psi: Array, t: float, dt: float) → Array[source]¶

Perform single SSFM time step.

Parameters:

psi – Wavefunction at time \(t\).
t – Current time (a.u.).
dt – Time step (a.u.).

Returns:

Wavefunction at time \(t + dt\).

schr.qm.estimate_stable_timestep(dx: float, mass: float = 1.0, hbar: float = 1.0, safety_factor: float = 0.5) → float[source]¶

Estimate stable timestep for RK4 solver.

For grid-based quantum simulations, RK4 has a stability limit related to the grid spacing. This function estimates a safe timestep using:

\[dt_{\text{max}} \approx C \frac{m \Delta x^2}{\hbar}\]

where C is a safety factor (default: 0.5).

Parameters:

dx – Grid spacing (a.u.).
mass – Particle mass (a.u., default: 1.0).
hbar – Reduced Planck constant (a.u., default: 1.0).
safety_factor – Safety factor C (default: 0.5 for conservative estimate).

Returns:

Estimated maximum stable timestep (a.u.).

Example

>>> dx = 0.1  # Grid spacing
>>> dt_max = estimate_stable_timestep(dx)
>>> print(f"Use dt < {dt_max:.4f}")
Use dt < 0.0050

Note

This is a conservative estimate. For strong potentials, you may need to reduce the timestep further. Always verify stability by monitoring wavefunction norm and energy conservation.

schr.qm.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.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 package¶

Submodules¶

Module contents¶

Schr

Navigation

Related Topics