core.models.uma.common.so3

core.models.uma.common.so3#

Copyright (c) Meta Platforms, Inc. and affiliates.

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 core.models.uma.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

lmax#
mmax#
res_size#
m_size#
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()

prepare_coefficient_idx()#

Construct a list of buffers

coefficient_idx(lmax: int, mmax: int)#
pre_compute_rotate_inv_rescale()#
__repr__()#
class core.models.uma.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

lmax#
mmax#
lat_resolution#
mapping#
rescale#
get_to_grid_mat(device=None)#
get_from_grid_mat(device=None)#
to_grid(embedding, lmax: int, mmax: int)#
from_grid(grid, lmax: int, mmax: int)#
Contents