enerzyme.models.layers.rbf.ExponentialGaussianRBFLayer#

class enerzyme.models.layers.rbf.ExponentialGaussianRBFLayer(num_rbf: int, no_basis_at_infinity: bool = False, init_alpha: float = 0.944863062918464, exp_weighting: bool = False, learnable_shape: bool = True, cutoff_sr: float = inf, cutoff_fn: Literal['polynomial', 'bump'] = 'polynomial', init_width_flavor: Literal['PhysNet', 'SpookyNet'] = 'PhysNet')[source]#

Bases: ExponentialRBF

__init__(num_rbf: int, no_basis_at_infinity: bool = False, init_alpha: float = 0.944863062918464, exp_weighting: bool = False, learnable_shape: bool = True, cutoff_sr: float = inf, cutoff_fn: Literal['polynomial', 'bump'] = 'polynomial', init_width_flavor: Literal['PhysNet', 'SpookyNet'] = 'PhysNet') → None[source]#

Radial basis functions based on exponential Gaussian functions given by:

g_i(x) = exp(-width_i*(exp(-alpha*x)-center_i)**2)

Params:#

num_basis_functions: Number of radial basis functions.

dtype: Data type of floating numbers.

no_basis_function_at_infinity: If True, no basis function is put at exp(-alpha*x) = 0, i.e. x = infinity.

init_alpha: Initial value for scaling parameter alpha (Default value corresponds to 0.5 1/Bohr converted to 1/Angstrom).

init_width_flavor: Initialization flavor for width of the exponentials. Options:

  • PhysNet: A constant number (2K^{-1}(1-exp(-cutoff)))^{-2}, where K is num_basis_functions

  • SpookyNet: A constant number K or K+1 (no_basis_function_at_infinity=True)

exp_weighting: If True, basis functions are weighted with a factor exp(-alpha*r).

learnable_shape: If True, centers and widths of exponentials are learnable.

cutoff: Short range cutoff threshold for radial base functions.

cutoff_fn: Short range cutoff function, whose are called by cutoff_fn(x, cutoff=cutoff) where x is the distance.

init_width_flavor: Initialization flavor for width of the exponentials. Options:

  • PhysNet: A constant number (2K^{-1}(1-exp(-cutoff)))^{-2}, where K is num_basis_functions [1].

  • SpookyNet: A constant number K or K+1 (no_basis_function_at_infinity=True).

References:#

[1] J. Chem. Theory Comput. 2019, 15, 3678−3693.

inner_fn(alphar: Tensor, expalphar: Tensor) → Tensor[source]#