core.models.gemnet_oc.layers.spherical_basis

core.models.gemnet_oc.layers.spherical_basis#

Copyright (c) Meta, Inc. and its affiliates. This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Classes#

CircularBasisLayer

2D Fourier Bessel Basis

SphericalBasisLayer

3D Fourier Bessel Basis

Module Contents#

class core.models.gemnet_oc.layers.spherical_basis.CircularBasisLayer(num_spherical: int, radial_basis: core.models.gemnet_oc.layers.radial_basis.RadialBasis, cbf: dict, scale_basis: bool = False)#

Bases: torch.nn.Module

2D Fourier Bessel Basis

Parameters:

  • num_spherical (int) – Number of basis functions. Controls the maximum frequency.

  • radial_basis (RadialBasis) – Radial basis function.

  • cbf (dict) – Name and hyperparameters of the circular basis function.

  • scale_basis (bool) – Whether to scale the basis values for better numerical stability.

radial_basis#
scale_basis#
forward(D_ca, cosφ_cab)#
class core.models.gemnet_oc.layers.spherical_basis.SphericalBasisLayer(num_spherical: int, radial_basis: core.models.gemnet_oc.layers.radial_basis.RadialBasis, sbf: dict, scale_basis: bool = False)#

Bases: torch.nn.Module

3D Fourier Bessel Basis

Parameters:

  • num_spherical (int) – Number of basis functions. Controls the maximum frequency.

  • radial_basis (RadialBasis) – Radial basis functions.

  • sbf (dict) – Name and hyperparameters of the spherical basis function.

  • scale_basis (bool) – Whether to scale the basis values for better numerical stability.

num_spherical#
radial_basis#
scale_basis#
forward(D_ca, cosφ_cab, θ_cabd)#