Rothe-JAX

JAX implementation of Gaussian-basis Rothe time stepping for the time-dependent Schrödinger equation (TDSE).

The wavefunction is expanded in a (generally non-orthogonal) Gaussian basis:

ψ(x, t) = Σₙ cₙ(t) · ∏ᵢ exp(−((aₙᵢ)² + i·bₙᵢ)(xᵢ − μₙᵢ)² + i·pₙᵢ·(xᵢ − μₙᵢ))

Each implicit Rothe step minimises the Rothe error functional to find the next-time Gaussian parameters and coefficients. When the per-step error exceeds a tolerance, a new Gaussian is added to the basis.

Installation

pip install -e .

For tests and plotting extras:

JAX backend: pyproject.toml lists plain jax, so pip will not install a GPU backend automatically. Install the variant that matches your hardware before running pip install -e .:

pip install "jax[cpu]"        # CPU
pip install "jax[cuda12]"     # NVIDIA GPU (CUDA 12)

float64 and complex128 precision are enabled automatically when you import rothe (via rothe/_config.py).

Running the bundled examples

All propagators require at minimum --epsilon (global Rothe-error budget, smaller = more accurate) and --regularization_lambda (Tikhonov strength for the linear coefficient solve).

1D Double Well — V = x⁴/4 − x²/2

No console entry point; run as a module:

python -m example_scripts.propagate_double_well \
  --num_gaussians 5 \
  --epsilon 1e-2 \
  --regularization_lambda 1e-8

Key options:

flag	default	meaning
--dt	0.1	time step
--t_final	20.0	end time
--t_start	—	resume from checkpoint at this time

Henon-Heiles

A D-dimensional chain of anharmonically coupled oscillators:

V = Σᵢ ½xᵢ² + Σᵢ (0.112 xᵢ²xᵢ₊₁ − 0.037 xᵢ³) + small quartic stabiliser

propagate-henon-heiles \
  --dim 2 \
  --num_gaussians 5 \
  --dt 0.01 \
  --epsilon 1e-3 \
  --regularization_lambda 1e-8

--dim is the number of dimensions. The propagation time is hardcoded to t = 100 (i.e. nsteps = int(100 / dt)).

Hydrogen — 3D Coulomb + optional laser field

propagate-hydrogen \
  --num_gauss_wavefunction 21 \
  --num_gauss_potential 21 \
  --epsilon 1e-2 \
  --regularization_lambda 1e-8

The Coulomb potential is pre-expanded in num_gauss_potential Gaussians (coefficients in gaussian_Coulomb/). Pre-computed ground-state initial states live in Hydrogen_ground_state_data/.

Additional flags:

flag	meaning
--dt 0.2j	imaginary-time propagation (ground-state search)
--frozen / --no-frozen	freeze ground-state basis; only newly added Gaussians are optimised
--num_dynamic 10	number of dynamic Gaussians per side when --frozen is active
--t_start <time>	resume from an existing checkpoint

---

Defining your own system

Three ingredients are needed:

1. A potential string describing V(x, t)
2. An initial parameter array of shape (n, D, 4)
3. Initial coefficients of shape (n,), dtype complex128

example_scripts/harmonic_oscillator.py is a fully annotated minimal example (1D harmonic oscillator). The sections below explain each ingredient.

The potential string DSL

Parsed by rothe.potential_parser.read_string. The first line must be dimension <D>. Then one or both of:

polynomial section — one monomial per line:

x0x0:         0.5          # 0.5 x₀²          (1D quadratic)
x0x0x1:       0.111        # coupling  x₀²·x₁
x0:           E*sin(w*t)   # time-dependent linear coupling

Powers are written by repeating xN (dimension index from 0). Coefficients are either float literals or sympy expressions in t for time-dependent potentials.

exponential section — one Gaussian term per line:

lincoeff, width, [mu0, mu1, ...]

Adds lincoeff * exp(-width² · Σ_d (x_d − mu_d)²). Used for Coulomb-type potentials approximated in a Gaussian basis.

constants section (optional, must be last) — named substitutions:

constants
E_0: 0.06
omega: 0.057

Then use E_0 and omega in the coefficient expressions above.

Both sections may appear together; the total potential is their sum.

Full example — 2D harmonic oscillator + a time-dependent laser:

POTENTIAL = """
dimension 2
polynomial
x0x0: 0.5
x1x1: 0.5
x0:   E_0*sin(omega*t)
constants
E_0:   0.06
omega: 0.057
"""

Gaussian parameter layout

params has shape (n, D, 4). The four slots per (Gaussian, dimension):

params[:, :, k]	name	meaning
0	a	Width; must be a > 0. Gaussian decays as exp(-a²(x−μ)²). Position std-dev: σ ≈ 1 / (√2 · a)
1	b	Imaginary part of the complex width α = a² + ib. Zero for a real envelope.
2	μ	Center position
3	p	Initial momentum

Common starting points:

import jax.numpy as jnp

# Single real Gaussian at rest at position x0
params = jnp.zeros((1, D, 4))
params = params.at[0, :, 0].set(a)      # width
params = params.at[0, d, 2].set(x0)     # center in dimension d

# n Gaussians spread over a grid, zero momentum
params = jnp.zeros((n, D, 4))
params = params.at[:, :, 0].set(a)
centers = jnp.linspace(x_min, x_max, n)
params = params.at[:, 0, 2].set(centers)

For the HO ground state exp(-x²/2): a = 1/√2 ≈ 0.707, b = μ = p = 0. For a displaced coherent state at x₀ with momentum p₀: set μ = x₀, p = p₀, keep a = 1/√2, b = 0.

Wiring to RotheSolver

from rothe.io import OutputConfig
from rothe.solver import RotheSolver, setUpRotheErrorAndGradient_jit
from rothe.systems.common import set_up_SHH2

# 1. Build the matrix-element callable from the potential string.
#    Signature: SHH2(t, params_ket, params_bra) -> (S, H, H²)
SHH2 = set_up_SHH2(potential_string=POTENTIAL_STRING, D=D)

# 2. Build JAX-compiled Rothe objective and objective+gradient callables.
rothe_error, rothe_vg = setUpRotheErrorAndGradient_jit()

# 3. Create and run the solver.
solver = RotheSolver(
    SHH2=SHH2,
    dt=dt,                          # float → real time; e.g. 0.1j → imaginary time
    t=0.0,
    epsilon=1e-3,                   # global Rothe-error budget
    regularization_lambda=1e-8,
    params_old=params_init,         # shape (n, D, 4)
    coeffs_old=coeffs_init,         # shape (n,), complex128
    rothe_grad_fn=rothe_vg,
    rothe_nograd=rothe_error,
    output_config=OutputConfig(name="my_sim", polynomial_string=POTENTIAL_STRING),
)

solver.propagate(nsteps, num_time_steps_total=nsteps)

Imaginary-time propagation (ground-state search): pass a purely imaginary dt, e.g. dt=0.1j. The solver renormalises after each step and prints the energy and variance, which converge to the ground-state values.

---

HDF5 output

Each run saves to wave_function_data/{name}.h5:

{name}.h5
├── attrs:  sim_name, dt, epsilon, regularization_lambda, created
├── polynomial_string              (scalar string dataset)
└── steps/
    ├── 0000/
    │   ├── params                 (n_total, D, 4)   float64
    │   ├── coeffs                 (n_total,)         complex128
    │   └── attrs: t, rothe_error, n_dynamic
    ├── 0001/
    │   └── ...
    └── ...

Optional per-step data (when polynomial_string is set in OutputConfig): dipole dataset and double_autocorrelation attribute.

The Gaussian count n_total can vary between steps — no padding.

Resuming a run:

python -m example_scripts.propagate_double_well ... --t_start 5.0

---

Testing

pytest

---

Repository layout

path	role
rothe/solver.py	RotheSolver — main loop, adaptive Gaussian addition
rothe/objective.py	Rothe objective, JAX-compiled value/gradient factory
rothe/wavefunction.py	Gaussian matrix elements (S, T, V, H, H²) and kinetic propagator
rothe/block_assembly.py	Block-matrix helpers (A, B, ρ blocks)
rothe/linear_solve.py	Regularised linear solve for coefficients
rothe/potential_parser.py	Potential string DSL parser
rothe/io.py	HDF5 save / load / resume helpers (OutputConfig)
rothe/systems/	System-specific setup (hydrogen, Henon-Heiles, shared helpers)
example_scripts/	Runnable propagation scripts + harmonic_oscillator.py template
tests/	pytest suite

Notes

• The first propagate() call triggers JAX compilation; subsequent steps are faster.
• scipy.optimize.minimize (L-BFGS-B for large bases, BFGS for small ones) drives the nonlinear parameter update on the host; GPU memory is used for matrix-element evaluations.
• Physical units and conventions are determined entirely by the potential string and the initial state you provide.