schr.qm.solvers module

Solvers for the time-dependent Schrödinger equation.

class schr.qm.solvers.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.solvers.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.solvers.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.solvers.compute_energy(psi: Array, hamiltonian: Hamiltonian, t: float, dx: float) → float[source]

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

Parameters:

  • psi – Wavefunction.

  • hamiltonian – Hamiltonian operator.

  • t – Time (a.u.).

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

Returns:

Energy expectation value (a.u.).

schr.qm.solvers.compute_norm(psi: Array, dx: float) → float[source]

Compute wavefunction norm \(\sqrt{\int|\psi|^2 dV}\).

Parameters:

  • psi – Wavefunction.

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

Returns:

Norm (should be 1 for normalized wavefunctions).

schr.qm.solvers.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.