Long-range

class metatrain.utils.long_range.LongRangeFeaturizer(hypers, feature_dim, neighbor_list_options)[source]

Bases: Module

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

Parameters:

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

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

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

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

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.

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