Metadata-Version: 2.4
Name: sun_diffusion
Version: 1.0.0
Summary: Diffusion models for SU(N) degrees of freedom.
Author-email: Gurtej Kanwar <gurtejkanwar@gmail.com>, Octavio Vega <vega00087@gmail.com>
License: MIT License
        
        Copyright (c) 2025 Octavio Vega
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
Project-URL: Repository, https://github.com/ovega14/sun_diffusion
Project-URL: Paper, https://arxiv.org/abs/2512.19877
Project-URL: Documentation, https://ovega14.github.io/sun_diffusion
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy<2,>=1.24
Requires-Dist: torch>=2.0.0
Provides-Extra: gpu
Requires-Dist: torch; extra == "gpu"
Provides-Extra: notebooks
Requires-Dist: jupyter; extra == "notebooks"
Requires-Dist: matplotlib; extra == "notebooks"
Requires-Dist: tqdm<5.0,>=4.65; extra == "notebooks"
Requires-Dist: ipywidgets==7.7.2; extra == "notebooks"
Provides-Extra: docs
Requires-Dist: sphinx>=7.0; extra == "docs"
Requires-Dist: furo; extra == "docs"
Requires-Dist: myst-parser; extra == "docs"
Requires-Dist: sphinx-autodoc-typehints; extra == "docs"
Requires-Dist: sphinx-copybutton; extra == "docs"
Dynamic: license-file

# sun_diffusion
[![arXiv](https://img.shields.io/badge/arXiv-2512.19877-b31b1b.svg)](https://arxiv.org/abs/2512.19877)
[![Docs](https://img.shields.io/badge/docs-GitHub%20Pages-blue)](https://ovega14.github.io/sun_diffusion/)
[![PyPI - Version](https://img.shields.io/pypi/v/sun_diffusion)](https://pypi.org/project/sun-diffusion/#description)
[![DOI](https://zenodo.org/badge/1046560489.svg)](https://doi.org/10.5281/zenodo.18419882)

Diffusion models for ${\rm SU}(N)$ degrees of freedom.

What this package does:
- Implements diffusion processes on the ${\rm SU}(N)$ Lie group manifold with
  - different noise schedules
  - customizable parameters
- Includes utilities for manipulating matrices (and their spectra) on ${\rm SU}(N)$ and $\mathfrak{su}(N)$, including
  - diagonalization
  - canonicalization
  - eigendecomposition
  - algebra-to-group projections (and vice versa)
  - irreducible representations
- Provides heat kernel evaluations, sampling routines, and score functions


## Installation
### CPU only (default)
```bash
pip install sun_diffusion
```

### GPU (CUDA) Users
Before installing, make sure to install a CUDA-enabled PyTorch compatible with your GPU.
For example, for CUDA 12.4:
```bash
pip install torch --index-url https://download.pytorch.org/whl/cu124
pip install sun_diffusion[gpu]
```
If PyTorch cannot detect a GPU, your code will fall back to CPU, or `set_device('cuda')` will raise an error. You can verify CUDA availability with a small example:
```python
from sun_diffusion.devices import set_device, summary, HAS_CUDA

# Check CUDA availability
print('CUDA available:', HAS_CUDA)
if HAS_CUDA:
    set_device('cuda', 0)
else:
    set_device('cpu')
print(summary())
```
```pycon
>>> CUDA available: True
>>> Using device: cuda:0 (NVIDIA GH200 120GB) with dtype: torch.float32
```
Further utilities for handling devices and dtypes can be found in the [`sun_diffusion.devices`](https://github.com/ovega14/sun_diffusion/blob/main/sun_diffusion/devices.py) module.

## Quickstart / Examples
**Note:** More in-depth examples can be found in the [`notebooks`](https://github.com/ovega14/sun_diffusion/blob/main/notebooks/) of this repository.

#### Physics Actions
This package allows one to define actions and evaluate them on batches of ${\rm SU}(N)$ configurations:
```python
from sun_diffusion.action import SUNToyAction
from sun_diffusion.sun import random_sun_element

# Create a toy action
action = SUNToyAction(beta=1.0)

# Random batch of SU(3) matrices
batch_size = 4
U = random_sun_element(batch_size, Nc=3)

# Evaluate the action
S = action(U)
print(S)
```
```pycon
>>> tensor([-0.0338, -0.0705, -0.5711, -0.7625])
```
See the [`sun_diffusion.action`](https://github.com/ovega14/sun_diffusion/blob/main/sun_diffusion/action.py) module for more details.

#### Diffusion Processes
This package also enables users to define diffusion processes on Euclidean space:
```python
import torch
from sun_diffusion.diffusion import VarianceExpandingDiffusion

batch_size = 512
x_0 = 0.1 * torch.randn((batch_size, 3))

# Diffuse x_0 -> x_1 on R^3
diffuser = VarianceExpandingDiffusion(kappa=1.1)
x_1 = diffuser(x_0, t=torch.ones(batch_size))
print('x_0 std =', x_0.std().item())
print('x_1 std =', x_1.std().item())
```
```pycon
>>> x_0 std = 0.10194464772939682
>>> x_1 std = 1.0630394220352173
```
as well as on the ${\rm SU}(N)$ manifold:
```python
from sun_diffusion.diffusion import PowerDiffusionSUN

batch_size = 512
U_0 = random_sun_element(batch_size, Nc=2, scale=0.1)

# Diffuse U_0 -> U_1 on SU(2)
diffuser = PowerDiffusionSUN(kappa=3.0, alpha=1.0)
U_1, xs, V = diffuser(U_0, torch.ones(batch_size))
```
See [`sun_diffusion.diffusion`](https://github.com/ovega14/sun_diffusion/blob/main/sun_diffusion/diffusion.py) for more diffusion processes and implementation details.

#### Heat Kernel and Group Algebra
Sampling from the ${\rm SU}(N)$ heat kernel over the diagonal subalgebra of eigenangles is also simple, and can easily be combined with this package's matrix algebra utilities to produce group elements:
```python
from sun_diffusion.heat import sample_sun_hk
from sun_diffusion.linalg import adjoint
from sun_diffusion.sun import random_un_haar_element, embed_diag, matrix_exp

# Batch of HK eigenangles
batch_size = 2
Nc = 2
xs = sample_sun_hk(batch_size, Nc=Nc, width=torch.rand(batch_size), n_iter=10)

# Promote to Algebra -> project to SU(N)
X = embed_diag(torch.tensor(xs))
V = random_un_haar_element(batch_size, Nc=Nc)
U = V @ matrix_exp(X) @ adjoint(V)
print(U)
```
```pycon
>>> tensor([[[ 0.8995-0.3124j,  0.0070+0.3055j],
         [-0.0070+0.3055j,  0.8995+0.3124j]],

        [[ 0.9748+0.1920j,  0.0399+0.1061j],
         [-0.0399+0.1061j,  0.9748-0.1920j]]])
```
See the [`sun_diffusion.heat`](https://github.com/ovega14/sun_diffusion/blob/main/sun_diffusion/heat.py), [`sun_diffusion.linalg`](https://github.com/ovega14/sun_diffusion/blob/main/sun_diffusion/linalg.py) and [`sun_diffusion.sun`](https://github.com/ovega14/sun_diffusion/blob/main/sun_diffusion/sun.py) modules for more.
 
## Conventions
Some notes on mathematical conventions and notation:

#### Exponential Map
In physics, the exponential map $\exp: \mathfrak{su}(N) \to {\rm SU}(N)$ is defined conventionally through 

$$U = \exp(A) := e^{iA}$$ 

where $A = A^\dagger$ is a *Hermitian* matrix. In the math literature, one absorbs the factor of $i$ into the matrix so that $A$ is *anti-Hermitian*. We adopt the **physicist's** convention, so our functions that map between group and algebra, namely `sun.matrix_exp()` and `sun.matrix_log()`, expect a Hermitian matrix as input/output, respectively.

#### Heat Equation
In mathematics, the Heat equation is often written simply as 

$$\partial_t p_t(U) = \Delta p_t(U)$$

where $\Delta$ is the Laplace-Beltrami operator on the manifold $\mathcal{M} \ni U$. But to define more expressive diffusion processes, one introduces a time-dependent diffusion-coefficient, denoted $g_t$, such that

$$\partial_t p(U) = \frac{g_t^2}{2}\Delta p_t(U).$$

The factor of 1/2 originates from not inserting a factor of $\sqrt{2}$ in the SDE that defines the diffusion process, which we adopt as well.
