Schr: GPU-Accelerated Quantum Mechanics and QED Simulator

Schr is a high-performance Python framework for simulating quantum mechanical and quantum electrodynamics (QED) systems using GPU acceleration via JAX. Designed for researchers in computational physics, quantum optics, and atomic physics, it provides numerically rigorous implementations of time-dependent Schrödinger equation solvers and field-theoretic methods.

Key Features

• GPU-Accelerated: Built on JAX for automatic differentiation and XLA compilation to GPUs/TPUs

• Spectral Methods: Split-step Fourier method for efficient time evolution with periodic boundary conditions

• QED Support: Fock space representation for photon fields and light-matter interaction

• Research-Grade Numerics: Absorbing boundary conditions, normalization preservation, and energy conservation

• Flexible Architecture: Modular design supporting 1D/2D/3D systems with arbitrary potentials

Quick Start

Installation

# Install uv
curl -LsSf https://astral.sh/uv/install.sh | sh

# Clone and sync
git clone https://github.com/sql-hkr/schr.git
cd schr
uv sync

For GPU support:

uv pip install "jax[cuda12]"

Basic Usage

Solve the time-dependent Schrödinger equation for a Gaussian wavepacket:

import jax.numpy as jnp
from schr.qm.hamiltonian import ParticleInPotential
from schr.qm.solvers import SplitStepFourier
from schr.qm.wavefunction import gaussian_wavepacket
from schr.utils.grid import create_grid_1d

# Create spatial grid
x, dx = create_grid_1d(-10.0, 10.0, 1024)

# Define free particle Hamiltonian
def potential(psi, t):
    return jnp.zeros_like(psi, dtype=jnp.float32)

hamiltonian = ParticleInPotential(
    potential=potential,
    grid_shape=(1024,),
    dx=dx,
    mass=1.0,
    hbar=1.0
)

# Initialize wavepacket
psi0 = gaussian_wavepacket(x, x0=0.0, p0=1.0, sigma=1.0)

# Time evolution
solver = SplitStepFourier(hamiltonian)
psi = psi0
for step in range(1000):
    psi = solver.step(psi, step * 0.01, dt=0.01)

Documentation Contents

User Guide

• Installation Guide
  • Requirements
  • Quick Install
  • GPU Support
  • Development Installation
  • Docker Installation
  • Platform-Specific Notes
  • Verifying Installation
  • Uninstallation
  • Next Steps
• Examples
  • 1D Examples
  • 2D Examples
  • QED Examples
  • Modifying Examples
  • Next Steps
• Theoretical Background
  • Quantum Mechanics Foundations
  • Numerical Methods
  • Boundary Conditions
  • Quantum Electrodynamics
  • Example Systems
  • Atomic Units
  • Numerical Considerations

API Reference

• src
  • schr package
• schr.core package
  • Submodules
  • Module contents
• schr.qm package
  • Submodules
  • Module contents
• schr.qed package
  • Submodules
  • Module contents
• schr.utils package
  • Submodules
  • Module contents

Additional Information

• Physical Units and Conversions
• Performance Optimization
• Contributing to Schr

Mathematical Foundation

Time-Dependent Schrödinger Equation

The package solves the TDSE in atomic units (\(\hbar = m_e = e = 1\)):

\[i\frac{\partial \psi(\mathbf{r}, t)}{\partial t} = \hat{H}\psi(\mathbf{r}, t) = \left[-\frac{1}{2m}\nabla^2 + V(\mathbf{r}, t)\right]\psi(\mathbf{r}, t)\]

where \(\psi(\mathbf{r}, t)\) is the wavefunction, \(m\) is the particle mass, and \(V(\mathbf{r}, t)\) is the time-dependent potential.

Split-Step Fourier Method

The primary solver implements the split-step Fourier method (SSFM), a pseudo-spectral technique:

\[\psi(t + \Delta t) \approx e^{-i\hat{V}\Delta t/2} \, e^{-i\hat{T}\Delta t} \, e^{-i\hat{V}\Delta t/2} \, \psi(t) + O(\Delta t^3)\]

Advantages:

• Second-order accuracy in time: local error \(O(\Delta t^3)\), global error \(O(\Delta t^2)\)

• Preserves unitarity and norm conservation

• Efficient FFT-based implementation: \(O(N \log N)\) per time step

• Handles arbitrary potentials without operator commutation requirements

Project Structure

The package is organized into several modules:

Core (schr.core)

Abstract base classes for Hamiltonians, Solvers, and Fields

Quantum Mechanics (schr.qm)

• hamiltonian: Hamiltonian operators (ParticleInPotential)

• solvers: Time evolution methods (SplitStepFourier, RungeKutta4)

• wavefunction: Wavefunction utilities and initial conditions

Quantum Electrodynamics (schr.qed)

• field: PhotonField representation in Fock space

• interaction: Light-matter interaction Hamiltonians

Utilities (schr.utils)

• grid: Grid generation for 1D/2D/3D systems

• fft: FFT utilities and momentum operators

• absorbing_boundary: Boundary condition implementations

• visualization: Plotting and animation tools

Physical Units

The package uses atomic units (Hartree atomic units) by default:

Quantity	Atomic Unit	SI Equivalent	Symbol
Length	Bohr radius	\(5.29 \times 10^{-11}\) m	\(a_0\)
Energy	Hartree	\(4.36 \times 10^{-18}\) J = 27.21 eV	\(E_h\)
Time	\(\hbar/E_h\)	\(2.42 \times 10^{-17}\) s = 24.2 as	—
Mass	Electron mass	\(9.11 \times 10^{-31}\) kg	\(m_e\)
Charge	Elementary charge	\(1.60 \times 10^{-19}\) C	\(e\)

In atomic units, \(\hbar = m_e = e = 1\).

Contributing

Contributions are welcome! Please see the Contributing to Schr guide for details on:

• Setting up a development environment

• Running tests

• Code style guidelines

• Submitting pull requests

Citation

If you use this software in your research, please cite:

@software{schr2025,
  author = {sql-hkr},
  title = {Schr: GPU-Accelerated Quantum Mechanics and QED Simulator},
  year = {2025},
  url = {https://github.com/sql-hkr/schr}
}

License

This project is licensed under the MIT License.

Indices and Tables

• Index

• Module Index

• Search Page