Tensor basis

Provides a tensor basis in which to predict spherical tensor targets, following the approach described in this work [1].

This hook creates a basis for each target and each angular channel (o3_lambda block) of the outputs. Then it asks for invariant coefficients to apply to the basis to produce the target. By using this hook one can:

  • Use an architecture that produces only scalar outputs, and still be able to predict tensorial targets.

  • Reduce the cost of equivariant models by asking them to produce only scalar outputs, and then use this hook to access the angular momentum channels of the target.

Installation

To install this hook along with the metatrain package, run:

pip install metatrain[hook-tensor_basis]

where the square brackets indicate that you want to install the optional dependencies required for the tensor_basis hook.

Hook hyperparameters

The default hyperparameters for this hook are:

tensor_basis:
  soap:
    max_angular: 6
    max_radial: 7
    cutoff:
      radius: 5.0
      width: 0.5
  inputs: null
  outputs: null

and here is the documentation for each hyperparameter:

Hypers.soap: SOAPConfig = {'cutoff': {'radius': 5.0, 'width': 0.5}, 'max_angular': 6, 'max_radial': 7}

Hyperparameters used to compute the spherical expansions from which the vector basis will be built. Higher angular momentum channels are built by augmenting the order of the vector basis.

Hypers.inputs: str | list | None = None

Name or names of the targets to use as invariant coefficients to apply to the tensor basis.

If None, they will be set as mtt::aux::scalars::{output_name.replace('mtt::', '')} for each output name.

Hypers.outputs: str | list | None = None

Name or names of the targets to predict through a tensor basis.

A separate tensor basis will be built for each target.

References