Batchwise optimizer performance improvements

docs/api/atomistic.rst

@@ -37,4 +37,5 @@ Output modules
     CoulombPotential
     DampedCoulombPotential
     EnergyCoulomb
-    EnergyEwald
+    EnergyEwald
+    HarmonicBond

examples/howtos/howto_priors.ipynb

@@ -0,0 +1,277 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "23179393",
+   "metadata": {},
+   "source": [
+    "# Adding a Prior to the Model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bb8b6bfc",
+   "metadata": {},
+   "source": [
+    "Sometimes the trained potential is not the whole energy: a restraint holding a bond at a\n",
+    "given length, a repulsive term the training data never covered, a bias pushing a\n",
+    "relaxation somewhere it would not go on its own. Such a term belongs to the *model*, not\n",
+    "to the calculator that evaluates it. Once it is part of the energy, the response module\n",
+    "differentiates it along with everything else, and the forces derived from it pick it up\n",
+    "on their own.\n",
+    "\n",
+    "SchNetPack's own priors are built exactly this way -- see `ZBLRepulsionEnergy`, the\n",
+    "Coulomb modules, or `HarmonicBond` in `schnetpack.atomistic`. The prior is an output\n",
+    "module that writes its energy under its own `output_key`, and an `Aggregation` module\n",
+    "sums the terms into the total energy."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2cbc4b6f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import numpy as np\n",
+    "from copy import deepcopy\n",
+    "\n",
+    "from ase.io import read\n",
+    "from ase.optimize import LBFGS\n",
+    "\n",
+    "import schnetpack as spk\n",
+    "from schnetpack import properties\n",
+    "from schnetpack.interfaces.ase_interface import SpkCalculator, AtomsConverter\n",
+    "from schnetpack.utils.compatibility import load_model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "24fc7a58",
+   "metadata": {},
+   "source": [
+    "We load the force field model and wrap it in a `SpkCalculator`, which evaluates one ASE `Atoms` object per call."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ec9f9939",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_path = \"../../tests/testdata/md_ethanol.model\"\n",
+    "\n",
+    "# set device\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "\n",
+    "# load model\n",
+    "model = load_model(model_path, device=device)\n",
+    "cutoff = model.representation.cutoff.item()\n",
+    "\n",
+    "ENERGY_UNIT = \"kcal/mol\"\n",
+    "POSITION_UNIT = \"Ang\"\n",
+    "FMAX = 0.001  # converged once no force exceeds this, in eV/Ang\n",
+    "MAX_STEPS = 1000  # give up after this many optimizer steps\n",
+    "\n",
+    "\n",
+    "def spk_calculator(model=model_path):\n",
+    "    return SpkCalculator(\n",
+    "        model=model,\n",
+    "        neighbor_list=spk.transform.MatScipyNeighborList(cutoff=cutoff),\n",
+    "        device=device,\n",
+    "        energy_unit=ENERGY_UNIT,\n",
+    "        position_unit=POSITION_UNIT,\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a2ac9298",
+   "metadata": {},
+   "source": [
+    "The demo below relaxes ethanol's C-O bond starting from a few different conformers."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fefd8058",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "input_structure_file = \"../../tests/testdata/ethanol_conformers.xyz\"\n",
+    "\n",
+    "# load the starting structures\n",
+    "conformers = read(input_structure_file, index=\":\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b520c0cb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def with_prior(model, prior, prior_key, energy_key=properties.energy):\n",
+    "    \"\"\"A copy of ``model`` whose energy carries an extra term, forces included.\n",
+    "\n",
+    "    The response module differentiates the total energy, so everything contributing to\n",
+    "    that energy has to run ahead of it: the prior and the ``Aggregation`` adding it to the\n",
+    "    model's own prediction are spliced in just before it.\n",
+    "    \"\"\"\n",
+    "    model = deepcopy(model)\n",
+    "    modules = list(model.output_modules)\n",
+    "\n",
+    "    response = next(\n",
+    "        idx\n",
+    "        for idx, module in enumerate(modules)\n",
+    "        if isinstance(module, (spk.atomistic.Forces, spk.atomistic.Response))\n",
+    "    )\n",
+    "    aggregation = spk.atomistic.Aggregation(\n",
+    "        keys=[energy_key, prior_key], output_key=energy_key\n",
+    "    )\n",
+    "    model.output_modules = nn.ModuleList(\n",
+    "        modules[:response] + [prior, aggregation] + modules[response:]\n",
+    "    )\n",
+    "\n",
+    "    # the model caches what its modules require and produce, so both have to be redone\n",
+    "    model.collect_derivatives()\n",
+    "    model.collect_outputs()\n",
+    "    return model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8c0e7d36",
+   "metadata": {},
+   "source": [
+    "Ethanol comes out of the file as `C C O H H H H H H`, so atoms `(0, 2)` are the C-O bond.\n",
+    "It relaxes to about 1.43 Å on its own; we restrain it to 1.8 Å and relax the same\n",
+    "structures twice, once with the plain model and once with the composed one, each\n",
+    "structure one at a time with ASE's `LBFGS`.\n",
+    "\n",
+    "`with_prior` re-runs `collect_derivatives` and `collect_outputs`, the two caches\n",
+    "`NeuralNetworkPotential` fills in at construction time. The second one is why\n",
+    "`energy_bond` turns up among the model outputs below, next to `energy` and `forces`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "94657210",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "BOND = (0, 2)  # the C-O bond of ethanol\n",
+    "BOND_KEY = \"energy_bond\"\n",
+    "\n",
+    "restrained_model = with_prior(\n",
+    "    model,\n",
+    "    spk.atomistic.HarmonicBond(\n",
+    "        atom_pair=BOND,\n",
+    "        bond_length=1.8,  # Angstrom, well beyond the equilibrium C-O distance\n",
+    "        force_constant=50.0,  # eV / Angstrom**2\n",
+    "        energy_unit=ENERGY_UNIT,\n",
+    "        position_unit=POSITION_UNIT,\n",
+    "        output_key=BOND_KEY,\n",
+    "    ),\n",
+    "    prior_key=BOND_KEY,\n",
+    ")\n",
+    "\n",
+    "print(\"output modules:\", [type(m).__name__ for m in restrained_model.output_modules])\n",
+    "print(\"model outputs: \", restrained_model.model_outputs)\n",
+    "\n",
+    "\n",
+    "def relaxed_bond_lengths(structures, model):\n",
+    "    \"\"\"Relax every structure with ``model``, one at a time, and report its C-O bond length.\"\"\"\n",
+    "    calculator = spk_calculator(model)\n",
+    "    lengths = []\n",
+    "    for structure in structures:\n",
+    "        atoms = structure.copy()\n",
+    "        atoms.calc = calculator\n",
+    "        LBFGS(atoms, logfile=None).run(fmax=FMAX, steps=MAX_STEPS)\n",
+    "        lengths.append(atoms.get_distance(*BOND))\n",
+    "    return np.array(lengths)\n",
+    "\n",
+    "\n",
+    "free_bonds = relaxed_bond_lengths(conformers, model_path)\n",
+    "restrained_bonds = relaxed_bond_lengths(conformers, restrained_model)\n",
+    "\n",
+    "print(f\"\\n{'':<14}{'C-O bond [Ang]':>16}\")\n",
+    "for label, bonds in [(\"unrestrained\", free_bonds), (\"restrained\", restrained_bonds)]:\n",
+    "    print(f\"{label:<14}{bonds.mean():>16.3f}\")\n",
+    "print(f\"{'target':<14}{1.8:>16.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "393bf175",
+   "metadata": {},
+   "source": [
+    "The restrained bond does not land exactly on the target, and it should not: the restraint\n",
+    "and the model pull against each other, and the relaxation stops where the two balance. At\n",
+    "that point the restraint still pulls outward with `2 * force_constant * (target - d)`, and\n",
+    "that is precisely what the model's own bond force cancels. A stiffer `force_constant`\n",
+    "moves the result closer to the target, a softer one leaves the model more say.\n",
+    "\n",
+    "The energy decomposition below makes the same point from the other side. `energy` is the\n",
+    "sum the forces are taken from, and `energy_bond` is the restraint's share of it. One\n",
+    "caveat when reading the numbers: the model's `AddOffsets` postprocessor runs after the\n",
+    "output modules, so the reported total carries the atomref/mean offset while the individual\n",
+    "terms do not."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b987ec9b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# a one-shot conversion this time, so a plain converter with a plain neighbor list --\n",
+    "# nothing is being propagated here, so there is no list worth keeping around\n",
+    "converter = AtomsConverter(\n",
+    "    neighbor_list=spk.transform.MatScipyNeighborList(cutoff=cutoff), device=device\n",
+    ")\n",
+    "plain_results = model(converter(deepcopy(conformers)))\n",
+    "prior_results = restrained_model(converter(deepcopy(conformers)))\n",
+    "\n",
+    "print(f\"{'structure':>10}{'energy_bond':>14}{'total energy':>14}   [{ENERGY_UNIT}]\")\n",
+    "for idx, (bond, total) in enumerate(\n",
+    "    zip(prior_results[BOND_KEY], prior_results[properties.energy])\n",
+    "):\n",
+    "    print(f\"{idx:>10}{bond.item():>14.2f}{total.item():>14.2f}\")\n",
+    "\n",
+    "# and the forces really changed, which is what drove the relaxation above\n",
+    "shift = (\n",
+    "    (prior_results[properties.forces] - plain_results[properties.forces]).abs().max()\n",
+    ")\n",
+    "print(f\"\\nlargest force change: {shift.item():.1f} {ENERGY_UNIT}/{POSITION_UNIT}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "schnetpack",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

pyproject.toml

@@ -53,6 +53,13 @@ dependencies = [
 [project.optional-dependencies]
 test = ["pytest", "pytest-datadir", "pytest-benchmark"]
 
+[tool.pytest.ini_options]
+markers = [
+    "benchmark_sweep: batch-size scaling benchmarks; run with -m benchmark_sweep",
+]
+# the sweep takes far too long for an ordinary test run, so it is opt-in
+addopts = "-m 'not benchmark_sweep'"
+
 [tool.setuptools]
 package-dir = { "" = "src" }
 script-files = [

src/schnetpack/atomistic/__init__.py

@@ -5,3 +5,4 @@
 from .electrostatic import *
 from .aggregation import *
 from .external_fields import *
+from .priors import *

src/schnetpack/atomistic/priors.py

@@ -0,0 +1,73 @@
+import torch
+import torch.nn as nn
+from typing import Dict, Tuple
+
+import schnetpack.properties as properties
+import schnetpack.units as spk_units
+
+__all__ = ["HarmonicBond"]
+
+
+class HarmonicBond(nn.Module):
+    """
+    A harmonic restraint on the distance between two atoms of every structure.
+
+    Follows the shape of schnetpack's other priors, e.g. ``ZBLRepulsionEnergy``: the
+    module contributes its own energy term under ``output_key`` and leaves the summing
+    to an ``Aggregation`` module, rather than writing into the model's energy directly.
+
+    The restraint is stated in eV and Angstrom and converted to whatever units the model
+    works in. ``atom_pair`` indexes into a single structure, so the two atoms have to
+    carry the same indices in every structure of the batch.
+
+    Args:
+        atom_pair: indices, within a single structure, of the two restrained atoms.
+        bond_length: equilibrium distance of the restraint, in Angstrom.
+        force_constant: spring constant of the restraint, in eV/Angstrom**2.
+        energy_unit (str/float): Energy unit the model works in.
+        position_unit (str/float): Unit used for distances by the model.
+        output_key (str): Key to which results will be stored.
+        trainable (bool): If set to true, bond length and force constant will be
+            optimized during training (default=False).
+    """
+
+    def __init__(
+        self,
+        atom_pair: Tuple[int, int],
+        bond_length: float,
+        force_constant: float,
+        energy_unit: str,
+        position_unit: str,
+        output_key: str,
+        trainable: bool = False,
+    ):
+        super().__init__()
+        self.output_key = output_key
+        self.model_outputs = [output_key]
+
+        self.register_buffer("atom_pair", torch.tensor(atom_pair, dtype=torch.long))
+        self.to_angstrom = spk_units.convert_units(position_unit, "Ang")
+        self.to_model_energy = spk_units.convert_units("eV", energy_unit)
+
+        self.bond_length = nn.Parameter(
+            torch.tensor(bond_length), requires_grad=trainable
+        )
+        self.force_constant = nn.Parameter(
+            torch.tensor(force_constant), requires_grad=trainable
+        )
+
+    def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        # the first atom of every structure, so that structures of differing size work
+        n_atoms = inputs[properties.n_atoms]
+        first = torch.cumsum(n_atoms, dim=0) - n_atoms
+
+        positions = inputs[properties.R]
+        Rij = (
+            positions[first + self.atom_pair[1]] - positions[first + self.atom_pair[0]]
+        )
+        distance = torch.norm(Rij, dim=-1) * self.to_angstrom
+
+        energy = self.force_constant * (distance - self.bond_length) ** 2
+        inputs[self.output_key] = energy * self.to_model_energy
+
+        return inputs