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

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 that alternates between position and momentum representations:

$$\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)$$

where $\hat{T} = -\nabla^2/(2m)$ is the kinetic energy operator and $\hat{V}$ is the potential energy operator.

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

Installation

Prerequisites

• Python >= 3.11
• uv (recommended package manager)
• CUDA-compatible GPU (optional, for GPU acceleration)
• CUDA Toolkit >= 11.8 (for GPU support)

Quick Install

# Install uv if you haven't already
curl -LsSf https://astral.sh/uv/install.sh | sh

# Clone the repository
git clone https://github.com/sql-hkr/schr.git
cd schr

# Sync dependencies and install the package
uv sync

GPU Support

For GPU acceleration, install JAX with CUDA support after syncing:

# Sync dependencies first
uv sync

# Then install JAX with CUDA support
# CUDA 12
uv pip install "jax[cuda12]"

# CUDA 11
uv pip install "jax[cuda11]"

See the JAX installation guide for detailed instructions.

Alternative: Using pip

If you prefer pip:

pip install -e .

Project Structure

schr/
├── src/schr/
│   ├── core/          # Abstract base classes
│   │   └── base.py    # Hamiltonian, Solver, Field
│   ├── qm/            # Quantum mechanics
│   │   ├── hamiltonian.py  # ParticleInPotential
│   │   ├── solvers.py      # SplitStepFourier, RungeKutta4
│   │   └── wavefunction.py # Wavefunction utilities
│   ├── qed/           # Quantum electrodynamics
│   │   ├── field.py        # PhotonField (Fock space)
│   │   └── interaction.py  # Light-matter coupling
│   └── utils/         # Numerical utilities
│       ├── grid.py             # Grid generation
│       ├── fft.py              # FFT utilities
│       ├── absorbing_boundary.py  # Boundary conditions
│       └── visualization.py    # Plotting tools
├── examples/          # Example simulations
├── tests/             # Unit tests
└── docs/              # Sphinx documentation

Physical Units

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

In atomic units, $\hbar = m_e = e = 1$.

Performance Considerations

GPU Acceleration

JAX automatically compiles numerical kernels to GPU. For optimal performance:

1. Use float32/complex64: Default precision balances accuracy and memory
2. JIT compilation: Use @jit decorator for hot loops
3. Batch operations: Vectorize over multiple initial conditions
4. Memory management: Monitor GPU memory with large grids

from jax import jit

@jit
def evolve_step(psi, t, dt):
    return solver.step(psi, t, dt)

# First call compiles, subsequent calls are fast
psi = evolve_step(psi, 0.0, 0.01)

Grid Size Selection

For spectral methods, choose grid sizes as powers of 2 for optimal FFT performance:

• 1D: 1024, 2048, 4096, ...
• 2D: 512×512, 1024×1024, 2048×2048, ...
• 3D: 128×128×128, 256×256×256, ...

Testing

Run the test suite:

# All tests
uv run pytest

# With coverage
uv run pytest --cov=schr

# Specific test file
uv run pytest tests/test_qm.py

Contributing

Contributions are welcome! Please see CONTRIBUTING.md for detailed guidelines on:

• Setting up your development environment
• Code style and standards
• Testing requirements
• Submitting pull requests

Quick start for contributors:

# Sync all dependencies (including dev)
uv sync --all-extras

# Run tests
uv run pytest

# Check code style
uv run ruff check src/

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 - see the LICENSE file for details.

Acknowledgments

• Built with JAX by Google Research
• Inspired by quantum optics research at [Your Institution]
• Special thanks to the computational physics community

Contact

• Author: sql-hkr
• Email: sql.hkr@gmail.com
• GitHub: @sql-hkr
• Issues: GitHub Issues

Note: This software is under active development. API stability is not guaranteed until version 1.0.0.