schr 0.1.0
Schr: GPU-Accelerated Quantum Mechanics and QED Simulator
Requires Python: >=3.11
Keywords: quantum mechanics, quantum electrodynamics, JAX, GPU, simulation, physics

# Schr: GPU-Accelerated Quantum Mechanics and QED Simulator

[![PyPI version](https://img.shields.io/pypi/v/schr)](https://pypi.org/project/schr/)
[![License](https://img.shields.io/github/license/sql-hkr/schr)](LICENSE)
[![Python versions](https://img.shields.io/pypi/pyversions/schr)](https://pypi.org/project/schr/)
[![CI](https://github.com/sql-hkr/schr/actions/workflows/ci.yml/badge.svg)](https://github.com/sql-hkr/schr/actions/workflows/ci.yml)

**Schr** is a high-performance Python framework for simulating quantum mechanical and quantum electrodynamics (QED) systems using GPU acceleration via [JAX](https://github.com/google/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](https://github.com/astral-sh/uv) (recommended package manager)
- CUDA-compatible GPU (optional, for GPU acceleration)
- CUDA Toolkit >= 11.8 (for GPU support)

### Quick Install

```bash
# 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:

```bash
# 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](https://jax.readthedocs.io/en/latest/installation.html) for detailed instructions.

### Alternative: Using pip

If you prefer pip:

```bash
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:

| Quantity | Atomic Unit | SI Equivalent |
|----------|-------------|---------------|
| Length   | Bohr radius $a_0$ | $5.29 \times 10^{-11}$ m |
| Energy   | Hartree $E_h$ | $4.36 \times 10^{-18}$ |
| Time     | $\hbar/E_h$ | $2.42 \times 10^{-17}$ |
| Mass     | Electron mass $m_e$ | $9.11 \times 10^{-31}$ kg |
| Charge   | Elementary charge $e$ | $1.60 \times 10^{-19}$ C |

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

```python
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:

```bash
# 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](CONTRIBUTING.md) for detailed guidelines on:

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

Quick start for contributors:

```bash
# 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:

```bibtex
@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](LICENSE) file for details.

## Acknowledgments

- Built with [JAX](https://github.com/google/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](https://github.com/sql-hkr)
- **Issues**: [GitHub Issues](https://github.com/sql-hkr/schr/issues)

---

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