### Register NequIP Extension Entry Point

Source: https://nequip.readthedocs.io/en/latest/dev/extension_packages/getting_started

This code snippet demonstrates how to register an extension package with the NequIP framework by defining an entry point in the pyproject.toml file. This allows NequIP to automatically discover and load custom functionalities from your package upon import.

```toml
[project.entry-points."nequip.extension"]
init_always = "your_package_name"

```

--------------------------------

### Install NequIP from PyPI

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/getting-started/install

Installs the NequIP package using pip from the Python Package Index. This is the simplest method for users who want to quickly start using NequIP.

```bash
pip install nequip
```

--------------------------------

### Run NequIP Unit Tests

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/getting-started/install

Installs the 'pytest' framework and runs the unit tests for NequIP to verify its functionality. This helps in diagnosing potential issues with the installation.

```bash
pip install pytest
pytest tests/unit/
```

--------------------------------

### Verify NequIP Installation with Tutorial Model

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/getting-started/install

Checks if the NequIP installation is working correctly by downloading tutorial data and training a model using the 'nequip-train' command with a tutorial configuration file.

```bash
cd configs
python get_tutorial_data.py
nequip-train -cn tutorial.yaml
```

--------------------------------

### Install NequIP from Source (Develop Branch)

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/getting-started/install

Clones the NequIP repository and checks out the develop branch before installing from source. This is for users who want to test the latest development features.

```bash
git clone https://github.com/mir-group/nequip.git
cd nequip
git checkout develop
pip install .
```

--------------------------------

### Install NequIP using pip or source

Source: https://context7.com/context7/nequip_readthedocs_io_en/llms.txt

Commands to install NequIP from PyPI or build from source. Requires PyTorch to be installed first. Supports installation from the latest release or the develop branch.

```bash
# Install PyTorch first (see pytorch.org for GPU-specific instructions)
pip install torch

# Install NequIP from PyPI
pip install nequip

# Or install from source (latest release)
git clone --depth 1 https://github.com/mir-group/nequip.git
cd nequip
pip install .

# Or install from develop branch
git clone https://github.com/mir-group/nequip.git
cd nequip
git checkout develop
pip install .

```

--------------------------------

### Install Weights and Biases Logger

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/getting-started/install

Installs the 'wandb' package, which is required to use the Weights and Biases logger with PyTorch Lightning for tracking training runs.

```bash
pip install wandb
```

--------------------------------

### Install NequIP from Source (Latest Release)

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/getting-started/install

Clones the NequIP repository for the latest release and installs it from source. This method is useful for developers or users who need the most recent stable version.

```bash
git clone --depth 1 https://github.com/mir-group/nequip.git
cd nequip
pip install .
```

--------------------------------

### Run NequIP Full Tests

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/getting-started/install

Executes all tests for NequIP, including unit and integration tests. This provides a comprehensive check of the installation and package stability.

```bash
pytest tests/
```

--------------------------------

### Verify NequIP Installation and Run Tests

Source: https://context7.com/context7/nequip_readthedocs_io_en/llms.txt

Steps to verify NequIP installation by running tutorial data and training, as well as executing unit and full test suites using pytest.

```bash
cd configs
python get_tutorial_data.py
nequip-train -cn tutorial.yaml

# Run unit tests
pip install pytest
pytest tests/unit/

# Run full test suite including integration tests
pytest tests/

```

--------------------------------

### Install cuEquivariance Libraries

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/accelerations/cuequivariance

Installs the necessary cuEquivariance libraries for PyTorch integration. These libraries provide GPU-accelerated tensor product operations.

```bash
pip install cuequivariance-torch cuequivariance-ops-torch-cu12
```

--------------------------------

### Train NequIP Model using nequip-train command

Source: https://context7.com/context7/nequip_readthedocs_io_en/llms.txt

Command-line examples for training NequIP models using `nequip-train`. Shows how to specify configuration files, paths, and continue training from checkpoints.

```bash
# Train with config file in current directory
nequip-train -cn config_name.yaml

# Train with config file in specific directory
nequip-train -cp /full/path/to/config/directory -cn config_name.yaml

# Continue training from checkpoint
nequip-train -cn config_name.yaml ++ckpt_path='path/to/checkpoint.ckpt'

```

--------------------------------

### SoftAdapt Callback Configuration Example (YAML)

Source: https://nequip.readthedocs.io/en/latest/_modules/nequip/train/callbacks/softadapt

Example configuration for the SoftAdapt callback in YAML format. This demonstrates how to set the `beta` hyperparameter, the `interval` for updates (epoch in this case), and the `frequency` of updates.

```yaml
callbacks:
  - _target_: nequip.train.callbacks.SoftAdapt
    beta: 1.1
    interval: epoch
    frequency: 5
```

--------------------------------

### Load Example Data from Package (Python)

Source: https://nequip.readthedocs.io/en/latest/_modules/nequip/model/saved_models/package

Loads example data from a `.nequip.zip` package file. This function uses `torch.package.PackageImporter` to access the `example_data.pkl` resource within the specified package.

```python
from nequip.data import AtomicDataDict
from nequip.model import data_dict_from_package

# Example usage:
package_path = "path/to/your/model.nequip.zip"
data_dict = data_dict_from_package(package_path=package_path)

# data_dict is of type AtomicDataDict.Type
```

--------------------------------

### Example: Combined Data Transforms Configuration

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/configuration/data

An example demonstrating the sequential application of both `ChemicalSpeciesToAtomTypeMapper` and `NeighborListTransform` in a NequIP configuration.

```yaml
transforms:
  - _target_: nequip.data.transforms.ChemicalSpeciesToAtomTypeMapper
    chemical_symbols: [C, H, O, Cu]
  - _target_: nequip.data.transforms.NeighborListTransform
    r_max: 5.0
```

--------------------------------

### Install Documentation Build Requirements with Pip

Source: https://nequip.readthedocs.io/en/latest/README

Installs the necessary Python packages required to build the NequIP documentation locally. These include Sphinx for documentation generation, and extensions like furo, myst_parser, and pygments-lammps for themes, markdown parsing, and syntax highlighting.

```bash
pip install sphinx sphinx_copybutton furo myst_parser pygments-lammps

```

--------------------------------

### NequIP Developer Guide - Framework Understanding

Source: https://nequip.readthedocs.io/en/latest/genindex

Information to help understand the internal workings of the NequIP framework.

```APIDOC
## Developer Guide: Understanding the NequIP Framework

### Description
This section provides insights into the architecture and design principles of the NequIP framework, intended for developers who want a deeper understanding of how NequIP works internally.

### Key Components
- **NequIP Workflows**: Describes the typical workflows involved in using NequIP, from data preparation to model deployment.
  - **Model Packaging**: Information on how NequIP models are packaged for distribution and inference.
- **Extension Packages**: Details on how to extend NequIP's functionality.
  - **Getting Started**: Initial steps for creating extension packages.
  - **Data Handling**: How to integrate custom data handling in extensions.
  - **Training Techniques**: Using custom training techniques within extensions.
  - **Custom Models**: Developing and integrating custom model architectures.

### Response
#### Success Response (200)
This section provides conceptual information for developers and does not represent an API endpoint.
```

--------------------------------

### NequIP Config: Simple Run Sequence

Source: https://nequip.readthedocs.io/en/latest/guide/configuration/config

Illustrates a basic configuration for the 'run' directive, specifying a sequence of tasks to be executed by 'nequip-train'. This example runs training followed by testing.

```yaml
run: [train, test]
```

--------------------------------

### NequIP Developer Guide - Contributing

Source: https://nequip.readthedocs.io/en/latest/genindex

Guidelines for contributing to the NequIP project.

```APIDOC
## Developer Guide: Contributing to NequIP

### Description
This section outlines the process and guidelines for developers who wish to contribute to the NequIP project.

### How to Contribute
Contributions are welcome and can include:
- Reporting bugs
- Suggesting new features
- Improving documentation
- Submitting code changes (bug fixes, new features)

### Contribution Workflow
1. **Fork the repository**
2. **Create a new branch** for your feature or bug fix.
3. **Make your changes** and ensure they follow the project's coding standards.
4. **Add tests** for your changes.
5. **Run tests** to ensure everything is working correctly.
6. **Submit a Pull Request** with a clear description of your changes.

### Coding Standards
- Adhere to Python style guides (e.g., PEP 8).
- Write clear and concise code with appropriate comments.
- Follow the existing code structure and patterns.

### Response
#### Success Response (200)
This section provides guidance for developers and does not represent an API endpoint.
```

--------------------------------

### NequIP Data Module Configuration Example

Source: https://nequip.readthedocs.io/en/latest/api/datamodule

Example configuration for NequIP model parameters, specifically for setting isolated atom energies per type. This is often recommended when using datasets like 3BPA to ensure accurate energy calculations.

```yaml
model:
  type_names: [C, H, N, O]
  per_type_energy_shifts:
    C: -1029.4889999855063
    H: -13.587222780835477
    N: -1484.9814568572233
    O: -2041.9816003861047
```

--------------------------------

### SLURM DDP Training Example with NequIP

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/accelerations/ddp_training

A minimal SLURM job submission script for initiating distributed data parallel (DDP) training with NequIP. This example configures a run across 2 nodes and 4 GPUs per node, demonstrating essential SLURM directives and the use of 'srun' to launch the training command with relevant trainer arguments passed from SLURM environment variables.

```bash
#!/bin/bash
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --gres=gpu:4
#SBATCH --cpus-per-task=8
... other slurm variables

# ... set up (e.g. module load, activate Python env, etc)

# ... cluster specific set up such as network interface 
# (e.g. MASTER_PORT, MASTER_ADDR, NCCL_SOCKET_IFNAME)

srun nequip-train -cn config.yaml ++trainer.num_nodes=${SLURM_NNODES}

```

--------------------------------

### Start WandB Sweep Agent with Sweep ID

Source: https://nequip.readthedocs.io/en/latest/guide/training-techniques/wandb_example

This command starts a Weights and Biases sweep agent using a specific sweep ID. The agent will connect to W&B to download the job defined by the sweep ID and begin execution. Replace `<sweep-ID>` with the actual ID obtained from a previous `wandb sweep` command or the W&B UI.

```bash
wandb agent <sweep-ID>

```

--------------------------------

### Initialize and Train with NequIPLightningModule

Source: https://nequip.readthedocs.io/en/latest/_modules/nequip/train/callbacks/loss_coeff_scheduler

This Python code initializes the learning rate scheduler with transition epochs and start epoch. It also captures the initial coefficients and asserts the start epoch and transition epoch values are positive. The on_train_epoch_start function handles coefficient interpolation during the transition period.

```Python
        self.transition_epochs = transition_epochs
        self.captured_initial_coeffs = None

        assert start_epoch >= 0, "Start epoch must be non-negative"
        assert transition_epochs > 0, "Transition epochs must be positive"

    def on_train_epoch_start(
        self,
        trainer: lightning.Trainer,
        pl_module: NequIPLightningModule,
    ) -> None:
        """""
        current_epoch = trainer.current_epoch

        # NOTE: initial coeffs captured should already be normalized
        # final coeffs were normalized at __init__

        if current_epoch == self.start_epoch and self.captured_initial_coeffs is None:
            # lazily capture the current coefficients when we start
            self.captured_initial_coeffs = {
                metric_name: metric_dict["coeff"]
                for metric_name, metric_dict in pl_module.loss.metrics.items()
                if metric_name in self.final_coeffs
            }
            # sanity check that all `final_coeffs` keys are present in the metrics
            assert set(self.final_coeffs.keys()) == set(
                self.captured_initial_coeffs.keys()
            ), (
                f"Mismatch between `final_coeffs` keys {set(self.final_coeffs.keys())} and available metrics {set(self.captured_initial_coeffs.keys())}"
            )

        if (
            self.start_epoch
            < current_epoch
            <= self.start_epoch + self.transition_epochs
        ):
            # linear interpolation during transition period
            assert self.captured_initial_coeffs is not None, (
                "Initial coefficients should have been captured"
            )
            epochs_into_transition = current_epoch - self.start_epoch
            alpha = epochs_into_transition / self.transition_epochs
            interpolated_coeffs = {}

            for key in self.final_coeffs.keys():
                initial_val = self.captured_initial_coeffs[key]
                final_val = self.final_coeffs[key]
                interpolated_coeffs[key] = initial_val + alpha * (
                    final_val - initial_val
                )

            pl_module.loss.set_coeffs(interpolated_coeffs)
```

--------------------------------

### NequIP Config: Run Sequence for Initial Evaluation

Source: https://nequip.readthedocs.io/en/latest/guide/configuration/config

Provides an example of a 'run' sequence that includes initial validation and testing before training, followed by post-training validation and testing. This is useful for assessing model performance at initialization.

```yaml
run: [val, test, train, val, test]
```

--------------------------------

### Get nequip-package command line options

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/getting-started/workflow

This command displays all available command-line options for the `nequip-package` tool. Further detailed help for specific subcommands like `build` and `info` can be accessed with their respective `-h` flags.

```bash
nequip-package -h
```

--------------------------------

### Launch WandB Sweep Agent

Source: https://nequip.readthedocs.io/en/latest/guide/training-techniques/wandb_example

This command initiates the Weights and Biases sweep agent, which will then fetch sweep configurations and start training runs based on the defined `sweep.yaml` file. Ensure that `sweep.yaml` is correctly configured and accessible.

```bash
wandb sweep sweep.yaml

```

--------------------------------

### YAML Configuration for Weighted Sum Coefficients

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/configuration/metrics

Example demonstrating how to set coefficients for different loss or metric terms, illustrating their normalization and impact on the 'weighted_sum'.

```yaml
coeffs:
  total_energy: 3.0
  forces: 1.0
```

```yaml
coeffs:
  total_energy_rmse: 1.0    # included in weighted_sum
  forces_rmse: 1.0          # included in weighted_sum
  total_energy_mae: null    # computed but not in weighted_sum
  forces_mae: null          # computed but not in weighted_sum
```

--------------------------------

### YAML Configuration for Loss and Metrics

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/configuration/metrics

Example of configuring NequIP's training module with energy and force loss, validation metrics, and using a weighted sum for early stopping and model checkpointing.

```yaml
# Define the monitored metric once for consistency
monitored_metric: val0_epoch/weighted_sum

training_module:
  _target_: nequip.train.EMALightningModule
  
  # Loss function
  loss:
    _target_: nequip.train.EnergyForceLoss
    coeffs:
      total_energy: 1.0
      forces: 1.0
  
  # Validation metrics - weighted_sum will be used for monitoring
  val_metrics:
    _target_: nequip.train.EnergyForceMetrics
    coeffs:
      total_energy_rmse: 1.0
      forces_rmse: 1.0
      total_energy_mae: null  # logged but not in weighted_sum
      forces_mae: null        # logged but not in weighted_sum

trainer:
  _target_: lightning.Trainer
  
  callbacks:
    # Early stopping using the monitored metric
    - _target_: lightning.pytorch.callbacks.EarlyStopping
      monitor: ${monitored_metric}
      patience: 20
      min_delta: 1e-3
    
    # Model checkpointing using the monitored metric  
    - _target_: lightning.pytorch.callbacks.ModelCheckpoint
      monitor: ${monitored_metric}
      filename: best

  # Learning rate scheduler using the monitored metric
  lr_scheduler:
    scheduler:
      _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
      factor: 0.6
      patience: 5
    monitor: ${monitored_metric}
```

--------------------------------

### NequIP Model Configuration

Source: https://nequip.readthedocs.io/en/latest/api/data_stats

Example of a NequIP model configuration in YAML, showing how to integrate dataset statistics like average number of neighbors and per-type energy shifts/scales.

```yaml
model:
  _target_: nequip.model.NequIPGNNEnergyModel

  # other model hyperparameters
  avg_num_neighbors: ${training_data_stats:num_neighbors_mean}
  per_type_energy_shifts: ${training_data_stats:per_atom_energy_mean}
  per_type_energy_scales: ${training_data_stats:total_energy_std}

```

--------------------------------

### NequIP Model Compilation with TorchScript (Bash)

Source: https://context7.com/context7/nequip_readthedocs_io_en/llms.txt

Commands for compiling a NequIP model using TorchScript for efficient inference. Examples are provided for both CPU and GPU compilation from different input formats (checkpoint or package).

```bash
# Compile for CPU inference
nequip-compile \
  path/to/best_model.ckpt \
  path/to/compiled_model.nequip.pth \
  --device cpu \
  --mode torchscript

# Compile for GPU inference
nequip-compile \
  path/to/model.nequip.zip \
  path/to/compiled_model.nequip.pth \
  --device cuda \
  --mode torchscript
```

--------------------------------

### Configure LinearLossCoefficientScheduler for Transitions

Source: https://nequip.readthedocs.io/en/latest/api/callbacks

The LinearLossCoefficientScheduler linearly interpolates loss coefficients from current values to specified final coefficients over a given number of epochs. It's stateful and captures coefficients at the start epoch. The example demonstrates transitioning to a 1:1:1 ratio for energy, forces, and stress over 200 epochs starting at epoch 100.

```yaml
callbacks:
  - _target_: nequip.train.callbacks.LinearLossCoefficientScheduler
    final_coeffs:
      per_atom_energy_mse: 1.0
      forces_mse: 1.0
      stress_mse: 1.0
    start_epoch: 100
    transition_epochs: 200
```

--------------------------------

### Install LAMMPS Python Interface

Source: https://nequip.readthedocs.io/en/latest/_sources/integrations/lammps/mliap

Installs the Python interface for the compiled LAMMPS build, allowing Python scripts to interact with LAMMPS, particularly for ML-IAP.

```bash
cd build-mliap
make install-python
```

--------------------------------

### Running nequip-train Command

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/reference/troubleshoot

Demonstrates the correct way to run the `nequip-train` command, addressing a common parsing error related to configuration files.

```bash
nequip-train -cn config.yaml
```

--------------------------------

### Install Dependencies for ML-IAP

Source: https://nequip.readthedocs.io/en/latest/_sources/integrations/lammps/mliap

Installs necessary Python dependencies including Cython and Cupy for GPU acceleration, required for building LAMMPS with ML-IAP support.

```bash
pip install cython==3.0.11 cupy-cuda12x
```

--------------------------------

### Install Dependencies for ML-IAP

Source: https://nequip.readthedocs.io/en/latest/integrations/lammps/mliap

Installs necessary Python packages for ML-IAP integration, including specific versions of cython and cupy. Ensure your CUDA version matches cupy-cuda12x.

```bash
pip install cython==3.0.11 cupy-cuda12x

```

--------------------------------

### Build HTML Documentation with Make

Source: https://nequip.readthedocs.io/en/latest/README

Executes the make command to build the HTML version of the NequIP documentation. The generated files will be placed in the \'_build\' directory, which can then be opened in a web browser.

```bash
make html

```

--------------------------------

### Build and Install LAMMPS

Source: https://nequip.readthedocs.io/en/latest/integrations/lammps/mliap

Builds the LAMMPS executable using the configured CMake options and then installs the Python interface for ML-IAP. This makes LAMMPS and its ML-IAP capabilities available for use.

```bash
cmake --build build-mliap -j 8
cd build-mliap
make install-python

```

--------------------------------

### NequIP Model Packaging (Bash)

Source: https://context7.com/context7/nequip_readthedocs_io_en/llms.txt

Commands for packaging a trained NequIP model checkpoint into a portable zip file. Includes options to build the package, inspect its metadata, and view help information.

```bash
# Build package from checkpoint
nequip-package build path/to/best_model.ckpt path/to/model.nequip.zip

# Inspect package metadata
nequip-package info path/to/model.nequip.zip

# Get help on packaging options
nequip-package -h
nequip-package build -h
```

--------------------------------

### Initialize NequIPCalculator

Source: https://nequip.readthedocs.io/en/latest/_modules/nequip/ase/nequip_calculator

Initializes the NequIPCalculator with a NequIP model, device, and unit conversion factors. The model must be in evaluation mode.

```python
calculator = NequIPCalculator(
    model=model,
    device=device,
    energy_units_to_eV=1.0,
    length_units_to_A=1.0,
    transforms=[],
    **kwargs
)
```

--------------------------------

### MetricsManager Configuration Example

Source: https://nequip.readthedocs.io/en/latest/api/metrics

This example demonstrates how to configure a custom MetricsManager in nequip to calculate weighted losses for energy and forces using MeanSquaredError. It shows the use of fields, coefficients, and nested metric definitions.

```yaml
_target_: nequip.train.MetricsManager
metrics:
    - name: per_atom_energy_mse
    field:
        _target_: nequip.data.PerAtomModifier
        field: total_energy
    coeff: 1
    metric:
        _target_: nequip.train.MeanSquaredError
    - name: forces_mse
    field: forces
    coeff: 1
    metric:
        _target_: nequip.train.MeanSquaredError

```

--------------------------------

### Build a packaged NequIP model using nequip-package

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/getting-started/workflow

This command builds a packaged NequIP model from a checkpoint file. The output must have the `.nequip.zip` extension. This process archives the model and its associated code, ensuring compatibility across NequIP updates.

```bash
nequip-package build path/to/ckpt_file path/to/packaged_model.nequip.zip
```

--------------------------------

### Initialize COLLDataModule and Prepare Data

Source: https://nequip.readthedocs.io/en/latest/_modules/nequip/data/datamodule/coll_datamodule

Initializes the COLLDataModule with dataset paths and prepares the data by downloading necessary files if they do not already exist. This function ensures that the training, validation, and test datasets are available in the specified directory.

```python
class COLLDataModule(ASEDataModule):
    """LightningDataModule for the COLL dataset from `<https://arxiv.org/abs/2011.14115>`_.

    Args:
        seed (int): data seed for reproducibility
        transforms (List[Callable]): list of data transforms
        data_source_dir (str): directory where dataset files will be downloaded to if not already present
    """

    def __init__(
        self,
        seed: int,
        transforms: List[Callable],
        data_source_dir: str,
        **kwargs,
    ):
        self.data_source_dir = data_source_dir

        train_file_path = os.path.join(data_source_dir, "coll_v1.2_AE_train.xyz")
        val_file_path = os.path.join(data_source_dir, "coll_v1.2_AE_val.xyz")
        test_file_path = os.path.join(data_source_dir, "coll_v1.2_AE_test.xyz")

        super().__init__(
            seed=seed,
            train_file_path=train_file_path,
            val_file_path=val_file_path,
            test_file_path=test_file_path,
            transforms=transforms,
            **kwargs,
        )

        self.train_file_path = train_file_path
        self.val_file_path = val_file_path
        self.test_file_path = test_file_path

    def prepare_data(self):
        """"""
        os.makedirs(self.data_source_dir, exist_ok=True)

        files_to_download = [
            (self.train_file_path, _URL_TRAIN, "training"),
            (self.val_file_path, _URL_VAL, "validation"),
            (self.test_file_path, _URL_TEST, "test"),
        ]

        for file_path, url, dataset_type in files_to_download:
            if not os.path.isfile(file_path):
                logger.info(f"Downloading {dataset_type} dataset to `{file_path}`")
                download_url(
                    url, self.data_source_dir, filename=os.path.basename(file_path)
                )
            else:
                logger.info(f"Using existing {dataset_type} data file `{file_path}`")

```

--------------------------------

### NequIPDataModule Initialization (Python)

Source: https://nequip.readthedocs.io/en/latest/_modules/nequip/data/datamodule/_base_datamodule

Initializes the NequIPDataModule, setting up datasets, dataloaders, and other configurations. It handles dataset conversion to lists, deep copying, and accounts for split datasets.

```python
import torch
from .. import AtomicDataDict
from nequip.utils.logger import RankedLogger

import lightning
from omegaconf import OmegaConf, DictConfig, ListConfig
from hydra.utils import instantiate
import copy
from typing import List, Dict, Any, Union, Optional

logger = RankedLogger(__name__, rank_zero_only=True)




[docs]
class NequIPDataModule(lightning.LightningDataModule):
    """
    Sanity checking is only performed at runtime -- ensure that the correct datasets are provided for the intended runs,
    which can be ``train``, ``val``, ``test``, and/or ``predict``.

    - ``train`` runs require ``train_dataset`` and ``val_dataset``
    - ``val`` runs require ``val_dataset``
    - ``test`` runs require ``test_dataset``
    - ``predict`` runs require ``predict_dataset``

    One can explicitly specify which ``train``, ``val``, ``test``, ``predict`` datasets to use, or randomly split a dataset to be used for any of those tasks with the ``split_dataset`` argument. These options are not mutually exclusive, e.g. if a single ``test_dataset`` is provided, and ``split_dataset`` is used to get another test set, there will now be two test sets (indexed by ``0`` and ``1``) used for testing. If ``test_dataset`` is a list, i.e. multiple test datasets are provided (e.g. if there are ``n`` test sets with indices ``0``, ``1``, ..., ``n - 1``) and multiple ``split_ataset`` is a list that contributes multiple test sets (say ``m`` such test sets are provided). There will be a total of ``m+n`` test sets, with the ones from ``test_dataset`` taking indices ``0``, ``1``, ..., ``n - 1`` and the ones from the ``split_dataset`` taking indices ``n``, ``n+1``, ..., ``n+m-1``.

    Args:
        seed (int): data seed for reproducibility
        train_dataset (Dict/List[Dict]): training dataset
        val_dataset (Dict/List[Dict]): validation dataset(s) (can provide multiple datasets in a list)
        test_dataset (Dict/List[Dict]): test dataset(s) (can provide multiple datasets in a list)
        predict_dataset (Dict/List[Dict]): prediction dataset(s) (can provide multiple datasets in a list)
        split_dataset (Dict/List[Dict]): dictionary with a ``dataset`` key, which defines the dataset and the keys ``train``, ``val``, ``test``, ``predict`` which represent the subsets to split ``dataset`` into and are either ``int`` s that sum up to the size of ``dataset`` or ``float`` s that sum up to 1 (at least 2, but not necessarily all of ``train``, ``val``, ``test``, ``predict`` must be provided if this option is used)
        train_dataloader (Dict): training ``DataLoader`` configuration dictionary
        val_dataloader (Dict): validation ``DataLoader`` configuration dictionary
        test_dataloader (Dict): testing ``DataLoader`` configuration dictionary
        predict_dataloader (Dict): prediction ``DataLoader`` configuration dictionary
        stats_manager (Dict): dictionary that can be instantiated into a :class:`~nequip.data.DataStatisticsManager` object
    """

    def __init__(
        self,
        seed: int,
        train_dataset: Optional[Union[Dict, List]] = [],
        val_dataset: Optional[Union[Dict, List]] = [],
        test_dataset: Optional[Union[Dict, List]] = [],
        predict_dataset: Optional[Union[Dict, List]] = [],
        split_dataset: Optional[Union[Dict, List]] = [],
        train_dataloader: Dict = {},
        val_dataloader: Dict = {},
        test_dataloader: Dict = {},
        predict_dataloader: Dict = {},
        stats_manager: Optional[Dict] = None,
    ):
        super().__init__()
        # internal logic follows lists in order of train, val, test, predict, split

        # == first convert all dataset configs to lists if not already lists ==
        dconfigs = []
        for dconfig in [
            train_dataset,
            val_dataset,
            test_dataset,
            predict_dataset,
            split_dataset,
        ]:
            # convert to primitives as later logic is based on types
            if isinstance(dconfig, DictConfig) or isinstance(dconfig, ListConfig):
                dconfig = OmegaConf.to_container(dconfig, resolve=True)
            assert isinstance(dconfig, dict) or isinstance(dconfig, list)
            # make deep copies of the dicts to avoid mutating them in case they are used outside (should be relatively cheap)
            if not isinstance(dconfig, list):
                dconfigs.append([copy.deepcopy(dconfig)])
            else:
                dconfigs.append(copy.deepcopy(dconfig))

        # == account for split datasets ==
        dataset_type_map = ["train", "val", "test", "predict"]  # index matches dconfig
        # loop over datasets to split
        for split_config in dconfigs[4]:
            split_dict = split_config.copy()
            dataset_to_split = split_dict.pop("dataset")
            assert all(

```

--------------------------------

### Initialize ConFIGLightningModule with Configuration

Source: https://nequip.readthedocs.io/en/latest/_modules/nequip/train/config

Initializes the ConFIGLightningModule, setting up parameters for gradient clipping, optimization strategy (least squares or pseudo-inverse), and normalization epsilon. It also handles PyTorch version-specific configurations for backward compatibility.

```python
def __init__(
    self,
    gradient_clip_val: Optional[float] = None,
    gradient_clip_algorithm: Optional[str] = None,
    lsqr: bool = True,
    norm_eps: float = 1e-8,
    **kwargs,
):
    # === torch>=2.6 requires the flag for compile with multiple backwards ===
    if _TORCH_GE_2_6:
        # relevant PyTorch commit: https://github.com/pytorch/pytorch/commit/87059d4547551f197731f5c084e3be6054797578
        # comments from PyTorch code:
        # This controls whether we collect donated buffer. This flag must be set
        # False if a user wants to retain_graph=True for backward.
        torch._functorch.config.donated_buffer = False

    super().__init__(**kwargs)
    # see https://lightning.ai/docs/pytorch/stable/common/optimization.html#id2
    self.automatic_optimization = False

    # === process model params ===
    # ingredients to do the reverse of torch.cat([torch.flatten(),...])
    param_dict = dict(self.model.named_parameters())
    self.ConFIG_model_param_names = list(param_dict.keys())
    self.ConFIG_param_numel_list = [
        param_dict[k].numel() for k in self.ConFIG_model_param_names
    ]
    self.ConFIG_param_batch_list = [
        0,
    ] + list(accumulate(self.ConFIG_param_numel_list))
    self.ConFIG_param_shape_list = [
        param_dict[k].shape for k in self.ConFIG_model_param_names
    ]
    del param_dict

    # === process loss functions ===
    assert len(self.loss) >= 1, (
        f"`ConFIGLightningModule` is used for cases with multiple loss components, but {len(self.loss)} found"
    )
    self.ConFIG_loss_component_keys: Dict[str, str] = {
        metric_name: f"train_loss_step{self.logging_delimiter}" + metric_name
        for metric_name in self.loss.keys()
    }

    # === method specific hyperparameters ===
    self.ConFIG_eps = norm_eps
    self.ConFIG_lsqr = lsqr

    # temporary narrow solution to accommodate only ReduceLROnPlateau LR scheduling
    if self.lr_scheduler_config is not None:
        scheduler = self.lr_scheduler_config["scheduler"]["_target_"]
        assert "ReduceLROnPlateau" in scheduler, (
            f"only `ReduceLROnPlateau` LR scheduler is usable with `ConFIGLightningModule` but found `{scheduler}`"
        )
        monitor = self.lr_scheduler_config["monitor"]
        assert "val" in monitor, (
            f"Only validation metrics can be monitored for LR scheduling with `ConFIGLightningModule`, but found {monitor}"
        )

```

--------------------------------

### Get Length of ASEDataset

Source: https://nequip.readthedocs.io/en/latest/_modules/nequip/data/dataset/ase_dataset

Returns the number of atomic structures loaded in the ASEDataset. This is determined by the number of entries in the internal `data_list`.

```python
def __len__(self) -> int:
        return len(self.data_list)

```

--------------------------------

### rMD17 DataModule Configuration (YAML)

Source: https://context7.com/context7/nequip_readthedocs_io_en/llms.txt

Configures the rMD17DataModule for loading benchmark datasets. This example shows configuration for the 'aspirin' molecule, similar to other datamodules, with specified root directory and transforms.

```yaml
# Revised MD17 dataset
data:
  _target_: nequip.data.datamodule.rMD17DataModule
  molecule_name: aspirin
  root: ./datasets/rmd17
  # ... rest of configuration similar to above
```

--------------------------------

### Fine-Tuning NequIP Models from Checkpoint (YAML)

Source: https://context7.com/context7/nequip_readthedocs_io_en/llms.txt

This YAML configuration outlines the process of fine-tuning a NequIP model by loading it from a checkpoint file. It specifies the training module, the checkpoint path, and allows for the definition of new loss, metrics, and optimizer settings for the continued training.

```yaml
training_module:
  _target_: nequip.train.EMALightningModule
  model:
    _target_: nequip.model.ModelFromCheckpoint
    checkpoint_path: /path/to/pretrained_checkpoint.ckpt
  # Define new loss, metrics, optimizer for fine-tuning
  loss:
    - _target_: nequip.train.SimplifiedPropertyWeightedLoss
      loss_coeffs:
        total_energy: 1.0
        forces: 50.0
  optimizer:
    _target_: torch.optim.Adam
    lr: 0.005

```

--------------------------------

### YAML: Coefficient Normalization Example

Source: https://nequip.readthedocs.io/en/latest/guide/configuration/metrics

Demonstrates how coefficients for loss or metric terms are automatically normalized to sum to 1, influencing the calculation of the 'weighted_sum' metric.

```yaml
coeffs:
  total_energy: 3.0
  forces: 1.0

# internally becomes: total_energy: 0.75, forces: 0.25
```

--------------------------------

### Run LAMMPS Simulation with Kokkos GPU

Source: https://nequip.readthedocs.io/en/latest/_sources/integrations/lammps/mliap

Example command to execute a LAMMPS simulation using the Kokkos backend for GPU acceleration, often used in conjunction with ML-IAP.

```bash
srun -n 1 /path/to/lammps/build/lmp -in in.lammps -k on g 1 -sf kk -pk kokkos newton on neigh half
```

--------------------------------

### NequIPCalculator Initialization - NequIP

Source: https://nequip.readthedocs.io/en/latest/api/ase

Initializes the NequIPCalculator for ASE, compatible with NequIP and Allegro models. It allows specifying the model, device, unit conversion factors, and data transforms. For optimal performance, it's recommended to use a compiled model.

```python
from nequip.ase import NequIPCalculator
from torch.nn import Module
from torch import device
from typing import List, Callable

# Example parameters (replace with your actual values)
model: Module
device: str | device = 'cpu'
energy_units_to_eV: float = 1.0
length_units_to_A: float = 1.0
transforms: List[Callable] = []

calculator = NequIPCalculator(
    model=model,
    device=device,
    energy_units_to_eV=energy_units_to_eV,
    length_units_to_A=length_units_to_A,
    transforms=transforms
)
```

--------------------------------

### Compile-time Model Modification - Bash

Source: https://nequip.readthedocs.io/en/latest/_sources/dev/extension_packages/models

Demonstrates how to apply model modifiers during the compilation phase using the `nequip-compile` command. This example enables the `enable_OpenEquivariance` modifier for accelerations.

```bash
nequip-compile model.ckpt compiled.pth --modifiers enable_OpenEquivariance
```

--------------------------------

### Get Field Type in Nequip

Source: https://nequip.readthedocs.io/en/latest/_modules/nequip/data/_key_registry

The `get_field_type` function determines the type ('graph', 'node', 'edge') of a given field name. It can optionally raise a `KeyError` if the field is not registered.

```python
def get_field_type(field: str, error_on_unregistered: bool = True) -> str:
    if field in _GRAPH_FIELDS:
        return "graph"
    elif field in _NODE_FIELDS:
        return "node"
    elif field in _EDGE_FIELDS:
        return "edge"
    else:
        if error_on_unregistered:
            raise KeyError(f"Unregistered field {field} found")
        else:
            return None
```

--------------------------------

### Initialize Nequip Project with Datasets

Source: https://nequip.readthedocs.io/en/latest/_modules/nequip/data/datamodule/_ase_datamodule

Initializes the Nequip project by passing dataset configurations and a split configuration to the parent class constructor. It assumes that dataset_configs is a list or tuple containing configurations for training, validation, testing, and prediction datasets, and split_config holds the split configuration.

```python
super().__init__(
            seed=seed,
            train_dataset=dataset_configs[0],
            val_dataset=dataset_configs[1],
            test_dataset=dataset_configs[2],
            predict_dataset=dataset_configs[3],
            split_dataset=split_config,
            **kwargs,
        )
```

--------------------------------

### Restarting NequIP Training from Checkpoint (Bash)

Source: https://nequip.readthedocs.io/en/latest/_sources/guide/getting-started/workflow

Shows the command-line interface for restarting a NequIP training session from a previously saved checkpoint. It uses Hydra's override syntax to specify the checkpoint path and configuration directory.

```bash
nequip-train -cp full/path/to/config/directory -cn config_name.yaml ++ckpt_path='path/to/ckpt_file'
```

--------------------------------

### Configure LossCoefficientMonitor for Logging Coefficients

Source: https://nequip.readthedocs.io/en/latest/api/callbacks

The LossCoefficientMonitor callback is used to monitor and log loss coefficients during training. The example shows how to configure it to log coefficients every 5 epochs.

```yaml
callbacks:
  - _target_: nequip.train.callbacks.LossCoefficientMonitor
    interval: epoch
    frequency: 5
```

--------------------------------

### Get Statistics from Data Source - Python

Source: https://nequip.readthedocs.io/en/latest/_modules/nequip/data/stats_manager

This method iterates through a data source, applies transformations, and computes statistics. It requires the `reset()` method to be called beforehand to clear any previous state.

```python
def get_statistics(self, data_source: Iterable[AtomicDataDict.Type]):
    """
    Remember to call reset before this is needed.

    Args:
        data_source (Iterable[AtomicDataDict]): iterable data source
    """
    for data in data_source:
        self(data)
    return self.compute()
```

--------------------------------

### Download and Extract Data - Python

Source: https://nequip.readthedocs.io/en/latest/_modules/nequip/data/datamodule/md22_datamodule

The `prepare_data` method checks if a data file exists at the specified path. If not, it downloads the data from a URL defined in `self.url_dict` and extracts it if it's a zip archive. This ensures the data is ready for use by the system.

```python
def prepare_data(self):
        """"""
        if not os.path.isfile(self.data_file_path):
            # download and unzip
            download_path = download_url(
                self.url_dict[self.dataset], self.data_source_dir
            )
            if download_path.endswith(".zip"):
                extract_zip(download_path, self.data_source_dir)
        else:
            logger.info(f"Using existing data file `{self.data_file_path}`")
```

--------------------------------

### Register External Libraries for Packaging

Source: https://nequip.readthedocs.io/en/latest/dev/dev

This Python function is used to register external libraries as dependencies for packaging NequIP models. It helps in managing dependencies correctly during the packaging process.

```python
from nequip.utils.packaging import register_libraries_as_external_for_packaging

# Example usage within a packaging script or context
register_libraries_as_external_for_packaging({
    "torch": "pytorch",
    "torch.nn": "pytorch",
    "torch.optim": "pytorch",
    "torch.utils.data": "pytorch",
    "numpy": "numpy",
    "ase": "ase",
    "scipy": "scipy",
    "numba": "numba",
    "jax": "jax",
    "jax.numpy": "jax",
    "jax.nn": "jax"
})
```

--------------------------------

### NequIP Model Compilation with AOTInductor (Bash)

Source: https://context7.com/context7/nequip_readthedocs_io_en/llms.txt

Commands for compiling a NequIP model using PyTorch 2.6+ AOTInductor for maximum performance and hardware-specific optimizations. This example shows compilation for ASE integration on CUDA.

```bash
# Compile for ASE integration
nequip-compile \
  path/to/model.nequip.zip \
  path/to/compiled_model.nequip.pt2 \
  --device cuda \
  --mode aotinductor \
  --target ase
```

--------------------------------

### Create ASE NequIPCalculator from Compiled Model

Source: https://nequip.readthedocs.io/en/latest/integrations/ase

Instantiates an ASE NequIPCalculator using a pre-compiled model file. The calculator is initialized with the path to the compiled model and the target device (e.g., 'cuda' or 'cpu').

```python
from nequip.ase import NequIPCalculator

calculator = NequIPCalculator.from_compiled_model(
    compile_path="path/to/compiled_model.nequip.pt2",
    device="cuda",  # or "cpu"
)

```

--------------------------------

### MetricsManager: Get Extra State

Source: https://nequip.readthedocs.io/en/latest/_modules/nequip/train/metrics_manager

Retrieves the extra state of the MetricsManager, including coefficient dictionaries and metric values for steps and epochs. This is useful for saving and resuming training.

```python
    def get_extra_state(self) -> None:
        """"""
        return {
            "coeff_dict": {k: v["coeff"] for k, v in self.metrics.items()},
            "metrics_values_step": self.metrics_values_step,
            "metrics_values_epoch": self.metrics_values_epoch,
        }
```

--------------------------------

### Configure EnergyOnlyMetrics for MAE/RMSE

Source: https://nequip.readthedocs.io/en/latest/api/metrics

This example demonstrates configuring EnergyOnlyMetrics for validation metrics. It specifies the target module and the coefficients for different energy metrics like total energy RMSE and per-atom energy RMSE.

```yaml
training_module:
  _target_: nequip.train.NequIPLightningModule

  val_metrics:
    _target_: nequip.train.EnergyOnlyMetrics
    coeffs:
      total_energy_rmse: 1.0
      per_atom_energy_rmse: null
      total_energy_mae: null
      per_atom_energy_mae: null

```

--------------------------------

### Configure EnergyForceStressLoss for Training

Source: https://nequip.readthedocs.io/en/latest/api/metrics

This configuration sets up the EnergyForceStressLoss wrapper for training, including coefficients for total energy, forces, and stress. It is integrated into the NequIPLightningModule for comprehensive loss management.

```yaml
training_module:
  _target_: nequip.train.NequIPLightningModule

  loss:
    _target_: nequip.train.EnergyForceStressLoss
    per_atom_energy: true
    coeffs:
      total_energy: 1.0
      forces: 1.0
      stress: 1.0

```