mattertune.normalization

Functions

compute_per_atom_references(dataset, ...[, ...])

compute_per_atom_references_cli_main(args, ...)

Classes

ComposeNormalizers(normalizers)

MeanStdNormalizerConfig(*, mean, std)

MeanStdNormalizerModule(config)

NormalizationContext(compositions)

The normalization context contains all the information required to normalize and denormalize the properties.

NormalizerConfigBase()

NormalizerModule(*args, **kwargs)

PerAtomReferencingNormalizerConfig(*, ...)

PerAtomReferencingNormalizerModule(config)

RMSNormalizerConfig(*, rms)

RMSNormalizerModule(config)

class mattertune.normalization.NormalizationContext(compositions)[source]

The normalization context contains all the information required to normalize and denormalize the properties. Currently, this only includes the compositions of the materials in the batch.

This flexibility allows for the “Normalizer” interface to be used for other types of normalization, beyond just simple mean and standard deviation normalization. For example, subtracting linear references from total energies can be implemented using this interface.

Parameters:

compositions (Tensor)

compositions: Tensor

The compositions should be provided as an integer tensor of shape (batch_size, num_elements), where each row (i.e., compositions[i]) corresponds to the composition vector of the i-th material in the batch.

The composition vector is a vector that maps each element to the number of atoms of that element in the material. For example, compositions[:, 1] corresponds to the number of Hydrogen atoms in each material in the batch, compositions[:, 2] corresponds to the number of Helium atoms, and so on.

__init__(compositions)
Parameters:

compositions (Tensor)

Return type:

None

class mattertune.normalization.NormalizerModule(*args, **kwargs)[source]
normalize(value, ctx)[source]

Normalizes the input tensor using the normalizer’s parameters and context.

Parameters:

  • value (torch.Tensor) – The input tensor to be normalized

  • ctx (NormalizationContext) – Context containing compositions information

Returns:

The normalized tensor

Return type:

torch.Tensor

denormalize(value, ctx)[source]

Denormalizes the input tensor using the normalizer’s parameters and context.

Parameters:

  • value (torch.Tensor) – The normalized tensor to be denormalized

  • ctx (NormalizationContext) – Context containing compositions information

Returns:

The denormalized tensor

Return type:

torch.Tensor

__init__(*args, **kwargs)
class mattertune.normalization.NormalizerConfigBase[source]
abstract create_normalizer_module()[source]
Return type:

NormalizerModule

class mattertune.normalization.MeanStdNormalizerConfig(*, mean, std)[source]
Parameters:

  • mean (float)

  • std (float)

mean: float

The mean of the property values.

std: float

The standard deviation of the property values.

class mattertune.normalization.MeanStdNormalizerModule(config)[source]
Parameters:

config (MeanStdNormalizerConfig)

mean: torch.Tensor
std: torch.Tensor
__init__(config)[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

config (MeanStdNormalizerConfig)

normalize(value, ctx)[source]

Normalizes the input tensor using the normalizer’s parameters and context.

Parameters:

  • value (torch.Tensor) – The input tensor to be normalized

  • ctx (NormalizationContext) – Context containing compositions information

Returns:

The normalized tensor

Return type:

torch.Tensor

denormalize(value, ctx)[source]

Denormalizes the input tensor using the normalizer’s parameters and context.

Parameters:

  • value (torch.Tensor) – The normalized tensor to be denormalized

  • ctx (NormalizationContext) – Context containing compositions information

Returns:

The denormalized tensor

Return type:

torch.Tensor

class mattertune.normalization.RMSNormalizerConfig(*, rms)[source]
Parameters:

rms (float)

rms: float

The root mean square of the property values.

class mattertune.normalization.RMSNormalizerModule(config)[source]
Parameters:

config (RMSNormalizerConfig)

rms: torch.Tensor
__init__(config)[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

config (RMSNormalizerConfig)

normalize(value, ctx)[source]

Normalizes the input tensor using the normalizer’s parameters and context.

Parameters:

  • value (torch.Tensor) – The input tensor to be normalized

  • ctx (NormalizationContext) – Context containing compositions information

Returns:

The normalized tensor

Return type:

torch.Tensor

denormalize(value, ctx)[source]

Denormalizes the input tensor using the normalizer’s parameters and context.

Parameters:

  • value (torch.Tensor) – The normalized tensor to be denormalized

  • ctx (NormalizationContext) – Context containing compositions information

Returns:

The denormalized tensor

Return type:

torch.Tensor

class mattertune.normalization.PerAtomReferencingNormalizerConfig(*, per_atom_references)[source]
Parameters:

per_atom_references (Mapping[int, float] | Sequence[float] | Path)

per_atom_references: Mapping[int, float] | Sequence[float] | Path

The reference values for each element.

  • If a dictionary is provided, it maps atomic numbers to reference values

  • If a list is provided, it’s a list of reference values indexed by atomic number

  • If a path is provided, it should point to a JSON file containing the references

create_normalizer_module()[source]
Return type:

NormalizerModule

class mattertune.normalization.PerAtomReferencingNormalizerModule(config)[source]
Parameters:

config (PerAtomReferencingNormalizerConfig)

references: torch.Tensor
__init__(config)[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

config (PerAtomReferencingNormalizerConfig)

normalize(value, ctx)[source]

Normalizes the input tensor using the normalizer’s parameters and context.

Parameters:

  • value (torch.Tensor) – The input tensor to be normalized

  • ctx (NormalizationContext) – Context containing compositions information

Returns:

The normalized tensor

Return type:

torch.Tensor

denormalize(value, ctx)[source]

Denormalizes the input tensor using the normalizer’s parameters and context.

Parameters:

  • value (torch.Tensor) – The normalized tensor to be denormalized

  • ctx (NormalizationContext) – Context containing compositions information

Returns:

The denormalized tensor

Return type:

torch.Tensor

class mattertune.normalization.ComposeNormalizers(normalizers)[source]
Parameters:

normalizers (Sequence[NormalizerModule])

__init__(normalizers)[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

normalizers (Sequence[NormalizerModule])

normalize(value, ctx)[source]

Normalizes the input tensor using the normalizer’s parameters and context.

Parameters:

  • value (torch.Tensor) – The input tensor to be normalized

  • ctx (NormalizationContext) – Context containing compositions information

Returns:

The normalized tensor

Return type:

torch.Tensor

denormalize(value, ctx)[source]

Denormalizes the input tensor using the normalizer’s parameters and context.

Parameters:

  • value (torch.Tensor) – The normalized tensor to be denormalized

  • ctx (NormalizationContext) – Context containing compositions information

Returns:

The denormalized tensor

Return type:

torch.Tensor

mattertune.normalization.compute_per_atom_references(dataset, property, reference_model, reference_model_kwargs={})[source]
Parameters:

  • dataset (Dataset[Atoms])

  • property (PropertyConfig)

  • reference_model (Literal['linear', 'ridge'])

  • reference_model_kwargs (dict[str, Any])

mattertune.normalization.compute_per_atom_references_cli_main(args, parser)[source]
Parameters:

  • args (Namespace)

  • parser (ArgumentParser)