MolecularDiffusion.modules.layers.esen.common.so3¶

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Classes¶

`CoefficientMapping`	Helper module for coefficients used to reshape l <--> m and to get coefficients of specific degree or order
`SO3_Grid`	Helper functions for grid representation of the irreps

Module Contents¶

class MolecularDiffusion.modules.layers.esen.common.so3.CoefficientMapping(lmax, mmax)¶

Bases: torch.nn.Module

Helper module for coefficients used to reshape l <–> m and to get coefficients of specific degree or order

Parameters:

(list (mmax_list) – int): List of maximum degree of the spherical harmonics
(list – int): List of maximum order of the spherical harmonics
use_rotate_inv_rescale (bool) – Whether to pre-compute inverse rotation rescale matrices

coefficient_idx(lmax: int, mmax: int)¶

complex_idx(m, lmax, m_complex, l_harmonic)¶: Add m_complex and l_harmonic to the input arguments since we cannot use self.m_complex.

pre_compute_coefficient_idx()¶: Pre-compute the results of coefficient_idx() and access them with prepare_coefficient_idx()

pre_compute_rotate_inv_rescale()¶

prepare_coefficient_idx()¶: Construct a list of buffers

lmax¶

m_size¶

mmax¶

res_size¶

class MolecularDiffusion.modules.layers.esen.common.so3.SO3_Grid(lmax: int, mmax: int, normalization: str = 'integral', resolution: int | None = None, rescale: bool = True)¶

Bases: torch.nn.Module

Helper functions for grid representation of the irreps

Parameters:

lmax (int) – Maximum degree of the spherical harmonics
mmax (int) – Maximum order of the spherical harmonics

from_grid(grid, lmax: int, mmax: int)¶

get_from_grid_mat(device=None)¶

get_to_grid_mat(device=None)¶

to_grid(embedding, lmax: int, mmax: int)¶

lat_resolution¶

lmax¶

mapping¶

mmax¶

rescale = True¶