sun-diffusion 0.1.0

Diffusion models for SU(N) degrees of freedom.

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.