sun_diffusion.heat module¶
Module for heat kernel evaluations and analytical score functions.
Notation: We use width to denote the standard deviation of the heat kernel instead of sigma, since this is used as a parameter label elsewhere. We often abbreviate the heat kernel to HK.
- sun_diffusion.heat.log_sun_hk(thetas, *, width, n_max=3, eig_meas=True)¶
Computes the \({\rm SU}(N)\) log HK over wrapped eigenangles.
- Return type:
Tensor
- sun_diffusion.heat.sun_hk(thetas, *, width, n_max=3, eig_meas=True)¶
Evaluates the \({\rm SU}(N)\) heat kernel on wrapped eigenangles.
Note
This function assumes the input only includes the \(N-1\) independent eigenangles.
- Parameters:
thetas (Tensor) – Wrapped eigenangles, shaped [B, Nc-1]
width (Tensor) – Standard deviation of the heat kernel, batched
n_max (int) – Max number of pre-image sum terms to include. Default: 3
eig_meas (bool) – Weather to include Haar measure term. Default: True
- Return type:
Tensor- Returns:
\({\rm SU}(N)\) heat kernel evaluated on the angles thetas
- sun_diffusion.heat.sun_dual_hk(thetas, *, width, max_weight=5, eig_meas=True)¶
Evaluates the \({\rm SU}(N)\) dual heat kernel over wrapped eigenangles as a character expansion.
Note
This function assumes the input only includes the \(N-1\) independent eigenangles.
- Parameters:
thetas (Tensor) – Wrapped eigenangles, shaped [B, Nc-1]
width (Tensor) – Standard deviation of the heat kernel, batched
max_weight (int) – Max total weight of irreps to sum. Default: 5
eig_meas (bool) – Whether to include Haar measure term. Default: True
- Return type:
Tensor- Returns:
\({\rm SU}(N)\) dual heat kernel evaluated on the angles thetas
- sun_diffusion.heat.sun_score_hk(thetas, *, width, n_max=3)¶
Computes the analytical score function for the wrapped \({\rm SU}(N)\) heat kernel over eigenangles thetas.
Note
This function assumes the input only includes the \(N-1\) independent eigenangles.
- Parameters:
thetas (Tensor) – Wrapped eigenangles, shaped [B, Nc-1]
width (Tensor) – Standard deviation of the heat kernel
n_max (int) – Max number of pre-image sum terms to include. Default: 3
- Return type:
Tensor- Returns:
Analytical gradient of the wrapped log HK
- sun_diffusion.heat.sun_score_dual_hk(thetas, *, width, max_weight=5)¶
Computes the analytical score function for the dual \({\rm SU}(N)\) heat kernel over eigenangles thetas.
Note
This function assumes the input only includes the \(N-1\) independent eigenangles.
- Parameters:
thetas (Tensor) – Wrapped eigenangles, shaped [B, Nc-1]
width (Tensor) – Standard deviation of the heat kernel, batched
max_weight (int) – Max total weight of irreps to sum. Default: 5
- Return type:
Tensor- Returns:
Analytical gradient of the character expansion HK
- sun_diffusion.heat.sun_score_hk_autograd(thetas, *, width, n_max=3)¶
Computes the score function for the wrapped \({\rm SU}(N)\) heat kernel via automatic differentiation of \(K\) in thetas, followed by division by \(K\).
This function uses autograd to compute the score indirectly as
\[s(\theta) = \frac{\nabla K(\theta)}{K(\theta)}\]- Parameters:
thetas (Tensor) – Wrapped eigenangles, shaped [B, Nc-1]
width (float) – Standard deviation of the heat kernel
n_max (int) – Max number of pre-image sum terms to include. Default: 3
- Return type:
Tensor- Returns:
Autograd derivative of the HK divided by \(K\)
- sun_diffusion.heat.sun_score_hk_autograd_v2(thetas, *, width, n_max=3)¶
Computes the score function for the wrapped \({\rm SU}(N)\) heat kernel via automatic differentiation of \(\log K\) in thetas.
This function uses autograd to compute the score directly as
\[s(\theta) = \nabla \log K(\theta)\]- Parameters:
thetas (Tensor) – Wrapped eigenangles, shaped [B, Nc-1]
width (float) – Standard deviation of the heat kernel
n_max (int) – Max number of pre-image sum terms to include. Default: 3
- Return type:
Tensor- Returns:
Autograd derivative of the HK log-density
- sun_diffusion.heat.sample_sun_hk(batch_size, Nc, *, width, n_iter=3, n_max=3)¶
Samples from the \({\rm SU}(N)\) heat kernel with importance sampling.
- Parameters:
batch_size (int) – Number of samples to generate
Nc (int) – Dimension of fundamental rep. of \({\rm SU}(N)\)
width (float)
n_iter (int) – Number of sampling iterations. Default: 3
n_max (int) – Max number of pre-image sum terms to include. Default: 3
- Returns:
Batch of eigenangles ([B, Nc]) from the heat kernel
- Return type:
xs (NDArray)