diff --git a/docs/api/atomistic.rst b/docs/api/atomistic.rst index 9ee244e8e..cdef5d3de 100644 --- a/docs/api/atomistic.rst +++ b/docs/api/atomistic.rst @@ -37,4 +37,5 @@ Output modules CoulombPotential DampedCoulombPotential EnergyCoulomb - EnergyEwald \ No newline at end of file + EnergyEwald + HarmonicBond \ No newline at end of file diff --git a/examples/howtos/howto_batchwise_relaxations.ipynb b/examples/howtos/howto_batchwise_relaxations.ipynb index 8c5309b11..4db5d42b9 100644 --- a/examples/howtos/howto_batchwise_relaxations.ipynb +++ b/examples/howtos/howto_batchwise_relaxations.ipynb @@ -1,13 +1,5 @@ { "cells": [ - { - "cell_type": "markdown", - "id": "2ebf008e", - "metadata": {}, - "source": [ - "## Batchwise structure optimization is deprecated and will be upgraded soon\n" - ] - }, { "cell_type": "markdown", "id": "d100d71c", @@ -21,26 +13,52 @@ "id": "cd51dfe1", "metadata": {}, "source": [ - "In this tutorial, we show how to use the ``ASEBatchwiseLBFGS``. It enables relaxation of structures in a batch-wise manner, i.e. it optimizes multiple structures in parallel. This is particularly useful, when many relatively similar structures (--> similar time until convergence) should be relaxed while requiring possibly short simulation time." + "In this tutorial, we show how to use the ``BatchwiseLBFGS``. It enables relaxation of structures in a batch-wise manner, i.e. it optimizes multiple structures in parallel. This is particularly useful, when many relatively similar structures (--> similar time until convergence) should be relaxed while requiring possibly short simulation time." ] }, { "cell_type": "code", "execution_count": null, "id": "68581ba5", - "metadata": {}, + "metadata": { + "execution": { + "iopub.execute_input": "2026-08-10T17:03:42.611231Z", + "iopub.status.busy": "2026-08-10T17:03:42.611127Z", + "iopub.status.idle": "2026-08-10T17:03:47.419502Z", + "shell.execute_reply": "2026-08-10T17:03:47.418697Z" + } + }, "outputs": [], "source": [ - "# import os\n", - "# import shutil\n", + "import os\n", + "import shutil\n", + "from copy import copy, deepcopy\n", + "\n", + "import torch\n", + "import torch.nn as nn\n", + "import numpy as np\n", + "from scipy.optimize import linear_sum_assignment\n", "\n", - "# import torch\n", - "# from ase.io import read\n", + "from ase.io import read\n", + "from ase.optimize import LBFGS\n", + "from ase.build import minimize_rotation_and_translation\n", + "from ase.visualize import view\n", "\n", - "# import schnetpack as spk\n", - "# from schnetpack import properties\n", - "# from schnetpack.interfaces.ase_interface import AtomsConverter\n", - "# from schnetpack.interfaces.batchwise_optimization import ASEBatchwiseLBFGS, BatchwiseCalculator" + "import schnetpack as spk\n", + "from schnetpack import properties\n", + "from schnetpack.interfaces.ase_interface import (\n", + " AtomsConverter,\n", + " SpkCalculator,\n", + " atoms_to_batch,\n", + " batch_to_atoms,\n", + ")\n", + "from schnetpack.interfaces.batchwise_optimization import (\n", + " BatchwiseCalculator,\n", + " BatchwiseLBFGS,\n", + ")\n", + "from schnetpack.interfaces.batchwise_trajectory import BatchwiseTrajectoryReader\n", + "from schnetpack.utils.compatibility import load_model\n", + "from schnetpack.datasets.rmd17 import rMD17" ] }, { @@ -48,42 +66,118 @@ "id": "6339e784", "metadata": {}, "source": [ - "First, we load the force field model that provides the forces for the relaxation process. Furthermore, we define the atoms converter, which is used to convert ase Atoms objects to SchNetPack input. Eventually the calculator is initialized. The latter provides the necessary functionality to load a model and calculates forces and energy for the respective structures. Please note that running batchwise relaxations is significantly faster on a cuda device." + "First, we load the force field model that provides the forces for the relaxation process. Furthermore, we define the atoms converter, which is used to convert ase Atoms objects to SchNetPack input.\n", + "\n", + "Please note: The benefit from using batch-wise relaxations highly depends on your hardware setup. It is significantly faster on a cuda device than cpu. A benchmarking test is implemented in `tests/interfaces/test_bw_benchmark.py`" ] }, { "cell_type": "code", "execution_count": null, "id": "7f1bd733", - "metadata": {}, + "metadata": { + "execution": { + "iopub.execute_input": "2026-08-10T17:03:47.421338Z", + "iopub.status.busy": "2026-08-10T17:03:47.420995Z", + "iopub.status.idle": "2026-08-10T17:03:47.545967Z", + "shell.execute_reply": "2026-08-10T17:03:47.545316Z" + } + }, "outputs": [], "source": [ - "# model_path = \"../../tests/testdata/md_ethanol.model\"\n", + "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", + "\n", + "cutoff = model.representation.cutoff.item()\n", + "CUTOFF_SKIN = 0.3 # Ang\n", + "\n", + "\n", + "def batch_neighbor_list():\n", + " \"\"\"Keeps the neighbor lists of the batch valid while the structures move.\n", + "\n", + " A structure's list is built out to ``cutoff + CUTOFF_SKIN`` and reused for as long\n", + " as no atom of it has drifted more than half the skin. During a relaxation the\n", + " structures change gradually, so most steps get away without rebuilding anything.\n", + " \"\"\"\n", + " return spk.transform.BatchNeighborList(\n", + " neighbor_list=spk.transform.MatScipyNeighborList(cutoff=cutoff),\n", + " cutoff_skin=CUTOFF_SKIN,\n", + " )\n", + "\n", + "\n", + "# settings shared by every relaxation in this notebook\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", + "NOISE = 0.05 # standard deviation of the Gaussian noise in Angstrom (added to atomic positions)\n", + "SEED = 0\n", + "\n", + "\n", + "def batchwise_calculator(model=model_path):\n", + " \"\"\"A calculator that evaluates a whole batch of structures in a single model call.\n", "\n", - "## set device\n", - "# device = torch.device(\"cpu\")\n", + " ``model`` is either a path or a ready model object -- passing a composed model is\n", + " how a prior gets added to the energy, see ``howto_priors.ipynb``.\n", + " \"\"\"\n", + " return BatchwiseCalculator(\n", + " model=model,\n", + " neighbor_list=batch_neighbor_list(),\n", + " device=device,\n", + " energy_unit=ENERGY_UNIT,\n", + " position_unit=POSITION_UNIT,\n", + " )\n", "\n", - "## load model\n", - "# model = torch.load(model_path, map_location=device, weights_only=False)\n", "\n", - "## define neighbor list\n", - "# cutoff = model.representation.cutoff.item()\n", - "# nbh_list=spk.transform.MatScipyNeighborList(cutoff=cutoff)\n", + "def ase_calculator():\n", + " \"\"\"A plain ase calculator, which evaluates one structure per call.\"\"\"\n", + " return SpkCalculator(\n", + " model=model_path,\n", + " neighbor_list=spk.transform.MatScipyNeighborList(cutoff=cutoff),\n", + " device=device,\n", + " energy_unit=ENERGY_UNIT,\n", + " position_unit=POSITION_UNIT,\n", + " )\n", "\n", - "## build atoms converter\n", - "# atoms_converter = AtomsConverter(\n", - "# neighbor_list=nbh_list,\n", - "# device=device,\n", - "# )\n", "\n", - "## build calculator\n", - "# calculator = BatchwiseCalculator(\n", - "# model=model_path,\n", - "# atoms_converter=atoms_converter,\n", - "# device=device,\n", - "# energy_unit=\"kcal/mol\",\n", - "# position_unit=\"Ang\",\n", - "# )" + "def batchwise_optimizer(atoms_list, fixed_atoms_mask=None, model=model_path, **kwargs):\n", + " \"\"\"A ``BatchwiseLBFGS`` set up to relax ``atoms_list`` as one batch.\n", + "\n", + " Builds its own calculator and converts the structures, so every call starts from a\n", + " clean state. ``atoms_to_batch`` only lays the structures out as tensors. The\n", + " calculator's neighbor list fills in the neighborhoods on every step. Remaining keyword \n", + " arguments are passed on to ``BatchwiseLBFGS`` (``trajectory``, ``logfile``, \n", + " ``trajectory_interval``, ...).\n", + " \"\"\"\n", + "\n", + " return BatchwiseLBFGS(\n", + " calculator=batchwise_calculator(model),\n", + " inputs=atoms_to_batch(deepcopy(atoms_list), device=device),\n", + " fixed_atoms_mask=fixed_atoms_mask,\n", + " **kwargs,\n", + " )\n", + "\n", + "\n", + "def make_batch(n_structures, seed=None):\n", + " \"\"\"``n_structures`` starting geometries, cycling through the conformers.\n", + "\n", + " Every copy gets Gaussian noise on its atomic positions, so the structures of a batch\n", + " are similar but not identical. Pass a ``seed`` to get a reproducible batch.\n", + " \"\"\"\n", + " rng = np.random.default_rng(seed)\n", + " batch = []\n", + " for idx in range(n_structures):\n", + " noisy_atoms = conformers[idx % len(conformers)].copy()\n", + " noisy_atoms.positions += rng.normal(\n", + " scale=NOISE, size=noisy_atoms.positions.shape\n", + " )\n", + " batch.append(noisy_atoms)\n", + " return batch" ] }, { @@ -98,13 +192,23 @@ "cell_type": "code", "execution_count": null, "id": "a3ebc6e9", - "metadata": {}, + "metadata": { + "execution": { + "iopub.execute_input": "2026-08-10T17:03:47.547441Z", + "iopub.status.busy": "2026-08-10T17:03:47.547300Z", + "iopub.status.idle": "2026-08-10T17:03:47.557537Z", + "shell.execute_reply": "2026-08-10T17:03:47.556821Z" + } + }, "outputs": [], "source": [ - "# input_structure_file = \"../../tests/testdata/md_ethanol.xyz\"\n", + "input_structure_file = \"../../tests/testdata/ethanol_conformers.xyz\"\n", + "\n", + "# load the initial structures\n", + "conformers = read(input_structure_file, index=\":\")\n", "\n", - "## load initial structures\n", - "# ats = read(input_structure_file, index=\":\")" + "n_replica = 5\n", + "ats = make_batch(n_replica * len(conformers), seed=SEED)" ] }, { @@ -112,21 +216,28 @@ "id": "5323eaa6", "metadata": {}, "source": [ - "For some systems it helps to fix the positions of certain atoms during the relaxation. This can be achieved by providing a mask of boolean entries to ``ASEBatchwiseLBFGS``. The mask is a list of $n_\\text{atoms}$ entries, indicating atoms, which positions are fixed during the relaxation. Here, we do not fix any atoms. Hence, the mask only contains ``True``." + "For some systems it helps to fix the positions of certain atoms during the relaxation. This can be achieved by providing a mask of boolean entries to ``BatchwiseLBFGS``. The mask is a list of $n_\\text{atoms}$ entries, indicating atoms, which positions are fixed during the relaxation. Here, we do not fix any atoms. Hence, the mask only contains ``True``." ] }, { "cell_type": "code", "execution_count": null, "id": "28e377f4", - "metadata": {}, + "metadata": { + "execution": { + "iopub.execute_input": "2026-08-10T17:03:47.559178Z", + "iopub.status.busy": "2026-08-10T17:03:47.559033Z", + "iopub.status.idle": "2026-08-10T17:03:47.561817Z", + "shell.execute_reply": "2026-08-10T17:03:47.561222Z" + } + }, "outputs": [], "source": [ - "## define structure mask for optimization (True for fixed, False for non-fixed)\n", - "# n_atoms = len(ats[0].get_atomic_numbers())\n", - "# single_structure_mask = [False for _ in range(n_atoms)]\n", - "## expand mask by number of input structures (fixed atoms are equivalent for all input structures)\n", - "# mask = single_structure_mask * len(ats)" + "# define structure mask for optimization (True for fixed, False for non-fixed)\n", + "n_atoms = len(ats[0].get_atomic_numbers())\n", + "single_structure_mask = [False for _ in range(n_atoms)]\n", + "# expand mask by number of input structures (fixed atoms are equivalent for all input structures)\n", + "mask = single_structure_mask * len(ats)" ] }, { @@ -141,63 +252,567 @@ "cell_type": "code", "execution_count": null, "id": "2532bb4a", + "metadata": { + "execution": { + "iopub.execute_input": "2026-08-10T17:03:47.563512Z", + "iopub.status.busy": "2026-08-10T17:03:47.563361Z", + "iopub.status.idle": "2026-08-10T17:03:51.238319Z", + "shell.execute_reply": "2026-08-10T17:03:51.237723Z" + } + }, + "outputs": [], + "source": [ + "results_dir = \"./howto_batchwise_relaxations_outputs\"\n", + "if not os.path.exists(results_dir):\n", + " os.makedirs(results_dir)\n", + "\n", + "# Initialize optimizer. The whole batch goes into a single HDF5 trajectory;\n", + "# trajectory_interval=1 records every step.\n", + "optimizer = batchwise_optimizer(\n", + " ats,\n", + " fixed_atoms_mask=mask,\n", + " trajectory=f\"{results_dir}/relax_traj.hdf5\",\n", + " trajectory_interval=1,\n", + ")\n", + "\n", + "# run optimization\n", + "optimizer.run(fmax=FMAX, steps=MAX_STEPS)\n", + "\n", + "# the optimizer works on tensors and hands them back; batch_to_atoms converts at the\n", + "# boundary, for the ase-based comparison further down\n", + "relaxed_batch, _ = optimizer.get_relaxation_results()\n", + "bw_atoms = batch_to_atoms(relaxed_batch)\n", + "\n", + "# release the trajectory file\n", + "optimizer.close()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cafaafac", + "metadata": { + "execution": { + "iopub.execute_input": "2026-08-10T17:03:51.239969Z", + "iopub.status.busy": "2026-08-10T17:03:51.239862Z", + "iopub.status.idle": "2026-08-10T17:03:51.379736Z", + "shell.execute_reply": "2026-08-10T17:03:51.379233Z" + } + }, + "outputs": [], + "source": [ + "# The trajectory holds the whole batch in one file. Every frame comes back as a\n", + "# complete input batch, so it converts to ase structures directly.\n", + "with BatchwiseTrajectoryReader(f\"{results_dir}/relax_traj.hdf5\") as traj:\n", + " print(f\"{traj.n_frames} frames of {traj.n_structures} structures\")\n", + " frames = [batch_to_atoms(traj.frame(idx)) for idx in range(traj.n_frames)]\n", + "\n", + "# the path the first structure of the batch took\n", + "view([frame[0] for frame in frames])" + ] + }, + { + "cell_type": "markdown", + "id": "ec2e6453", "metadata": {}, + "source": [ + "Now we run the sequential relaxations" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2b942191", + "metadata": { + "execution": { + "iopub.execute_input": "2026-08-10T17:03:51.381381Z", + "iopub.status.busy": "2026-08-10T17:03:51.381228Z", + "iopub.status.idle": "2026-08-10T17:04:01.243163Z", + "shell.execute_reply": "2026-08-10T17:04:01.242585Z" + } + }, "outputs": [], "source": [ - "# results_dir = \"./howto_batchwise_relaxations_outputs\"\n", - "# if not os.path.exists(results_dir):\n", - "# os.makedirs(results_dir)\n", + "# a plain ase calculator, to relax the same structures one at a time\n", + "calculator = ase_calculator()\n", + "\n", + "seq_atoms = []\n", + "for at_idx, at in enumerate(deepcopy(ats)):\n", + "\n", + " at.calc = calculator\n", + " lbfgs = LBFGS(at, trajectory=f\"{results_dir}/relax_traj_seq_{at_idx}.xyz\")\n", + " lbfgs.run(fmax=FMAX, steps=MAX_STEPS)\n", + "\n", + " seq_atoms.append(at)" + ] + }, + { + "cell_type": "markdown", + "id": "49f7706f", + "metadata": {}, + "source": [ + "## Comparing batch-wise and sequential relaxation\n", + "\n", + "* **RMSD** between the two relaxed geometries\n", + "* **Potential energy**, evaluated for both sets of structures with the *same* `SpkCalculator`\n", "\n", - "## Initialize optimizer\n", - "# optimizer = ASEBatchwiseLBFGS(\n", - "# calculator=calculator,\n", - "# atoms=ats,\n", - "# trajectory=\"./howto_batchwise_relaxations_outputs/relax_traj\",\n", - "# )\n", + "Two effects make a plain RMSD misleading. A relaxation may translate and rotate the molecule as a whole -- removed with ASE's `minimize_rotation_and_translation`. And identical atoms are indistinguishable: ethanol's CH$_3$ group rotates freely, so a 120° methyl rotation returns the same physical structure with the three hydrogens relabelled. Pairing atoms by index would report a large deviation for two structures that are in fact identical.\n", "\n", - "## run optimization\n", - "# optimizer.run(fmax=0.0005, steps=1000)" + "We therefore also report a **permutation-matched RMSD**, minimised over rigid-body motion *and* relabellings of like atoms (Hungarian assignment via `scipy.optimize.linear_sum_assignment`, alternated with the superposition until the assignment stops changing)." ] }, { "cell_type": "code", "execution_count": null, - "id": "fb369782", + "id": "2c217726", + "metadata": { + "execution": { + "iopub.execute_input": "2026-08-10T17:04:01.244667Z", + "iopub.status.busy": "2026-08-10T17:04:01.244554Z", + "iopub.status.idle": "2026-08-10T17:04:01.631505Z", + "shell.execute_reply": "2026-08-10T17:04:01.630706Z" + } + }, + "outputs": [], + "source": [ + "def rmsd(P, Q):\n", + " return np.sqrt(((P - Q) ** 2).sum() / len(P))\n", + "\n", + "\n", + "def matched_rmsd(atoms, reference, max_iter=20):\n", + " \"\"\"RMSD minimized over rigid-body motion *and* relabellings of atoms with identical type.\n", + "\n", + " Alternates the Kabsch superposition with an optimal assignment of identical\n", + " elements until the permutation stops changing. Returns (rmsd, permutation),\n", + " where permutation[i] is the atom of `atoms` matched to atom i of `reference`.\n", + " \"\"\"\n", + " ref_pos = reference.get_positions()\n", + " symbols = np.array(reference.get_chemical_symbols())\n", + " groups = [np.where(symbols == s)[0] for s in sorted(set(symbols))]\n", + "\n", + " perm = np.arange(len(atoms))\n", + " for _ in range(max_iter):\n", + " current = atoms.copy()\n", + " current.positions = atoms.get_positions()[perm]\n", + " minimize_rotation_and_translation(reference, current)\n", + " pos = current.get_positions()\n", + "\n", + " new_perm = perm.copy()\n", + " for idx in groups:\n", + " cost = ((pos[idx][:, None, :] - ref_pos[None, idx, :]) ** 2).sum(-1)\n", + " row, col = linear_sum_assignment(cost)\n", + " new_perm[idx[col]] = perm[idx[row]]\n", + "\n", + " if np.array_equal(new_perm, perm):\n", + " return rmsd(pos, ref_pos), perm\n", + " perm = new_perm\n", + "\n", + " raise RuntimeError(\"permutation did not converge\")\n", + "\n", + "\n", + "def evaluate(atoms_list):\n", + " \"\"\"Energies and maximum force components, all from the same SpkCalculator.\"\"\"\n", + " energies, fmax = [], []\n", + " for a in atoms_list:\n", + " a = a.copy()\n", + " a.calc = calculator\n", + " energies.append(a.get_potential_energy())\n", + " fmax.append(np.abs(a.get_forces()).max())\n", + " return np.array(energies), np.array(fmax)\n", + "\n", + "\n", + "e_batch, fmax_batch = evaluate(bw_atoms)\n", + "e_seq, fmax_seq = evaluate(seq_atoms)\n", + "\n", + "rmsd_raw = np.array(\n", + " [rmsd(b.positions, s.positions) for b, s in zip(bw_atoms, seq_atoms)]\n", + ")\n", + "rmsd_matched = np.array([matched_rmsd(b, s)[0] for b, s in zip(bw_atoms, seq_atoms)])\n", + "\n", + "header = (\n", + " f\"{'i':>2} {'RMSD raw':>10} {'RMSD matched':>13} \"\n", + " f\"{'E_batch [eV]':>14} {'E_seq [eV]':>14} {'dE [meV]':>10}\"\n", + ")\n", + "print(header)\n", + "print(\"-\" * len(header))\n", + "for i in range(len(bw_atoms)):\n", + " print(\n", + " f\"{i:>2} {rmsd_raw[i]:>10.2e} {rmsd_matched[i]:>13.2e} \"\n", + " f\"{e_batch[i]:>14.6f} {e_seq[i]:>14.6f} {(e_batch[i] - e_seq[i]) * 1e3:>10.2e}\"\n", + " )\n", + "\n", + "dE = e_batch - e_seq\n", + "print()\n", + "print(\n", + " f\"RMSD matched : mean {rmsd_matched.mean():.2e} A, max {rmsd_matched.max():.2e} A\"\n", + ")\n", + "print(f\"RMSD raw : mean {rmsd_raw.mean():.2e} A, max {rmsd_raw.max():.2e} A\")\n", + "print(\n", + " f\"dE : mean {dE.mean() * 1e3:.2e} meV, max|dE| {np.abs(dE).max() * 1e3:.2e} meV\"\n", + ")\n", + "print(\n", + " f\"energy spread within batch-wise results : {(e_batch.max() - e_batch.min()) * 1e3:.2e} meV\"\n", + ")\n", + "print(\n", + " f\"energy spread within sequential results : {(e_seq.max() - e_seq.min()) * 1e3:.2e} meV\"\n", + ")\n", + "print(\n", + " f\"max force component: batch-wise {fmax_batch.max():.2e}, sequential {fmax_seq.max():.2e} eV/A\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "c8bbca69", "metadata": {}, + "source": [ + "Sanity check of RMSD matched" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dd9771b3", + "metadata": { + "execution": { + "iopub.execute_input": "2026-08-10T17:04:01.633416Z", + "iopub.status.busy": "2026-08-10T17:04:01.633243Z", + "iopub.status.idle": "2026-08-10T17:04:01.666150Z", + "shell.execute_reply": "2026-08-10T17:04:01.665616Z" + } + }, "outputs": [], "source": [ - "# if os.path.exists(results_dir):\n", - "# shutil.rmtree(results_dir)" + "reference = seq_atoms[0]\n", + "symbols = np.array(reference.get_chemical_symbols())\n", + "positions = reference.get_positions()\n", + "print(\"atom order:\", reference.get_chemical_symbols())\n", + "\n", + "# a methyl group is the three hydrogens sharing one carbon, so find each hydrogen's\n", + "# nearest heavy atom rather than trusting the order the structures happen to come in\n", + "hydrogens = np.where(symbols == \"H\")[0]\n", + "heavy = np.where(symbols != \"H\")[0]\n", + "host = heavy[\n", + " np.linalg.norm(positions[hydrogens, None] - positions[None, heavy], axis=-1).argmin(\n", + " axis=1\n", + " )\n", + "]\n", + "(methyl_carbon,) = [a for a in set(host) if list(host).count(a) == 3]\n", + "methyl = hydrogens[host == methyl_carbon]\n", + "\n", + "# relabel the methyl hydrogens cyclically: this is the very same physical structure,\n", + "# only the three identical hydrogens have swapped names\n", + "order = np.arange(len(reference))\n", + "order[methyl] = np.roll(methyl, 1)\n", + "\n", + "relabelled = reference.copy()\n", + "relabelled.positions = reference.get_positions()[order]\n", + "\n", + "e_ref, e_relabelled = evaluate([reference, relabelled])[0]\n", + "r, perm = matched_rmsd(relabelled, reference)\n", + "\n", + "print(f\"methyl hydrogens : {methyl.tolist()} (on C{int(methyl_carbon)})\")\n", + "print(f\"energy difference : {(e_relabelled - e_ref) * 1e3:.2e} meV\")\n", + "print(f\"RMSD raw : {rmsd(relabelled.get_positions(), positions):.2e} A\")\n", + "print(f\"RMSD matched : {r:.2e} A\")\n", + "print(f\"recovered permutation : {perm}\")" ] }, { "cell_type": "markdown", - "id": "1f579c83", + "id": "126425ef", "metadata": {}, "source": [ - "Optimzed structures (in the form of ASE `Atoms`) and properties can be obtained with the `get_relaxation_results` function." + "## How the relaxation time scales with the batch size\n", + "\n", + "Everything above compares the *result* of the two relaxations. What the batch-wise optimizer is actually for is the *cost*: the whole batch is concatenated into a single graph and pushed through the model in one forward pass, so all structures are evaluated in parallel and take their LBFGS step together.\n", + "\n", + "That parallelism is where the time goes. A single ethanol molecule is nowhere near large enough to occupy a gpu -- one forward pass keeps a small fraction of the cores busy, and the step time is set by kernel launch latency and idle hardware rather than by the amount of arithmetic. Relaxing structures one at a time repeats that under-utilised call once per structure and per step. Batching raises the work per call instead of the number of calls, so until the device is saturated the extra structures are close to free.\n", + "\n", + "The sweep below times both optimizers over a range of batch sizes. A few points keep the numbers honest:\n", + "\n", + "* Only the relaxation itself is timed. Building the calculator, converting the structures and constructing the optimizer all happen before the clock starts.\n", + "* Every batch size is repeated a few times, each repeat with a different noise seed, and both optimizers relax the *identical* batch in each repeat (paired samples). The confidence interval therefore covers machine noise as well as how much the starting geometries happen to vary in difficulty.\n", + "* The batch-wise optimizer stops when the *last* structure of the batch has converged, so a batch is only as fast as its slowest member. This is why the method pays off for batches of similar structures.\n", + "\n", + "How much there is to gain is a property of the device. On cpu there is far less parallelism to exploit, so batching mostly saves the per-call overhead and the gain settles at a modest constant factor. On cuda the flat part of the curve extends much further, which is where this optimizer earns its keep." ] }, { "cell_type": "code", "execution_count": null, - "id": "78f81235", + "id": "b3dcf7a3", + "metadata": { + "execution": { + "iopub.execute_input": "2026-08-10T17:04:01.667482Z", + "iopub.status.busy": "2026-08-10T17:04:01.667372Z", + "iopub.status.idle": "2026-08-10T17:08:46.061092Z", + "shell.execute_reply": "2026-08-10T17:08:46.060350Z" + } + }, + "outputs": [], + "source": [ + "import time\n", + "from scipy import stats\n", + "\n", + "batch_sizes = [1, 2, 4, 8, 16, 32]\n", + "n_repeats = 5\n", + "\n", + "\n", + "def time_batchwise(batch):\n", + " \"\"\"Wall time of relaxing the whole batch in parallel. The setup is not timed.\"\"\"\n", + " optimizer = batchwise_optimizer(batch, logfile=None)\n", + "\n", + " start = time.perf_counter()\n", + " optimizer.run(fmax=FMAX, steps=MAX_STEPS)\n", + " elapsed = time.perf_counter() - start\n", + "\n", + " # a run that is fast because it never converged is not a faster run\n", + " assert optimizer.nsteps < MAX_STEPS, \"batch-wise run hit the step limit\"\n", + " return elapsed\n", + "\n", + "\n", + "def time_sequential(batch):\n", + " \"\"\"Wall time of relaxing the same structures one after another.\"\"\"\n", + " calculator = ase_calculator()\n", + " structures = deepcopy(batch) # LBFGS relaxes the Atoms objects in place\n", + "\n", + " start = time.perf_counter()\n", + " steps = []\n", + " for structure in structures:\n", + " structure.calc = calculator\n", + " lbfgs = LBFGS(structure, logfile=None)\n", + " lbfgs.run(fmax=FMAX, steps=MAX_STEPS)\n", + " steps.append(lbfgs.nsteps)\n", + " elapsed = time.perf_counter() - start\n", + "\n", + " assert max(steps) < MAX_STEPS, \"a sequential run hit the step limit\"\n", + " return elapsed\n", + "\n", + "\n", + "# The first relaxation pays for lazy torch initialisation, which would otherwise land\n", + "# entirely on the smallest batch size and fake a slow start.\n", + "time_batchwise(make_batch(2, seed=999))\n", + "time_sequential(make_batch(2, seed=999))\n", + "\n", + "timings = {\n", + " \"batch-wise\": np.zeros((len(batch_sizes), n_repeats)),\n", + " \"sequential\": np.zeros((len(batch_sizes), n_repeats)),\n", + "}\n", + "\n", + "for size_idx, n_structures in enumerate(batch_sizes):\n", + " for repeat in range(n_repeats):\n", + " # both optimizers relax the identical batch, so the repeats are paired\n", + " batch = make_batch(n_structures, seed=repeat)\n", + " timings[\"batch-wise\"][size_idx, repeat] = time_batchwise(batch)\n", + " timings[\"sequential\"][size_idx, repeat] = time_sequential(batch)\n", + "\n", + " print(\n", + " f\"batch size {n_structures:>3}: \"\n", + " f\"batch-wise {timings['batch-wise'][size_idx].mean():6.2f} s, \"\n", + " f\"sequential {timings['sequential'][size_idx].mean():6.2f} s\",\n", + " flush=True,\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "b91ce79d", "metadata": {}, + "source": [ + "The mean and the confidence interval over the repeats. With only a handful of repeats the interval is a Student-t interval, which is why it is noticeably wider than the standard error: increase `n_repeats` to tighten it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a7973259", + "metadata": { + "execution": { + "iopub.execute_input": "2026-08-10T17:08:46.062907Z", + "iopub.status.busy": "2026-08-10T17:08:46.062736Z", + "iopub.status.idle": "2026-08-10T17:08:46.069338Z", + "shell.execute_reply": "2026-08-10T17:08:46.068820Z" + } + }, + "outputs": [], + "source": [ + "confidence = 0.95\n", + "\n", + "\n", + "def mean_ci(samples, confidence=confidence):\n", + " \"\"\"Mean and half-width of the Student-t confidence interval, along the repeats.\"\"\"\n", + " mean = samples.mean(axis=1)\n", + " half_width = stats.sem(samples, axis=1) * stats.t.ppf(\n", + " 0.5 + confidence / 2, samples.shape[1] - 1\n", + " )\n", + " return mean, half_width\n", + "\n", + "\n", + "mean_bw, ci_bw = mean_ci(timings[\"batch-wise\"])\n", + "mean_seq, ci_seq = mean_ci(timings[\"sequential\"])\n", + "\n", + "header = (\n", + " f\"{'batch size':>10} {'batch-wise [s]':>22} {'sequential [s]':>22} {'speedup':>9}\"\n", + ")\n", + "print(header)\n", + "print(\"-\" * len(header))\n", + "for i, n_structures in enumerate(batch_sizes):\n", + " print(\n", + " f\"{n_structures:>10} \"\n", + " f\"{mean_bw[i]:>13.2f} +- {ci_bw[i]:<6.2f} \"\n", + " f\"{mean_seq[i]:>13.2f} +- {ci_seq[i]:<6.2f} \"\n", + " f\"{mean_seq[i] / mean_bw[i]:>8.2f}x\"\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b59c90de", + "metadata": { + "execution": { + "iopub.execute_input": "2026-08-10T17:08:46.070842Z", + "iopub.status.busy": "2026-08-10T17:08:46.070707Z", + "iopub.status.idle": "2026-08-10T17:08:46.313782Z", + "shell.execute_reply": "2026-08-10T17:08:46.313247Z" + } + }, "outputs": [], "source": [ - "## get list of optimized structures and properties\n", - "# opt_atoms, opt_props = optimizer.get_relaxation_results()\n", + "import matplotlib.pyplot as plt\n", "\n", - "# for oatoms in opt_atoms:\n", - "# print(oatoms.get_positions())\n", + "SURFACE = \"#fcfcfb\"\n", + "INK = \"#0b0b0b\"\n", + "INK_SECONDARY = \"#52514e\"\n", + "MUTED = \"#898781\"\n", + "GRID = \"#e1e0d9\"\n", + "AXIS = \"#c3c2b7\"\n", + "SERIES = {\"batch-wise\": \"#2a78d6\", \"sequential\": \"#eb6834\"}\n", "\n", - "# print(opt_props)" + "fig, ax = plt.subplots(figsize=(7.2, 4.8), dpi=130)\n", + "fig.patch.set_facecolor(SURFACE)\n", + "ax.set_facecolor(SURFACE)\n", + "\n", + "for label, (mean, half_width) in {\n", + " \"batch-wise\": (mean_bw, ci_bw),\n", + " \"sequential\": (mean_seq, ci_seq),\n", + "}.items():\n", + " color = SERIES[label]\n", + " # confidence interval as a band, the mean as the line on top of it\n", + " ax.fill_between(\n", + " batch_sizes,\n", + " mean - half_width,\n", + " mean + half_width,\n", + " color=color,\n", + " alpha=0.18,\n", + " linewidth=0,\n", + " zorder=2,\n", + " )\n", + " ax.plot(\n", + " batch_sizes,\n", + " mean,\n", + " color=color,\n", + " linewidth=2,\n", + " marker=\"o\",\n", + " markersize=6,\n", + " markeredgecolor=SURFACE,\n", + " markeredgewidth=2,\n", + " label=label,\n", + " zorder=3,\n", + " )\n", + " # label the series at its end, so the curves can be read without the legend\n", + " ax.annotate(\n", + " label,\n", + " (batch_sizes[-1], mean[-1]),\n", + " textcoords=\"offset points\",\n", + " xytext=(10, 0),\n", + " va=\"center\",\n", + " color=color,\n", + " fontsize=10,\n", + " )\n", + "\n", + "# both scalings are power laws, which read as straight lines on log-log axes\n", + "ax.set_xscale(\"log\", base=2)\n", + "ax.set_yscale(\"log\", base=2)\n", + "ax.set_xticks(batch_sizes)\n", + "ax.set_xticklabels(batch_sizes)\n", + "ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda y, _: f\"{y:g}\"))\n", + "ax.minorticks_off()\n", + "ax.set_xlim(batch_sizes[0] / 1.3, batch_sizes[-1] * 2.6)\n", + "\n", + "ax.grid(True, which=\"major\", color=GRID, linewidth=0.8)\n", + "ax.set_axisbelow(True)\n", + "for side in (\"top\", \"right\"):\n", + " ax.spines[side].set_visible(False)\n", + "for side in (\"left\", \"bottom\"):\n", + " ax.spines[side].set_color(AXIS)\n", + "ax.tick_params(colors=MUTED, labelsize=9, length=0)\n", + "\n", + "ax.set_xlabel(\n", + " \"batch size (number of structures)\", color=INK_SECONDARY, fontsize=10, labelpad=8\n", + ")\n", + "ax.set_ylabel(\"relaxation time [s]\", color=INK_SECONDARY, fontsize=10, labelpad=8)\n", + "ax.set_title(\n", + " \"Batch-wise relaxation pulls ahead as the batch grows\",\n", + " color=INK,\n", + " fontsize=13,\n", + " loc=\"left\",\n", + " pad=34,\n", + ")\n", + "ax.text(\n", + " 0.0,\n", + " 1.015,\n", + " f\"wall time to relax the whole batch to fmax = {FMAX} eV/A on {device.type},\\n\"\n", + " f\"mean of {n_repeats} runs, band is the {confidence:.0%} confidence interval\",\n", + " transform=ax.transAxes,\n", + " color=MUTED,\n", + " fontsize=9,\n", + " va=\"bottom\",\n", + " linespacing=1.6,\n", + ")\n", + "legend = ax.legend(frameon=False, loc=\"upper left\", fontsize=10)\n", + "for text in legend.get_texts():\n", + " text.set_color(INK_SECONDARY)\n", + "\n", + "fig.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "f4199a77", + "metadata": {}, + "source": [ + "Both curves grow with the batch size, but with different slopes. The sequential loop is linear: every structure costs a full relaxation of its own, so doubling the batch doubles the time. The batch-wise curve is much shallower, because the added structures ride along inside the same kernels instead of adding calls of their own.\n", + "\n", + "The two curves cross at a small batch size. Below the crossing a single structure is relaxed slightly faster by plain ase `LBFGS` -- with one structure there is no parallelism to exploit, and the batch machinery only adds bookkeeping. Above it the gap widens, and the speedup column in the table above is the ratio of the two curves.\n", + "\n", + "Once the batch is large enough to saturate the device, the batch-wise curve turns linear as well: from there on extra structures really do cost extra arithmetic, and the slope approaches the sequential one, offset by whatever factor the parallel evaluation won. Where that knee sits depends on the hardware, the model and the size of the structures, and it lies much further out on a gpu than on cpu -- so it is worth re-running this sweep with `batch_sizes` extended for your own setup." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b1c69f79", + "metadata": { + "execution": { + "iopub.execute_input": "2026-08-10T17:08:46.315391Z", + "iopub.status.busy": "2026-08-10T17:08:46.315250Z", + "iopub.status.idle": "2026-08-10T17:08:46.318036Z", + "shell.execute_reply": "2026-08-10T17:08:46.317583Z" + } + }, + "outputs": [], + "source": [ + "if os.path.exists(results_dir):\n", + " shutil.rmtree(results_dir)" ] }, { "cell_type": "code", "execution_count": null, - "id": "11776e3c", + "id": "b8d83cdc", "metadata": {}, "outputs": [], "source": [] @@ -205,7 +820,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "schnetpack", "language": "python", "name": "python3" }, @@ -219,7 +834,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.12" + "version": "3.13.13" } }, "nbformat": 4, diff --git a/examples/howtos/howto_priors.ipynb b/examples/howtos/howto_priors.ipynb new file mode 100644 index 000000000..6fd8ceb66 --- /dev/null +++ b/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 +} diff --git a/pyproject.toml b/pyproject.toml index f8d691dc9..3230ee5cb 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/src/schnetpack/atomistic/__init__.py b/src/schnetpack/atomistic/__init__.py index 67144181b..5199ca932 100644 --- a/src/schnetpack/atomistic/__init__.py +++ b/src/schnetpack/atomistic/__init__.py @@ -5,3 +5,4 @@ from .electrostatic import * from .aggregation import * from .external_fields import * +from .priors import * diff --git a/src/schnetpack/atomistic/priors.py b/src/schnetpack/atomistic/priors.py new file mode 100644 index 000000000..4cd5273da --- /dev/null +++ b/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 diff --git a/src/schnetpack/data/loader.py b/src/schnetpack/data/loader.py index 56f8d95f3..21018a619 100644 --- a/src/schnetpack/data/loader.py +++ b/src/schnetpack/data/loader.py @@ -1,13 +1,13 @@ import torch from torch.utils.data import DataLoader -from typing import Optional, Sequence +from typing import Dict, List, Optional, Sequence from torch.utils.data import Dataset, Sampler from torch.utils.data.dataloader import _collate_fn_t, _T_co import schnetpack.properties as structure -__all__ = ["AtomsLoader"] +__all__ = ["AtomsLoader", "split_batch"] def _atoms_collate_fn(batch): @@ -58,6 +58,73 @@ def _atoms_collate_fn(batch): return coll_batch +#: entries that describe the structures themselves, as opposed to their neighborhoods. +#: These are what a neighbor list needs, and what :func:`split_batch` splits by default. +_STRUCTURE_KEYS = ( + structure.n_atoms, + structure.Z, + structure.R, + structure.cell, + structure.pbc, +) + + +def split_batch( + inputs: Dict[str, torch.Tensor], keys: Optional[Sequence[str]] = None +) -> List[Dict[str, torch.Tensor]]: + """Split a collated batch back into one input dictionary per structure. + + The inverse of :func:`_atoms_collate_fn` for the structure-defining entries: atom-wise + entries are cut along the atom axis at the ``n_atoms`` boundaries, structure-wise ones + are indexed. Every structure gets its position in the batch as ``properties.idx``, the + sample index a per-sample transform keys its caches by, so a batch that carries none -- + one read back from a trajectory, say -- still splits into usable samples. + + Neighbor lists are deliberately not split. They are the one thing that cannot be + recovered by cutting: the pair indices are shifted into batch-global numbering, and any + caller splitting a batch is about to rebuild them anyway. + + Args: + inputs: collated input batch. + keys: entries to split in addition to ``n_atoms``, ``Z``, ``R``, ``cell`` and + ``pbc``. An entry whose first dimension matches the number of atoms in the + batch is treated as atom-wise, anything else as structure-wise. The two only + coincide when every structure holds a single atom, and there the two readings + cut at the same places anyway. + + Returns: + list(dict(str, torch.Tensor)): one input dictionary per structure, in batch order. + """ + n_atoms = inputs[structure.n_atoms] + n_structures = n_atoms.shape[0] + n_total_atoms = int(n_atoms.sum()) + + offsets = torch.cat( + [torch.zeros(1, dtype=n_atoms.dtype, device=n_atoms.device), n_atoms.cumsum(0)] + ).tolist() + + split_keys = list(_STRUCTURE_KEYS) + [ + key for key in (keys or ()) if key not in _STRUCTURE_KEYS + ] + + samples = [] + for idx in range(n_structures): + sample = {structure.idx: torch.tensor([idx])} + for key in split_keys: + if key not in inputs: + continue + value = inputs[key] + if key == structure.n_atoms: + sample[key] = value[idx : idx + 1] + elif value.shape[0] == n_total_atoms: + sample[key] = value[offsets[idx] : offsets[idx + 1]] + else: + sample[key] = value[idx : idx + 1] + samples.append(sample) + + return samples + + class AtomsLoader(DataLoader): """Data loader for subclasses of ASEAtomsData""" diff --git a/src/schnetpack/interfaces/__init__.py b/src/schnetpack/interfaces/__init__.py index 90301e69f..7e77036d4 100644 --- a/src/schnetpack/interfaces/__init__.py +++ b/src/schnetpack/interfaces/__init__.py @@ -1,2 +1 @@ from .ase_interface import * -from .batchwise_optimization import * diff --git a/src/schnetpack/interfaces/ase_interface.py b/src/schnetpack/interfaces/ase_interface.py index 957491f29..8d60e71d5 100644 --- a/src/schnetpack/interfaces/ase_interface.py +++ b/src/schnetpack/interfaces/ase_interface.py @@ -45,7 +45,14 @@ log = logging.getLogger(__name__) -__all__ = ["SpkCalculator", "AseInterface", "AtomsConverter", "SpkEnsembleCalculator"] +__all__ = [ + "SpkCalculator", + "AseInterface", + "AtomsConverter", + "SpkEnsembleCalculator", + "atoms_to_batch", + "batch_to_atoms", +] class AtomsConverterError(Exception): @@ -83,7 +90,6 @@ def __init__( stored to the input batch. """ - self.neighbor_list = deepcopy(neighbor_list) self.device = device self.dtype = dtype self.additional_inputs = additional_inputs or {} @@ -129,17 +135,7 @@ def __call__(self, atoms: List[Atoms] or Atoms): inputs_batch = [] for at_idx, at in enumerate(atoms): - - inputs = { - properties.n_atoms: torch.tensor([at.get_global_number_of_atoms()]), - properties.Z: torch.from_numpy(at.get_atomic_numbers()), - properties.R: torch.from_numpy(at.get_positions()), - properties.cell: torch.from_numpy(at.get_cell().array).view(-1, 3, 3), - properties.pbc: torch.from_numpy(at.get_pbc()).view(-1, 3), - } - - # specify sample index - inputs.update({properties.idx: torch.tensor([at_idx])}) + inputs = _atoms_to_inputs(at, at_idx) # add additional inputs (specified in AtomsConverter __init__) inputs.update(self.additional_inputs) @@ -156,6 +152,86 @@ def __call__(self, atoms: List[Atoms] or Atoms): return inputs +def _atoms_to_inputs(atoms: Atoms, idx: int = 0) -> Dict[str, torch.Tensor]: + """The tensors describing a single ase structure, before any transform has run.""" + return { + properties.n_atoms: torch.tensor([atoms.get_global_number_of_atoms()]), + properties.Z: torch.from_numpy(atoms.get_atomic_numbers()), + properties.R: torch.from_numpy(atoms.get_positions()), + properties.cell: torch.from_numpy(atoms.get_cell().array).view(-1, 3, 3), + properties.pbc: torch.from_numpy(atoms.get_pbc()).view(-1, 3), + properties.idx: torch.tensor([idx]), + } + + +def atoms_to_batch( + atoms: Union[List[Atoms], Atoms], + device: Union[str, torch.device] = "cpu", + dtype: torch.dtype = torch.float32, +) -> Dict[str, torch.Tensor]: + """Turn ase structures into a schnetpack input batch, neighbor lists aside. + + The inverse of :func:`batch_to_atoms`, and what code working on batches of tensors + needs at its entry: the batch-wise optimizer takes it from here, and its calculator's + :class:`~schnetpack.transform.BatchNeighborList` fills in the neighborhoods on every + step. Use :class:`AtomsConverter` instead when the batch has to be complete right + away, e.g. to call a model on it directly. + + Args: + atoms: a list of ase structures, or a single one. + device: device the batch is placed on. + dtype: float precision of the batch. + + Returns: + dict[str, torch.Tensor]: input batch holding positions, atomic numbers, cells, + pbc and the atom counts. + """ + if isinstance(atoms, Atoms): + atoms = [atoms] + + cast = CastTo32() if dtype == torch.float32 else CastTo64() + if dtype not in (torch.float32, torch.float64): + raise AtomsConverterError(f"Unrecognized precision {dtype}") + + inputs = _atoms_collate_fn( + [cast(_atoms_to_inputs(at, at_idx)) for at_idx, at in enumerate(atoms)] + ) + return {key: value.to(device) for key, value in inputs.items()} + + +def batch_to_atoms(inputs: Dict[str, torch.Tensor]) -> List[Atoms]: + """Turn a schnetpack input batch back into ase structures. + + The inverse of :meth:`AtomsConverter.__call__`, and the one place in schnetpack + that does this conversion -- code that works on batches of tensors (the batch-wise + optimizer, for instance) can stay free of ase and call this at its boundary. + + Args: + inputs: input batch. Only positions, atomic numbers, cells, pbc and the atom + counts are read, so a batch straight out of an optimizer works. + + Returns: + list(ase.Atoms): one structure per entry of the batch, in batch order. + """ + n_atoms = inputs[properties.n_atoms].detach().cpu() + positions = inputs[properties.R].detach().cpu().numpy() + atomic_numbers = inputs[properties.Z].detach().cpu().numpy() + cells = inputs[properties.cell].detach().cpu().numpy().reshape(-1, 3, 3) + pbc = inputs[properties.pbc].detach().cpu().numpy().reshape(-1, 3) + + # split rather than slice by a fixed width, so ragged batches work too + offsets = np.pad(np.cumsum(n_atoms.numpy()), (1, 0)) + return [ + Atoms( + positions=positions[offsets[idx] : offsets[idx + 1]], + numbers=atomic_numbers[offsets[idx] : offsets[idx + 1]], + cell=cells[idx], + pbc=pbc[idx], + ) + for idx in range(len(n_atoms)) + ] + + class SpkCalculatorError(Exception): pass diff --git a/src/schnetpack/interfaces/batchwise_optimization.py b/src/schnetpack/interfaces/batchwise_optimization.py deleted file mode 100644 index a5745d742..000000000 --- a/src/schnetpack/interfaces/batchwise_optimization.py +++ /dev/null @@ -1,916 +0,0 @@ -from copy import deepcopy -import os -import pickle -import time - -import ase -import numpy as np -from math import sqrt -from os.path import isfile - -from ase.optimize.optimize import Dynamics -from ase.parallel import world, barrier -from ase.io import write -from ase import Atoms - -from typing import Dict, Optional, List, Tuple - -import torch -from torch import nn -from schnetpack.units import convert_units -from schnetpack.interfaces.ase_interface import AtomsConverter - -__all__ = [ - "ASEBatchwiseLBFGS", - "BatchwiseCalculator", - "BatchwiseEnsembleCalculator", - "NNEnsemble", -] - - -class AtomsConverterError(Exception): - pass - - -class NNEnsemble(nn.Module): - # TODO: integrate this into EnsembleCalculator directly - def __init__(self, models: nn.ModuleList, properties: List[str]): - super(NNEnsemble, self).__init__() - self.models = models - if type(properties) == str: - properties = [properties] - self.properties = properties - - def setup(self, stage: Optional[str] = None) -> None: - for model in self.models: - model.setup(stage) - - def forward( - self, - x: Dict, - ) -> Tuple: - results = {} - for p in self.properties: - results[p] = [] - - inputs = deepcopy(x) - for model in self.models: - x = deepcopy(inputs) - predictions = model(x) - for prop, values in predictions.items(): - if prop in self.properties: - results[prop].append(values) - - means = {} - stds = {} - for prop, values in results.items(): - stacked_values = torch.stack(values) - means[prop] = stacked_values.mean(dim=0) - stds[prop] = stacked_values.std(dim=0) - - return means, stds - - -class BatchwiseCalculator: - """ - Calculator for neural network models for batchwise optimization. - """ - - def __init__( - self, - model: nn.Module or str, - atoms_converter: AtomsConverter, - device: str or torch.device = "cpu", - auxiliary_output_modules: Optional[List] = None, - energy_key: str = "energy", - force_key: str = "forces", - stress_key: Optional[str] = None, - energy_unit: str = "eV", - position_unit: str = "Ang", - dtype: torch.dtype = torch.float32, - ): - """ - model: - path to trained model or trained model - - atoms_converter: - Class used to convert ase Atoms objects to schnetpack input - - device: - device used for calculations (default="cpu") - - auxiliary_output_modules: - auxiliary module to manipulate output properties (e.g., prior energy or forces) - - energy_key: - name of energies in model (default="energy") - - force_key: - name of forces in model (default="forces") - - stress_key: - name of stress in model (default=None) - - energy_unit: - energy units used by model (default="eV") - - position_unit: - position units used by model (default="Angstrom") - - dtype: - required data type for the model input (default: torch.float32) - """ - - self.results = None - self.atoms = None - - if type(device) == str: - device = torch.device(device) - self.device = device - self.dtype = dtype - self.atoms_converter = atoms_converter - self.auxiliary_output_modules = auxiliary_output_modules or [] - - self.energy_key = energy_key - self.force_key = force_key - self.stress_key = stress_key - - # set up basic conversion factors - self.energy_conversion = convert_units(energy_unit, "eV") - self.position_conversion = convert_units(position_unit, "Angstrom") - - # Unit conversion to default ASE units - self.property_units = { - self.energy_key: self.energy_conversion, - self.force_key: self.energy_conversion / self.position_conversion, - } - if self.stress_key is not None: - self.property_units[self.stress_key] = ( - self.energy_conversion / self.position_conversion**3 - ) - - # load model from path if needed - if type(model) == str: - model = self._load_model(model) - - self._initialize_model(model) - - def _load_model(self, model: str) -> nn.Module: - return torch.load(model, map_location="cpu", weights_only=False).to( - torch.float64 - ) - - def _initialize_model(self, model: nn.Module) -> None: - for auxiliary_output_module in self.auxiliary_output_modules: - model.output_modules.insert(1, auxiliary_output_module) - self.model = model.eval() - self.model.to(device=self.device, dtype=self.dtype) - - def _requires_calculation(self, property_keys: List[str], atoms: List[ase.Atoms]): - if self.results is None: - return True - for name in property_keys: - if name not in self.results: - return True - if len(self.atoms) != len(atoms): - return True - for atom, atom_ref in zip(atoms, self.atoms): - if atom != atom_ref: - return True - - def get_forces( - self, atoms: List[ase.Atoms], fixed_atoms_mask: Optional[List[int]] = None - ) -> np.array: - """ - atoms: - - fixed_atoms_mask: - list of indices corresponding to atoms with positions fixed in space. - """ - if self._requires_calculation( - property_keys=[self.energy_key, self.force_key], atoms=atoms - ): - self.calculate(atoms) - f = self.results[self.force_key] - if fixed_atoms_mask is not None: - f[fixed_atoms_mask] = 0.0 - return f - - def get_potential_energy(self, atoms: List[ase.Atoms]) -> float: - if self._requires_calculation(property_keys=[self.energy_key], atoms=atoms): - self.calculate(atoms) - return self.results[self.energy_key] - - def calculate(self, atoms: List[ase.Atoms]) -> None: - property_keys = list(self.property_units.keys()) - inputs = self.atoms_converter(atoms) - model_results = self.model(inputs) - - results = {} - # store model results in calculator - for prop in property_keys: - if prop in model_results: - results[prop] = ( - model_results[prop].detach().cpu().numpy() - * self.property_units[prop] - ) - else: - raise AtomsConverterError( - "'{:s}' is not a property of your model. Please " - "check the model " - "properties!".format(prop) - ) - - self.results = results - self.atoms = atoms.copy() - - -class BatchwiseEnsembleCalculator(BatchwiseCalculator): - """ - Calculator for ensemble of neural network models for batchwise optimization. - """ - - # TODO: inherit from SpkEnsembleCalculator - def __init__( - self, - model: str or nn.ModuleList, - atoms_converter: AtomsConverter, - device: str or torch.device = "cpu", - auxiliary_output_modules: Optional[List[nn.Module]] = None, - energy_key: str = "energy", - force_key: str = "forces", - stress_key: Optional[str] = None, - energy_unit: str = "eV", - position_unit: str = "Ang", - dtype: torch.dtype = torch.float32, - ): - """ - model: - Directory of trained models or module list of trained models - - atoms_converter: - Class used to convert ase Atoms objects to schnetpack input - - device: - device used for calculations (default="cpu") - - auxiliary_output_modules: - auxiliary module to manipulate output properties (e.g., prior energy or forces) - - energy_key: - name of energies in model (default="energy") - - force_key: - name of forces in model (default="forces") - - energy_unit: - energy units used by model (default="eV") - - stress_key: - name of stress in model (default=None) - - position_unit: - position units used by model (default="Angstrom") - - dtype: - required data type for the model input (default: torch.float32) - """ - super(BatchwiseEnsembleCalculator, self).__init__( - model=model, - atoms_converter=atoms_converter, - device=device, - auxiliary_output_modules=auxiliary_output_modules, - energy_key=energy_key, - force_key=force_key, - stress_key=stress_key, - energy_unit=energy_unit, - position_unit=position_unit, - dtype=dtype, - ) - - def _load_model(self, model: str) -> nn.ModuleList: - # get model paths - model_names = os.listdir(model) - model_paths = [os.path.join(model, model_name) for model_name in model_names] - - # create module list - models = torch.nn.ModuleList() - for m_path in model_paths: - m = torch.load( - os.path.join(m_path, "best_model"), - map_location="cpu", - weights_only=False, - ).to(torch.float64) - models.append(m) - - return models - - def _initialize_model(self, model: nn.ModuleList) -> None: - # add auxiliary output modules - for m in model: - for auxiliary_output_module in self.auxiliary_output_modules: - m.output_modules.insert(1, auxiliary_output_module) - - # initialize ensemble - ensemble = NNEnsemble(models=model, properties=list(self.property_units.keys())) - self.model = ensemble.eval().to(device=self.device, dtype=self.dtype) - - def calculate(self, atoms: List[ase.Atoms]) -> None: - property_keys = list(self.property_units.keys()) - inputs = self.atoms_converter(atoms) - model_results, stds = self.model(inputs) - - results = {} - # store model uncertainties in calculator - for prop in property_keys: - if prop in model_results: - results["{}_uncertainty".format(prop)] = ( - stds[prop].detach().cpu().numpy() * self.property_units[prop] - ) - - # store model results in calculator - for prop in property_keys: - if prop in model_results: - results[prop] = ( - model_results[prop].detach().cpu().numpy() - * self.property_units[prop] - ) - else: - raise AtomsConverterError( - "'{:s}' is not a property of your model. Please " - "check the model " - "properties!".format(prop) - ) - - self.results = results - self.atoms = atoms.copy() - - -class BatchwiseDynamics(Dynamics): - """Base-class for batch-wise MD and structure optimization classes.""" - - def __init__( - self, - calculator: BatchwiseCalculator, - atoms: List[Atoms], - logfile: str, - trajectory: Optional[str], - append_trajectory: bool = False, - master: Optional[bool] = None, - log_every_step: bool = False, - fixed_atoms_mask: Optional[List[int]] = None, - ): - """Structure dynamics object. - - Parameters: - - calculator: - This calculator provides properties such as forces and energy, which can be used for MD simulations or - relaxations - - atoms: - The Atoms objects to relax. - - restart: - Filename for restart file. Default value is *None*. - - logfile: - If *logfile* is a string, a file with that name will be opened. - Use '-' for stdout. - - trajectory: - Attach trajectory object. If *trajectory* is a string a - Trajectory will be constructed. Use *None* for no - trajectory. - - append_trajectory: - Appended to the trajectory file instead of overwriting it. - - master: - Defaults to None, which causes only rank 0 to save files. If - set to true, this rank will save files. - - log_every_step: - set to True to log Dynamics after each step (default=False) - - fixed_atoms: - list of indices corresponding to atoms with positions fixed in space. - """ - super().__init__( - atoms=atoms, - logfile=logfile, - trajectory=trajectory, - append_trajectory=append_trajectory, - master=master, - ) - - self.calculator = calculator - self.trajectory = trajectory - self.log_every_step = log_every_step - self.fixed_atoms_mask = fixed_atoms_mask - self.n_configs = len(self.atoms) - self.n_atoms = len(self.atoms[0]) - - def irun(self): - # compute initial structure and log the first step - self.calculator.get_forces(self.atoms, fixed_atoms_mask=self.fixed_atoms_mask) - - # yield the first time to inspect before logging - yield False - - if self.nsteps == 0: - self.log() - pass - - # run the algorithm until converged or max_steps reached - while not self.converged() and self.nsteps < self.max_steps: - - # compute the next step - self.step() - self.nsteps += 1 - - # let the user inspect the step and change things before logging - # and predicting the next step - yield False - - # log the step - if self.log_every_step: - self.log() - - # log last step - self.log() - - # finally check if algorithm was converged - yield self.converged() - - def run(self) -> bool: - """Run dynamics algorithm. - - This method will return when the forces on all individual - atoms are less than *fmax* or when the number of steps exceeds - *steps*.""" - - for converged in BatchwiseDynamics.irun(self): - pass - return converged - - -class BatchwiseOptimizer(BatchwiseDynamics): - """Base-class for all structure optimization classes.""" - - # default maxstep for all optimizers - defaults = {"maxstep": 0.2} - - def __init__( - self, - calculator: BatchwiseCalculator, - atoms: List[Atoms], - restart: Optional[bool] = None, - logfile: Optional[str] = None, - trajectory: Optional[str] = None, - master: Optional[str] = None, - append_trajectory: bool = False, - log_every_step: bool = False, - fixed_atoms_mask: Optional[List[int]] = None, - ): - """Structure optimizer object. - - Parameters: - - calculator: - This calculator provides properties such as forces and energy, which can be used for MD simulations or - relaxations - - atoms: - The Atoms objects to relax. - - restart: - Filename for restart file. Default value is *None*. - - logfile: - If *logfile* is a string, a file with that name will be opened. - Use '-' for stdout. - - trajectory: - Attach trajectory object. If *trajectory* is a string a - Trajectory will be constructed. Use *None* for no - trajectory. - - master: - Defaults to None, which causes only rank 0 to save files. If - set to true, this rank will save files. - - append_trajectory: - Appended to the trajectory file instead of overwriting it. - - log_every_step: - set to True to log Dynamics after each step (default=False) - - fixed_atoms: - list of indices corresponding to atoms with positions fixed in space. - """ - BatchwiseDynamics.__init__( - self, - calculator=calculator, - atoms=atoms, - logfile=logfile, - trajectory=trajectory, - master=master, - append_trajectory=append_trajectory, - log_every_step=log_every_step, - fixed_atoms_mask=fixed_atoms_mask, - ) - - self.restart = restart - - # initialize attribute - self.fmax = None - - if restart is None or not isfile(restart): - self.initialize() - else: - self.read() - barrier() - - def todict(self) -> Dict: - description = { - "type": "optimization", - "optimizer": self.__class__.__name__, - } - return description - - def initialize(self): - pass - - def irun(self, fmax: float = 0.05, steps: Optional[int] = None): - """call Dynamics.irun and keep track of fmax""" - self.fmax = fmax - if steps: - self.max_steps = steps - return BatchwiseDynamics.irun(self) - - def run(self, fmax: float = 0.05, steps: Optional[int] = None): - """call Dynamics.run and keep track of fmax""" - self.fmax = fmax - if steps: - self.max_steps = steps - return BatchwiseDynamics.run(self) - - def converged(self, forces: Optional[np.array] = None) -> bool: - """Did the optimization converge?""" - if forces is None: - forces = self.calculator.get_forces( - self.atoms, fixed_atoms_mask=self.fixed_atoms_mask - ) - # todo: maybe np.linalg.norm? - return (forces**2).sum(axis=1).max() < self.fmax**2 - - def log(self, forces: Optional[np.array] = None) -> None: - if forces is None: - forces = self.calculator.get_forces( - self.atoms, fixed_atoms_mask=self.fixed_atoms_mask - ) - fmax = sqrt((forces**2).sum(axis=1).max()) - T = time.localtime() - if self.logfile is not None: - name = self.__class__.__name__ - if self.nsteps == 0: - args = (" " * len(name), "Step", "Time", "fmax") - msg = "%s %4s %8s %12s\n" % args - self.logfile.write(msg) - - args = (name, self.nsteps, T[3], T[4], T[5], fmax) - msg = "%s: %3d %02d:%02d:%02d %12.4f\n" % args - self.logfile.write(msg) - - self.logfile.flush() - - if self.trajectory is not None: - for struc_idx, at in enumerate(self.atoms): - # store in trajectory - write( - self.trajectory + "_{}.xyz".format(struc_idx), - at, - format="extxyz", - append=False if self.nsteps == 0 else True, - ) - - def get_relaxation_results(self) -> Tuple[Atoms, Dict]: - self.calculator.get_forces(self.atoms) - return self.atoms, self.calculator.results - - def dump(self, data): - if world.rank == 0 and self.restart is not None: - with open(self.restart, "wb") as fd: - pickle.dump(data, fd, protocol=2) - - def load(self): - with open(self.restart, "rb") as fd: - return pickle.load(fd) - - -class ASEBatchwiseLBFGS(BatchwiseOptimizer): - """Limited memory BFGS optimizer. - - LBFGS optimizer that allows for relaxation of multiple structures in parallel. This optimizer is an - extension/adaptation of the ase.optimize.LBFGS optimizer particularly designed for batch-wise relaxation - of atomic structures. The inverse Hessian is approximated for each sample separately, which allows for - optimizing batches of different structures/compositions. - - """ - - def __init__( - self, - calculator: BatchwiseCalculator, - atoms: List[Atoms], - restart: Optional[bool] = None, - logfile: str = "-", - trajectory: Optional[str] = None, - maxstep: Optional[float] = None, - memory: int = 100, - damping: float = 1.0, - alpha: float = 70.0, - use_line_search: bool = False, - master: Optional[str] = None, - log_every_step: bool = False, - fixed_atoms_mask: Optional[List[int]] = None, - verbose: bool = False, - ): - """Parameters: - - calculator: - This calculator provides properties such as forces and energy, which can be used for MD simulations or - relaxations - - atoms: - The Atoms objects to relax. - - restart: - Pickle file used to store vectors for updating the inverse of - Hessian matrix. If set, file with such a name will be searched - and information stored will be used, if the file exists. - - logfile: - If *logfile* is a string, a file with that name will be opened. - Use '-' for stdout. - - trajectory: - Pickle file used to store trajectory of atomic movement. - - maxstep: - How far is a single atom allowed to move. This is useful for DFT - calculations where wavefunctions can be reused if steps are small. - Default is 0.2 Angstrom. - - memory: - Number of steps to be stored. Default value is 100. Three numpy - arrays of this length containing floats are stored. - - damping: - The calculated step is multiplied with this number before added to - the positions. - - alpha: - Initial guess for the Hessian (curvature of energy surface). A - conservative value of 70.0 is the default, but number of needed - steps to converge might be less if a lower value is used. However, - a lower value also means risk of instability. - - use_line_search: - Not implemented yet. - - master: - Defaults to None, which causes only rank 0 to save files. If - set to true, this rank will save files. - - log_every_step: - set to True to log Dynamics after each step (default=False) - - fixed_atoms: - list of indices corresponding to atoms with positions fixed in space. - """ - - BatchwiseOptimizer.__init__( - self, - calculator=calculator, - atoms=atoms, - restart=restart, - logfile=logfile, - trajectory=trajectory, - master=master, - log_every_step=log_every_step, - fixed_atoms_mask=fixed_atoms_mask, - ) - - if maxstep is not None: - self.maxstep = maxstep - else: - self.maxstep = self.defaults["maxstep"] - - if self.maxstep > 1.0: - raise ValueError( - "You are using a much too large value for " - + "the maximum step size: %.1f Angstrom" % maxstep - ) - - self.memory = memory - # Initial approximation of inverse Hessian 1./70. is to emulate the - # behaviour of BFGS. Note that this is never changed! - self.H0 = 1.0 / alpha - self.damping = damping - self.use_line_search = use_line_search - self.p = None - self.function_calls = 0 - self.force_calls = 0 - self.n_normalizations = 0 - - self.verbose = verbose - - if use_line_search: - raise NotImplementedError("Lines search has not been implemented yet") - - def initialize(self) -> None: - """Initialize everything so no checks have to be done in step""" - self.iteration = 0 - self.s = [] - self.y = [] - # Store also rho, to avoid calculating the dot product again and - # again. - self.rho = [] - - self.r0 = None - self.f0 = None - self.e0 = None - self.task = "START" - self.load_restart = False - - def read(self) -> None: - """Load saved arrays to reconstruct the Hessian""" - ( - self.iteration, - self.s, - self.y, - self.rho, - self.r0, - self.f0, - self.e0, - self.task, - ) = self.load() - self.load_restart = True - - def step(self, f: np.array = None) -> None: - """Take a single step - - Use the given forces, update the history and calculate the next step -- - then take it""" - - if f is None: - f = self.calculator.get_forces( - self.atoms, fixed_atoms_mask=self.fixed_atoms_mask - ) - - # check if updates for respective structures are required - q_euclidean = -f.reshape(self.n_configs, -1, 3) - squared_max_forces = (q_euclidean**2).sum(axis=-1).max(axis=-1) - configs_mask = squared_max_forces < self.fmax**2 - mask = ( - configs_mask[:, None] - .repeat(q_euclidean.shape[1], 0) - .repeat(q_euclidean.shape[2], 1) - ) - r = np.zeros((self.n_atoms * self.n_configs, 3), dtype=np.float64) - for config_idx, at in enumerate(self.atoms): - first_idx = config_idx * self.n_atoms - last_idx = config_idx * self.n_atoms + self.n_atoms - r[first_idx:last_idx] = at.get_positions() - - self.update(r, f, self.r0, self.f0) - - s = self.s - y = self.y - rho = self.rho - H0 = self.H0 - - loopmax = np.min([self.memory, self.iteration]) - a = np.empty( - ( - loopmax, - self.n_configs, - 1, - 1, - ), - dtype=np.float64, - ) - - # ## The algorithm itself: - q = -f.reshape(self.n_configs, 1, -1) - for i in range(loopmax - 1, -1, -1): - a[i] = rho[i] * np.matmul(s[i], np.transpose(q, axes=(0, 2, 1))) - q -= a[i] * y[i] - - z = H0 * q - - for i in range(loopmax): - b = rho[i] * np.matmul(y[i], np.transpose(z, axes=(0, 2, 1))) - z += s[i] * (a[i] - b) - - p = -z.reshape((-1, 3)) - self.p = np.where(mask, np.zeros_like(p), p) - # ## - - g = -f - if self.use_line_search is True: - e = self.func(r) - self.line_search(r, g, e) - dr = (self.alpha_k * self.p).reshape(self.n_atoms * self.n_configs, -1) - else: - self.force_calls += 1 - self.function_calls += 1 - dr = self.determine_step(self.p) * self.damping - - # update positions - pos_updated = r + dr - - # create new list of ase Atoms objects with updated positions - ats = [] - for config_idx, at in enumerate(self.atoms): - first_idx = config_idx * self.n_atoms - last_idx = config_idx * self.n_atoms + self.n_atoms - at = Atoms( - positions=pos_updated[first_idx:last_idx], - numbers=self.atoms[config_idx].get_atomic_numbers(), - ) - at.pbc = self.atoms[config_idx].pbc - at.cell = self.atoms[config_idx].cell - ats.append(at) - self.atoms = ats - - self.iteration += 1 - self.r0 = r - self.f0 = -g - self.dump( - ( - self.iteration, - self.s, - self.y, - self.rho, - self.r0, - self.f0, - self.e0, - self.task, - ) - ) - - def determine_step(self, dr: np.array) -> np.array: - """Determine step to take according to maxstep - - Normalize all steps as the largest step. This way - we still move along the eigendirection. - """ - steplengths = (dr**2).sum(-1) ** 0.5 - # check if any step in entire batch is greater than maxstep - if np.max(steplengths) >= self.maxstep: - # rescale steps for each config separately - for config_idx in range(self.n_configs): - first_idx = config_idx * self.n_atoms - last_idx = config_idx * self.n_atoms + self.n_atoms - longest_step = np.max(steplengths[first_idx:last_idx]) - if longest_step >= self.maxstep: - if self.verbose: - print("normalized integration step") - self.n_normalizations += 1 - dr[first_idx:last_idx] *= self.maxstep / longest_step - return dr - - def update(self, r: np.array, f: np.array, r0: np.array, f0: np.array) -> None: - """Update everything that is kept in memory - - This function is mostly here to allow for replay_trajectory. - """ - if self.iteration > 0: - s0 = r.reshape(self.n_configs, 1, -1) - r0.reshape(self.n_configs, 1, -1) - self.s.append(s0) - - # We use the gradient which is minus the force! - y0 = f0.reshape(self.n_configs, 1, -1) - f.reshape(self.n_configs, 1, -1) - self.y.append(y0) - - rho0 = np.ones((self.n_configs, 1, 1), dtype=np.float64) - for config_idx in range(self.n_configs): - ys0 = np.dot(y0[config_idx, 0], s0[config_idx, 0]) - if ys0 > 1e-8: - rho0[config_idx, 0, 0] = 1.0 / ys0 - self.rho.append(rho0) - - if self.iteration > self.memory: - self.s.pop(0) - self.y.pop(0) - self.rho.pop(0) - - def func(self, x): - """Objective function for use of the optimizers""" - raise NotImplementedError("func not implemented yet") - - def line_search(self, r, g, e): - self.alpha_k = None - raise NotImplementedError("LineSearch not implemented yet") diff --git a/src/schnetpack/md/neighborlist_md.py b/src/schnetpack/md/neighborlist_md.py index 9c2010d97..a62ed95b4 100644 --- a/src/schnetpack/md/neighborlist_md.py +++ b/src/schnetpack/md/neighborlist_md.py @@ -1,10 +1,13 @@ +from typing import Dict + import torch -import torch.nn as nn -from schnetpack.transform import NeighborListTransform, CollectAtomTriples from schnetpack.data.loader import _atoms_collate_fn -from typing import List, Dict -from schnetpack import properties +from schnetpack.transform import ( + BatchNeighborList, + CollectAtomTriples, + NeighborListTransform, +) __all__ = ["NeighborListMD"] @@ -13,6 +16,9 @@ class NeighborListMD: """ Wrapper for neighbor list transforms to make them suitable for molecular dynamics simulations. Introduces handling of multiple replicas and a cutoff shell (buffer region) to avoid recomputations of the neighbor list in every step. + + The work is done by :class:`~schnetpack.transform.BatchNeighborList`, shared with batchwise structure relaxation + framework. """ def __init__( @@ -39,63 +45,11 @@ def __init__( self.requires_triples = requires_triples self._collate = collate_fn - # Build neighbor list transform - self.transform = [base_nbl(self.cutoff_full)] - - if self.requires_triples: - self.transform.append(CollectAtomTriples()) - - self.transform = nn.Sequential(*self.transform) - - # Previous cells and positions for determining update - self.previous_positions = None - self.previous_cells = None - self.molecular_indices = None - - def _update_required( - self, - positions: torch.tensor, - cells: torch.tensor, - idx_m: torch.tensor, - n_molecules: int, - ): - """ - Use displacement and cell changes to determine, whether an update of the neighbor list is necessary. - - Args: - positions (torch.Tensor): Atom positions. - cells (torch.Tensor): Simulation cells. - idx_m (torch.Tensor): Molecular indices. - n_molecules (int): Number of molecules in simulation - - Returns: - bool: Udate is required. - """ - - if self.previous_positions is None: - # Everything needs to be updated - update_required = torch.ones(n_molecules, device=idx_m.device).bool() - elif n_molecules != len(self.molecular_indices): - self.molecular_indices = None - update_required = torch.ones(n_molecules, device=idx_m.device).bool() - else: - # Check for changes is positions - update_positions = ( - torch.norm(self.previous_positions - positions, dim=1) - > 0.5 * self.cutoff_shell - ).float() - - # Map to individual molecules - update_required = torch.zeros(n_molecules, device=idx_m.device).float() - update_required = update_required.index_add( - 0, idx_m, update_positions - ).bool() - - # Check for cell changes (is no cells are required, this will always be zero) - update_cells = torch.any((self.previous_cells != cells).view(-1, 9), dim=1) - update_required = torch.logical_or(update_required, update_cells) - - return update_required + self.neighbor_list = BatchNeighborList( + neighbor_list=base_nbl(cutoff), + cutoff_skin=cutoff_shell, + transforms=[CollectAtomTriples()] if requires_triples else None, + ) def get_neighbors(self, inputs: Dict[str, torch.Tensor]): """ @@ -105,130 +59,7 @@ def get_neighbors(self, inputs: Dict[str, torch.Tensor]): inputs (dict(str, torch.Tensor)): input batch. Returns: - torch.tensor: indices of neighbors. + dict(str, torch.Tensor): indices of neighbors, and nothing else -- the caller + merges them into the batch it already holds. """ - # TODO: check consistent wrapping - atom_types = inputs[properties.Z] - positions = inputs[properties.R] - n_atoms = inputs[properties.n_atoms] - idx_m = inputs[properties.idx_m] - cells = inputs[properties.cell] - pbc = inputs[properties.pbc] - - n_molecules = n_atoms.shape[0] - - # Check which molecular environments need to be updated - update_required = self._update_required(positions, cells, idx_m, n_molecules) - - if torch.any(update_required): - # if updated, store current positions and cells for future comparisons - self.previous_positions = positions.clone() - self.previous_cells = cells.clone() - - # Split everything into individual structures - input_batch = self._split_batch( - atom_types, positions, n_atoms, cells, pbc, n_molecules - ) - - # Set batch construct - if self.molecular_indices is None: - self.molecular_indices = [{} for _ in range(n_molecules)] - - # Check which molecule needs to be updated and compute neighborhoods - for idx in range(n_molecules): - if update_required[idx]: - # Get neighbors and if necessary triple indices - self.molecular_indices[idx] = self.transform(input_batch[idx]) - - # Remove superfluous entries before aggregation - del self.molecular_indices[idx][properties.R] - del self.molecular_indices[idx][properties.Z] - del self.molecular_indices[idx][properties.cell] - del self.molecular_indices[idx][properties.pbc] - - neighbor_idx = self._collate(self.molecular_indices) - # Remove n_atoms - del neighbor_idx[properties.n_atoms] - - # Move everything to correct device - neighbor_idx = {p: neighbor_idx[p].to(positions.device) for p in neighbor_idx} - - # filter out all pairs in the buffer zone - neighbor_idx = self._filter_indices(positions, neighbor_idx) - - return neighbor_idx - - def _filter_indices( - self, positions: torch.Tensor, neighbor_idx: Dict[str, torch.Tensor] - ) -> Dict[str, torch.Tensor]: - """ - Routine for filtering out pair indices and offets due to the buffer region, which would otherwise slow down - the calculators. - - Args: - positions (torch.Tensor): Tensor of the Cartesian atom positions. - neighbor_idx (dict(str, torch.Tensor)): Dictionary containing pair indices and offets - - Returns: - dict(str, torch.Tensor): Dictionary containing updated pair indices and offets - """ - offsets = neighbor_idx[properties.offsets] - idx_i = neighbor_idx[properties.idx_i] - idx_j = neighbor_idx[properties.idx_j] - - Rij = positions[idx_j] - positions[idx_i] + offsets - d_ij = torch.linalg.norm(Rij, dim=1) - d_ij_filter = d_ij <= self.cutoff - - neighbor_idx[properties.idx_i] = neighbor_idx[properties.idx_i][d_ij_filter] - neighbor_idx[properties.idx_j] = neighbor_idx[properties.idx_j][d_ij_filter] - neighbor_idx[properties.offsets] = neighbor_idx[properties.offsets][ - d_ij_filter, : - ] - - return neighbor_idx - - @staticmethod - def _split_batch( - atom_types: torch.Tensor, - positions: torch.Tensor, - n_atoms: torch.Tensor, - cells: torch.Tensor, - pbc: torch.Tensor, - n_molecules: int, - ) -> List[Dict[str, torch.tensor]]: - """ - Split the tensors containing molecular information into the different molecules for neighbor list computation. - Args: - atom_types (torch.Tensor): Atom type tensor. - positions (torch.Tensor): Atomic positions. - n_atoms (torch.Tensor): Number of atoms in each molecule. - cells (torch.Tensor): Simulation cells. - pbc (torch.Tensor): Periodic boundary conditions used for each molecule. - n_molecules (int): Number of molecules. - - Returns: - list(dict(str, torch.Tensor))): List of input dictionaries for each molecule. - """ - input_batch = [] - - idx_c = 0 - for idx_mol in range(n_molecules): - curr_n_atoms = n_atoms[idx_mol] - inputs = { - properties.n_atoms: torch.tensor([curr_n_atoms]).cpu(), - properties.Z: atom_types[idx_c : idx_c + curr_n_atoms].cpu(), - properties.R: positions[idx_c : idx_c + curr_n_atoms].cpu(), - } - - if cells is None: - inputs[properties.cell] = None - inputs[properties.pbc] = None - else: - inputs[properties.cell] = cells[idx_mol].cpu() - inputs[properties.pbc] = pbc[idx_mol].cpu() - - idx_c += curr_n_atoms - input_batch.append(inputs) - - return input_batch + return self.neighbor_list.neighbors(inputs) diff --git a/src/schnetpack/relax/__init__.py b/src/schnetpack/relax/__init__.py new file mode 100644 index 000000000..621d14d7c --- /dev/null +++ b/src/schnetpack/relax/__init__.py @@ -0,0 +1 @@ +from .batchwise_optimization import * diff --git a/src/schnetpack/relax/batchwise_optimization.py b/src/schnetpack/relax/batchwise_optimization.py new file mode 100644 index 000000000..f11e846cc --- /dev/null +++ b/src/schnetpack/relax/batchwise_optimization.py @@ -0,0 +1,681 @@ +""" +Batch-wise structure relaxation for SchNetPack models. + +``BatchwiseLBFGS`` relaxes a whole batch of structures in parallel, keeping one inverse +Hessian approximation per structure so that batches of differing compositions can be +optimized together. + +The batch is a SchNetPack input dictionary of torch tensors and stays that way for the +entire run -- this module contains no ase code at all. Trajectories go to a single +buffered HDF5 file (see :mod:`schnetpack.interfaces.batchwise_trajectory`), and callers +convert at the boundary with :func:`~schnetpack.interfaces.ase_interface.atoms_to_batch` +on the way in and :func:`~schnetpack.interfaces.ase_interface.batch_to_atoms` on the way +out. + +Note: + ``BatchwiseEnsembleCalculator`` and ``NNEnsemble`` have not been migrated to the + tensor-based calculator contract. ``BatchwiseEnsembleCalculator.calculate`` still + expects a list of ``ase.Atoms`` and returns numpy arrays, so the inherited + ``get_forces(inputs)`` raises. They are kept for backwards compatibility, are not + exercised by any test, and are the only ase-shaped thing left in this module. +""" + +import os +import sys +import time +from abc import ABC, abstractmethod +from contextlib import ExitStack +from copy import deepcopy +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union + +import torch +from torch import nn + +from schnetpack import properties +from schnetpack.relax.batchwise_trajectory import BatchwiseTrajectoryWriter +from schnetpack.units import convert_units +from schnetpack.utils.compatibility import load_model + +from schnetpack.transform import BatchNeighborList + +if TYPE_CHECKING: # import only for type checking, so the runtime stays ase-free + from ase import Atoms + +__all__ = [ + "BatchwiseLBFGS", + "BatchwiseCalculator", + "BatchwiseCalculatorError", + "BatchwiseEnsembleCalculator", + "BatchwiseOptimizer", + "NNEnsemble", +] + + +class BatchwiseCalculatorError(Exception): + pass + + +class NNEnsemble(nn.Module): + # TODO: integrate this into EnsembleCalculator directly + def __init__(self, models: nn.ModuleList, properties: List[str]): + super(NNEnsemble, self).__init__() + self.models = models + if isinstance(properties, str): + properties = [properties] + self.properties = properties + + def setup(self, stage: Optional[str] = None) -> None: + for model in self.models: + model.setup(stage) + + def forward(self, x: Dict) -> Tuple: + results = {p: [] for p in self.properties} + + inputs = deepcopy(x) + for model in self.models: + predictions = model(deepcopy(inputs)) + for prop, values in predictions.items(): + if prop in self.properties: + results[prop].append(values) + + means = {} + stds = {} + for prop, values in results.items(): + stacked_values = torch.stack(values) + means[prop] = stacked_values.mean(dim=0) + stds[prop] = stacked_values.std(dim=0) + + return means, stds + + +class BatchwiseCalculator: + """Evaluates a SchNetPack model on a whole batch of structures at once. + + Results are cached and only recomputed when the structure actually changed, so a + relaxation pays for exactly one model call per step even though the optimizer asks + for forces several times per step. + + Args: + model: trained model, or a path to one. The calculator evaluates it as it is + given -- to add a prior to the energy, compose it into the model's output + modules beforehand (see ``examples/howtos/howto_batchwise_relaxations.ipynb``). + neighbor_list: keeps the batch's neighbor lists valid as the structures move. + Most steps reuse the previous list rather than rebuilding it, which is a + large part of why relaxing a batch pays off. + device: device the model runs on. + energy_key, force_key, stress_key: names of these properties in the model. + ``stress_key=None`` disables stress. + energy_unit, position_unit: units the model works in. Results are converted to + ase units (eV, Angstrom). + dtype: precision of the model input. + """ + + def __init__( + self, + model: Union[nn.Module, str], + neighbor_list: BatchNeighborList, + device: Union[str, torch.device] = "cpu", + energy_key: str = "energy", + force_key: str = "forces", + stress_key: Optional[str] = None, + energy_unit: str = "eV", + position_unit: str = "Ang", + dtype: torch.dtype = torch.float32, + ): + self.results = None + self.device = torch.device(device) if isinstance(device, str) else device + self.dtype = dtype + self.neighbor_list = neighbor_list + + self.energy_key = energy_key + self.force_key = force_key + self.stress_key = stress_key + + # unit conversion to default ase units + energy_conversion = convert_units(energy_unit, "eV") + position_conversion = convert_units(position_unit, "Angstrom") + self.property_units = { + energy_key: energy_conversion, + force_key: energy_conversion / position_conversion, + } + if stress_key is not None: + self.property_units[stress_key] = energy_conversion / position_conversion**3 + + if isinstance(model, str): + model = self._load_model(model) + self._initialize_model(model) + + # the structure self.results was computed for + self._cached_structure = None + + def _load_model(self, model: str) -> nn.Module: + return load_model(model, device="cpu").to(torch.float64) + + def _initialize_model(self, model: nn.Module) -> None: + self.model = model.eval() + self.model.to(device=self.device, dtype=self.dtype) + + #: input entries that decide whether a cached result is still valid + _structure_keys = (properties.R, properties.cell, properties.pbc) + + def _structure_id(self, inputs: Dict[str, torch.Tensor]) -> Tuple: + """Fingerprint of the structure the inputs describe. + + Pairs each tensor with ``_version``, the counter autograd bumps on in-place + mutation, so both a rebound entry and an updated one are noticed. Unlike an + element-wise comparison this never synchronizes with the device, which matters + because it is checked on every force request. Keeping the tensors themselves in + the fingerprint also keeps them alive, so a freed tensor cannot be mistaken for + the cached one. + """ + return tuple( + (inputs[key], inputs[key]._version) for key in self._structure_keys + ) + + def _requires_calculation( + self, property_keys: List[str], inputs: Dict[str, torch.Tensor] + ) -> bool: + if self.results is None or self._cached_structure is None: + return True + if any(name not in self.results for name in property_keys): + return True + return any( + inputs[key] is not tensor or tensor._version != version + for (tensor, version), key in zip( + self._cached_structure, self._structure_keys + ) + ) + + def get_forces( + self, + inputs: Dict[str, torch.Tensor], + fixed_atoms_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Forces on every atom of the batch, in eV/Angstrom. + + Args: + inputs: schnetpack input batch. + fixed_atoms_mask: boolean mask selecting the atoms to return forces for. + Defaults to all of them. + """ + if self._requires_calculation([self.energy_key, self.force_key], inputs): + self.calculate(inputs) + forces = self.results[self.force_key] + return forces if fixed_atoms_mask is None else forces[fixed_atoms_mask] + + def get_potential_energy(self, inputs: Dict[str, torch.Tensor]) -> torch.Tensor: + """Potential energy of every structure of the batch, in eV.""" + if self._requires_calculation([self.energy_key], inputs): + self.calculate(inputs) + return self.results[self.energy_key] + + def calculate(self, inputs: Dict[str, torch.Tensor]) -> None: + structure_id = self._structure_id(inputs) + + # Shallow copy: the update replaces entries (neighbor lists, casts) and must not + # write them back into the caller's batch, but the tensors themselves are only + # read, so there is no reason to copy them all -- deep copying the batch would + # mean copying the neighbor list on every single step. + # Positions are the exception: the model marks them as requiring grad to get the + # forces, and would otherwise turn the optimizer's live positions into a leaf + # variable that can no longer be updated in place. + inputs = dict(inputs) + inputs[properties.R] = inputs[properties.R].detach().clone() + inputs = self.neighbor_list.update(inputs) + + model_results = self.model(inputs) + + results = {} + for prop, unit in self.property_units.items(): + if prop not in model_results: + raise BatchwiseCalculatorError( + f"'{prop}' is not a property of your model. " + "Please check the model properties!" + ) + results[prop] = model_results[prop].detach() * unit + + self.results = results + self._cached_structure = structure_id + + +class BatchwiseEnsembleCalculator(BatchwiseCalculator): + """Calculator for an ensemble of models, reporting per-property uncertainties. + + Warning: + Not migrated to the tensor-based calculator contract, see the module docstring. + ``calculate`` below still expects a list of ``ase.Atoms`` and returns numpy, so + the inherited ``get_forces(inputs)`` raises. + + Args: + model: directory of trained models, or a module list of them. Remaining + arguments are those of :class:`BatchwiseCalculator`. + """ + + # TODO: inherit from SpkEnsembleCalculator + def _load_model(self, model: str) -> nn.ModuleList: + models = torch.nn.ModuleList() + for model_name in os.listdir(model): + models.append( + load_model( + os.path.join(model, model_name, "best_model"), device="cpu" + ).to(torch.float64) + ) + return models + + def _initialize_model(self, model: nn.ModuleList) -> None: + ensemble = NNEnsemble(models=model, properties=list(self.property_units.keys())) + self.model = ensemble.eval().to(device=self.device, dtype=self.dtype) + + def calculate(self, atoms: List["Atoms"]) -> None: + from schnetpack.interfaces.ase_interface import atoms_to_batch + + inputs = self.neighbor_list.update( + atoms_to_batch(atoms, device=self.device, dtype=self.dtype) + ) + model_results, stds = self.model(inputs) + + results = {} + for prop, unit in self.property_units.items(): + if prop not in model_results: + raise BatchwiseCalculatorError( + f"'{prop}' is not a property of your model. " + "Please check the model properties!" + ) + results[prop] = model_results[prop].detach().cpu().numpy() * unit + results[f"{prop}_uncertainty"] = stds[prop].detach().cpu().numpy() * unit + + self.results = results + self.atoms = atoms.copy() + + +class BatchwiseOptimizer(ABC): + """Drives a batch of structures downhill until every one of them is relaxed. + + Subclasses supply :meth:`step`; everything else -- the run loop, the convergence + criterion, the text log and the HDF5 trajectory -- lives here. + + Positions are updated in place, so ``inputs`` holds the relaxed structures once the + run is over. + + Args: + calculator: provides the forces and energies driving the relaxation. + inputs: schnetpack input batch holding the structures to relax. All structures + must have the same number of atoms. + logfile: text progress log. A path, ``"-"`` for stdout, or ``None`` for no log. + log_interval: how often to write a log line. See below. + trajectory: path of the HDF5 trajectory to write, or ``None`` for none. + trajectory_interval: how often to write a trajectory frame. See below. + store_forces: store the forces of every trajectory frame, not just the + positions. Doubles the file size. + fixed_atoms_mask: boolean mask over all atoms in the batch, True for atoms + whose positions are held fixed in space. + max_steps: step limit used when ``run`` is called without one. + + Both intervals count optimizer steps: ``0`` writes only the first and last frame, + ``1`` writes every step, ``n`` writes every nth. The first and last frame are + always written whatever the interval, and the two intervals are independent. + """ + + def __init__( + self, + calculator: BatchwiseCalculator, + inputs: Dict[str, torch.Tensor], + logfile: Optional[str] = None, + log_interval: int = 1, + trajectory: Optional[str] = None, + trajectory_interval: int = 0, + store_forces: bool = False, + fixed_atoms_mask: Optional[List[bool]] = None, + max_steps: int = 100_000, + ): + self.calculator = calculator + self.inputs = inputs + self.nsteps = 0 + self.max_steps = max_steps + self.fmax = None + # per-structure squared max force for the forces the last convergence check + # saw, reused by step() and the loggers so the reduction is not repeated + self._max_sq_force_per_config = None + + n_atoms = inputs[properties.n_atoms] + self.n_configs = n_atoms.shape[0] + if not bool((n_atoms == n_atoms[0]).all()): + raise ValueError( + "batch-wise optimization requires all structures in the batch to have " + f"the same number of atoms, got {n_atoms.tolist()}" + ) + self.n_atoms = int(n_atoms[0]) + + # kept as a float mask rather than an index: zeroing the displacement of fixed + # atoms is equivalent to dropping them from the optimization (their history + # contributions are identically zero) and needs no device synchronization, + # which boolean-mask indexing would force on every step + n_total_atoms = self.n_configs * self.n_atoms + device = inputs[properties.R].device + if fixed_atoms_mask is None: + self.free_atoms = torch.ones( + (n_total_atoms, 1), dtype=torch.float64, device=device + ) + else: + fixed = torch.as_tensor(fixed_atoms_mask, dtype=torch.bool).view(-1, 1) + if fixed.shape[0] != n_total_atoms: + raise ValueError( + f"fixed_atoms_mask has {fixed.shape[0]} entries, expected one per " + f"atom in the batch ({n_total_atoms})" + ) + self.free_atoms = (~fixed).to(dtype=torch.float64, device=device) + + self._closer = ExitStack() + self.log_interval = log_interval + if logfile is None: + self.logfile = None + elif logfile == "-": + self.logfile = sys.stdout + else: + self.logfile = self._closer.enter_context( + open(logfile, "a", encoding="utf-8") + ) + + self.trajectory = trajectory + self.trajectory_interval = trajectory_interval + self.store_forces = store_forces + self._writer = None + + # ------------------------------------------------------------------ run loop + + @abstractmethod + def step(self) -> None: + """Move every structure of the batch one step downhill.""" + + def max_squared_force_per_config(self, forces: torch.Tensor) -> torch.Tensor: + """Largest squared force norm within each structure of the batch. + + Fixed atoms are excluded -- their residual force says nothing about whether the + free atoms have relaxed. + """ + squared = forces.view(self.n_configs, self.n_atoms, 3).pow(2).sum(-1) + squared = squared * self.free_atoms.view(self.n_configs, self.n_atoms) + return squared.max(-1).values + + def converged(self, forces: Optional[torch.Tensor] = None) -> bool: + """Is every structure of the batch relaxed to within ``fmax``?""" + if forces is None: + forces = self.calculator.get_forces(self.inputs) + self._max_sq_force_per_config = self.max_squared_force_per_config(forces) + return bool((self._max_sq_force_per_config.max() < self.fmax**2).item()) + + def irun(self, fmax: float = 0.05, steps: Optional[int] = None): + """Drive the relaxation step by step, yielding after every step. + + The final value yielded is whether the batch converged. + """ + self.fmax = fmax + if steps is not None: + self.max_steps = steps + + converged = False + while True: + # one model call per iteration; the loggers and step() below reuse it + converged = self.converged() + final = converged or self.nsteps >= self.max_steps + self._write_frame(final=final) + if final: + break + + self.step() + self.nsteps += 1 + # let the caller inspect the step before the next one is computed + yield False + + yield converged + + def run(self, fmax: float = 0.05, steps: Optional[int] = None) -> bool: + """Relax until converged or ``steps`` steps have been taken. + + Returns whether the maximum force on every free atom dropped below ``fmax``. + """ + converged = False + for converged in self.irun(fmax=fmax, steps=steps): + pass + return converged + + def get_relaxation_results( + self, + ) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]: + """The relaxed batch and the model results for it. + + Both are the live tensors, at full precision and on the calculator's device. + Use :func:`schnetpack.interfaces.ase_interface.batch_to_atoms` on the batch to + get ``ase.Atoms`` back. + """ + self.calculator.get_forces(self.inputs) + return self.inputs, self.calculator.results + + # ------------------------------------------------------------------- logging + + def _is_due(self, interval: int, final: bool) -> bool: + if final or self.nsteps == 0: + return True + return interval > 0 and self.nsteps % interval == 0 + + def _write_frame(self, final: bool = False) -> None: + """Write a log line and a trajectory frame, if this step calls for them.""" + log_now = self.logfile is not None and self._is_due(self.log_interval, final) + trajectory_now = self.trajectory is not None and self._is_due( + self.trajectory_interval, final + ) + if not (log_now or trajectory_now): + # nothing to write, so do not pay for the transfer off the device + return + + forces = self.calculator.get_forces(self.inputs) + if self._max_sq_force_per_config is None: + self._max_sq_force_per_config = self.max_squared_force_per_config(forces) + + if log_now: + self._write_log_line(self._max_sq_force_per_config.max().sqrt().item()) + if trajectory_now: + self._write_trajectory_frame(forces) + + def _write_log_line(self, fmax: float) -> None: + name = self.__class__.__name__ + if self.nsteps == 0: + header = (" " * len(name), "Step", "Time", "fmax") + self.logfile.write("%s %4s %8s %12s\n" % header) + + clock = time.localtime() + line = (name, self.nsteps, clock[3], clock[4], clock[5], fmax) + self.logfile.write("%s: %3d %02d:%02d:%02d %12.4f\n" % line) + self.logfile.flush() + + def _write_trajectory_frame(self, forces: torch.Tensor) -> None: + if self._writer is None: + self._writer = self._closer.enter_context( + BatchwiseTrajectoryWriter( + self.trajectory, + atomic_numbers=self.inputs[properties.Z], + pbc=self.inputs[properties.pbc], + store_forces=self.store_forces, + attrs={"optimizer": self.__class__.__name__, "fmax": self.fmax}, + ) + ) + self._writer.write( + step=self.nsteps, + positions=self.inputs[properties.R], + cell=self.inputs[properties.cell], + energy=self.calculator.get_potential_energy(self.inputs), + forces=forces if self.store_forces else None, + converged=self._max_sq_force_per_config < self.fmax**2, + ) + + def close(self) -> None: + """Close the log file and the trajectory.""" + self._closer.close() + self._writer = None + + def __enter__(self) -> "BatchwiseOptimizer": + return self + + def __exit__(self, *args) -> None: + self.close() + + +class BatchwiseLBFGS(BatchwiseOptimizer): + """Limited memory BFGS, relaxing a batch of structures in parallel. + + An adaptation of ``ase.optimize.LBFGS`` for batch-wise relaxation: the inverse + Hessian is approximated for each structure separately, so batches of different + structures and compositions can be optimized together. + + Args: + maxstep: how far a single atom is allowed to move in one step, in Angstrom. + Each structure of the batch is rescaled on its own. + memory: number of steps of history kept for the two-loop recursion. + damping: the calculated step is multiplied by this before it is taken. + alpha: initial guess for the curvature of the energy surface. The conservative + default of 70.0 emulates BFGS; a lower value may converge in fewer steps at + the cost of stability. + device: device the L-BFGS bookkeeping runs on (default: cpu). The two-loop + recursion is bound by kernel launches rather than arithmetic -- it was + measured 5-7x slower on cuda than on cpu for batches up to 256 structures + of 1000 atoms -- so the default is cpu regardless of where the model runs. + Worth re-measuring before overriding for much larger batches. + + Remaining keyword arguments are those of :class:`BatchwiseOptimizer`. + """ + + #: how far a single atom may move in one step, in Angstrom + default_maxstep = 0.2 + + def __init__( + self, + calculator: BatchwiseCalculator, + inputs: Dict[str, torch.Tensor], + maxstep: Optional[float] = None, + memory: int = 100, + damping: float = 1.0, + alpha: float = 70.0, + device: Optional[Union[str, torch.device]] = None, + **kwargs, + ): + super().__init__(calculator=calculator, inputs=inputs, **kwargs) + + self.maxstep = self.default_maxstep if maxstep is None else maxstep + if self.maxstep > 1.0: + raise ValueError( + "You are using a much too large value for the maximum step size: " + f"{self.maxstep:.1f} Angstrom" + ) + + self.memory = memory + # Initial approximation of the inverse Hessian, 1./70. to emulate the behaviour + # of BFGS. Note that this is never changed! + self.H0 = 1.0 / alpha + self.damping = damping + self.device = ( + torch.device(device) if device is not None else torch.device("cpu") + ) + # same mask as self.free_atoms, but on the device the recursion runs on + self._free_atoms_opt = self.free_atoms.to(self.device) + + self.iteration = 0 + self.s = [] + self.y = [] + # rho is stored alongside, to avoid calculating the dot product again and again + self.rho = [] + self.r0 = None + self.f0 = None + + def step( + self, f: Optional[torch.Tensor] = None, normalize_step: bool = True + ) -> None: + """Update the history, compute the next step, and take it.""" + if f is None: + f = self.calculator.get_forces(self.inputs) + # forces on fixed atoms are zeroed rather than removed: their history + # contributions vanish, so the recursion below is unchanged by them + f = f.to(device=self.device, dtype=torch.float64) * self._free_atoms_opt + + # structures that already meet the force criterion must not be moved further. + # f is masked, so the fixed atoms drop out of the reduction on their own + if self._max_sq_force_per_config is None: + max_sq_force = f.view(self.n_configs, self.n_atoms, 3).pow(2).sum(-1) + max_sq_force = max_sq_force.max(-1).values + else: + max_sq_force = self._max_sq_force_per_config.to(self.device) + relaxed = max_sq_force < self.fmax**2 + + r = self.inputs[properties.R].to(device=self.device, dtype=torch.float64) + self.update(r, f, self.r0, self.f0) + + loopmax = min(self.memory, self.iteration) + a = torch.empty( + (loopmax, self.n_configs, 1), dtype=torch.float64, device=self.device + ) + + # ## The algorithm itself: + q = -f.view(self.n_configs, -1) + for i in range(loopmax - 1, -1, -1): + a[i] = self.rho[i] * (self.s[i] * q).sum(-1, keepdim=True) + q -= a[i] * self.y[i] + + z = self.H0 * q + + for i in range(loopmax): + b = self.rho[i] * (self.y[i] * z).sum(-1, keepdim=True) + z += self.s[i] * (a[i] - b) + + p = -z.view(self.n_configs, self.n_atoms, 3) + # broadcast rather than materialize a full-size boolean mask + p = p * (~relaxed).view(-1, 1, 1) + # ## + + dr = self.determine_step(p) if normalize_step else p.view(-1, 3) + dr = dr * self.damping + + self.inputs[properties.R] += dr.to( + device=self.inputs[properties.R].device, + dtype=self.inputs[properties.R].dtype, + ) + # the forces the cached reduction belongs to are stale now + self._max_sq_force_per_config = None + + self.iteration += 1 + self.r0 = r + self.f0 = f + + def determine_step(self, dr: torch.Tensor) -> torch.Tensor: + """Scale the step down to ``maxstep``, each structure on its own. + + Every atom of a structure is scaled by the same factor, so the step still + points along the eigendirection. + """ + dr = dr.view(self.n_configs, self.n_atoms, 3) + longest_step = dr.pow(2).sum(-1).sqrt().max(dim=1, keepdim=True).values + # clamp instead of branching: structures below maxstep are scaled by 1 + scale = (self.maxstep / longest_step).clamp(max=1.0) + return (dr * scale.unsqueeze(-1)).view(-1, 3) + + def update( + self, + r: torch.Tensor, + f: torch.Tensor, + r0: Optional[torch.Tensor], + f0: Optional[torch.Tensor], + ) -> None: + """Append the latest position and gradient difference to the history.""" + if self.iteration > 0: + s0 = (r - r0).view(self.n_configs, -1) + self.s.append(s0) + + # we use the gradient, which is minus the force + y0 = (f0 - f).view(self.n_configs, -1) + self.y.append(y0) + + ys0 = (y0 * s0).sum(-1, keepdim=True) + self.rho.append(torch.where(ys0 > 1e-8, 1.0 / ys0, torch.zeros_like(ys0))) + + if self.iteration > self.memory: + self.s.pop(0) + self.y.pop(0) + self.rho.pop(0) diff --git a/src/schnetpack/relax/batchwise_trajectory.py b/src/schnetpack/relax/batchwise_trajectory.py new file mode 100644 index 000000000..fd6f0a934 --- /dev/null +++ b/src/schnetpack/relax/batchwise_trajectory.py @@ -0,0 +1,302 @@ +"""HDF5 trajectories for batch-wise structure relaxations. + +A batch-wise relaxation propagates ``n_structures`` structures of equal size in +lockstep, so every frame is a dense ``(n_structures, n_atoms, ...)`` block. That is +stored here as named, directly sliceable datasets:: + + / attrs: n_structures, n_atoms, optimizer, fmax, + energy_unit, position_unit, schnetpack_version + /atomic_numbers (n_structures, n_atoms) + /pbc (n_structures, 3) + /steps (n_frames,) optimizer step per frame + /positions (n_frames, n_structures, n_atoms, 3) + /cell (n_frames, n_structures, 3, 3) + /energy (n_frames, n_structures) + /converged (n_frames, n_structures) + /forces (n_frames, n_structures, n_atoms, 3) only if store_forces + +so that, e.g., ``reader.positions[:, 3]`` is the whole path of structure 3 without +touching the rest of the file. + +This is deliberately not the layout ``schnetpack.md`` writes. That one carries a +replica axis, packs positions, energies, cells and stresses into a single flat array +decoded by hard-coded offsets, and needs the masses and time step of an MD run -- +none of which a relaxation has. +""" + +from typing import Dict, Optional, Union + +import h5py +import numpy as np +import torch + +from schnetpack import properties + +__all__ = ["BatchwiseTrajectoryWriter", "BatchwiseTrajectoryReader"] + + +def _schnetpack_version() -> str: + """Recorded so a trajectory can be traced back to the code that wrote it.""" + from schnetpack import __version__ + + return __version__ + + +def _to_numpy(value: Union[torch.Tensor, np.ndarray], dtype) -> np.ndarray: + if isinstance(value, torch.Tensor): + value = value.detach().cpu().numpy() + return np.asarray(value, dtype=dtype) + + +class BatchwiseTrajectoryWriter: + """Buffered writer for the layout described in the module docstring. + + Frames are accumulated in memory and written one slab at a time, so a relaxation + pays for ``n_frames / buffer_size`` HDF5 writes rather than one per frame. The + frame datasets grow as the run goes on, since a relaxation does not know how many + steps it will take. + + Args: + filename: path of the HDF5 file to create. + atomic_numbers: (n_structures, n_atoms) or (n_structures * n_atoms,). + pbc: (n_structures, 3) periodic boundary conditions. + store_forces: also store the forces of every frame. Off by default -- they are + as large as the positions, and the forces of the final frame come back from + the optimizer anyway. + buffer_size: frames held in memory before a write, doubling as the HDF5 chunk + size along the frame axis. ``None`` derives it from ``max_buffer_bytes``, + which keeps chunks near the size HDF5 is happy with whatever the batch + size -- 64 frames of a 256 x 1000 atom batch would be a 200 MB chunk, two + orders of magnitude above the recommended maximum and a 200 MB read for + any caller that wanted a single frame. One frame is the floor, so a batch + whose single frame already exceeds the budget still gets a chunk that big. + max_buffer_bytes: byte budget used to derive ``buffer_size``. + precision: 32 or 64, the float precision of the stored data. + attrs: extra root attributes, e.g. the optimizer name and the force criterion. + """ + + def __init__( + self, + filename: str, + atomic_numbers: Union[torch.Tensor, np.ndarray], + pbc: Union[torch.Tensor, np.ndarray], + store_forces: bool = False, + buffer_size: Optional[int] = None, + precision: int = 32, + attrs: Optional[Dict] = None, + max_buffer_bytes: int = 8 << 20, + ): + if precision not in (32, 64): + raise ValueError(f"precision must be 32 or 64, got {precision}") + self.dtype = np.float32 if precision == 32 else np.float64 + self.store_forces = store_forces + + pbc = _to_numpy(pbc, bool).reshape(-1, 3) + self.n_structures = pbc.shape[0] + atomic_numbers = _to_numpy(atomic_numbers, np.int32).reshape( + self.n_structures, -1 + ) + self.n_atoms = atomic_numbers.shape[1] + + self.file = h5py.File(filename, "w") + self.file.attrs["n_structures"] = self.n_structures + self.file.attrs["n_atoms"] = self.n_atoms + self.file.attrs["energy_unit"] = "eV" + self.file.attrs["position_unit"] = "Ang" + self.file.attrs["schnetpack_version"] = _schnetpack_version() + for key, value in (attrs or {}).items(): + self.file.attrs[key] = value + + self.file.create_dataset("atomic_numbers", data=atomic_numbers) + self.file.create_dataset("pbc", data=pbc) + + # (name, per-frame shape, dtype) + specs = [ + ("steps", (), np.int32), + ("positions", (self.n_structures, self.n_atoms, 3), self.dtype), + ("cell", (self.n_structures, 3, 3), self.dtype), + ("energy", (self.n_structures,), self.dtype), + ("converged", (self.n_structures,), bool), + ] + if store_forces: + specs.append(("forces", (self.n_structures, self.n_atoms, 3), self.dtype)) + + if buffer_size is None: + frame_bytes = sum( + int(np.prod(shape, dtype=np.int64)) * np.dtype(dtype).itemsize + for _, shape, dtype in specs + ) + buffer_size = int(np.clip(max_buffer_bytes // max(frame_bytes, 1), 1, 64)) + self.buffer_size = buffer_size + + self.datasets = {} + self.buffers = {} + for name, shape, dtype in specs: + self.datasets[name] = self.file.create_dataset( + name, + shape=(0,) + shape, + maxshape=(None,) + shape, + chunks=(buffer_size,) + shape, + dtype=dtype, + ) + self.buffers[name] = np.zeros((buffer_size,) + shape, dtype=dtype) + + self.n_frames = 0 + self._buffered = 0 + + def write( + self, + step: int, + positions: torch.Tensor, + cell: torch.Tensor, + energy: Optional[torch.Tensor] = None, + forces: Optional[torch.Tensor] = None, + converged: Optional[torch.Tensor] = None, + ) -> None: + """Append one frame. Tensors may live on any device. + + ``positions`` and ``forces`` are taken flat, ``(n_structures * n_atoms, 3)``, + the way the optimizer holds them. + """ + shape = (self.n_structures, self.n_atoms, 3) + frame = { + "steps": step, + "positions": _to_numpy(positions, self.dtype).reshape(shape), + "cell": _to_numpy(cell, self.dtype).reshape(self.n_structures, 3, 3), + "energy": ( + np.zeros(self.n_structures, self.dtype) + if energy is None + else _to_numpy(energy, self.dtype).reshape(self.n_structures) + ), + "converged": ( + np.zeros(self.n_structures, bool) + if converged is None + else _to_numpy(converged, bool).reshape(self.n_structures) + ), + } + if self.store_forces: + frame["forces"] = ( + np.zeros(shape, self.dtype) + if forces is None + else _to_numpy(forces, self.dtype).reshape(shape) + ) + + for name, value in frame.items(): + self.buffers[name][self._buffered] = value + self._buffered += 1 + self.n_frames += 1 + + if self._buffered == self.buffer_size: + self.flush() + + def flush(self) -> None: + """Write the buffered frames out and empty the buffer.""" + if self._buffered == 0: + return + start = self.n_frames - self._buffered + for name, dataset in self.datasets.items(): + dataset.resize(self.n_frames, axis=0) + dataset[start : self.n_frames] = self.buffers[name][: self._buffered] + self._buffered = 0 + self.file.flush() + + def close(self) -> None: + if self.file: + self.flush() + self.file.close() + self.file = None + + def __enter__(self) -> "BatchwiseTrajectoryWriter": + return self + + def __exit__(self, *args) -> None: + self.close() + + +class BatchwiseTrajectoryReader: + """Read-only view on a file written by ``BatchwiseTrajectoryWriter``. + + The frame arrays are exposed as h5py datasets rather than numpy arrays, so + ``reader.positions[:, 3]`` reads one structure's path off disk without + materializing the whole trajectory. + """ + + _frame_keys = ("steps", "positions", "cell", "energy", "converged", "forces") + + def __init__(self, filename: str): + self.file = h5py.File(filename, "r") + self.n_structures = int(self.file.attrs["n_structures"]) + self.n_atoms = int(self.file.attrs["n_atoms"]) + self.n_frames = self.file["positions"].shape[0] + + def __getattr__(self, name: str): + # datasets are reached as attributes: reader.positions, reader.energy, ... + if name in self._frame_keys or name in ("atomic_numbers", "pbc"): + try: + return self.__dict__["file"][name] + except KeyError: + raise AttributeError( + f"'{name}' was not stored in this trajectory" + ) from None + raise AttributeError(name) + + @property + def has_forces(self) -> bool: + return "forces" in self.file + + def frame(self, index: int = -1) -> Dict[str, torch.Tensor]: + """A single frame as a schnetpack input batch, defaulting to the last one. + + The result is a complete batch, so it feeds straight into + :func:`schnetpack.interfaces.ase_interface.batch_to_atoms` or back into an + optimizer to continue from that frame. Energies, forces and the convergence + flags of the frame come along under their usual property keys. + """ + n_structures, n_atoms = self.n_structures, self.n_atoms + batch = { + properties.n_atoms: torch.full((n_structures,), n_atoms, dtype=torch.long), + properties.idx_m: torch.repeat_interleave( + torch.arange(n_structures), n_atoms + ), + properties.Z: torch.as_tensor( + self.file["atomic_numbers"][:], dtype=torch.long + ).view(-1), + properties.pbc: torch.as_tensor(self.file["pbc"][:], dtype=torch.bool), + properties.R: torch.as_tensor(self.file["positions"][index]).view(-1, 3), + properties.cell: torch.as_tensor(self.file["cell"][index]), + } + batch[properties.energy] = torch.as_tensor(self.file["energy"][index]) + batch["converged"] = torch.as_tensor(self.file["converged"][index]) + batch["step"] = int(self.file["steps"][index]) + if self.has_forces: + batch[properties.forces] = torch.as_tensor(self.file["forces"][index]).view( + -1, 3 + ) + return batch + + def structure(self, index: int) -> Dict[str, np.ndarray]: + """The whole path of a single structure, as numpy arrays over frames. + + Unlike :meth:`frame` this is not a batch -- it spans frames rather than + structures, and is meant for analysis and plotting. + """ + structure = { + "atomic_numbers": self.file["atomic_numbers"][index], + "pbc": self.file["pbc"][index], + "steps": self.file["steps"][:], + } + for key in ("positions", "cell", "energy", "converged", "forces"): + if key in self.file: + structure[key] = self.file[key][:, index] + return structure + + def close(self) -> None: + if self.file: + self.file.close() + self.file = None + + def __enter__(self) -> "BatchwiseTrajectoryReader": + return self + + def __exit__(self, *args) -> None: + self.close() diff --git a/src/schnetpack/transform/__init__.py b/src/schnetpack/transform/__init__.py index 5076539ee..314035c1b 100644 --- a/src/schnetpack/transform/__init__.py +++ b/src/schnetpack/transform/__init__.py @@ -10,5 +10,6 @@ from .atomistic import * from .casting import * from .neighborlist import * +from .batch_neighborlist import * from .response import * from .base import * diff --git a/src/schnetpack/transform/batch_neighborlist.py b/src/schnetpack/transform/batch_neighborlist.py new file mode 100644 index 000000000..ffbe61a1d --- /dev/null +++ b/src/schnetpack/transform/batch_neighborlist.py @@ -0,0 +1,300 @@ +"""Neighbor lists for a batch of structures that keeps moving. + +A relaxation or an MD run evaluates the same structures over and over, moving them a +little each step. Rebuilding their neighbor lists every time is wasteful, so +:class:`BatchNeighborList` builds them out to ``cutoff + cutoff_skin`` and reuses them for +as long as no atom of a structure has drifted more than half the skin. Each step the +cached list is restricted to the pairs actually within the cutoff, which is a single +masked gather on the device the batch already lives on. + +This is the batch-level counterpart of +:class:`~schnetpack.transform.SkinNeighborList`, which does the same thing one sample at a +time while a batch is being assembled. Unlike the transforms it is not a +:class:`~schnetpack.transform.Transform`: it holds state, and it works on a collated batch +rather than on a single sample. +""" + +from typing import Dict, List, Optional, Sequence, Union + +import torch + +from schnetpack import properties +from schnetpack.data.loader import _atoms_collate_fn, split_batch + +from .base import Transform + +__all__ = ["BatchNeighborList"] + +#: entries of a cached list that are pair-indexed, i.e. shrink when the skin is pruned +_PAIR_KEYS = ( + properties.idx_i, + properties.idx_j, + properties.offsets, + properties.lidx_i, + properties.lidx_j, +) + +#: entries that index into the pair arrays and have to be renumbered along with them +_TRIPLE_KEYS = (properties.idx_j_triples, properties.idx_k_triples) + + +class BatchNeighborList: + """Keeps the neighbor lists of a batch valid while its structures move. + + The batch is a schnetpack input dictionary and stays on its device throughout, except + on the steps that have to rebuild: the neighbor list implementations run on cpu, so + the structures needing a new list -- and only those -- are shipped over and back. + + ``idx_i``, ``idx_j`` and ``offsets`` are what this produces; the ``Rij`` that can come + with them is a by-product of the pruning, and models recompute it with their + ``PairwiseDistances`` input module anyway. + + Note: + The offsets of a pair list are cell shifts, so pruning takes ``R[j] - R[i] + + offsets`` for the pair vector -- the convention ``PairwiseDistances`` and every + neighbor list in schnetpack share. Atoms that have wandered outside their cell + break it, here and in the model alike, so a structure with a cell has to stay in + it. Structures without one are unaffected. + + Args: + neighbor_list: the neighbor list transform to build with, e.g. + :class:`~schnetpack.transform.MatScipyNeighborList`. Its cutoff is read off and + then widened by ``cutoff_skin`` -- the transform is modified in place, the way + :class:`~schnetpack.transform.SkinNeighborList` does it. + cutoff_skin: an atom may drift half of this before its structure's list is + rebuilt. A wider skin means fewer rebuilds and more pairs to carry along. + transforms: transforms applied to each structure after its neighbor list is built, + e.g. :class:`~schnetpack.transform.CollectAtomTriples`. + device: device the returned entries live on. Defaults to following the batch, + which is what a caller propagating structures on a device wants. + dtype: float precision of the returned entries. Defaults to following the batch. + additional_inputs: entries added to every structure before the transforms run, for + transforms that need them. + """ + + def __init__( + self, + neighbor_list: Transform, + cutoff_skin: float = 0.3, + transforms: Optional[Union[Transform, List[Transform]]] = None, + device: Optional[Union[str, torch.device]] = None, + dtype: Optional[torch.dtype] = None, + additional_inputs: Optional[Dict[str, torch.Tensor]] = None, + ): + self._device = torch.device(device) if isinstance(device, str) else device + self._dtype = dtype + self.additional_inputs = additional_inputs or {} + + self.cutoff = neighbor_list._cutoff + self.cutoff_skin = cutoff_skin + # build out to cutoff + skin, so a list stays usable while the atoms move + neighbor_list._cutoff = self.cutoff + cutoff_skin + + if transforms is None: + transforms = [] + elif not isinstance(transforms, list): + transforms = [transforms] + + if dtype not in (None, torch.float32, torch.float64): + raise ValueError(f"Unrecognized precision {dtype}") + + self.transforms: List[Transform] = [neighbor_list] + transforms + + # resolved per call, from the batch, unless they were pinned in the constructor + self.device = self._device or torch.device("cpu") + self.dtype = self._dtype or torch.float32 + + self.reset() + + def reset(self) -> None: + """Forget every cached list, so that the next call rebuilds from scratch.""" + #: the cutoff+skin list of each structure, with the structure it was built for + self._references: Dict[int, Dict[str, torch.Tensor]] = {} + #: those lists concatenated into batch numbering, on device, ready to be pruned + self._cache: Optional[Dict[str, torch.Tensor]] = None + + # -------------------------------------------------------------- public interface + + def update(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """The batch with its neighbor lists refreshed for the current positions. + + Entries the caller put in the batch are carried over untouched; only the + neighborhood entries are replaced. + """ + return {**inputs, **self.neighbors(inputs, with_distances=True)} + + def neighbors( + self, inputs: Dict[str, torch.Tensor], with_distances: bool = False + ) -> Dict[str, torch.Tensor]: + """Just the neighborhood entries for the structures in ``inputs``. + + Deliberately free of positions, atomic numbers, cells and atom counts, so that a + caller holding the live batch can ``update()`` its dictionary with the result + without overwriting the very tensors it is propagating. + + Args: + inputs: input batch. Only the structure-defining entries are read. + with_distances: also return ``Rij`` for the surviving pairs. + """ + positions = inputs[properties.R] + self.device = self._device or positions.device + self.dtype = self._dtype or ( + positions.dtype if positions.is_floating_point() else torch.float32 + ) + + stale = self._stale_structures(inputs) + if stale: + self._rebuild(inputs, stale) + return self._prune(inputs, with_distances=with_distances) + + # -------------------------------------------------------------- rebuild or reuse + + def _stale_structures(self, inputs: Dict[str, torch.Tensor]) -> List[int]: + """Which structures have moved far enough to need a new list. + + Answered per structure and on the device the batch already lives on, against the + positions each structure's own list was built for. A structure drifting just under + the threshold must not have its budget reset because another structure in the + batch was rebuilt -- which is what measuring against one batch-wide reference, + refreshed whenever anything rebuilds, would do. + """ + n_atoms = inputs[properties.n_atoms] + n_structures = int(n_atoms.shape[0]) + + if self._cache is None or len(self._references) != n_structures: + # a batch of a different size is a different batch; nothing may be reused + self.reset() + return list(range(n_structures)) + + positions = inputs[properties.R] + if positions.shape[0] != self._cache[properties.R].shape[0]: + self.reset() + return list(range(n_structures)) + + squared_drift = ( + (self._cache[properties.R] - positions).pow(2).sum(-1).to(torch.float64) + ) + idx_m = torch.repeat_interleave( + torch.arange(n_structures, device=positions.device), n_atoms + ) + per_structure = torch.zeros( + n_structures, dtype=torch.float64, device=positions.device + ).scatter_reduce_(0, idx_m, squared_drift, reduce="amax", include_self=False) + + stale = per_structure >= 0.25 * self.cutoff_skin**2 + stale |= ~( + torch.isclose( + self._cache[properties.cell], + inputs[properties.cell].view(n_structures, 3, 3).to(self.dtype), + ) + .view(n_structures, -1) + .all(-1) + ) + stale |= ( + self._cache[properties.pbc] != inputs[properties.pbc].view(n_structures, 3) + ).any(-1) + + return torch.nonzero(stale).view(-1).tolist() + + def _rebuild(self, inputs: Dict[str, torch.Tensor], stale: Sequence[int]) -> None: + """Build fresh cutoff+skin lists for the given structures, and cache them. + + Straight off the batch: the structures are cut out of it as tensors and handed to + the neighbor list transforms, which read positions, atomic numbers, cells and pbc + out of an input dictionary and never need ``ase.Atoms``. + """ + samples = split_batch(inputs) + + for idx in stale: + sample = {key: value.cpu() for key, value in samples[idx].items()} + sample.update(self.additional_inputs) + for transform in self.transforms: + sample = transform(sample) + + # the positions a list was built for are the reference the drift check + # measures against, so they are kept alongside it. Kept on cpu, where they + # were built and where they are concatenated; only the concatenated result + # is worth moving to the device. + self._references[idx] = { + key: value for key, value in sample.items() if key != properties.idx + } + + self._cache = self._collate_references(len(samples)) + + def _collate_references(self, n_structures: int) -> Dict[str, torch.Tensor]: + """Concatenate the per structure lists into batch numbering, once per rebuild. + + The pair indices of a structure count from its own first atom, so they have to be + shifted by everything ahead of it -- which is exactly what the collate function + the data loader uses does, triples included. It runs on cpu, where the lists were + built; what lands on the device is the finished, concatenated cache, which every + step from here to the next rebuild reads without touching the cpu again. + """ + collated = _atoms_collate_fn( + [self._references[idx] for idx in range(n_structures)] + ) + collated = {key: value.to(self.device) for key, value in collated.items()} + collated[properties.R] = collated[properties.R].to(self.dtype) + collated[properties.offsets] = collated[properties.offsets].to(self.dtype) + collated[properties.cell] = collated[properties.cell].view(n_structures, 3, 3) + collated[properties.pbc] = collated[properties.pbc].view(n_structures, 3) + return collated + + def _prune( + self, inputs: Dict[str, torch.Tensor], with_distances: bool = False + ) -> Dict[str, torch.Tensor]: + """Restrict the cached cutoff+skin lists to the pairs within the cutoff. + + The whole batch at once and on its own device -- the counterpart of + :meth:`SkinNeighborList._remove_neighbors_in_skin`, which does the same thing one + structure at a time on the cpu. + """ + cache = self._cache + idx_i, idx_j = cache[properties.idx_i], cache[properties.idx_j] + offsets = cache[properties.offsets] + + positions = inputs[properties.R] + Rij = positions[idx_j] - positions[idx_i] + offsets + within_cutoff = Rij.pow(2).sum(-1) <= self.cutoff**2 + + neighbors = { + key: cache[key][within_cutoff] for key in _PAIR_KEYS if key in cache + } + if with_distances: + neighbors[properties.Rij] = Rij[within_cutoff] + + if properties.idx_i_triples in cache: + neighbors.update(self._prune_triples(cache, within_cutoff)) + + return neighbors + + @staticmethod + def _prune_triples( + cache: Dict[str, torch.Tensor], within_cutoff: torch.Tensor + ) -> Dict[str, torch.Tensor]: + """Renumber the triples onto the pairs that survived the pruning. + + ``idx_j_triples`` and ``idx_k_triples`` index into the pair arrays, so dropping + pairs without renumbering would leave them pointing at the wrong pairs, or past + the end of the array altogether. Triples with a leg that did not survive are + dropped. + """ + renumbered = torch.full( + within_cutoff.shape, + -1, + dtype=torch.long, + device=within_cutoff.device, + ) + renumbered[within_cutoff] = torch.arange( + int(within_cutoff.sum()), device=within_cutoff.device + ) + + legs = [cache[key] for key in _TRIPLE_KEYS] + keep = renumbered[legs[0]] >= 0 + for leg in legs[1:]: + keep &= renumbered[leg] >= 0 + + triples = {properties.idx_i_triples: cache[properties.idx_i_triples][keep]} + for key, leg in zip(_TRIPLE_KEYS, legs): + triples[key] = renumbered[leg[keep]] + return triples diff --git a/src/schnetpack/transform/neighborlist.py b/src/schnetpack/transform/neighborlist.py index dc52e75f5..a94c1ae45 100644 --- a/src/schnetpack/transform/neighborlist.py +++ b/src/schnetpack/transform/neighborlist.py @@ -348,14 +348,18 @@ def _remove_neighbors_in_skin( self, inputs: Dict[str, torch.Tensor], ) -> Dict[str, torch.Tensor]: + """Restrict the cutoff+skin list to the pairs within the actual cutoff. + + Rebinds rather than mutating in place, so the unpruned list that ``_build`` + handed to ``previous_inputs`` -- the one a later step reuses -- stays intact. + """ Rij = inputs[properties.Rij] idx_i = inputs[properties.idx_i] idx_j = inputs[properties.idx_j] offsets = inputs[properties.offsets] - rij = torch.norm(inputs[properties.Rij], dim=-1) - cidx = torch.nonzero(rij <= self.cutoff).squeeze(-1) + cidx = torch.nonzero(Rij.pow(2).sum(-1) <= self.cutoff**2).squeeze(-1) inputs[properties.Rij] = Rij[cidx] inputs[properties.idx_i] = idx_i[cidx] @@ -370,33 +374,35 @@ def _update(self, inputs): # get sample index sample_idx = inputs[properties.idx].item() - # check if previous neighbor list exists and make sure that this is not the - # first update step - if sample_idx in self.previous_inputs.keys(): + # check if previous neighbor list exists + if sample_idx in self.previous_inputs: + # load previous inputs previous_inputs = self.previous_inputs[sample_idx] + # extract previous structure - previous_positions = np.array(previous_inputs[properties.R], copy=True) - previous_cell = np.array( - previous_inputs[properties.cell].view(3, 3), copy=True - ) - previous_pbc = np.array(previous_inputs[properties.pbc], copy=True) + previous_positions = previous_inputs[properties.R] + previous_cell = previous_inputs[properties.cell].view(3, 3) + previous_pbc = previous_inputs[properties.pbc] + # extract current structure positions = inputs[properties.R] cell = inputs[properties.cell].view(3, 3) pbc = inputs[properties.pbc] - # check if structure change is sufficiently small to reuse previous neighbor - # list + + # check if structure change is sufficiently small to reuse previous neighbor list if ( - (previous_pbc == pbc.numpy()).any() - and (previous_cell == cell.numpy()).any() - and ((previous_positions - positions.numpy()) ** 2).sum(1).max() + torch.equal(previous_pbc, pbc) + and torch.allclose(previous_cell, cell) + and torch.max( + torch.sum(torch.square(previous_positions - positions), dim=-1) + ).item() < 0.25 * self.cutoff_skin**2 ): - # reuse previous neighbor list - inputs[properties.idx_i] = previous_inputs[properties.idx_i].clone() - inputs[properties.idx_j] = previous_inputs[properties.idx_j].clone() - inputs[properties.offsets] = previous_inputs[properties.offsets].clone() + inputs[properties.idx_i] = previous_inputs[properties.idx_i] + inputs[properties.idx_j] = previous_inputs[properties.idx_j] + inputs[properties.offsets] = previous_inputs[properties.offsets] + return False, inputs # build new neighbor list @@ -410,15 +416,17 @@ def _build(self, inputs): for nbh_transform in self.nbh_transforms: inputs = nbh_transform(inputs) - # store new reference conformation and remove old one + # store new reference conformation and remove old one. This runs from _update, + # i.e. before forward prunes the skin away, so what is stored is the full + # cutoff+skin list -- the one a later step can reuse. sample_idx = inputs[properties.idx].item() stored_inputs = { - properties.R: inputs[properties.R].detach().clone(), - properties.cell: inputs[properties.cell].detach().clone(), - properties.pbc: inputs[properties.pbc].detach().clone(), - properties.idx_i: inputs[properties.idx_i].detach().clone(), - properties.idx_j: inputs[properties.idx_j].detach().clone(), - properties.offsets: inputs[properties.offsets].detach().clone(), + properties.R: inputs[properties.R], + properties.cell: inputs[properties.cell], + properties.pbc: inputs[properties.pbc], + properties.idx_i: inputs[properties.idx_i], + properties.idx_j: inputs[properties.idx_j], + properties.offsets: inputs[properties.offsets], } self.previous_inputs.update({sample_idx: stored_inputs}) @@ -568,23 +576,22 @@ def __init__(self, selection_name: str): self.selection_name = selection_name super().__init__() - def forward( - self, - inputs: Dict[str, torch.Tensor], - ) -> Dict[str, torch.Tensor]: + def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + filtered_out_indices = inputs[self.selection_name] + + # filter out pairs where both atoms are contained in filtered_out_indices + at_i_is_not_filtered_out = torch.isin( + inputs[properties.idx_i], filtered_out_indices, invert=True + ) + at_j_is_not_filtered_out = torch.isin( + inputs[properties.idx_j], filtered_out_indices, invert=True + ) + + mask = at_i_is_not_filtered_out | at_j_is_not_filtered_out - n_neighbors = inputs[properties.idx_i].shape[0] - slab_indices = inputs[self.selection_name].tolist() - kept_nbh_indices = [] - for nbh_idx in range(n_neighbors): - i = inputs[properties.idx_i][nbh_idx].item() - j = inputs[properties.idx_j][nbh_idx].item() - if i not in slab_indices or j not in slab_indices: - kept_nbh_indices.append(nbh_idx) - - inputs[properties.idx_i] = inputs[properties.idx_i][kept_nbh_indices] - inputs[properties.idx_j] = inputs[properties.idx_j][kept_nbh_indices] - inputs[properties.offsets] = inputs[properties.offsets][kept_nbh_indices] + inputs[properties.idx_i] = inputs[properties.idx_i][mask] + inputs[properties.idx_j] = inputs[properties.idx_j][mask] + inputs[properties.offsets] = inputs[properties.offsets][mask] return inputs diff --git a/src/scripts/spkconvert b/src/scripts/spkconvert index 2ee9f5824..e24ffa177 100755 --- a/src/scripts/spkconvert +++ b/src/scripts/spkconvert @@ -26,10 +26,10 @@ if __name__ == "__main__": parser.add_argument( "--expand_property_dims", default=[], - nargs='+', + nargs="+", help="Expanding the first dimension of the given property " - "(required for example for old FieldSchNet datasets). " - "Add property names here in the form 'property1 property2 property3'", + "(required for example for old FieldSchNet datasets). " + "Add property names here in the form 'property1 property2 property3'", ) args = parser.parse_args() with connect(args.data_path) as db: @@ -80,4 +80,4 @@ if __name__ == "__main__": data[p] = np.expand_dims(v, 0) else: data[p] = v - db.update(i + 1, data=data) \ No newline at end of file + db.update(i + 1, data=data) diff --git a/tests/interfaces/__init__.py b/tests/interfaces/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/interfaces/test_bw_benchmark.py b/tests/interfaces/test_bw_benchmark.py new file mode 100644 index 000000000..92d07a6e3 --- /dev/null +++ b/tests/interfaces/test_bw_benchmark.py @@ -0,0 +1,65 @@ +"""Wall clock scaling of ``BatchwiseLBFGS`` against sequential ase ``LBFGS``. + +Nothing here asserts on time, since wall clock thresholds are machine dependent. + +These are deselected by default. Run them with:: + pytest tests/interfaces -m benchmark_sweep --benchmark-group-by=param --benchmark-time-unit=ms +""" + +import pytest + +from .test_bw_optimizer import ( + FMAX, + MAX_STEPS, + build_batchwise_optimizer, + make_structures, + relax_sequential, + spk_calculator, +) + +# ensure that batch size appears in ascending order +N_VALUES = [pytest.param(n, id=f"{n:02d}") for n in (1, 3, 9, 18)] +ROUNDS = 3 + + +@pytest.mark.benchmark_sweep +@pytest.mark.parametrize("n_structures", N_VALUES) +def test_batchwise_relaxation(benchmark, n_structures): + structures = make_structures(n_structures) + built = [] + + def setup(): + # Ensure fresh optimizer for each for every round. + # Re-run before every round and left out of the timing. + optimizer = build_batchwise_optimizer(structures) + built.append(optimizer) + return (optimizer,), {} + + benchmark.pedantic( + lambda optimizer: optimizer.run(fmax=FMAX, steps=MAX_STEPS), + setup=setup, + rounds=ROUNDS, + iterations=1, + ) + + # a run that is fast because it never converged is not a faster run + assert built[-1].nsteps < MAX_STEPS + + +@pytest.mark.benchmark_sweep +@pytest.mark.parametrize("n_structures", N_VALUES) +def test_sequential_relaxation(benchmark, n_structures): + structures = make_structures(n_structures) + results = [] + + def setup(): + # Sequential relaxation copies the structures itself. + # Only the calculator needs renewing. + return (structures, spk_calculator()), {} + + def run(atoms_list, calculator): + results.append(relax_sequential(atoms_list, calculator)) + + benchmark.pedantic(run, setup=setup, rounds=ROUNDS, iterations=1) + + assert max(results[-1].steps) < MAX_STEPS diff --git a/tests/interfaces/test_bw_optimizer.py b/tests/interfaces/test_bw_optimizer.py new file mode 100644 index 000000000..79c7929f4 --- /dev/null +++ b/tests/interfaces/test_bw_optimizer.py @@ -0,0 +1,270 @@ +"""Compare ``BatchwiseLBFGS`` against a sequential loop over ase ``LBFGS``. + +The comparison is on the outcome of the relaxation only: both optimizers must land in +the same minima, and a structure must relax the same way alone as inside a batch. Wall +clock timing lives in ``test_bw_vs_sequ_benchmark.py``, which measures it with +pytest-benchmark. +""" + +import os +from copy import deepcopy +from dataclasses import dataclass +from typing import List + +import numpy as np +import pytest +import torch +from ase import Atoms +from ase.io import read +from ase.optimize import LBFGS + +import schnetpack as spk +from schnetpack import properties +from schnetpack.interfaces.ase_interface import ( + SpkCalculator, + atoms_to_batch, + batch_to_atoms, +) +from schnetpack.relax.batchwise_optimization import ( + BatchwiseCalculator, + BatchwiseLBFGS, +) +from schnetpack.utils.compatibility import load_model + +TESTDATA = os.path.join(os.path.dirname(__file__), "..", "testdata") +MODEL_PATH = os.path.join(TESTDATA, "md_ethanol.model") +STRUCTURE_PATH = os.path.join(TESTDATA, "ethanol_conformers.xyz") + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +SEED = 0 +N_STRUCTURES = 9 # multiple of number of conformers +NOISE = 0.05 # added to atomic positions +FMAX = 0.01 +MAX_STEPS = 200 +CUTOFF_SKIN = 0.3 # Ang, reuse the neighbor list while atoms move less than half +ENERGY_UNIT = "kcal/mol" +POSITION_UNIT = "Ang" + + +@dataclass +class RelaxationResult: + """Outcome of relaxing a batch of structures, however it was relaxed.""" + + atoms: List[Atoms] + steps: List[int] # optimizer steps, per structure for the sequential run + + +def _neighbor_list(): + """The plain neighbor list, as the sequential ase reference uses it.""" + model = load_model(MODEL_PATH, device=DEVICE) + return spk.transform.MatScipyNeighborList(cutoff=model.representation.cutoff.item()) + + +def _batch_neighbor_list(cutoff_skin: float = CUTOFF_SKIN): + """The neighbor list the batch-wise optimizer gets. + + It reuses the previous list while no atom of a structure has moved more than half + the skin, which is what makes the batch-wise path worth using. The sequential ase + reference gets a plain list instead -- ase rebuilds per structure anyway. + """ + return spk.transform.BatchNeighborList( + neighbor_list=_neighbor_list(), cutoff_skin=cutoff_skin + ) + + +def spk_calculator(): + return SpkCalculator( + model=MODEL_PATH, + neighbor_list=_neighbor_list(), + device=DEVICE, + energy_unit=ENERGY_UNIT, + position_unit=POSITION_UNIT, + ) + + +def make_structures(n_structures: int = N_STRUCTURES) -> List[Atoms]: + """A batch seeded by cycling the ethanol conformers, + and adding noise to atomic positions. + """ + conformers = read(STRUCTURE_PATH, index=":") + rng = np.random.default_rng(SEED) + structures = [] + for idx in range(n_structures): + noisy = conformers[idx % len(conformers)].copy() + noisy.positions += rng.normal(scale=NOISE, size=noisy.positions.shape) + structures.append(noisy) + return structures + + +def build_batchwise_optimizer(atoms_list: List[Atoms]) -> BatchwiseLBFGS: + """Everything needed to relax a batch, short of actually running it. + + Kept separate from the run so the benchmark can time only the relaxation. + """ + calculator = BatchwiseCalculator( + model=MODEL_PATH, + neighbor_list=_batch_neighbor_list(), + device=DEVICE, + energy_unit=ENERGY_UNIT, + position_unit=POSITION_UNIT, + ) + inputs = atoms_to_batch(deepcopy(atoms_list), device=DEVICE) + + n_atoms = len(atoms_list[0]) + return BatchwiseLBFGS( + calculator=calculator, + inputs=inputs, + logfile=None, + fixed_atoms_mask=[False] * (n_atoms * len(atoms_list)), + ) + + +def relax_batchwise(atoms_list: List[Atoms]) -> RelaxationResult: + """Relax a batch of structures in parallel with ``BatchwiseLBFGS``.""" + optimizer = build_batchwise_optimizer(atoms_list) + optimizer.run(fmax=FMAX, steps=MAX_STEPS) + + # the optimizer hands back tensors; ase structures are a boundary conversion + relaxed, _ = optimizer.get_relaxation_results() + return RelaxationResult(atoms=batch_to_atoms(relaxed), steps=[optimizer.nsteps]) + + +def relax_sequential(atoms_list: List[Atoms], calculator) -> RelaxationResult: + """Relax the structures one at a time, the way ase would normally be used.""" + atoms, steps = [], [] + # LBFGS relaxes in place, so the caller's structures must not be handed over + for structure in deepcopy(atoms_list): + structure.calc = calculator + optimizer = LBFGS(structure, logfile=None) + optimizer.run(fmax=FMAX, steps=MAX_STEPS) + + atoms.append(structure) + steps.append(optimizer.nsteps) + + return RelaxationResult(atoms=atoms, steps=steps) + + +@pytest.fixture(scope="module") +def shared_calculator(): + """One calculator used to score every relaxed structure, whatever produced it.""" + return spk_calculator() + + +@pytest.fixture(scope="module") +def initial_structures(): + return make_structures() + + +@pytest.fixture(scope="module") +def sequential_result(initial_structures): + return relax_sequential(initial_structures, spk_calculator()) + + +@pytest.fixture(scope="module") +def batchwise_result(initial_structures): + return relax_batchwise(initial_structures) + + +@pytest.fixture(scope="module") +def single_structure_atoms(initial_structures): + """Every structure relaxed on its own, but still through the batch-wise optimizer.""" + return [relax_batchwise([structure]).atoms[0] for structure in initial_structures] + + +def evaluate(atoms_list: List[Atoms], calculator): + """Energies and max force norms, all from the same calculator.""" + energies, fmax = [], [] + for structure in atoms_list: + structure = structure.copy() + structure.calc = calculator + energies.append(structure.get_potential_energy()) + fmax.append(np.sqrt((structure.get_forces() ** 2).sum(axis=1).max())) + return np.array(energies), np.array(fmax) + + +def test_both_optimizers_converge( + sequential_result, batchwise_result, shared_calculator +): + """Neither run may be compared against a reference that never converged.""" + assert max(sequential_result.steps) < MAX_STEPS, ( + "sequential relaxation hit the step limit, so it is not a valid reference: " + f"steps={sequential_result.steps}" + ) + assert batchwise_result.steps[0] < MAX_STEPS + + _, fmax_seq = evaluate(sequential_result.atoms, shared_calculator) + _, fmax_batch = evaluate(batchwise_result.atoms, shared_calculator) + + assert fmax_seq.max() <= FMAX, f"sequential did not reach fmax: {fmax_seq}" + assert fmax_batch.max() <= FMAX, f"batch-wise did not reach fmax: {fmax_batch}" + + +def test_batchwise_reaches_same_minima( + sequential_result, batchwise_result, shared_calculator +): + """Both optimizers must land in the same basin. + + Energy is the right quantity to compare: it is invariant under the rigid body + motion and the relabelling of identical atoms that a relaxation may introduce, so + no structural alignment is needed. The conformers seeding this batch settle into + two basins 2.7 meV apart, comfortably above the tolerance below, so a structure + ending up in the wrong one would still be caught. + """ + energies_seq, _ = evaluate(sequential_result.atoms, shared_calculator) + energies_batch, _ = evaluate(batchwise_result.atoms, shared_calculator) + + np.testing.assert_allclose(energies_batch, energies_seq, atol=1e-3) + + +def test_batch_size_invariance(batchwise_result, single_structure_atoms): + """Relaxing a structure alone or inside a batch must give the same result.""" + for idx, (in_batch, alone) in enumerate( + zip(batchwise_result.atoms, single_structure_atoms) + ): + deviation = np.abs(in_batch.get_positions() - alone.get_positions()).max() + assert deviation < 1e-4, ( + f"structure {idx} relaxed differently inside a batch of {N_STRUCTURES} " + f"than on its own: max deviation {deviation:.2e} Ang" + ) + + +@pytest.mark.parametrize("trajectory_interval", [0, 1]) +def test_forces_are_computed_once_per_step( + initial_structures, tmp_path, trajectory_interval +): + """The convergence check, the frame written from it, and the step share one call. + + The calculator caches its results and decides whether they are still valid from + the identity and mutation counter of the position tensor. If that check ever goes + wrong in the conservative direction, relaxations silently cost twice as much. + Writing a frame on every step must not cost a second call either, which is why + ``trajectory_interval=1`` is covered here too. + """ + optimizer = build_batchwise_optimizer(initial_structures[:3]) + optimizer.trajectory = str(tmp_path / "relax.hdf5") + optimizer.trajectory_interval = trajectory_interval + calculate = optimizer.calculator.calculate + calls = [] + + def counting_calculate(inputs): + calls.append(None) + return calculate(inputs) + + optimizer.calculator.calculate = counting_calculate + optimizer.run(fmax=FMAX, steps=MAX_STEPS) + optimizer.close() + + # one for the initial forces, one per step taken + assert len(calls) == optimizer.nsteps + 1 + + +def test_cached_forces_are_dropped_when_the_positions_move(initial_structures): + """The other direction: a structure that changed must not return stale forces.""" + optimizer = build_batchwise_optimizer(initial_structures[:2]) + calculator, inputs = optimizer.calculator, optimizer.inputs + + before = calculator.get_forces(inputs).clone() + assert torch.equal(calculator.get_forces(inputs), before), "cache should have hit" + + inputs[properties.R] += 0.1 + assert not torch.equal(calculator.get_forces(inputs), before) diff --git a/tests/interfaces/test_bw_optimizer_units.py b/tests/interfaces/test_bw_optimizer_units.py new file mode 100644 index 000000000..00ca2f386 --- /dev/null +++ b/tests/interfaces/test_bw_optimizer_units.py @@ -0,0 +1,318 @@ +"""Behaviour of ``BatchwiseLBFGS`` that does not need a trained model. + +``test_bw_optimizer.py`` checks that a relaxation lands where ase's ``LBFGS`` lands. +The mechanics around it -- what comes back out of the batch, which atoms are allowed +to move, what an invalid batch does -- are cheaper and clearer to pin down against an +analytic potential, which is what this module does. +""" + +from typing import Dict, List, Optional + +import numpy as np +import pytest +import torch +from ase import Atoms + +from schnetpack import properties +from schnetpack.interfaces.ase_interface import atoms_to_batch, batch_to_atoms +from schnetpack.relax.batchwise_optimization import BatchwiseLBFGS +from schnetpack.relax.batchwise_trajectory import BatchwiseTrajectoryReader + + +class HarmonicCalculator: + """Every atom is pulled towards the origin by a spring. + + Stands in for a ``BatchwiseCalculator``: the optimizer only ever asks it for + forces, and the minimum is known exactly. + """ + + def __init__(self, spring_constant: float = 1.0): + self.spring_constant = spring_constant + self.device = torch.device("cpu") + self.results = {} + + def get_forces( + self, + inputs: Dict[str, torch.Tensor], + fixed_atoms_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + forces = -self.spring_constant * inputs[properties.R] + self.results = { + "energy": self.get_potential_energy(inputs), + "forces": forces, + } + if fixed_atoms_mask is not None: + forces = forces[fixed_atoms_mask] + return forces + + def get_potential_energy(self, inputs: Dict[str, torch.Tensor]) -> torch.Tensor: + # one energy per structure, the way a BatchwiseCalculator reports them + per_atom = 0.5 * self.spring_constant * inputs[properties.R].pow(2).sum(-1) + counts = inputs[properties.n_atoms].tolist() + return torch.stack([chunk.sum() for chunk in torch.split(per_atom, counts)]) + + +def make_inputs( + n_atoms_per_config: List[int], + cell: Optional[np.ndarray] = None, + seed: int = 0, +) -> Dict[str, torch.Tensor]: + """A schnetpack input batch, built without a converter or a neighbor list.""" + n_configs = len(n_atoms_per_config) + n_total = sum(n_atoms_per_config) + rng = np.random.default_rng(seed) + + cells = torch.zeros(n_configs, 3, 3, dtype=torch.float32) + if cell is not None: + cells[:] = torch.as_tensor(cell, dtype=torch.float32) + + return { + properties.n_atoms: torch.tensor(n_atoms_per_config), + properties.R: torch.tensor( + rng.normal(scale=1.0, size=(n_total, 3)), dtype=torch.float32 + ), + properties.Z: torch.full((n_total,), 6, dtype=torch.long), + properties.cell: cells, + properties.pbc: torch.full((n_configs, 3), cell is not None), + properties.idx_m: torch.repeat_interleave( + torch.arange(n_configs), torch.tensor(n_atoms_per_config) + ), + } + + +def make_optimizer(inputs, **kwargs) -> BatchwiseLBFGS: + kwargs.setdefault("logfile", None) + return BatchwiseLBFGS(calculator=HarmonicCalculator(), inputs=inputs, **kwargs) + + +def test_full_cell_survives_the_round_trip(): + """A triclinic cell must come back as it went in, off-diagonal entries included.""" + cell = np.array([[4.0, 0.0, 0.0], [1.5, 3.5, 0.0], [0.5, 1.0, 5.0]]) + optimizer = make_optimizer(make_inputs([4, 4], cell=cell)) + + relaxed, _ = optimizer.get_relaxation_results() + + for structure in batch_to_atoms(relaxed): + np.testing.assert_allclose(structure.cell[:], cell, atol=1e-5) + assert structure.pbc.all() + + +def test_no_mask_matches_an_all_free_mask(): + """``fixed_atoms_mask=None`` must mean the same as "no atom is fixed".""" + without = make_optimizer(make_inputs([5, 5])) + without.run(fmax=1e-3, steps=20) + + explicit = make_optimizer(make_inputs([5, 5]), fixed_atoms_mask=[False] * 10) + explicit.run(fmax=1e-3, steps=20) + + np.testing.assert_allclose( + without.inputs[properties.R].numpy(), + explicit.inputs[properties.R].numpy(), + atol=1e-6, + ) + + +def test_fixed_atoms_do_not_move(): + inputs = make_inputs([5, 5]) + initial = inputs[properties.R].clone() + fixed = [True, False, False, False, False] * 2 + + optimizer = make_optimizer(inputs, fixed_atoms_mask=fixed) + optimizer.run(fmax=1e-3, steps=20) + + moved = (inputs[properties.R] - initial).abs().max(dim=1).values + assert torch.allclose(moved[torch.tensor(fixed)], torch.zeros(2)) + assert (moved[~torch.tensor(fixed)] > 1e-3).all() + + +def test_fixed_atoms_are_left_out_of_the_convergence_check(): + """A fixed atom far from the minimum must not keep the relaxation running.""" + inputs = make_inputs([2]) + # the fixed atom carries a large force that can never be relaxed away + inputs[properties.R] = torch.tensor([[50.0, 0.0, 0.0], [0.1, 0.0, 0.0]]) + + optimizer = make_optimizer(inputs, fixed_atoms_mask=[True, False]) + assert optimizer.run(fmax=0.05, steps=50) + + +def test_ragged_batches_are_rejected(): + with pytest.raises(ValueError, match="same number of atoms"): + make_optimizer(make_inputs([3, 4])) + + +def test_mask_length_is_checked(): + with pytest.raises(ValueError, match="one per atom"): + make_optimizer(make_inputs([3, 3]), fixed_atoms_mask=[False] * 5) + + +def test_run_without_a_step_limit_still_takes_steps(): + """``max_steps`` used to default to 0, so ``run(fmax=...)`` did nothing.""" + optimizer = make_optimizer(make_inputs([4])) + + assert optimizer.run(fmax=1e-3) + assert optimizer.nsteps > 0 + + +def test_run_reports_failure_when_the_step_limit_is_hit(): + optimizer = make_optimizer(make_inputs([4])) + + assert not optimizer.run(fmax=1e-12, steps=2) + assert optimizer.nsteps == 2 + + +def test_logfile_none_disables_logging(): + optimizer = make_optimizer(make_inputs([3]), logfile=None) + assert optimizer.logfile is None + optimizer.run(fmax=1e-3, steps=5) + + +def test_trajectory_holds_one_frame_per_step(tmp_path): + """``trajectory_interval=1`` records every state the batch passed through.""" + path = str(tmp_path / "relax.hdf5") + optimizer = make_optimizer( + make_inputs([3, 3]), trajectory=path, trajectory_interval=1 + ) + optimizer.run(fmax=1e-3, steps=5) + optimizer.close() + + with BatchwiseTrajectoryReader(path) as traj: + # the initial state plus one frame after each step + assert traj.n_frames == optimizer.nsteps + 1 + assert traj.n_structures == 2 + assert traj.n_atoms == 3 + assert traj.positions.shape == (traj.n_frames, 2, 3, 3) + assert list(traj.steps) == list(range(traj.n_frames)) + # a relaxation towards the origin has to shrink the positions monotonically + assert np.abs(traj.positions[-1]).max() < np.abs(traj.positions[0]).max() + assert not traj.has_forces + + +def test_trajectory_interval_zero_keeps_only_the_endpoints(tmp_path): + path = str(tmp_path / "relax.hdf5") + optimizer = make_optimizer(make_inputs([2]), trajectory=path, trajectory_interval=0) + optimizer.run(fmax=1e-3, steps=20) + optimizer.close() + + with BatchwiseTrajectoryReader(path) as traj: + assert list(traj.steps) == [0, optimizer.nsteps] + + +def test_trajectory_interval_n_always_includes_first_and_last(tmp_path): + path = str(tmp_path / "relax.hdf5") + optimizer = make_optimizer(make_inputs([2]), trajectory=path, trajectory_interval=3) + optimizer.run(fmax=1e-6, steps=10) + optimizer.close() + + with BatchwiseTrajectoryReader(path) as traj: + steps = list(traj.steps) + assert steps[0] == 0 and steps[-1] == optimizer.nsteps + assert set(steps) == {0, optimizer.nsteps} | { + s for s in range(optimizer.nsteps + 1) if s % 3 == 0 + } + + +def test_trajectory_can_store_forces_and_energies(tmp_path): + path = str(tmp_path / "relax.hdf5") + inputs = make_inputs([2]) + optimizer = make_optimizer( + inputs, trajectory=path, trajectory_interval=1, store_forces=True + ) + optimizer.run(fmax=1e-3, steps=10) + optimizer.close() + + with BatchwiseTrajectoryReader(path) as traj: + assert traj.has_forces + # the harmonic calculator pulls every atom towards the origin + np.testing.assert_allclose( + traj.forces[-1].reshape(-1, 3), + -traj.positions[-1].reshape(-1, 3), + atol=1e-5, + ) + # and the energy falls as the structures relax + assert traj.energy[-1].sum() < traj.energy[0].sum() + # every structure ends up flagged converged + assert traj.converged[-1].all() + assert not traj.converged[0].any() + + +def test_no_trajectory_is_written_without_a_path(tmp_path): + optimizer = make_optimizer(make_inputs([2]), trajectory=None) + optimizer.run(fmax=1e-3, steps=5) + optimizer.close() + assert list(tmp_path.iterdir()) == [] + + +def test_optimizer_is_a_context_manager(tmp_path): + path = str(tmp_path / "relax.hdf5") + with make_optimizer( + make_inputs([2]), trajectory=path, trajectory_interval=1 + ) as optimizer: + optimizer.run(fmax=1e-3, steps=5) + + with BatchwiseTrajectoryReader(path) as traj: + assert traj.n_frames == optimizer.nsteps + 1 + + +def test_relaxation_finds_the_analytic_minimum(): + inputs = make_inputs([6, 6]) + optimizer = make_optimizer(inputs) + + assert optimizer.run(fmax=1e-4, steps=100) + np.testing.assert_allclose( + inputs[properties.R].numpy(), np.zeros((12, 3)), atol=1e-4 + ) + + +def test_batch_to_atoms_handles_ragged_batches(): + """The helper is the inverse of atoms_to_batch, so it must not assume equal sizes. + + The optimizer itself rejects ragged batches, but the helper is also used on batches + that never went through it. + """ + inputs = { + properties.n_atoms: torch.tensor([2, 3]), + properties.R: torch.arange(15, dtype=torch.float32).view(5, 3), + properties.Z: torch.tensor([1, 6, 8, 1, 1]), + properties.cell: torch.zeros(2, 3, 3), + properties.pbc: torch.tensor([[True] * 3, [False] * 3]), + } + + structures = batch_to_atoms(inputs) + + assert [len(s) for s in structures] == [2, 3] + assert list(structures[0].numbers) == [1, 6] + assert list(structures[1].numbers) == [8, 1, 1] + np.testing.assert_allclose(structures[1].positions, np.arange(6, 15).reshape(3, 3)) + assert structures[0].pbc.all() and not structures[1].pbc.any() + + +def test_batch_to_atoms_keeps_the_full_cell(): + """Off-diagonal cell entries must survive; taking the diagonal was a real bug.""" + cell = np.array([[4.0, 0.0, 0.0], [1.5, 3.5, 0.0], [0.5, 1.0, 5.0]]) + inputs = make_inputs([3], cell=cell) + + np.testing.assert_allclose(batch_to_atoms(inputs)[0].cell[:], cell, atol=1e-5) + + +def test_atoms_to_batch_round_trips(): + """The two boundary conversions are each other's inverse, ragged batches included.""" + structures = [ + Atoms( + numbers=[1, 6], + positions=np.arange(6).reshape(2, 3) * 0.5, + cell=np.eye(3) * 4, + pbc=True, + ), + Atoms( + numbers=[8, 1, 1], positions=np.arange(9).reshape(3, 3) * 0.25, pbc=False + ), + ] + + recovered = batch_to_atoms(atoms_to_batch(structures, dtype=torch.float64)) + + assert [len(s) for s in recovered] == [2, 3] + for original, structure in zip(structures, recovered): + assert list(original.numbers) == list(structure.numbers) + np.testing.assert_allclose(original.positions, structure.positions, atol=1e-6) + np.testing.assert_allclose(original.cell[:], structure.cell[:], atol=1e-6) + assert (original.pbc == structure.pbc).all() diff --git a/tests/interfaces/test_bw_trajectory.py b/tests/interfaces/test_bw_trajectory.py new file mode 100644 index 000000000..e81c8af50 --- /dev/null +++ b/tests/interfaces/test_bw_trajectory.py @@ -0,0 +1,212 @@ +"""Round trip through ``BatchwiseTrajectoryWriter`` / ``BatchwiseTrajectoryReader``. + +The optimizer's own trajectory tests cover which frames get written. These cover the +storage itself: that the values survive, that buffering is invisible to the reader, and +that the optional datasets really are optional. +""" + +import numpy as np +import pytest +import torch + +from schnetpack import properties +from schnetpack.interfaces.ase_interface import batch_to_atoms +from schnetpack.relax.batchwise_trajectory import ( + BatchwiseTrajectoryReader, + BatchwiseTrajectoryWriter, +) + +N_STRUCTURES = 3 +N_ATOMS = 4 + + +def make_frame(step: int, rng: np.random.Generator) -> dict: + return { + "step": step, + "positions": torch.tensor( + rng.normal(size=(N_STRUCTURES * N_ATOMS, 3)), dtype=torch.float64 + ), + "cell": torch.tensor( + rng.normal(size=(N_STRUCTURES, 3, 3)), dtype=torch.float64 + ), + "energy": torch.tensor(rng.normal(size=N_STRUCTURES), dtype=torch.float64), + "forces": torch.tensor( + rng.normal(size=(N_STRUCTURES * N_ATOMS, 3)), dtype=torch.float64 + ), + "converged": torch.tensor([step % 2 == 0] * N_STRUCTURES), + } + + +def write_frames(path, n_frames, buffer_size=64, store_forces=True, precision=32): + rng = np.random.default_rng(0) + frames = [make_frame(step, rng) for step in range(n_frames)] + with BatchwiseTrajectoryWriter( + path, + atomic_numbers=torch.full((N_STRUCTURES, N_ATOMS), 6), + pbc=torch.zeros(N_STRUCTURES, 3, dtype=torch.bool), + store_forces=store_forces, + buffer_size=buffer_size, + precision=precision, + attrs={"optimizer": "BatchwiseLBFGS", "fmax": 0.05}, + ) as writer: + for frame in frames: + writer.write(**frame) + return frames + + +def test_values_survive_the_round_trip(tmp_path): + path = str(tmp_path / "traj.hdf5") + frames = write_frames(path, n_frames=5) + + with BatchwiseTrajectoryReader(path) as traj: + assert traj.n_frames == 5 + assert traj.n_structures == N_STRUCTURES + assert traj.n_atoms == N_ATOMS + assert list(traj.steps) == [0, 1, 2, 3, 4] + + for index, frame in enumerate(frames): + expected = frame["positions"].numpy().reshape(N_STRUCTURES, N_ATOMS, 3) + np.testing.assert_allclose(traj.positions[index], expected, rtol=1e-6) + np.testing.assert_allclose( + traj.cell[index], frame["cell"].numpy(), rtol=1e-6 + ) + np.testing.assert_allclose( + traj.energy[index], frame["energy"].numpy(), rtol=1e-6 + ) + assert (traj.converged[index] == frame["converged"].numpy()).all() + + +@pytest.mark.parametrize("n_frames", [1, 3, 4, 5, 9]) +def test_buffering_is_invisible(tmp_path, n_frames): + """Frame counts either side of a buffer boundary must read back identically.""" + path = str(tmp_path / "traj.hdf5") + frames = write_frames(path, n_frames=n_frames, buffer_size=4) + + with BatchwiseTrajectoryReader(path) as traj: + assert traj.n_frames == n_frames + assert list(traj.steps) == list(range(n_frames)) + np.testing.assert_allclose( + traj.positions[-1], + frames[-1]["positions"].numpy().reshape(N_STRUCTURES, N_ATOMS, 3), + rtol=1e-6, + ) + + +def test_forces_are_optional(tmp_path): + path = str(tmp_path / "traj.hdf5") + write_frames(path, n_frames=3, store_forces=False) + + with BatchwiseTrajectoryReader(path) as traj: + assert not traj.has_forces + with pytest.raises(AttributeError, match="not stored"): + traj.forces + + +def test_metadata_is_stored(tmp_path): + path = str(tmp_path / "traj.hdf5") + write_frames(path, n_frames=2) + + with BatchwiseTrajectoryReader(path) as traj: + assert traj.file.attrs["optimizer"] == "BatchwiseLBFGS" + assert traj.file.attrs["fmax"] == pytest.approx(0.05) + assert traj.file.attrs["energy_unit"] == "eV" + assert traj.file.attrs["schnetpack_version"] + assert traj.atomic_numbers.shape == (N_STRUCTURES, N_ATOMS) + assert traj.pbc.shape == (N_STRUCTURES, 3) + + +def test_a_frame_reads_back_as_an_input_batch(tmp_path): + """``frame`` must produce a batch that batch_to_atoms and an optimizer accept.""" + path = str(tmp_path / "traj.hdf5") + frames = write_frames(path, n_frames=6) + + with BatchwiseTrajectoryReader(path) as traj: + last = traj.frame() + + assert last[properties.n_atoms].tolist() == [N_ATOMS] * N_STRUCTURES + assert last[properties.idx_m].tolist() == sum( + ([idx] * N_ATOMS for idx in range(N_STRUCTURES)), [] + ) + assert last[properties.Z].shape == (N_STRUCTURES * N_ATOMS,) + assert last[properties.cell].shape == (N_STRUCTURES, 3, 3) + np.testing.assert_allclose( + last[properties.R].numpy(), frames[-1]["positions"].numpy(), rtol=1e-6 + ) + np.testing.assert_allclose( + last[properties.forces].numpy(), frames[-1]["forces"].numpy(), rtol=1e-6 + ) + assert last["step"] == 5 + + # and it really is convertible, off-diagonal cells included + structures = batch_to_atoms(last) + assert [len(s) for s in structures] == [N_ATOMS] * N_STRUCTURES + np.testing.assert_allclose( + structures[1].cell[:], frames[-1]["cell"].numpy()[1], rtol=1e-6 + ) + + +def test_structure_gives_one_path_across_frames(tmp_path): + path = str(tmp_path / "traj.hdf5") + write_frames(path, n_frames=6) + + with BatchwiseTrajectoryReader(path) as traj: + path_of_one = traj.structure(1) + + assert path_of_one["positions"].shape == (6, N_ATOMS, 3) + assert path_of_one["energy"].shape == (6,) + assert path_of_one["atomic_numbers"].shape == (N_ATOMS,) + assert list(path_of_one["steps"]) == list(range(6)) + + +def test_double_precision_is_kept(tmp_path): + path = str(tmp_path / "traj.hdf5") + frames = write_frames(path, n_frames=2, precision=64) + + with BatchwiseTrajectoryReader(path) as traj: + assert traj.positions.dtype == np.float64 + np.testing.assert_allclose( + traj.positions[0], + frames[0]["positions"].numpy().reshape(N_STRUCTURES, N_ATOMS, 3), + rtol=1e-12, + ) + + +def test_unknown_precision_is_rejected(tmp_path): + with pytest.raises(ValueError, match="precision"): + BatchwiseTrajectoryWriter( + str(tmp_path / "traj.hdf5"), + atomic_numbers=torch.full((N_STRUCTURES, N_ATOMS), 6), + pbc=torch.zeros(N_STRUCTURES, 3, dtype=torch.bool), + precision=16, + ) + + +@pytest.mark.parametrize( + "n_structures, n_atoms, expected", + [(3, 4, 64), (256, 1000, 2), (2000, 5000, 1)], +) +def test_buffer_size_follows_the_batch_size(tmp_path, n_structures, n_atoms, expected): + """A fixed 64-frame buffer would be a 200 MB chunk for a large batch. + + HDF5 wants chunks well under a megabyte, and a chunk that large also means any + caller reading a single frame pays for all 64. + """ + writer = BatchwiseTrajectoryWriter( + str(tmp_path / "traj.hdf5"), + atomic_numbers=torch.full((n_structures, n_atoms), 6), + pbc=torch.zeros(n_structures, 3, dtype=torch.bool), + ) + assert writer.buffer_size == expected + assert writer.datasets["positions"].chunks[0] == expected + writer.close() + + +def test_explicit_buffer_size_wins(tmp_path): + writer = BatchwiseTrajectoryWriter( + str(tmp_path / "traj.hdf5"), + atomic_numbers=torch.full((N_STRUCTURES, N_ATOMS), 6), + pbc=torch.zeros(N_STRUCTURES, 3, dtype=torch.bool), + buffer_size=7, + ) + assert writer.buffer_size == 7 + writer.close() diff --git a/tests/testdata/ethanol_conformers.xyz b/tests/testdata/ethanol_conformers.xyz new file mode 100644 index 000000000..0fe4a5269 --- /dev/null +++ b/tests/testdata/ethanol_conformers.xyz @@ -0,0 +1,33 @@ +9 +Properties=species:S:1:pos:R:3:forces:R:3 pbc="F F F" +C 0.01578971 -0.56446709 -0.00000001 -0.00002002 0.00001423 0.00000181 +C -1.26713065 0.24594442 -0.00000002 0.00000786 0.00000108 0.00000501 +O 1.11974167 0.35543938 0.00000000 0.00002332 0.00000199 0.00000997 +H 0.05397592 -1.21669037 0.89324391 -0.00000852 -0.00000068 -0.00000179 +H 0.05397594 -1.21669037 -0.89324391 -0.00000156 -0.00000385 -0.00000500 +H -1.31873531 0.88816843 0.89054363 0.00000525 -0.00000507 -0.00001478 +H -1.31873532 0.88816841 -0.89054361 0.00000132 -0.00001167 0.00001647 +H -2.14185171 -0.41976893 -0.00000001 -0.00000767 -0.00000318 -0.00000339 +H 1.94206975 -0.16220386 -0.00000000 0.00000012 0.00000726 -0.00000824 +9 +Properties=species:S:1:pos:R:3:forces:R:3 pbc="F F F" +C 0.01965182 -0.70813042 -0.09759178 -0.00002002 0.00001423 0.00000181 +C -1.16025535 0.25079963 0.00901837 0.00000786 0.00000108 0.00000501 +O 1.27313737 -0.04311800 -0.30964909 0.00002332 0.00000199 0.00000997 +H 0.06899438 -1.35753705 0.79724365 -0.00000852 -0.00000068 -0.00000179 +H -0.10017085 -1.36746848 -0.97024937 -0.00000156 -0.00000385 -0.00000500 +H -1.06768884 0.90263318 0.89225952 0.00000525 -0.00000507 -0.00001478 +H -1.22323851 0.88920776 -0.88346033 0.00000132 -0.00001167 0.00001647 +H -2.10457979 -0.30586752 0.10726536 -0.00000767 -0.00000318 -0.00000339 +H 1.43324969 0.53738080 0.45516371 0.00000012 0.00000726 -0.00000824 +9 +Properties=species:S:1:pos:R:3:forces:R:3 pbc="F F F" +C 0.01920773 -0.70873886 0.09553212 -0.00002002 0.00001423 0.00000181 +C -1.15993787 0.25152224 -0.00745704 0.00000786 0.00000108 0.00000501 +O 1.27330394 -0.04549091 0.30947937 0.00002332 0.00000199 0.00000997 +H -0.10090774 -1.37094682 0.96597319 -0.00000852 -0.00000068 -0.00000179 +H 0.06774765 -1.35512450 -0.80153001 -0.00000156 -0.00000385 -0.00000500 +H -1.22214065 0.88692654 0.88721863 0.00000525 -0.00000507 -0.00001478 +H -1.06708642 0.90629581 -0.88849282 0.00000132 -0.00001167 0.00001647 +H -2.10475240 -0.30402572 -0.10732862 -0.00000767 -0.00000318 -0.00000339 +H 1.43366581 0.53748222 -0.45339501 0.00000012 0.00000726 -0.00000824 diff --git a/tests/transform/__init__.py b/tests/transform/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/transform/test_batch_neighborlist.py b/tests/transform/test_batch_neighborlist.py new file mode 100644 index 000000000..5f5856d90 --- /dev/null +++ b/tests/transform/test_batch_neighborlist.py @@ -0,0 +1,335 @@ +"""Behaviour of ``BatchNeighborList``, the neighbor lists of a batch that keeps moving. + +The reference throughout is a plain neighbor list built from scratch for the current +positions: whatever the module decides to reuse, what reaches the model has to be the +same set of pairs it would have got from a fresh build. +""" + +from typing import Dict, List + +import numpy as np +import pytest +import torch +from ase import Atoms + +import schnetpack as spk +from schnetpack import properties +from schnetpack.data.loader import _atoms_collate_fn, split_batch +from schnetpack.interfaces.ase_interface import AtomsConverter, atoms_to_batch + +CUTOFF = 3.0 +CUTOFF_SKIN = 0.5 +N_ATOMS = 12 + + +def make_structures( + pbc: bool, n_structures: int = 3, seed: int = 1, cell: bool = True +) -> List[Atoms]: + """Structures of ``N_ATOMS`` carbons scattered through an 8 Angstrom box. + + ``cell=False`` leaves them without one, for the tests that move a structure far from + where it started: the offsets of a pair list are cell shifts, so an atom outside its + cell breaks ``R[j] - R[i] + offsets`` -- the convention ``PairwiseDistances`` and + every neighbor list in schnetpack share. + """ + rng = np.random.default_rng(seed) + return [ + Atoms( + numbers=[6] * N_ATOMS, + positions=rng.uniform(0.0, 8.0, size=(N_ATOMS, 3)), + cell=8.0 * np.eye(3) if cell else None, + pbc=pbc, + ) + for _ in range(n_structures) + ] + + +def make_batch_neighbor_list(**kwargs) -> spk.transform.BatchNeighborList: + return spk.transform.BatchNeighborList( + neighbor_list=spk.transform.MatScipyNeighborList(cutoff=CUTOFF), + cutoff_skin=CUTOFF_SKIN, + dtype=torch.float64, + **kwargs, + ) + + +def batch_of(structures: List[Atoms]) -> Dict[str, torch.Tensor]: + """The structures as a batch, neighbor lists left to the module under test.""" + return atoms_to_batch(structures, dtype=torch.float64) + + +def freshly_built(structures: List[Atoms], **kwargs) -> Dict[str, torch.Tensor]: + """The same structures through a plain neighbor list, built from scratch.""" + return AtomsConverter( + neighbor_list=spk.transform.MatScipyNeighborList(cutoff=CUTOFF), + dtype=torch.float64, + **kwargs, + )(structures) + + +def neighbor_pairs(inputs: Dict[str, torch.Tensor]) -> np.ndarray: + """The neighbor list of a batch, in an order-independent, comparable form.""" + pairs = np.column_stack( + [ + inputs[properties.idx_i].cpu().numpy(), + inputs[properties.idx_j].cpu().numpy(), + inputs[properties.offsets].cpu().numpy().round(6), + ] + ) + return pairs[np.lexsort(pairs.T[::-1])] + + +def displaced(structures: List[Atoms], scale: float, seed: int = 2) -> List[Atoms]: + rng = np.random.default_rng(seed) + moved = [] + for structure in structures: + structure = structure.copy() + structure.positions += rng.normal(scale=scale, size=(N_ATOMS, 3)) + moved.append(structure) + return moved + + +@pytest.mark.parametrize("pbc", [False, True], ids=["free", "periodic"]) +@pytest.mark.parametrize( + "displacement, path", + # half the skin is the threshold, so these pick the reuse and the rebuild branch + [(0.1 * CUTOFF_SKIN, "reuse"), (2.0 * CUTOFF_SKIN, "rebuild")], + ids=["reuse", "rebuild"], +) +def test_update_never_returns_pairs_beyond_the_cutoff(pbc, displacement, path): + """Reusing a list must not leak the skin into what the model sees. + + The list is built out to ``cutoff + cutoff_skin`` so that it stays valid while the + atoms move; the pairs beyond the cutoff have to be dropped again before the batch + reaches the model, on the step that rebuilds the list and on every step that reuses + it. Checked against a plain neighbor list built for the same positions. + """ + structures = make_structures(pbc) + neighbor_list = make_batch_neighbor_list() + neighbor_list.update(batch_of(structures)) + + moved = displaced(structures, displacement) + updated = neighbor_list.update(batch_of(moved)) + + np.testing.assert_allclose( + neighbor_pairs(updated), neighbor_pairs(freshly_built(moved)), atol=1e-6 + ) + + # and the branch under test really is the one that ran + rebuilt = torch.allclose( + neighbor_list._references[0][properties.R], + torch.from_numpy(moved[0].positions), + ) + assert rebuilt == (path == "rebuild") + + +def test_update_accepts_a_batch_without_a_sample_index(): + """A batch read back from a trajectory has no sample index; it must still update. + + ``BatchwiseTrajectoryReader.frame`` stores no ``idx``, so a relaxation resumed from + a frame would otherwise fail on its very first step. + """ + structures = make_structures(pbc=False) + inputs = batch_of(structures) + del inputs[properties.idx] + + updated = make_batch_neighbor_list().update(inputs) + + assert updated[properties.idx_i].shape == updated[properties.idx_j].shape + assert len(updated[properties.idx_i]) > 0 + + +def test_only_the_structures_that_moved_are_rebuilt(): + """Each structure carries its own list and its own reference positions. + + A batch is rebuilt structure by structure, so one structure walking away must not + cost its neighbors in the batch a rebuild -- nor touch the positions their own + lists were measured against. + """ + structures = make_structures(pbc=False) + neighbor_list = make_batch_neighbor_list() + neighbor_list.update(batch_of(structures)) + references = { + idx: neighbor_list._references[idx][properties.R].clone() for idx in range(3) + } + + moved = [structure.copy() for structure in structures] + moved[1].positions += 2.0 * CUTOFF_SKIN + + assert neighbor_list._stale_structures(batch_of(moved)) == [1] + + neighbor_list.update(batch_of(moved)) + for idx in (0, 2): + assert torch.allclose( + references[idx], neighbor_list._references[idx][properties.R] + ) + assert not torch.allclose(references[1], neighbor_list._references[1][properties.R]) + + +def test_drift_is_measured_against_the_positions_each_list_was_built_for(): + """A structure creeping along must not have its budget reset by the others. + + Its drift is measured against the positions its own list was built for, so it is + rebuilt once the total displacement since then passes half the skin -- however many + steps that took, and whatever the rest of the batch did in the meantime. + """ + structures = make_structures(pbc=False, cell=False) + neighbor_list = make_batch_neighbor_list() + + current = [structure.copy() for structure in structures] + neighbor_list.update(batch_of(current)) + + # a quarter of the threshold per step for the creeper, well past it for its neighbor, + # so that the other structure forces a rebuild on nearly every step + step = 0.125 * CUTOFF_SKIN + for _ in range(8): + current[0].positions[:, 0] += step + current[2].positions += 2.0 * CUTOFF_SKIN + + updated = neighbor_list.update(batch_of(current)) + np.testing.assert_allclose( + neighbor_pairs(updated), neighbor_pairs(freshly_built(current)), atol=1e-6 + ) + + +def test_a_batch_of_a_different_size_is_not_reused(): + """Cached lists are keyed by position in the batch, so a new batch starts over.""" + neighbor_list = make_batch_neighbor_list() + neighbor_list.update(batch_of(make_structures(pbc=False, n_structures=3))) + + structures = make_structures(pbc=False, n_structures=5, seed=7) + updated = neighbor_list.update(batch_of(structures)) + + np.testing.assert_allclose( + neighbor_pairs(updated), neighbor_pairs(freshly_built(structures)), atol=1e-6 + ) + + +def test_a_changed_cell_forces_a_rebuild(): + """Positions can sit still while the cell moves the periodic images.""" + structures = make_structures(pbc=True) + neighbor_list = make_batch_neighbor_list() + neighbor_list.update(batch_of(structures)) + + squeezed = [structure.copy() for structure in structures] + for structure in squeezed: + structure.set_cell(6.5 * np.eye(3)) + + updated = neighbor_list.update(batch_of(squeezed)) + + np.testing.assert_allclose( + neighbor_pairs(updated), neighbor_pairs(freshly_built(squeezed)), atol=1e-6 + ) + + +def test_caller_entries_survive_a_rebuild(): + """``update`` refreshes the neighborhoods and leaves everything else alone. + + A batch carries more than its structures -- energies, convergence flags, whatever + the caller put there -- and dropping those on the steps that happen to rebuild is + the kind of asymmetry that only shows up much later. + """ + structures = make_structures(pbc=False) + neighbor_list = make_batch_neighbor_list() + neighbor_list.update(batch_of(structures)) + + inputs = batch_of(displaced(structures, 2.0 * CUTOFF_SKIN)) + inputs["converged"] = torch.tensor([True, False, True]) + inputs[properties.energy] = torch.zeros(3, dtype=torch.float64) + + updated = neighbor_list.update(inputs) + + assert updated["converged"].tolist() == [True, False, True] + assert properties.energy in updated + assert properties.Z in updated and properties.n_atoms in updated + + +def test_neighbors_returns_no_structure_entries(): + """The MD calculator merges the result into the batch it is propagating. + + Handing back positions -- copies, at that -- would overwrite the live ones. + """ + structures = make_structures(pbc=False) + + neighbors = make_batch_neighbor_list().neighbors(batch_of(structures)) + + for key in (properties.R, properties.Z, properties.cell, properties.pbc): + assert key not in neighbors + assert properties.idx_i in neighbors and properties.offsets in neighbors + + +def test_triples_are_renumbered_onto_the_pruned_pairs(): + """``idx_j_triples`` indexes pairs, so pruning pairs has to renumber the triples. + + Left alone they would point at the wrong pairs, or past the end of the array. + """ + structures = make_structures(pbc=False) + neighbor_list = make_batch_neighbor_list( + transforms=spk.transform.CollectAtomTriples() + ) + + updated = neighbor_list.update(batch_of(structures)) + reference = freshly_built(structures, transforms=spk.transform.CollectAtomTriples()) + + n_pairs = len(updated[properties.idx_i]) + for key in (properties.idx_j_triples, properties.idx_k_triples): + assert int(updated[key].max()) < n_pairs + + # the triples are the same ones a fresh build finds, read through the pair arrays + def triple_atoms(inputs): + idx_j = inputs[properties.idx_j] + triples = np.column_stack( + [ + inputs[properties.idx_i_triples].cpu().numpy(), + idx_j[inputs[properties.idx_j_triples]].cpu().numpy(), + idx_j[inputs[properties.idx_k_triples]].cpu().numpy(), + ] + ) + return triples[np.lexsort(triples.T[::-1])] + + np.testing.assert_array_equal(triple_atoms(updated), triple_atoms(reference)) + + +def test_split_batch_inverts_the_collate_function(): + """Cutting a batch apart and putting it back must give the batch back.""" + structures = make_structures(pbc=True, n_structures=4) + inputs = batch_of(structures) + + recollated = _atoms_collate_fn(split_batch(inputs)) + + for key in (properties.n_atoms, properties.Z, properties.R, properties.pbc): + assert torch.equal(recollated[key], inputs[key]) + assert torch.allclose(recollated[properties.cell], inputs[properties.cell]) + + +def test_split_batch_handles_ragged_and_single_atom_structures(): + """Atom-wise and structure-wise entries are told apart by more than their length.""" + structures = [ + Atoms(numbers=[1], positions=[[0.0, 0.0, 0.0]], cell=np.eye(3) * 4, pbc=True), + Atoms(numbers=[8, 1], positions=[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]), + Atoms(numbers=[2], positions=[[2.0, 0.0, 0.0]]), + ] + samples = split_batch(atoms_to_batch(structures, dtype=torch.float64)) + + assert [int(s[properties.n_atoms]) for s in samples] == [1, 2, 1] + assert [s[properties.idx].item() for s in samples] == [0, 1, 2] + for sample, structure in zip(samples, structures): + assert sample[properties.cell].shape == (1, 3, 3) + assert sample[properties.pbc].shape == (1, 3) + np.testing.assert_allclose( + sample[properties.R].numpy(), structure.positions, atol=1e-6 + ) + + +def test_split_batch_of_single_atom_structures_round_trips(): + """The one case where atom count and structure count coincide.""" + structures = [ + Atoms(numbers=[1], positions=[[float(i), 0.0, 0.0]], cell=np.eye(3) * 4) + for i in range(3) + ] + inputs = atoms_to_batch(structures, dtype=torch.float64) + + recollated = _atoms_collate_fn(split_batch(inputs)) + + assert torch.equal(recollated[properties.Z], inputs[properties.Z]) + assert torch.allclose(recollated[properties.cell], inputs[properties.cell])