### YAML Example: Full Training Configuration with Metrics

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/configuration/metrics.md

A comprehensive configuration example demonstrating the setup of loss functions, validation metrics, and trainer callbacks using nequip and PyTorch Lightning. It highlights the use of 'weighted_sum' for monitoring and conditioning training.

```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}
```

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

### Install torch-sim from Source

Source: https://github.com/mir-group/nequip/blob/main/docs/integrations/torchsim.md

Install torch-sim by cloning the repository and installing it in editable mode. This is the recommended installation method.

```bash
git clone https://github.com/TorchSim/torch-sim.git
cd torch-sim
pip install -e .
```

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

### Verify Installation with Tutorial Model

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/getting-started/install.md

Run the tutorial configuration to ensure the package is correctly installed and functional.

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

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

### Install NequIP from Source

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/getting-started/install.md

Methods for installing from the latest release or the develop branch.

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

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

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

### Install NequIP via PyPI

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/getting-started/install.md

Standard installation method using the Python package manager.

```bash
pip install nequip
```

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

### Install documentation dependencies

Source: https://github.com/mir-group/nequip/blob/main/docs/README.md

Install the required Python packages for building the documentation.

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

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

### NequIP3BPADataModule Configuration Example

Source: https://github.com/mir-group/nequip/blob/main/docs/api/datamodule.md

Example configuration for isolated atom energies within the NequIP3BPADataModule. Ensure these values match the 'iso_atoms.xyz' file.

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

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

### Initiate model training

Source: https://github.com/mir-group/nequip/blob/main/configs/README.md

Starts the training process using the specified YAML configuration file.

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

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

### Install CuEquivariance dependencies

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/accelerations/cuequivariance.md

Install the required PyTorch and CUDA-specific CuEquivariance packages.

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

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

### Build and Install LAMMPS

Source: https://github.com/mir-group/nequip/blob/main/docs/integrations/lammps/mliap.md

Compile the LAMMPS binary and install the Python interface.

```bash
cmake --build build-mliap -j 8
```

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

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

### Install and Run Pre-commit Hooks

Source: https://github.com/mir-group/nequip/blob/main/docs/dev/contributing.md

Install pre-commit to automatically run code formatting and linting checks locally before each commit. This ensures consistency with CI checks.

```bash
pip install pre-commit
pre-commit install
```

```bash
pre-commit run --all-files
```

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

### Run Unit and Integration Tests

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/getting-started/install.md

Execute test suites to validate the installation, including GPU-accelerated tests if hardware is available.

```default
pip install pytest
pytest tests/unit/
```

```default
pytest tests/
```

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

### SLURM Example for DDP Training

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/accelerations/ddp_training.md

A minimal SLURM job submission script for distributed data parallel training using 2 nodes and 4 GPUs per node. Ensure cluster-specific setup for network interfaces and environment variables is performed.

```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}

```

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

### Configure SoftAdapt Callback

Source: https://github.com/mir-group/nequip/blob/main/docs/api/callbacks.md

Example configuration for the SoftAdapt callback to update loss coefficients every 5 epochs.

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

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

### Install ML-IAP Dependencies

Source: https://github.com/mir-group/nequip/blob/main/docs/integrations/lammps/mliap.md

Install the necessary Python packages for the ML-IAP build process.

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

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

### Basic NequIP Training

Source: https://context7.com/mir-group/nequip/llms.txt

Use this command to start training a NequIP model with a specified configuration file. Ensure the YAML file is correctly formatted.

```bash
nequip-train --config-name tutorial.yaml
```

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

### Configure LinearLossCoefficientScheduler Callback

Source: https://github.com/mir-group/nequip/blob/main/docs/api/callbacks.md

Examples for linear interpolation of loss coefficients, including single-stage and multi-stage transitions.

```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
```

```yaml
callbacks:
  # First transition: current -> 1:5:1 from epoch 50-150
  - _target_: nequip.train.callbacks.LinearLossCoefficientScheduler
    final_coeffs:
      per_atom_energy_mse: 1.0
      forces_mse: 5.0
      stress_mse: 1.0
    start_epoch: 50
    transition_epochs: 100
  # Second transition: current -> 1:1:1 from epoch 200-400
  - _target_: nequip.train.callbacks.LinearLossCoefficientScheduler
    final_coeffs:
      per_atom_energy_mse: 1.0
      forces_mse: 1.0
      stress_mse: 1.0
    start_epoch: 200
    transition_epochs: 200
```

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

### Calculate energy-volume curve

Source: https://github.com/mir-group/nequip/blob/main/docs/integrations/ase.md

Example demonstrating the use of the NequIP calculator to compute energies and forces for an energy-volume curve.

```python
from ase.build import bulk
import numpy as np
import matplotlib.pyplot as plt
from nequip.integrations.ase import NequIPCalculator
import torch

# Initialize the nequip calculator
calculator = NequIPCalculator.from_compiled_model(
    compile_path="path/to/compiled_model.nequip.pt2",
    chemical_species_to_atom_type_map={"Si": "Si"},
    device="cuda" if torch.cuda.is_available() else "cpu",
)  # use GPUs if available

# Range of scaling factors for lattice constant
scaling_factors = np.linspace(0.95, 1.05, 10)
volumes = []
energies = []
```

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

### Configure Adam Optimizer

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/configuration/config.md

Example configuration for the Adam optimizer using PyTorch. Ensure `torch.optim.Adam` is available.

```yaml
optimizer:
  _target_: torch.optim.Adam
  lr: 0.01
```

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

### Configure LossCoefficientScheduler Callback

Source: https://github.com/mir-group/nequip/blob/main/docs/api/callbacks.md

Example configuration for scheduling specific loss coefficients at defined epoch milestones.

```yaml
callbacks:
  - _target_: nequip.train.callbacks.LossCoefficientScheduler
    schedule:
      100:
        per_atom_energy_mse: 1.0
        forces_mse: 5.0
      200:
        per_atom_energy_mse: 5.0
        forces_mse: 1.0
```

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

### Start a W&B Sweep Agent

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/training-techniques/wandb_example.md

Run this command to start an agent that will execute the sweep jobs. Replace `<sweep-ID>` with the actual sweep ID.

```bash
wandb agent <sweep-ID>
```

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

### YAML Example: Coefficients for Loss and Metrics

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/configuration/metrics.md

Example of setting coefficients for loss and metric terms. Coefficients are normalized to sum to 1 and affect the 'weighted_sum'. Metrics with null coefficients are computed but excluded from 'weighted_sum'.

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

```yaml
coeffs:
  total_energy_rmse: 1.0
  forces_rmse: 1.0
  total_energy_mae: null
  forces_mae: null
```

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

### Combine Multiple Transforms

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/configuration/data.md

Example showing the sequential application of species mapping and neighbor list transforms.

```yaml
transforms:
  - _target_: nequip.data.transforms.ChemicalSpeciesToAtomTypeMapper
    model_type_names: ${model_type_names}
  - _target_: nequip.data.transforms.NeighborListTransform
    r_max: 5.0
```

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

### Configure MetricsManager via YAML

Source: https://github.com/mir-group/nequip/blob/main/docs/api/metrics.md

Example configuration for a MetricsManager instance equivalent to an EnergyForceLoss, defining metrics for energy and forces.

```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
```

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

### Configure Tensorboard Logger

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/configuration/config.md

Example configuration for integrating Tensorboard logging via PyTorch Lightning within the trainer section.

```yaml
logger:
  _target_: lightning.pytorch.loggers.TensorBoardLogger
  # The run name in tensorboard can be, for example, inherited from Hydra.
  version: ${hydra:job.name}
  # By default (not if overridden) Hydra will make `./outputs` and put various runs at `./outputs/{name}`.
  # Here we add an additional `./outputs/tensorboard_logs` within which logs will be stored _across_ runs.
  save_dir: outputs/tensorboard_logs
```

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

### Configuring SimpleDDPStrategy in YAML

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/accelerations/ddp_training.md

Example of how to specify NequIP's SimpleDDPStrategy in a YAML configuration file for distributed training, particularly when using train-time compilation.

```yaml
trainer:
  _target_: lightning.Trainer
  # other trainer arguments
  devices: ${num_devices}
  num_nodes: ${num_nodes}
  strategy:
    _target_: nequip.train.SimpleDDPStrategy

```

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

### Dynamic Batch Size Configuration with Environment Variables

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/accelerations/ddp_training.md

Example of using Omegaconf's variable interpolation and environment variable resolver to dynamically set the per-rank batch size based on the effective global batch size and the SLURM world size.

```yaml
batch_size: ${int_div:${effective_global_batch_size},${oc.env:SLURM_NTASKS}}

```

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

### Register Libraries for Packaging

Source: https://github.com/mir-group/nequip/blob/main/docs/dev/understanding_nequip/workflows/packaging.md

Use this function to register libraries as 'external' or 'mocked' for packaging. External libraries are not included in the package file, requiring them to be installed in the target environment. Mocked libraries allow packaging but will raise an error if used at runtime. Avoid registering NequIP extension packages as external.

```python
from nequip.scripts._package_utils import register_libraries_as_external_for_packaging

# Example: Registering 'openequivariance' as external and 'matplotlib' as mock
register_libraries_as_external_for_packaging(extern_modules=['openequivariance'], mock_modules=['matplotlib'])
```

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

### Start a W&B Sweep

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/training-techniques/wandb_example.md

Use this command to initiate a Weights and Biases sweep based on a sweep configuration file.

```bash
wandb sweep sweep.yaml
```

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

### Lazy Import for Optional Dependencies

Source: https://github.com/mir-group/nequip/blob/main/docs/dev/understanding_nequip/workflows/packaging.md

When dealing with optional dependencies that might not be installed at package-time, use lazy imports to prevent packaging failures. This example shows a 'bad' pattern with a top-level import that would fail if the dependency is not present.

```python
# bad: top-level import
from openequivariance import TensorProductConv  # fails if not installed
```

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

### Build documentation

Source: https://github.com/mir-group/nequip/blob/main/docs/README.md

Execute the build process to generate HTML documentation files.

```bash
make html
```

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

### Run NequIP training command

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/reference/troubleshoot.md

The standard command to initiate training using a configuration file.

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

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

### Initialize NequIPCalculator

Source: https://github.com/mir-group/nequip/blob/main/docs/integrations/ase.md

Create an ASE calculator instance from a compiled NequIP model file.

```python
from nequip.integrations.ase import NequIPCalculator

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

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

### Download tutorial training data

Source: https://github.com/mir-group/nequip/blob/main/configs/README.md

Executes the script to fetch the fcu.xyz dataset from Materials Cloud.

```bash
python get_tutorial_data.py
```

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

### NequIP3BPADataModule Initialization

Source: https://github.com/mir-group/nequip/blob/main/docs/api/datamodule.md

Initialize the NequIP3BPADataModule with specified parameters. The data_source_dir should point to the dataset location or download directory.

```python
NequIP3BPADataModule(seed=123, transforms=[...], train_val_split=[0.9, 0.1], data_source_dir='/path/to/data')
```

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

### Initialize NequIP ASE Calculator

Source: https://context7.com/mir-group/nequip/llms.txt

Use the NequIPCalculator to interface with ASE for atomistic simulations.

```python
from nequip.integrations.ase import NequIPCalculator

calc = NequIPCalculator.from_compiled_model(
    model_path="nequip.net:foundation/nequip-mptrj-omat:latest",
    device="cuda",
    transforms=[...],  # Check model documentation for required transforms
)
```

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

### NumNeighbors

Source: https://github.com/mir-group/nequip/blob/main/docs/api/data_modifiers.md

Get number of neighbors from an AtomicDataDict.

```APIDOC
## NumNeighbors

### Description
Get number of neighbors from an `AtomicDataDict`.

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Request Example
None

### Response
#### Success Response (200)
None

#### Response Example
None
```

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

### Train a NequIP model

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/getting-started/workflow.md

Execute the training process using a specified configuration directory and file name.

```bash
nequip-train -cp full/path/to/config/directory -cn config_name.yaml
```

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

### EdgeLengths

Source: https://github.com/mir-group/nequip/blob/main/docs/api/data_modifiers.md

Get edge lengths from an AtomicDataDict.

```APIDOC
## EdgeLengths

### Description
Get edge lengths from an `AtomicDataDict`.

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Request Example
None

### Response
#### Success Response (200)
None

#### Response Example
None
```

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

### LAMMPS Script Configuration

Source: https://github.com/mir-group/nequip/blob/main/docs/integrations/lammps/mliap.md

Example configuration for using a NequIP model within a LAMMPS input script.

```lammps
# general settings
units         metal
boundary      p p p
atom_style    atomic
atom_modify   map yes
newton        on

# data
read_data     system.data
replicate     2 2 2

# set up potential
pair_style    mliap unified output.nequip.lmp.pt 0
pair_coeff    * * H O
```

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

### Package Models with nequip-package

Source: https://context7.com/mir-group/nequip/llms.txt

Create and inspect portable model archives for deployment.

```bash
# Build a packaged model from checkpoint
nequip-package build /path/to/best.ckpt output_model.nequip.zip

# Inspect package metadata
nequip-package info output_model.nequip.zip

# Get available modifiers in package
nequip-package info output_model.nequip.zip --get-modifiers

# Output in YAML format
nequip-package info output_model.nequip.zip --yaml
```

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

### Run NequIP training with Hydra configuration

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/reference/troubleshoot.md

Correct command-line usage for specifying a configuration file with Hydra.

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

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

### NequIPCalculator Initialization

Source: https://github.com/mir-group/nequip/blob/main/docs/api/ase.md

Initialize the NequIPCalculator with a NequIP framework model and specify device and unit conversion factors.

```APIDOC
## NequIPCalculator

### Description
NequIP framework ASE Calculator compatible with NequIP and Allegro models.

### Method
__init__

### Parameters
- **model** (Module) - Required - A model in the NequIP framework.
- **device** (str | torch.device) - Required - Device for model evaluation (e.g., 'cpu' or 'cuda').
- **energy_units_to_eV** (float) - Optional - Energy conversion factor (default 1.0).
- **length_units_to_A** (float) - Optional - Length units conversion factor (default 1.0).
- **transforms** (List[Callable]) - Optional - List of data transforms.

### Implemented Properties
- **energy** (list[str])
- **energies** (list[str])
- **forces** (list[str])
- **stress** (list[str])
- **free_energy** (list[str])

### WARNING
If running MD with custom species, ensure correct masses are set for ASE.
```

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

### Create NequIPTorchSimCalc Instance

Source: https://github.com/mir-group/nequip/blob/main/docs/integrations/torchsim.md

Instantiate the `NequIPTorchSimCalc` using a pre-compiled model file. Specify the device to be used for calculations.

```python
from nequip.integrations.torchsim import NequIPTorchSimCalc

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

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

### Understanding the NequIP Framework

Source: https://github.com/mir-group/nequip/blob/main/docs/dev/dev.md

An overview of the NequIP framework, detailing its workflows, particularly model packaging.

```APIDOC
## Understanding the NequIP Framework

### Description
Explores the core NequIP framework and its workflows, with a focus on model packaging.

### Workflows
- Model Packaging
  - Overview
  - Developer Notes
    - Package Format Versioning
    - Dependency Management
      - `register_libraries_as_external_for_packaging()`
    - Repackaging Support
    - Sharp Edges with Model Modifiers
```

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

### Build NequIP Models Programmatically

Source: https://context7.com/mir-group/nequip/llms.txt

Instantiate models using either custom architecture parameters or predefined presets.

```python
model = NequIPGNNModel(
    seed=123,
    model_dtype="float32",
    type_names=["C", "H", "O"],
    r_max=5.0,

    # Architecture parameters
    num_layers=4,           # Number of interaction blocks
    l_max=2,                # Maximum rotation order (l=1 faster, l=2 more accurate)
    parity=True,            # Include odd-parity features
    num_features=32,        # Feature multiplicity per irrep

    # Radial network
    radial_mlp_depth=2,
    radial_mlp_width=64,
    num_bessels=8,
    polynomial_cutoff_p=6,

    # Normalization
    avg_num_neighbors=15.0,

    # Per-type energy parameters (from dataset statistics)
    per_type_energy_scales={"C": 1.2, "H": 0.8, "O": 1.5},
    per_type_energy_shifts={"C": -38.0, "H": -0.5, "O": -75.0},
)
```

```python
from nequip.model import PresetNequIPGNNModel

model_large = PresetNequIPGNNModel(
    preset="L",  # Large preset: 6 layers, l_max=3
    seed=123,
    type_names=["Si", "O"],
    r_max=6.0,
    avg_num_neighbors=20.0,
)
```

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

### Configure PerTypeScaleShift modifications

Source: https://github.com/mir-group/nequip/blob/main/docs/api/nn.md

Example YAML configuration for updating per-type atomic energy shifts during model fine-tuning.

```yaml
shifts:
  C: 1.23
  H: 0.12
  O: 2.13
```

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

### Fine-tune Pre-trained Models

Source: https://context7.com/mir-group/nequip/llms.txt

Load and modify pre-trained models for fine-tuning using YAML configuration or programmatic setup.

```yaml
# config.yaml for fine-tuning
training_module:
  model:
    # Use ModelFromCheckpoint to load pre-trained weights
    _target_: nequip.model.saved_models.ModelFromCheckpoint
    ckpt_path: /path/to/pretrained/best.ckpt
    compile_mode: compile

    # Apply modifiers for fine-tuning
    modifiers:
      - modifier: PerTypeScaleShift
        per_type_energy_scales: ${training_data_stats:per_type_forces_rms}
        per_type_energy_shifts:
          C: -38.1  # New isolated atom energies for target data
          H: -0.5
          O: -75.2
```

```python
# Programmatic fine-tuning setup
from nequip.model.saved_models import load_saved_model
from nequip.model.modify_utils import modify

# Load base model
model = load_saved_model(
    "pretrained.nequip.zip",
    compile_mode="eager",
)

# Apply modifications for fine-tuning
model = modify(model, [
    {
        "modifier": "PerTypeScaleShift",
        "per_type_energy_shifts": {"C": -38.0, "H": -0.5, "O": -75.0},
    }
])
```

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

### Initialize ASEDataModule

Source: https://context7.com/mir-group/nequip/llms.txt

Manage data loading, splitting, and statistics calculation using ASE-compatible datasets.

```python
from nequip.data.datamodule import ASEDataModule

# Create data module with split from single file
datamodule = ASEDataModule(
    seed=42,

    # Split a single dataset file
    split_dataset={
        "file_path": "training_data.xyz",  # ASE-readable format
        "train": 0.8,   # 80% for training
        "val": 0.1,     # 10% for validation
        "test": 0.1,    # 10% for testing
    },

    # Data transforms
    transforms=[
        {"_target_": "nequip.data.transforms.ChemicalSpeciesToAtomTypeMapper",
         "model_type_names": ["H", "C", "O"]},
        {"_target_": "nequip.data.transforms.NeighborListTransform",
         "r_max": 5.0},
    ],

    # DataLoader configurations
    train_dataloader={
        "_target_": "torch.utils.data.DataLoader",
        "batch_size": 8,
        "shuffle": True,
        "num_workers": 4,
    },
    val_dataloader={
        "_target_": "torch.utils.data.DataLoader",
        "batch_size": 16,
    },

    # Statistics manager for model initialization
    stats_manager={
        "_target_": "nequip.data.CommonDataStatisticsManager",
        "type_names": ["H", "C", "O"],
        "dataloader_kwargs": {"batch_size": 32},
    },
)

# Get dataset statistics
stats = datamodule.get_statistics(dataset="train")
# Returns: {"num_neighbors_mean", "per_atom_energy_mean", "forces_rms", "per_type_forces_rms"}
```

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

### Integrate NequIP with ASE

Source: https://context7.com/mir-group/nequip/llms.txt

Use compiled models as ASE calculators for molecular dynamics and geometry optimization.

```python
from ase import Atoms
from ase.build import molecule, bulk
from ase.optimize import BFGS
from ase.md.langevin import Langevin
from ase import units

from nequip.integrations.ase import NequIPCalculator
from nequip.data.transforms import ChemicalSpeciesToAtomTypeMapper, NeighborListTransform

# Load compiled model into calculator
calc = NequIPCalculator.from_compiled_model(
    model_path="compiled_model.nequip.pt2",
    device="cuda",
    transforms=[
        ChemicalSpeciesToAtomTypeMapper(model_type_names=["C", "H", "O"]),
        NeighborListTransform(r_max=5.0),
    ],
)

# Create atoms and attach calculator
atoms = molecule("H2O")
atoms.calc = calc

# Get energy and forces
energy = atoms.get_potential_energy()  # in eV
forces = atoms.get_forces()             # in eV/Angstrom
print(f"Energy: {energy:.4f} eV")
print(f"Max force: {forces.max():.4f} eV/Ang")

# Geometry optimization
opt = BFGS(atoms)
opt.run(fmax=0.01)  # Optimize until max force < 0.01 eV/Ang

# Molecular dynamics with Langevin thermostat
atoms = bulk("Cu", cubic=True) * (3, 3, 3)
atoms.calc = calc

dyn = Langevin(
    atoms,
    timestep=1.0 * units.fs,
    temperature_K=300,
    friction=0.01,
)
dyn.run(steps=1000)
```

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

### Configure Multi-GPU Training

Source: https://context7.com/mir-group/nequip/llms.txt

Set up distributed training using PyTorch Lightning strategies.

```yaml
# config.yaml for multi-GPU training
trainer:
  _target_: lightning.Trainer
  accelerator: gpu
  devices: 4           # Number of GPUs
  strategy: ddp        # Distributed Data Parallel
  max_epochs: 500

  # Or use NequIP's SimpleDDPStrategy for better performance
  # strategy:
  #   _target_: nequip.train.SimpleDDPStrategy
```

```bash
# Run with multiple GPUs
nequip-train --config-name config.yaml

# Explicit GPU selection
CUDA_VISIBLE_DEVICES=0,1,2,3 nequip-train --config-name config.yaml
```

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

### Load Model from Package Configuration

Source: https://github.com/mir-group/nequip/blob/main/docs/api/save_model.md

Use this configuration to load a NequIP model from a packaged zip file. The `package_path` should point to the `.nequip.zip` file. Using absolute paths is recommended.

```yaml
model:
  _target_: nequip.model.ModelFromPackage
  package_path: path/to/pkg
  compile_mode: eager/compile
```

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

### nequip.model.PresetNequIPGNNModel

Source: https://github.com/mir-group/nequip/blob/main/docs/api/nequip_model.md

Builds a NequIPGNNModel using predefined architecture presets (S, M, L, XL).

```APIDOC
## nequip.model.PresetNequIPGNNModel

### Description
A wrapper for NequIPGNNModel that injects preset hyperparameters. Users can override preset defaults by providing explicit arguments.

### Parameters
- **preset** (str) - Required - One of 'S', 'M', 'L', 'XL'
- **kwargs** (dict) - Optional - Arguments to override preset defaults

### Preset Configurations
- **S**: {'num_layers': 2, 'l_max': 1, 'num_features': [128, 64]}
- **M**: {'num_layers': 4, 'l_max': 2, 'num_features': [128, 64, 32]}
- **L**: {'num_layers': 6, 'l_max': 3, 'num_features': [128, 64, 32, 32]}
- **XL**: {'num_layers': 6, 'l_max': 4, 'num_features': [320, 96, 64, 32, 32]}
```

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

### Configure ReduceLROnPlateau Scheduler

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/configuration/config.md

Example configuration for the ReduceLROnPlateau learning rate scheduler. This scheduler monitors a specified metric and reduces the learning rate when the metric has stopped improving.

```yaml
lr_scheduler:
  scheduler:
    _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
    factor: 0.6
    patience: 5
    threshold: 0.2
    min_lr: 1e-6
  monitor: val0_epoch/weighted_sum
  interval: epoch
  frequency: 1
```

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

### Restart training from a checkpoint

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/getting-started/workflow.md

Use the ++ckpt_path override to resume training from a specific checkpoint file while maintaining the original configuration.

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

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

### Prepare NequIP Models

Source: https://github.com/mir-group/nequip/blob/main/docs/integrations/lammps/mliap.md

Use the command-line tool to convert NequIP models into the format required by LAMMPS ML-IAP.

```bash
# Using local checkpoint or package file
nequip-prepare-lmp-mliap \
  ckpt_file_or_package_file \
  output.nequip.lmp.pt \
  --modifiers modifier_to_apply

# Using model from nequip.net
nequip-prepare-lmp-mliap \
  nequip.net:group-name/model-name:version \
  output.nequip.lmp.pt \
  --modifiers modifier_to_apply
```

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

### Resume NequIP Training from Checkpoint

Source: https://context7.com/mir-group/nequip/llms.txt

Resume a previously interrupted training session by specifying the path to the last checkpoint file. This allows for continued training without losing progress.

```bash
nequip-train --config-name tutorial.yaml ckpt_path=/path/to/last.ckpt
```

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

### nequip.model.NequIPGNNModel

Source: https://github.com/mir-group/nequip/blob/main/docs/api/nequip_model.md

Initializes a NequIP GNN model capable of predicting energies, forces, and stresses.

```APIDOC
## nequip.model.NequIPGNNModel

### Description
Constructs a NequIP GNN model. This model can be configured for various depths, feature multiplicities, and rotation orders.

### Parameters
- **seed** (int) - Optional - Seed for reproducibility
- **model_dtype** (str) - Optional - float32 or float64
- **r_max** (float) - Optional - Cutoff radius
- **num_layers** (int) - Optional - Number of interaction blocks (default 4)
- **l_max** (int) - Optional - Maximum rotation order (default 1)
- **parity** (bool) - Optional - Include odd mirror parity features (default True)
- **num_features** (int/List[int]) - Optional - Multiplicity of features (default 32)
- **radial_mlp_depth** (int) - Optional - Number of radial layers (default 1)
- **radial_mlp_width** (int) - Optional - Number of hidden neurons in radial function (default 128)
- **do_derivatives** (bool) - Optional - Compute forces and stresses via autograd (default True)
```

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

### Import NequIPGNNModel and model_builder

Source: https://context7.com/mir-group/nequip/llms.txt

Imports the necessary classes for building NequIP models. `NequIPGNNModel` is the core model, and `model_builder` is a utility for constructing models.

```python
from nequip.model import NequIPGNNModel, model_builder
```

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

### ConFIGLightningModule

Source: https://github.com/mir-group/nequip/blob/main/docs/api/lightning_module.md

A LightningModule implementation for the Conflict-free inverse gradient (ConFIG) approach to multitask learning.

```APIDOC
## ConFIGLightningModule

### Description
Implements the Conflict-free inverse gradient (ConFIG) approach for multitask learning. Arguments are identical to NequIPLightningModule with additional parameters for gradient handling.

### Parameters
- **gradient_clip_val** (float) - Optional - Value for gradient clipping.
- **gradient_clip_algorithm** (str) - Optional - Algorithm for gradient clipping.
- **lsqr** (bool) - Optional - Whether to use least squares method (default True).
- **norm_eps** (float) - Optional - Epsilon for normalization (default 1e-08).
```

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

### COLLDataModule Initialization

Source: https://github.com/mir-group/nequip/blob/main/docs/api/datamodule.md

Initialize the COLLDataModule. The data_source_dir is where the COLL dataset will be downloaded or is already located.

```python
COLLDataModule(seed=123, transforms=[...], data_source_dir='/path/to/data')
```

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

### NequIP Training with Custom Config Path

Source: https://context7.com/mir-group/nequip/llms.txt

Train a NequIP model using a configuration file located in a custom directory. This is useful for organizing multiple configuration files.

```bash
nequip-train --config-path /path/to/configs --config-name my_config.yaml
```

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

### Load compiled model in ASE

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/accelerations/cuequivariance.md

Import the CuEquivariance library before initializing the NequIPCalculator with a compiled model.

```python
import cuequivariance_torch
from nequip.integrations.ase import NequIPCalculator

# Load the compiled model
calc = NequIPCalculator.from_compiled_model(
    "path/to/compiled_model.nequip.pt2/pth",
    device="cuda"
)

# Use with ASE atoms object
atoms.calc = calc
energy = atoms.get_potential_energy()
forces = atoms.get_forces()
```

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

### Inspect packaged model metadata

Source: https://github.com/mir-group/nequip/blob/main/docs/guide/getting-started/workflow.md

Retrieve information about a previously packaged NequIP model.

```bash
nequip-package info path/to/pkg_file.nequip.zip
```

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

### NequIPDataModule Initialization

Source: https://github.com/mir-group/nequip/blob/main/docs/api/datamodule.md

Initializes the NequIPDataModule with datasets, dataloaders, and statistics management configuration.

```APIDOC
## CLASS NequIPDataModule

### Description
Initializes the data module for managing training, validation, testing, and prediction datasets. It supports explicit dataset definition or random splitting.

### Parameters
- **seed** (int) - Required - Data seed for reproducibility
- **train_dataset** (Dict/List) - Optional - Training dataset
- **val_dataset** (Dict/List) - Optional - Validation dataset(s)
- **test_dataset** (Dict/List) - Optional - Test dataset(s)
- **predict_dataset** (Dict/List) - Optional - Prediction dataset(s)
- **split_dataset** (Dict/List) - Optional - Configuration for splitting a dataset into subsets
- **train_dataloader** (Dict) - Optional - Training DataLoader configuration
- **val_dataloader** (Dict) - Optional - Validation DataLoader configuration
- **test_dataloader** (Dict) - Optional - Testing DataLoader configuration
- **predict_dataloader** (Dict) - Optional - Prediction DataLoader configuration
- **stats_manager** (Dict) - Optional - Configuration for DataStatisticsManager
```