### Installation and imports

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/qm_fitting.md

Initial setup for the Espaloma environment.

```python

```

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

### Install espaloma with mamba

Source: https://github.com/choderalab/espaloma/blob/main/docs/install.md

Use mamba to create a new environment and install espaloma version 0.3.2. Mamba is recommended for its speed.

```bash
mamba create --name espaloma -c conda-forge "espaloma=0.3.2"
```

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

### Install Conda and Dependencies

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/qm_fitting.md

Sets up the environment by installing condacolab and necessary scientific libraries via mamba.

```python
! pip install -q condacolab
import condacolab
condacolab.install()
```

```python
%%capture
! mamba install --yes --strict-channel-priority --channel jaimergp/label/unsupported-cudatoolkit-shim --channel omnia --channel omnia/label/cuda100 --channel dglteam --channel numpy openmm openmmtools openmmforcefields rdkit openff-toolkit dgl-cuda10.0 qcportal
```

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

### Install Conda environment

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Initializes the Conda environment within a Google Colab session.

```python
# install conda
! pip install -q condacolab
import condacolab
condacolab.install()
```

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

### Install dependencies via Mamba

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Installs required scientific and machine learning libraries using Mamba.

```python
%%capture
! mamba install --yes --strict-channel-priority --channel jaimergp/label/unsupported-cudatoolkit-shim --channel omnia --channel omnia/label/cuda100 --channel dglteam --channel numpy openmm openmmtools openmmforcefields rdkit openff-toolkit dgl-cuda10.0 qcportal
```

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

### Install Espaloma and Dependencies

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/typing.md

Installs Espaloma and its dependencies using Conda. This process includes updating Conda, creating a new environment, and cloning the Espaloma repository. It's recommended to run this in a Linux environment with GPU support.

```bash
%%capture
! wget -c https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
! bash Miniconda3-latest-Linux-x86_64.sh -b -f -p /usr/local
! conda config --add channels conda-forge --add channels omnia --add channels omnia/label/cuda100 --add channels dglteam
! conda update --yes --all
! conda create --yes -n openmm python=3.6 numpy openmm openmmtools rdkit openforcefield==0.7.0 dgl-cuda10.0 qcportal
! git clone https://github.com/choderalab/espaloma.git
```

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

### Install Sphinx dependencies

Source: https://github.com/choderalab/espaloma/blob/main/docs/README.md

Install the required Sphinx and ReadTheDocs theme packages via Conda.

```bash
conda install sphinx sphinx_rtd_theme
```

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

### Install espaloma with Mamba

Source: https://github.com/choderalab/espaloma/blob/main/README.md

Installs espaloma version 0.3.2 using mamba, a faster alternative to conda. Ensure you have mamba installed and configured with the conda-forge channel.

```bash
$ mamba create --name espaloma -c conda-forge "espaloma=0.3.2"
```

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

### Install Condacolab

Source: https://github.com/choderalab/espaloma/blob/main/docs/qm_fitting.md

Installs the condacolab package for managing conda environments in Google Colab.

```python
! pip install -q condacolab
import condacolab
condacolab.install()
```

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

### Compile documentation

Source: https://github.com/choderalab/espaloma/blob/main/docs/README.md

Execute the Makefile to generate static HTML documentation.

```bash
make html
```

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

### Orchestrate Training Experiments with Train

Source: https://context7.com/choderalab/espaloma/llms.txt

Demonstrates setting up and running a training experiment using `esp.app.experiment.Train`, including model, data, metrics, optimizer, and epoch configuration.

```python
import espaloma as esp
import torch

# Setup model, data, and training
representation = esp.nn.Sequential(
    layer=esp.nn.layers.dgl_legacy.gn("SAGEConv"),
    config=[128, "relu", 128, "relu", 128, "relu"],
)
readout = esp.nn.readout.janossy.JanossyPooling(
    in_features=128,
    config=[128, "relu", 128, "relu"],
    out_features={2: {"k": 1, "eq": 1}, 3: {"k": 1, "eq": 1}, 4: {"k": 6}},
)
model = torch.nn.Sequential(
    representation, readout,
    esp.mm.geometry.GeometryInGraph(),
    esp.mm.energy.EnergyInGraph(),
)

# Load and split dataset
dataset = esp.data.dataset.GraphDataset.load("./dataset")
dataset.shuffle(seed=2666)
ds_train, ds_valid, ds_test = dataset.split([8, 1, 1])

# Configure training
train_experiment = esp.app.experiment.Train(
    net=model,
    data=ds_train.view("graph", batch_size=32),
    metrics=[esp.metrics.GraphMetric(
        base_metric=torch.nn.MSELoss(),
        between=["u", "u_ref"],
        level="g"
    )],
    optimizer=lambda net: torch.optim.Adam(net.parameters(), lr=1e-4),
    n_epochs=100,
    record_interval=10,
    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
)

# Run training
trained_model = train_experiment.train()

# Access training states for checkpointing
states = train_experiment.states  # Dict of {epoch: state_dict}
```

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

### Load and Prepare Dataset

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/qm_fitting.md

Configures the dataset name, downloads the data, and splits it into training, validation, and test sets.

```python
dataset_name = "gen2"
# dataset_name = "pepconf"
# dataset_name = "vehicle"
# dataset_name = "phalkethoh"
```

```python
%%capture
! wget "data.wangyq.net/esp_dataset/"$dataset_name".zip"
! unzip $dataset_name".zip"
```

```python
ds = esp.data.dataset.GraphDataset.load(dataset_name)
ds.shuffle(seed=2666)
ds_tr, ds_vl, ds_te = ds.split([8, 1, 1])
```

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

### esp.graphs.deploy.openmm_system_from_graph()

Source: https://context7.com/choderalab/espaloma/llms.txt

Converts an Espaloma-parameterized graph into an OpenMM System.

```APIDOC
## esp.graphs.deploy.openmm_system_from_graph(graph, forcefield, charge_method)

### Description
Generates an OpenMM System object from a parameterized molecular graph.

### Parameters
- **graph** (Graph) - Required - The Espaloma graph object.
- **forcefield** (string) - Optional - The base force field to use.
- **charge_method** (string) - Optional - Method for charge assignment ('nn', 'am1-bcc', 'from-molecule').

### Request Example
```python
openmm_system = esp.graphs.deploy.openmm_system_from_graph(
    molecule_graph,
    forcefield="openff_unconstrained-2.1.1",
    charge_method="nn"
)
```
```

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

### Configure Loss Function and Optimizer

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/qm_fitting.md

Sets up the MSE loss function for energy prediction and the Adam optimizer.

```python
loss_fn = esp.metrics.GraphMetric(
        base_metric=torch.nn.MSELoss(), # use mean-squared error loss
        between=['u', "u_ref"],         # between predicted and QM energies
        level="g", # compare on graph level
)
```

```python
optimizer = torch.optim.Adam(espaloma_model.parameters(), 1e-4)
```

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

### Import Espaloma

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Configures the system path and imports the Espaloma library.

```python
import torch
import sys
sys.path.append("/content/espaloma")
import espaloma as esp
```

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

### Initialize Loss Tracking

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Creates lists to store training and validation losses.

```python
loss_tr = []
loss_vl = []
```

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

### Build and Run Espaloma Experiment

Source: https://context7.com/choderalab/espaloma/llms.txt

Defines a sequential model, loads and splits a dataset, configures a training and testing experiment, and runs the experiment. The model includes SAGEConv layers, Janossy Pooling, and geometry/energy modules. The experiment is configured with specified datasets, metrics, optimizer, epochs, and device.

```python
model = torch.nn.Sequential(
    esp.nn.Sequential(
        layer=esp.nn.layers.dgl_legacy.gn("SAGEConv"),
        config=[128, "relu", 128, "relu", 128, "relu"],
    ),
    esp.nn.readout.janossy.JanossyPooling(
        in_features=128,
        config=[128, "relu", 128, "relu"],
        out_features={2: {"k": 1, "eq": 1}, 3: {"k": 1, "eq": 1}, 4: {"k": 6}},
    ),
    esp.mm.geometry.GeometryInGraph(),
    esp.mm.energy.EnergyInGraph(),
)

ds = esp.data.dataset.GraphDataset.load("./qm_dataset")
ds.shuffle(seed=42)
ds_train, ds_valid, ds_test = ds.split([8, 1, 1])

experiment = esp.app.experiment.TrainAndTest(
    net=model,
    ds_tr=ds_train.view("graph", batch_size=16),
    ds_te=ds_test.view("graph", batch_size=16),
    ds_vl=ds_valid.view("graph", batch_size=16),
    metrics_tr=[
        esp.metrics.GraphMetric(
            base_metric=torch.nn.MSELoss(),
            between=["u", "u_ref"],
            level="g"
        )
    ],
    metrics_te=[
        esp.metrics.GraphMetric(
            base_metric=esp.metrics.rmse,
            between=["u", "u_ref"],
            level="g"
        )
    ],
    optimizer=lambda net: torch.optim.Adam(net.parameters(), lr=1e-4),
    n_epochs=1000,
    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
)

results = experiment.run()
print(f"Test RMSE: {results['test']}")
print(f"Train RMSE: {results['train']}")
print(f"Validation RMSE: {results['validate']}")
```

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

### Model Download Utilities with esp.get_model_path()

Source: https://context7.com/choderalab/espaloma/llms.txt

Downloads pre-trained model files to a specified directory with progress tracking and caching. Allows downloading the 'latest' version or a specific version, with options to overwrite existing files. The downloaded model can then be loaded using torch.load().

```python
import espaloma as esp

# Download latest model to default directory
model_path = esp.get_model_path(
    model_dir=".espaloma/",
    version="latest",
    disable_progress_bar=False,
    overwrite=False
)

# Download specific version
model_path = esp.get_model_path(
    model_dir="./models/",
    version="0.3.2",
    overwrite=True  # Re-download even if exists
)

print(f"Model saved to: {model_path}")

# Load downloaded model
import torch
model = torch.load(model_path, map_location="cpu")
model.eval()
```

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

### Manage Datasets with GraphDataset

Source: https://context7.com/choderalab/espaloma/llms.txt

Demonstrates dataset management using `esp.data.dataset.GraphDataset`, including creation from molecules, loading, shuffling, splitting, subsampling, applying transformations, and creating DataLoaders.

```python
import espaloma as esp
from openff.toolkit.topology import Molecule

# Create dataset from SMILES or Molecule objects
smiles_list = ["CCO", "CC=O", "CC(=O)O", "CCCC", "c1ccccc1"]
molecules = [Molecule.from_smiles(s) for s in smiles_list]
dataset = esp.data.dataset.GraphDataset(molecules)

# Load pre-computed dataset
dataset = esp.data.dataset.GraphDataset.load("./my_dataset")

# Dataset operations
dataset.shuffle(seed=42)
ds_train, ds_valid, ds_test = dataset.split([8, 1, 1])  # 80/10/10 split

# Subsample for quick experiments
small_dataset = dataset.subsample(ratio=0.1, seed=42)

# Apply transformations
dataset.apply(lambda g: some_transform(g), in_place=True)

# Create DataLoader view for training
dataloader = dataset.view(
    collate_fn="graph",
    batch_size=32,
    shuffle=True
)

# Save dataset
dataset.save("./processed_dataset")
```

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

### Deploy OpenMM Systems from Graphs

Source: https://context7.com/choderalab/espaloma/llms.txt

Convert parameterized molecular graphs into OpenMM systems using various charge assignment methods.

```python
import espaloma as esp
import torch
from openff.toolkit.topology import Molecule

# Complete workflow: molecule -> graph -> parameters -> OpenMM system
molecule = Molecule.from_smiles("CN1C=NC2=C1C(=O)N(C(=O)N2C)C")
molecule_graph = esp.Graph(molecule)

# Load and apply model to assign parameters
espaloma_model = esp.get_model("latest")
espaloma_model(molecule_graph.heterograph)

# Create OpenMM system with different charge methods
# Using neural network predicted charges (default)
openmm_system = esp.graphs.deploy.openmm_system_from_graph(
    molecule_graph,
    forcefield="openff_unconstrained-2.1.1",
    charge_method="nn"
)

# Using AM1-BCC charges
openmm_system_am1bcc = esp.graphs.deploy.openmm_system_from_graph(
    molecule_graph,
    forcefield="openff_unconstrained-2.1.1",
    charge_method="am1-bcc"
)

# Using charges from the original molecule
openmm_system_from_mol = esp.graphs.deploy.openmm_system_from_graph(
    molecule_graph,
    charge_method="from-molecule"
)
```

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

### Download and Unzip Dataset

Source: https://github.com/choderalab/espaloma/blob/main/docs/qm_fitting.md

Downloads the specified dataset from a URL and unzips the contents. Ensure the dataset name variable is set correctly before running.

```python
%%capture
! wget "data.wangyq.net/esp_dataset/"$dataset_name".zip"
! unzip $dataset_name".zip"
```

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

### Load and split dataset

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Loads the dataset, shuffles it, and partitions it into training, validation, and test sets.

```python
ds = esp.data.dataset.GraphDataset.load("phalkethoh")
ds.shuffle(seed=2666)
ds_tr, ds_vl, ds_te = ds.split([8, 1, 1])
```

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

### Load Pre-trained Espaloma Models

Source: https://context7.com/choderalab/espaloma/llms.txt

Download and load pre-trained models from GitHub or local files. Local files require calling eval() after loading.

```python
import espaloma as esp
import torch

# Load the latest pre-trained espaloma model
espaloma_model = esp.get_model("latest")

# Load a specific version
espaloma_model = esp.get_model("0.3.2")

# Loading from a local .pt file requires calling eval()
espaloma_model = torch.load("espaloma-0.3.2.pt")
espaloma_model.eval()
```

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

### Deploying Espaloma force field from latest model

Source: https://github.com/choderalab/espaloma/blob/main/docs/deploy.md

Uses the Open Force Field toolkit to create a graph and apply the latest pretrained Espaloma model for OpenMM system generation.

```python
# imports
import os
import torch
import espaloma as esp

# define or load a molecule of interest via the Open Force Field toolkit
from openff.toolkit.topology import Molecule
molecule = Molecule.from_smiles("CN1C=NC2=C1C(=O)N(C(=O)N2C)C")

# create an Espaloma Graph object to represent the molecule of interest
molecule_graph = esp.Graph(molecule)

# load pretrained model
espaloma_model = esp.get_model("latest")

# apply a trained espaloma model to assign parameters
espaloma_model(molecule_graph.heterograph)

# create an OpenMM System for the specified molecule
openmm_system = esp.graphs.deploy.openmm_system_from_graph(molecule_graph)
```

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

### Reproduce Paper Environment (Linux-64)

Source: https://github.com/choderalab/espaloma/blob/main/scripts/perses-benchmark/README.md

This command allows you to recreate the exact Conda environment used in the paper for Linux-64 systems.

```bash
conda env create -n espaloma-perses -f espaloma-perses.export.yaml
```

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

### Define Training Metrics with GraphMetric

Source: https://context7.com/choderalab/espaloma/llms.txt

Shows how to define training metrics using `esp.metrics.GraphMetric` for graph-level or node-level attributes, including energy loss, force loss, and bond parameter metrics.

```python
import espaloma as esp
import torch

# Energy loss: MSE between predicted and reference energies
energy_loss = esp.metrics.GraphMetric(
    base_metric=torch.nn.MSELoss(),
    between=["u", "u_ref"],
    level="g"  # graph level
)

# Force loss: RMSE on gradients
force_loss = esp.metrics.GraphMetric(
    base_metric=esp.metrics.rmse,
    between=["u_prime", "u_ref_prime"],
    level="n1"  # atom level
)

# Bond parameter metrics
bond_k_metric = esp.metrics.GraphMetric(
    base_metric=torch.nn.MSELoss(),
    between=["k", "k_ref"],
    level="n2"  # bond level
)

# Derivative metrics for force matching
gradient_metric = esp.metrics.GraphHalfDerivativeMetric(
    base_metric=torch.nn.MSELoss(),
    input_level="g",
    input_name="u",
    target_prime_level="n1",
    target_prime_name="u_ref_prime",
    weight=1.0
)

# Center energies for conformational analysis
centered_loss = esp.metrics.center(torch.nn.L1Loss())
```

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

### Create training dataloader

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Constructs a dataloader for the training set with a specified batch size.

```python
ds_tr_loader = ds_tr.view(batch_size=100, shuffle=True)
```

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

### Select Dataset Name

Source: https://github.com/choderalab/espaloma/blob/main/docs/qm_fitting.md

Defines the name of the dataset to be loaded. Available options include 'gen2', 'pepconf', 'vehicle', and 'phalkethoh'.

```python
dataset_name = "gen2"
# dataset_name = "pepconf"
# dataset_name = "vehicle"
# dataset_name = "phalkethoh"
```

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

### Import Core Libraries

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/typing.md

Imports essential Python libraries for the experiment, including PyTorch for deep learning, DGL for graph neural networks, and NumPy for numerical operations.

```python
import torch
import dgl
import numpy as np
```

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

### Clone and Import Espaloma

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/qm_fitting.md

Clones the repository and imports the core Espaloma module.

```python
! git clone https://github.com/choderalab/espaloma.git
```

```python
import torch
import sys
sys.path.append("/content/espaloma")
import espaloma as esp
```

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

### Class and Method Documentation Template

Source: https://github.com/choderalab/espaloma/blob/main/docs/_templates/custom-class-template.rst

This documentation outlines the structure for documenting classes, methods, and attributes within the Espaloma project.

```APIDOC
## Class Documentation Structure

### Description
This template defines how classes, their methods, and attributes are documented within the Espaloma project using Sphinx-style directives.

### Components
- **autoclass**: Used to document the class, including members and inherited members.
- **automethod**: Used to document specific methods, starting with the constructor (__init__).
- **autosummary**: Used to generate tables for methods and attributes.

### Usage
- **Methods**: Listed under the 'Methods' rubric.
- **Attributes**: Listed under the 'Attributes' rubric.
```

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

### Loading a local Espaloma model file

Source: https://github.com/choderalab/espaloma/blob/main/docs/deploy.md

Loads a model from a local .pt file and sets it to evaluation mode for inference.

```python
# load local pretrained model
espaloma_model = torch.load("espaloma-0.3.2.pt")
espaloma_model.eval()
```

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

### Create and Manage Molecular Graphs

Source: https://context7.com/choderalab/espaloma/llms.txt

Initialize graph objects from SMILES or OpenFF molecules and access underlying DGL representations or save/load state.

```python
import espaloma as esp
from openff.toolkit.topology import Molecule

# Create graph from SMILES string (caffeine)
graph_from_smiles = esp.Graph("CN1C=NC2=C1C(=O)N(C(=O)N2C)C")

# Create graph from OpenFF Molecule object
molecule = Molecule.from_smiles("CCO")  # ethanol
graph_from_mol = esp.Graph(molecule)

# Access underlying graph representations
heterograph = graph_from_smiles.heterograph  # DGL heterogeneous graph
homograph = graph_from_smiles.homograph      # DGL homogeneous graph
mol = graph_from_smiles.mol                  # OpenFF Molecule

# Save and load graphs
graph_from_smiles.save("./my_molecule")
loaded_graph = esp.Graph.load("./my_molecule")
```

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

### Download and extract dataset

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Fetches the PhAlkEthOH dataset archive and extracts it.

```python
%%capture
! wget http://data.wangyq.net/esp_dataset/phalkethoh_mm_small.zip
! unzip phalkethoh_mm_small.zip
```

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

### Split and Batch Dataset

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/typing.md

Splits the dataset into training, testing, and validation sets with an 80:10:10 ratio. It then prepares these sets for model training by batching them with a specified batch size and enabling shuffling for the training set.

```python
ds_tr, ds_te, ds_vl = ds.split([8, 1, 1])
```

```python
ds_tr = ds_tr.view('graph', batch_size=100, shuffle=True)
ds_te = ds_te.view('graph', batch_size=100)
ds_vl = ds_vl.view('graph', batch_size=100)
```

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

### esp.get_model()

Source: https://context7.com/choderalab/espaloma/llms.txt

Downloads and loads pre-trained Espaloma models for force field parameter prediction.

```APIDOC
## esp.get_model(version)

### Description
Loads a pre-trained Espaloma model from GitHub releases or local files.

### Parameters
- **version** (string) - Required - The version string (e.g., 'latest', '0.3.2') or path to a local .pt file.

### Request Example
```python
import espaloma as esp
espaloma_model = esp.get_model("latest")
```
```

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

### Analyze benchmark results

Source: https://github.com/choderalab/espaloma/blob/main/scripts/perses-benchmark/tyk2/espaloma-0.2.2/README.md

Generate CSV output and free energy plots for a specific target protein using the analysis script.

```python
python benchmark_analysis.py --target [protein-name]
```

```python
python benchmark_analysis.py --target tyk2
```

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

### Execute Training Loop

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Iterates through the dataset to update model weights and save checkpoints.

```python
for idx_epoch in range(10000):
    for g in ds_tr_loader:
        optimizer.zero_grad()
        if torch.cuda.is_available():
            g = g.to("cuda:0")
        g = espaloma_model(g)
        loss = loss_fn(g)
        loss.backward()
        optimizer.step()
        torch.save(espaloma_model.state_dict(), "%s.th" % idx_epoch)
```

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

### Initialize Inspection Metric

Source: https://github.com/choderalab/espaloma/blob/main/docs/qm_fitting.md

Sets up a metric for inspecting model performance, using L1Loss centered for evaluation. This metric will be used to calculate performance on training and validation sets.

```python
inspect_metric = esp.metrics.center(torch.nn.L1Loss()) # use mean-squared error loss
```

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

### Create Espaloma-Perses Conda Environment

Source: https://github.com/choderalab/espaloma/blob/main/scripts/perses-benchmark/README.md

Use this command to create a Conda environment for espaloma and perses using the specified YAML file.

```bash
conda env create -n espaloma-perses -f espaloma-perses.yaml
```

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

### Move Data to GPU

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Transfers validation and training graphs to GPU.

```python
if torch.cuda.is_available():
    g_vl = g_vl.to("cuda:0")
    g_tr = g_tr.to("cuda:0")
```

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

### Clone Espaloma repository

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Downloads the Espaloma source code from GitHub.

```python
! git clone https://github.com/choderalab/espaloma.git
```

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

### Define and Assemble Espaloma Model

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Configures the JanossyPooling readout and composes the full end-to-end model with geometry and energy components.

```python
readout = esp.nn.readout.janossy.JanossyPooling(
    in_features=128, config=[128, "relu", 128, "relu", 128, "relu"],
    out_features={              # define modular MM parameters Espaloma will assign
        1: {"e": 1, "s": 1}, # atom hardness and electronegativity
        2: {"log_coefficients": 2}, # bond linear combination, enforce positive
        3: {"log_coefficients": 2}, # angle linear combination, enforce positive
        4: {"k": 6}, # torsion barrier heights (can be positive or negative)
    },
)

espaloma_model = torch.nn.Sequential(
                 representation, readout, esp.nn.readout.janossy.ExpCoefficients(),
                 esp.mm.geometry.GeometryInGraph(),
                 esp.mm.energy.EnergyInGraph(),
                 esp.mm.energy.EnergyInGraph(suffix="_ref"),
                 esp.nn.readout.charge_equilibrium.ChargeEquilibrium(),
)
```

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

### Inspect and Visualize Performance

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/qm_fitting.md

Evaluates the model on training and validation sets and plots the loss curves.

```python
inspect_metric = esp.metrics.center(torch.nn.L1Loss()) # use mean-squared error loss
```

```python
loss_tr = []
loss_vl = []
```

```python
with torch.no_grad():
    for idx_epoch in range(10000):
        espaloma_model.load_state_dict(
            torch.load("%s.th" % idx_epoch)
        )

        # training set performance
        u = []
        u_ref = []
        for g in ds_tr:
            if torch.cuda.is_available():
                g.heterograph = g.heterograph.to("cuda:0")
            espaloma_model(g.heterograph)
            u.append(g.nodes['g'].data['u'])
            u_ref.append(g.nodes['g'])
        u = torch.cat(u, dim=0)
        u_ref = torch.cat(u_ref, dim=0)
        loss_tr.append(inspect_metric(u, u_ref))


        # validation set performance
        u = []
        u_ref = []
        for g in ds_vl:
            if torch.cuda.is_available():
                g.heterograph = g.heterograph.to("cuda:0")
            espaloma_model(g.heterograph)
            u.append(g.nodes['g'].data['u'])
            u_ref.append(g.nodes['g'])
        u = torch.cat(u, dim=0)
        u_ref = torch.cat(u_ref, dim=0)
        loss_vl.append(inspect_metric(u, u_ref))
```

```python
import numpy as np
loss_tr = np.array(loss_tr) * 627.5
loss_vl = np.array(loss_vl) * 627.5
```

```python
from matplotlib import pyplot as plt
plt.plot(loss_tr, label="train")
plt.plot(loss_vl, label="valid")
plt.yscale("log")
plt.legend()
```

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

### Train the Model

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/qm_fitting.md

Executes the training loop over the dataset and saves model checkpoints.

```python
for idx_epoch in range(10000):
    for g in ds_tr:
        optimizer.zero_grad()
        if torch.cuda.is_available():
            g.heterograph = g.heterograph.to("cuda:0")
        g = espaloma_model(g.heterograph)
        loss = loss_fn(g)
        loss.backward()
        optimizer.step()
    torch.save(espaloma_model.state_dict(), "%s.th" % idx_epoch)
```

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

### Run single benchmark edge

Source: https://github.com/choderalab/espaloma/blob/main/scripts/perses-benchmark/tyk2/espaloma-0.2.2/README.md

Execute a specific transformation edge for a target protein. Calculations can be resumed by re-running the same command if they fail.

```python
python run_benchmarks.py --target [protein-name] --edge [edge-index]
```

```python
# Set up and run edge 6
python run_benchmarks.py --target tyk2 --edge 6
```

```python
# Resume failed edge 6
python run_benchmarks.py --target tyk2 --edge 6
```

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

### Load and Split Dataset

Source: https://github.com/choderalab/espaloma/blob/main/docs/qm_fitting.md

Loads a molecular dataset using GraphDataset, shuffles it, and then splits it into training, validation, and testing sets with specified proportions.

```python
ds = esp.data.dataset.GraphDataset.load(dataset_name)
ds.shuffle(seed=2666)
ds_tr, ds_vl, ds_te = ds.split([8, 1, 1])
```

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

### Run LSF batch job

Source: https://github.com/choderalab/espaloma/blob/main/scripts/perses-benchmark/tyk2/espaloma-0.2.2/README.md

Submit a benchmark job to an LSF batch scheduler using the provided template.

```bash
bsub < LSF-job-template.sh
```

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

### Plot Training Results

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Visualizes training and validation loss curves using matplotlib.

```python
from matplotlib import pyplot as plt
plt.plot(loss_tr, label="train")
plt.plot(loss_vl, label="valid")
plt.yscale("log")
plt.legend()
```

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

### Evaluate Model Checkpoints

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Loads saved model states and calculates losses for training and validation sets.

```python
for idx_epoch in range(10000):
    espaloma_model.load_state_dict(
        torch.load("%s.th" % idx_epoch)
    )

    espaloma_model(g_tr)
    loss_tr.append(inspect_metric(g_tr).item())

    espaloma_model(g_vl)
    loss_vl.append(inspect_metric(g_vl).item())
```

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

### esp.Graph()

Source: https://context7.com/choderalab/espaloma/llms.txt

Creates a graph representation of a molecule for GNN processing.

```APIDOC
## esp.Graph(molecule)

### Description
Initializes an Espaloma Graph object from a SMILES string or an OpenFF Molecule object.

### Parameters
- **molecule** (string/Molecule) - Required - Input molecule as a SMILES string or OpenFF Molecule object.

### Request Example
```python
import espaloma as esp
graph = esp.Graph("CN1C=NC2=C1C(=O)N(C(=O)N2C)C")
```
```

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

### Move Model to GPU

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Transfers the model to CUDA device if available.

```python
if torch.cuda.is_available():
    espaloma_model = espaloma_model.cuda()
```

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

### Define Readout and Full Model

Source: https://github.com/choderalab/espaloma/blob/main/docs/qm_fitting.md

Defines the readout module for generating molecular mechanics parameters and composes all Espaloma stages into an end-to-end model. The readout module is configured to output parameters for atom, bond, angle, and torsion potentials.

```python
readout = esp.nn.readout.janossy.JanossyPooling(
    in_features=128, config=[128, "relu", 128, "relu", 128, "relu"],
    out_features={
        1: {"e": 1, "s": 1}, # atom hardness and electronegativity
        2: {"log_coefficients": 2}, # bond linear combination, enforce positive
        3: {"log_coefficients": 2}, # angle linear combination, enforce positive
        4: {"k": 6}, # torsion barrier heights (can be positive or negative)
    },
)

espaloma_model = torch.nn.Sequential(
                 representation, readout, esp.nn.readout.janossy.ExpCoefficients(),
                 esp.mm.geometry.GeometryInGraph(),
                 esp.mm.energy.EnergyInGraph(),
)
```

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

### Define Optimizer

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Initializes the Adam optimizer for the model parameters.

```python
optimizer = torch.optim.Adam(espaloma_model.parameters(), 1e-4)
```

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

### Load Molecular Dataset

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/typing.md

Downloads and loads the ZINC molecular dataset if it does not already exist. This dataset is used for training and evaluating the atom typing model.

```python
import os
if not os.path.exists("zinc"):
    os.system("wget data.wangyq.net/esp_datasets/zinc")
ds = esp.data.dataset.GraphDataset.load("zinc")
```

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

### Configure HTML Static Path in Sphinx

Source: https://github.com/choderalab/espaloma/blob/main/docs/_templates/README.md

Set the path to the templates directory in your Sphinx `conf.py` file. This allows Sphinx to find and use your custom templates.

```python
html_static_path = ['_templates']
```

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

### Configure static templates path in Sphinx

Source: https://github.com/choderalab/espaloma/blob/main/docs/_static/README.md

Defines the directory for custom static files relative to the conf.py file.

```python
templates_path = ['_static']
```

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

### Calculate Training and Validation Loss

Source: https://github.com/choderalab/espaloma/blob/main/docs/qm_fitting.md

Iterates through epochs to load saved model states, calculate performance metrics (L1 loss) on both training and validation datasets, and stores these losses for later analysis.

```python
loss_tr = []
loss_vl = []

with torch.no_grad():
    for idx_epoch in range(10000):
        espaloma_model.load_state_dict(
            torch.load("%s.th" % idx_epoch)
        )

        # training set performance
        u = []
        u_ref = []
        for g in ds_tr:
            if torch.cuda.is_available():
                g.heterograph = g.heterograph.to("cuda:0")
            espaloma_model(g.heterograph)
            u.append(g.nodes['g'].data['u'])
            u_ref.append(g.nodes['g'])
        u = torch.cat(u, dim=0)
        u_ref = torch.cat(u_ref, dim=0)
        loss_tr.append(inspect_metric(u, u_ref))


        # validation set performance
        u = []
        u_ref = []
        for g in ds_vl:
            if torch.cuda.is_available():
                g.heterograph = g.heterograph.to("cuda:0")
            espaloma_model(g.heterograph)
            u.append(g.nodes['g'].data['u'])
            u_ref.append(g.nodes['g'])
        u = torch.cat(u, dim=0)
        u_ref = torch.cat(u_ref, dim=0)
        loss_vl.append(inspect_metric(u, u_ref))
```

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

### Parameter Readout with Janossy Pooling

Source: https://context7.com/choderalab/espaloma/llms.txt

Aggregate atom representations into bond, angle, and torsion parameters.

```python
import espaloma as esp
import torch
```

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

### esp.nn.Sequential()

Source: https://context7.com/choderalab/espaloma/llms.txt

Constructs the message-passing GNN for atom-level latent representations.

```APIDOC
## esp.nn.Sequential(layer, config, feature_units, input_units)

### Description
Defines the Stage I graph neural network architecture.

### Parameters
- **layer** (object) - Required - The GNN layer type (e.g., SAGEConv).
- **config** (list) - Required - List of hidden units and activation functions.
- **feature_units** (int) - Optional - Input atom feature dimension.
- **input_units** (int) - Optional - Initial embedding dimension.

### Request Example
```python
representation = esp.nn.Sequential(
    layer=esp.nn.layers.dgl_legacy.gn("SAGEConv"),
    config=[128, "relu", 128, "relu"]
)
```
```

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

### Build Neural Network Models

Source: https://context7.com/choderalab/espaloma/llms.txt

Construct message-passing GNNs for atom-level latent representations using various layer types.

```python
import espaloma as esp
import torch

# Define Stage I: Graph neural network for atom embeddings
# Uses SAGEConv layers with 128 hidden units and ReLU activations
representation = esp.nn.Sequential(
    layer=esp.nn.layers.dgl_legacy.gn("SAGEConv"),
    config=[128, "relu", 128, "relu", 128, "relu"],
    feature_units=117,  # input atom features
    input_units=128     # initial embedding dimension
)

# Alternative architectures using different GNN layers
representation_gat = esp.nn.Sequential(
    layer=esp.nn.layers.dgl_legacy.gn("GATConv"),
    config=[64, "relu", 64, "relu", 64, "relu"],
)

representation_gin = esp.nn.Sequential(
    layer=esp.nn.layers.dgl_legacy.gn("GINConv"),
    config=[256, "relu", 256, "relu"],
)
```

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

### Extract dataset batches

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Retrieves the full training and validation sets as single batches.

```python
g_tr = next(iter(ds_tr.view(batch_size=len(ds_tr))))
g_vl = next(iter(ds_vl.view(batch_size=len(ds_vl))))
```

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

### Define Espaloma Model Architecture

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/qm_fitting.md

Constructs the graph representation and readout stages, then composes them into a full end-to-end model.

```python
representation = esp.nn.Sequential(
    layer=esp.nn.layers.dgl_legacy.gn("SAGEConv"), # use SAGEConv implementation in DGL
    config=[128, "relu", 128, "relu", 128, "relu"], # 3 layers, 128 units, ReLU activation
)
```

```python
readout = esp.nn.readout.janossy.JanossyPooling(
    in_features=128, config=[128, "relu", 128, "relu", 128, "relu"],
    out_features={              # define modular MM parameters Espaloma will assign
        1: {"e": 1, "s": 1}, # atom hardness and electronegativity
        2: {"log_coefficients": 2}, # bond linear combination, enforce positive
        3: {"log_coefficients": 2}, # angle linear combination, enforce positive
        4: {"k": 6}, # torsion barrier heights (can be positive or negative)
    },
)

espaloma_model = torch.nn.Sequential(
                 representation, readout, esp.nn.readout.janossy.ExpCoefficients(),
                 esp.mm.geometry.GeometryInGraph(),
                 esp.mm.energy.EnergyInGraph(),
)
```

```python
if torch.cuda.is_available():
    espaloma_model = espaloma_model.cuda()
```

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

### Define Espaloma Representation Layer

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Initializes the graph neural network representation stage using SAGEConv layers.

```python
representation = esp.nn.Sequential(
    layer=esp.nn.layers.dgl_legacy.gn("SAGEConv"), # use SAGEConv implementation in DGL
    config=[128, "relu", 128, "relu", 128, "relu"], # 3 layers, 128 units, ReLU activation
)
```

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

### Define Graph Neural Network Model

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/typing.md

Defines a graph neural network model using SAGEConv layers. The model consists of a sequential representation with three layers, each having 128 units and ReLU activation, followed by a node typing readout layer.

```python
# define a layer
layer = esp.nn.layers.dgl_legacy.gn("SAGEConv")

# define a representation
representation = esp.nn.Sequential(
        layer,
        [128, "relu", 128, "relu", 128, "relu"],
)

# define a readout
readout = esp.nn.readout.node_typing.NodeTyping(
        in_features=128,
        n_classes=100
)

net = torch.nn.Sequential(
    representation,
    readout
)
```

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

### Energy Calculation with esp.mm.energy.energy_in_graph()

Source: https://context7.com/choderalab/espaloma/llms.txt

Computes molecular mechanics energy from graph parameters including bonded and nonbonded terms. Ensure geometry is computed first and the model has been applied to assign parameters. Access computed energies via node data keys like 'u', 'u_n2', 'u_n3', 'u_n4'.

```python
import espaloma as esp
import torch

# After applying model to assign parameters
molecule_graph = esp.Graph("CCO")
espaloma_model = esp.get_model("latest")
espaloma_model(molecule_graph.heterograph)

# Compute geometry (bond lengths, angles, dihedrals)
esp.mm.geometry.geometry_in_graph(molecule_graph.heterograph)

# Compute energy with all bonded terms
esp.mm.energy.energy_in_graph(
    molecule_graph.heterograph,
    suffix="",
    terms=["n2", "n3", "n4", "n4_improper"]
)

# Access computed energies
total_energy = molecule_graph.heterograph.nodes["g"].data["u"]
bond_energy = molecule_graph.heterograph.nodes["g"].data["u_n2"]
angle_energy = molecule_graph.heterograph.nodes["g"].data["u_n3"]
torsion_energy = molecule_graph.heterograph.nodes["g"].data["u_n4"]

# Use as PyTorch module in model pipeline
energy_module = esp.mm.energy.EnergyInGraph(
    terms=["n2", "n3", "n4"]
)
```

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

### Train the GNN Model

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/typing.md

Trains the defined graph neural network model using the Adam optimizer for 3000 epochs. In each epoch, it iterates through the training dataset, performs a forward pass, calculates the loss, and updates the model's weights.

```python
# define optimizer
optimizer = torch.optim.Adam(net.parameters(), 1e-5)

# train the model
for _ in range(3000):
    for g in ds_tr:
        optimizer.zero_grad()
        net(g.heterograph)
        loss = loss_fn(g.heterograph)
        loss.backward()
        optimizer.step()
```

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

### Define Inspection Metric

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Sets up an L1 loss metric for evaluating model performance.

```python
inspect_metric = esp.metrics.GraphMetric(
        base_metric=torch.nn.L1Loss(), # use mean-squared error loss
        between=['u', "u_ref"],         # between predicted and QM energies
        level="g", # compare on graph level
)
```

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

### Define Loss Function

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Sets up the Mean Squared Error loss for comparing predicted and reference energies at the graph level.

```python
loss_fn = esp.metrics.GraphMetric(
        base_metric=torch.nn.MSELoss(), # use mean-squared error loss
        between=['u', "u_ref"],         # between predicted and QM energies
        level="g", # compare on graph level
)
```

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

### Convert and Scale Losses

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/mm_fitting_small.md

Converts loss lists to numpy arrays and scales by 627.5.

```python
import numpy as np
loss_tr = np.array(loss_tr) * 627.5
loss_vl = np.array(loss_vl) * 627.5
```

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

### Geometry Computation with esp.mm.geometry.geometry_in_graph()

Source: https://context7.com/choderalab/espaloma/llms.txt

Calculates internal coordinates (bond lengths, angles, dihedrals) from Cartesian coordinates. Requires coordinates to be set in the 'xyz' node data. The computed values are stored in 'x' data for nodes 'n2', 'n3', and 'n4'. Can be used as a PyTorch module.

```python
import espaloma as esp
import torch
import numpy as np

# Create graph with coordinates
molecule_graph = esp.Graph("CCO")

# Set atomic coordinates (in angstroms)
coords = torch.tensor([
    [-1.5, 0.0, 0.0],  # C
    [0.0, 0.0, 0.0],   # C
    [0.75, 1.2, 0.0],  # O
    # ... hydrogen positions
], dtype=torch.float32)

molecule_graph.heterograph.nodes["n1"].data["xyz"] = coords

# Compute all geometric quantities
esp.mm.geometry.geometry_in_graph(molecule_graph.heterograph)

# Access computed values
bond_lengths = molecule_graph.heterograph.nodes["n2"].data["x"]  # in distance units
bond_angles = molecule_graph.heterograph.nodes["n3"].data["x"]   # in radians
torsion_angles = molecule_graph.heterograph.nodes["n4"].data["x"] # in radians

# Use as module in differentiable pipeline
geometry_module = esp.mm.geometry.GeometryInGraph()
```

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

### Define Loss Function

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/typing.md

Initializes the TypingAccuracy metric as the loss function for training the atom typing model. This metric is used to evaluate the accuracy of the predicted atom types.

```python
loss_fn = esp.metrics.TypingAccuracy()
```

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

### Define Janossy Pooling Readout and Compose Model

Source: https://context7.com/choderalab/espaloma/llms.txt

Defines Stage II-III Janossy pooling for parameter prediction and composes a full Espaloma model including representation, readout, and geometry/energy modules.

```python
readout = esp.nn.readout.janossy.JanossyPooling(
    in_features=128,
    config=[128, "relu", 128, "relu", 128, "relu"],
    out_features={
        1: {"e": 1, "s": 1, "q": 1},  # atom: hardness, electronegativity, charge
        2: {"k": 1, "eq": 1},          # bond: force constant, equilibrium length
        3: {"k": 1, "eq": 1},          # angle: force constant, equilibrium angle
        4: {"k": 6},                    # proper torsion: 6 barrier heights
    },
)

# For improper torsions, use dedicated pooling
improper_readout = esp.nn.readout.janossy.JanossyPoolingImproper(
    in_features=128,
    config=[128, "relu", 128, "relu"],
    out_features={"k": 6}
)

# Compose full espaloma model
espaloma_model = torch.nn.Sequential(
    representation,
    readout,
    improper_readout,
    esp.mm.geometry.GeometryInGraph(),
    esp.mm.energy.EnergyInGraph(),
)
```

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

### Assign Legacy Atom Typing

Source: https://github.com/choderalab/espaloma/blob/main/docs/experiments/typing.md

Applies the GAFF-1.81 force field to assign legacy atom types to the molecules in the dataset. This operation modifies the dataset in place.

```python
typing = esp.graphs.legacy_force_field.LegacyForceField('gaff-1.81')
ds.apply(typing, in_place=True) # this modify the original data
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.