RandomFeaturesHead

class franken.rf.heads.RandomFeaturesHead(input_dim, num_random_features=2**10, num_species=None, chemically_informed_ratio=None)

Bases: Module

Base class for random-feature heads

Parameters:

  • input_dim (int) – Dimension of the input features.

  • num_random_features (int) – The number of random features to use in the feature mapping. Defaults to \(2^{10} = 1024\).

  • num_species (int | None) – The number of chemical species for which the kernel is computed. This parameter is relevant for systems with multiple chemical species. Defaults to None.

  • chemically_informed_ratio (float | None) – The relative weight of chemically-informed kernels with respect to the all-species kernel. Ignored if num_species is None. Defaults to None.

init_args()

Returns the arguments needed to re-initialize this class.

species_scatter_sum(Z, atomic_numbers=None, batch_ids=None)

Average features across all atoms in one or more configurations.

Depending on the configuration of random features, this function will either perform a simple average, or will use a chemically-informed averaging method where features are averaged within each atomic type and concatenated across atomic types.

Returns a feature per-system feature map of shape [n_systems, total_feature_dim].