core.models.gemnet_gp.layers.radial_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

PolynomialEnvelope	Polynomial envelope function that ensures a smooth cutoff.
ExponentialEnvelope	Exponential envelope function that ensures a smooth cutoff,
SphericalBesselBasis	1D spherical Bessel basis
BernsteinBasis	Bernstein polynomial basis,
RadialBasis

Module Contents

class core.models.gemnet_gp.layers.radial_basis.PolynomialEnvelope(exponent: int)

Bases: torch.nn.Module

Polynomial envelope function that ensures a smooth cutoff.

Parameters:

exponent (int) – Exponent of the envelope function.

p

a: float

b: float

c: float

forward(d_scaled: torch.Tensor) → torch.Tensor

class core.models.gemnet_gp.layers.radial_basis.ExponentialEnvelope

Bases: torch.nn.Module

Exponential envelope function that ensures a smooth cutoff, as proposed in Unke, Chmiela, Gastegger, Schütt, Sauceda, Müller 2021. SpookyNet: Learning Force Fields with Electronic Degrees of Freedom and Nonlocal Effects

forward(d_scaled) → torch.Tensor

class core.models.gemnet_gp.layers.radial_basis.SphericalBesselBasis(num_radial: int, cutoff: float)

Bases: torch.nn.Module

1D spherical Bessel basis

Parameters:

• num_radial (int) – Controls maximum frequency.

• cutoff (float) – Cutoff distance in Angstrom.

norm_const

frequencies

forward(d_scaled)

class core.models.gemnet_gp.layers.radial_basis.BernsteinBasis(num_radial: int, pregamma_initial: float = 0.45264)

Bases: torch.nn.Module

Bernstein polynomial basis, as proposed in Unke, Chmiela, Gastegger, Schütt, Sauceda, Müller 2021. SpookyNet: Learning Force Fields with Electronic Degrees of Freedom and Nonlocal Effects

Parameters:

• num_radial (int) – Controls maximum frequency.

• pregamma_initial (float) – Initial value of exponential coefficient gamma. Default: gamma = 0.5 * a_0**-1 = 0.94486, inverse softplus -> pregamma = log e**gamma - 1 = 0.45264

pregamma

softplus

forward(d_scaled) → torch.Tensor

class core.models.gemnet_gp.layers.radial_basis.RadialBasis(num_radial: int, cutoff: float, rbf: dict[str, str] | None = None, envelope: dict[str, str | int] | None = None)

Bases: torch.nn.Module

Parameters:

• num_radial (int) – Controls maximum frequency.

• cutoff (float) – Cutoff distance in Angstrom.

• rbf (dict = {"name": "gaussian"}) – Basis function and its hyperparameters.

• envelope (dict = {"name": "polynomial", "exponent": 5}) – Envelope function and its hyperparameters.

inv_cutoff

forward(d)