jax_md_mod.model.layers.RadialBesselLayer

jax_md_mod.model.layers.RadialBesselLayer#

class RadialBesselLayer(cutoff, num_radial=16, envelope_p=6, name='BesselRadial')[source]#

Radial Bessel Function (RBF) representation of pairwise distances.

Variables:

freq_init – RBFFrequencyInitializer

Methods

__init__(cutoff, num_radial=16, envelope_p=6, name='BesselRadial')[source]#

Initializes the RBF layer.

Parameters:

  • cutoff – Radial cut-off distance

  • num_radial – Number of radial Bessel embedding functions

  • envelope_p – Power of envelope polynomial

  • name – Name of RBF layer

__call__(distances)[source]#

Returns the RBF embeddings of edge distances.

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