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
~~~~~~~~~~~~

.. code-block:: bash

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

.. code-block:: bash

   uv pip install "jax[cuda12]"

Basic Usage
~~~~~~~~~~~

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

.. code-block:: python

   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
----------------------

.. toctree::
   :maxdepth: 2
   :caption: User Guide

   installation
   examples
   theory

.. toctree::
   :maxdepth: 2
   :caption: API Reference

   api/modules
   api/schr.core
   api/schr.qm
   api/schr.qed
   api/schr.utils

.. toctree::
   :maxdepth: 1
   :caption: Additional Information

   physical_units
   performance
   contributing

Mathematical Foundation
-----------------------

Time-Dependent Schrödinger Equation
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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

.. math::

   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 :math:`\psi(\mathbf{r}, t)` is the wavefunction, :math:`m` is the particle mass, and :math:`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:

.. math::

   \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 :math:`O(\Delta t^3)`, global error :math:`O(\Delta t^2)`
* Preserves unitarity and norm conservation
* Efficient FFT-based implementation: :math:`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:

.. list-table::
   :header-rows: 1
   :widths: 20 25 35 20

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

In atomic units, :math:`\hbar = m_e = e = 1`.

Contributing
------------

Contributions are welcome! Please see the :doc:`contributing` 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:

.. code-block:: 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.

Indices and Tables
------------------

* :ref:`genindex`
* :ref:`modindex`
* :ref:`search`