FeatureScaler

class franken.rf.scaler.FeatureScaler(input_dim, statistics, scale_by_Z, num_species)

Bases: Module

Mean and standard deviation scaler for GNN features.

Can be initialized from a franken.rf.scaler.Statistics instance, and when called will scale GNN features. Supports both global and per-species normalization.

Parameters:

• input_dim (int) – Dimension of the GNN features

• statistics (Statistics | None) – Instance of franken.rf.scaler.Statistics from which the feature mean and standard deviation can be fetched. If set to None, the class will be initialized to perform no normalization. You can modify the statistics by calling franken.rf.scaler.FeatureScaler.set_from_statistics().

• scale_by_Z (bool) – Whether to scale per-species or globally.

• num_species (int) – The number of distinct species in the data.

forward(descriptors, atomic_numbers=None)

Scale the given features to have zero-mean and unit standard deviation.

Parameters:

• descriptors (torch.Tensor) – GNN features

• atomic_numbers (Tensor | None) – Atomic numbers for each atom. This can be left to None unless the feature-scaler has been configured to perform per-species normalization. Defaults to None.

Returns:

Normalized GNN features

Return type:

torch.Tensor

set_from_statistics(statistics)

Set the mean and standard deviation statistics for scaling.

Parameters:

statistics (Statistics) – Instance of franken.rf.scaler.Statistics from which the feature mean and standard deviation can be fetched.