SchNetPack - Deep Neural Networks for Atomistic Systems

The schnetpack.transform module provides a suite of transformation classes used to manipulate atomistic data (coordinates, species, etc.) during data loading or preprocessing. These transformations are typically applied to datasets to ensure physical consistency or to prepare data for specific model architectures.

Atomistic Transformations

These classes modify the physical properties of the system:

• AddOffsets: Adds offsets to coordinates.
• RemoveOffsets: Removes offsets from coordinates.
• SubtractCenterOfMass: Translates the system so the center of mass is at the origin.
• SubtractCenterOfGeometry: Translates the system so the geometric center is at the origin.

Casting Transformations

These classes manage data types (precision) for tensors:

• CastMap: Maps specific keys to specific types.
• CastTo32: Casts tensors to float32.
• CastTo64: Casts tensors to float64.

Neighbor List Transformations

These classes handle the construction and filtering of neighbor lists, which are critical for message-passing neural networks:

• MatScipyNeighborList: Uses SciPy for neighbor list construction.
• ASENeighborList: ASE-based neighbor list.
• VesinNeighborList: Vesin-based neighbor list.
• TorchNeighborList: PyTorch-native neighbor list.
• CachedNeighborList: A neighbor list that caches results.
• CountNeighbors: Counts the number of neighbors for each atom.
• FilterNeighbors: Filters neighbors based on specific criteria.
• WrapPositions: Wraps positions according to periodic boundary conditions.
• CollectAtomTriples: Collects atom triples for specific interaction types.

SchNetPack uses hierarchical Hydra configuration files in YAML format to define models and tasks. These configurations can be modified dynamically via command-line arguments.

Key features include:

• Hierarchical Structure: Configs are organized into groups (e.g., data, model, task) that can be nested.
• Interpolation: Use ${group.variable} to reference values from other parts of the configuration (e.g., ${globals.cutoff}).
• Resolvers: Use ${resolver:arguments} to evaluate functions during config construction. Built-in resolvers include ${hydra:runtime.cwd} for the current working directory and custom resolvers like ${uuid:1} for generating cached unique identifiers.
• Object Instantiation: The special key _target_ specifies the class to be instantiated, while other keys provide the arguments for that class's __init__ method.

# Example command to run a training experiment using a specific config group
$ spktrain experiment=qm9_atomwise

SchNetPack uses a hierarchical configuration system (via Hydra) to assemble experiments. You can define an experiment config that overrides defaults for models, datasets, and tasks.

Key concepts:

• # @package _global_: Placing this at the top of your config file ensures the configuration is placed at the base level of the hierarchy, allowing it to directly override train.yaml.
• defaults: Used to select and override predefined configuration groups (e.g., model, data, task).
• _target_: Specifies the Python class/module to instantiate for a specific component.
• Interpolation (${...}): Allows referencing variables defined elsewhere in the config (e.g., ${globals.property}).

An experiment config typically defines:

1. Defaults: Which model and data modules to use.
2. Globals: Shared variables like cutoff, lr, or the target property.
3. Data Transforms: A list of pre-processing steps (e.g., SubtractCenterOfMass, RemoveOffsets, CastTo32).
4. Model Architecture: Output modules (e.g., Atomwise) and post-processors (e.g., AddOffsets).
5. Task Definition: The loss function (e.g., torch.nn.MSELoss) and metrics (e.g., MeanAbsoluteError).

# @package _global_

defaults:
  - override /model: nnp
  - override /data: qm9

run.path: runs/qm9_${globals.property}

globals:
  cutoff: 5.
  lr: 5e-4
  property: energy_U0

data:
  transforms:
    - _target_: schnetpack.transform.SubtractCenterOfMass
    - _target_: schnetpack.transform.RemoveOffsets
      property: ${globals.property}
      remove_atomrefs: True
      remove_mean: True
    - _target_: schnetpack.transform.MatScipyNeighborList
      cutoff: ${globals.cutoff}
    - _target_: schnetpack.transform.CastTo32

model:
  output_modules:
    - _target_: schnetpack.atomistic.Atomwise
      output_key: ${globals.property}
      n_in: ${model.representation.n_atom_basis}
      aggregation_mode: sum
  postprocessors:
    - _target_: schnetpack.transform.CastTo64
    - _target_: schnetpack.transform.AddOffsets
      property: ${globals.property}
      add_mean: True
      add_atomrefs: True

task:
  outputs:
    - _target_: schnetpack.task.ModelOutput
      name: ${globals.property}
      loss_fn:
        _target_: torch.nn.MSELoss
      metrics:
        mae:
          _target_: torchmetrics.regression.MeanAbsoluteError
        mse:
          _target_: torchmetrics.regression.MeanSquaredError
      loss_weight: 1.

SchNetPack training runs are based on a train.yaml file which sets up default configuration groups. Instead of writing a full configuration from scratch, you can use these groups to predefine templates.

Default config groups in train.yaml include:

Default configurations are located in src/schnetpack/configs.