pylibcugraphops.dimenet.radial_basis_fwd

pylibcugraphops.dimenet.radial_basis_fwd = <nanobind.nb_func object>

Computes the forward pass for Dimenet++ radial basis features.

radial_basis_fwd(
    output_rbf: device array, output_sbf_rad: device array,
    input_vector: device array, input_w: device array
)

We define the following dimensions:

• n_spherical: number of spherical basis functions. Must be 7.

• n_radial: number of radial basis functions. Must be 6.

• n_dist: number of distance dimensions. Must be 1.

Parameters:

output_rbfdevice array type

Device array containing the output radial basis features. Dimension is assumed to be [#edges, n_radial].

output_sbf_raddevice array type

Device array containing the output radial basis part of spherical basis features. Dimension is assumed to be [#edges, n_spherical * n_radial].

input_vectordevice array type

Device array containing the vector values (position of input node - position of output node + edge offsets) for each edge. Dimension is assumed to be [#edges, n_dist].

input_wdevice array type

Device array containing the input frequencies. Dimension is assumed to be [n_radial].