Documentation¶

sun_diffusion¶

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

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:

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:

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

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

Quickstart / Examples¶

Note: More in-depth examples can be found in the notebooks of this repository.

Physics Actions¶

This package allows one to define actions and evaluate them on batches of \({\rm SU}(N)\) configurations:

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)

>>> tensor([-0.0338, -0.0705, -0.5711, -0.7625])

See the sun_diffusion.action module for more details.

Diffusion Processes¶

This package also enables users to define diffusion processes on Euclidean space:

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

>>> x_0 std = 0.10194464772939682
>>> x_1 std = 1.0630394220352173

as well as on the \({\rm SU}(N)\) manifold:

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

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)

>>> 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, sun_diffusion.linalg and sun_diffusion.sun 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.

API Reference¶

Modules: