jax_md_mod.model.layers.SphericalBesselLayer

jax_md_mod.model.layers.SphericalBesselLayer#

class SphericalBesselLayer(r_cutoff, num_spherical, num_radial, envelope_p=6, name='BesselSpherical')[source]#

Spherical Bessel Function (SBF) representation of angular triplets.

Methods

__init__(r_cutoff, num_spherical, num_radial, envelope_p=6, name='BesselSpherical')[source]#

Initializes the SBF layer.

Parameters:

  • r_cutoff – Radial cut-off

  • num_spherical – Number of spherical Bessel embedding functions

  • num_radial – Number of radial Bessel embedding functions

  • envelope_p – Power of envelope polynomial

  • name – Name of SBF layer

__call__(pair_distances, angles, angle_mask, expand_to_kj)[source]#

Returns the SBF embeddings of angular triplets.

params_dict()

Returns parameters keyed by name for this module and submodules.

state_dict()

Returns state keyed by name for this module and submodules.
Contents