Long-range

class metatrain.utils.long_range.LongRangeFeaturizer(hypers: Dict, feature_dim: int, neighbor_list_options: NeighborListOptions)[source]

Bases: Module

A class to compute long-range features starting from short-range features.

Parameters:

  • hypers (Dict) – Dictionary containing the hyperparameters for the long-range featurizer.

  • feature_dim (int) – The dimension of the short-range features (which also corresponds to the number of long-range features that will be returned).

  • neighbor_list_options (NeighborListOptions) – A NeighborListOptions object containing the neighbor list information for the short-range model.

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

ewald_calculator

Calculator to compute the long-range electrostatic potential using the Ewald summation method.

p3m_calculator

Calculator to compute the long-range electrostatic potential using the P3M method.

use_ewald

If True, use the Ewald summation method instead of the P3M method for periodic systems during training.

direct_calculator

Calculator for the electrostatic potential in non-periodic systems.

neighbor_list_options

Neighbor list information for the short-range model.

charges_map

Map the short-range features to atomic charges.

forward(systems: List[System], features: Tensor, neighbor_distances: Tensor) → Tensor[source]

Compute the long-range features for a list of systems.

Parameters:

  • systems (List[System]) – A list of System objects for which to compute the long-range features. Each system must contain a neighbor list consistent with the neighbor list options used to create the class.

  • features (Tensor) – A tensor of short-range features for the systems.

  • neighbor_distances (Tensor) – A tensor of neighbor distances for the systems, which must be consistent with the neighbor list options used to create the class.

Returns:

A tensor of long-range features for the systems.

Return type:

Tensor

class metatrain.utils.long_range.DummyLongRangeFeaturizer[source]

Bases: Module

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

forward(systems: List[System], features: Tensor, neighbor_distances: Tensor) → Tensor[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:
Return type:

Tensor