{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e27adf39",
   "metadata": {},
   "source": [
    "# ANI-1, block by block: reproducing the original implementation with `xnn`\n",
    "\n",
    "[ANI-1 (Smith, Isayev & Roitberg, *Chem. Sci.* **8**, 3192, 2017)](https://doi.org/10.1039/C6SC05720A)\n",
    "is a High-Dimensional Neural Network Potential whose per-atom descriptor is the\n",
    "**Atomic Environment Vector (AEV)**: element-resolved radial *and* angular\n",
    "symmetry functions (paper eqns 2–4). One neural network per element maps the AEV\n",
    "to an atomic energy; the atomic energies (plus per-element self energies) sum to\n",
    "the total energy.\n",
    "\n",
    "The reference implementation is [**aiqm/torchani**](https://github.com/aiqm/torchani)\n",
    "(the PyTorch ANI the paper's authors maintain). This notebook shows, block by\n",
    "block, that `xnn` reproduces torchani **element-for-element**:\n",
    "\n",
    "| block | what we check |\n",
    "|-------|---------------|\n",
    "| cutoff | `xnn` `CosineCutoff` vs `torchani.aev.cutoff_cosine` |\n",
    "| AEV (ANI-1x, 384) | `AEV.ani1x()` vs `torchani.AEVComputer` with the ANI-1x grid |\n",
    "| AEV (ANI-1, 768)  | `AEV.ani1()` vs `torchani.AEVComputer` with the paper's grid |\n",
    "| element nets + E, F | transplant torchani's **pretrained** ANI-1x weights → energies & forces |\n",
    "| ensemble | 8-model ANI-1x ensemble mean |\n",
    "\n",
    "Everything runs in `float64`. Run this notebook with the **`xnn-ani`** kernel\n",
    "(`torch` + `torchani==2.2.4` + `xnn`); see `pyproject.toml`'s `[ani]` extra.\n",
    "\n",
    "> **Conventions note.** torchani follows NeuroChem, not the paper's eqn 3\n",
    "> literally: it multiplies the radial term by `0.25` and scales `cos(θ)` by\n",
    "> `0.95` inside `acos` (to keep the gradient finite at ±1). `xnn` bakes both in\n",
    "> by default (`radial_prefactor=0.25`, `angular_cos_factor=0.95`), so the AEVs\n",
    "> match exactly. Set them to `1.0` for the literal Behler-Parrinello form."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da211960",
   "metadata": {},
   "source": [
    "## 0. Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0b063c19",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T06:17:29.051742Z",
     "iopub.status.busy": "2026-07-20T06:17:29.051438Z",
     "iopub.status.idle": "2026-07-20T06:17:32.395494Z",
     "shell.execute_reply": "2026-07-20T06:17:32.394514Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch 2.5.1+cpu | torchani 2.2.4\n"
     ]
    }
   ],
   "source": [
    "import warnings; warnings.filterwarnings(\"ignore\")\n",
    "import math\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchani\n",
    "\n",
    "torch.set_default_dtype(torch.float64)\n",
    "\n",
    "from xnn.dnn.models.ani import ANI\n",
    "from xnn.dnn.featurizers import AEV\n",
    "from xnn.dnn.featurizers.aev import _even_shifts, _angle_shifts\n",
    "from xnn.common.featurizers import CosineCutoff\n",
    "from xnn.common.data import AtomicGraph\n",
    "from xnn.common.data.neighborlist import build_neighbor_list\n",
    "from xnn.common.models.outputs import ForceStressOutput\n",
    "\n",
    "print(\"torch\", torch.__version__, \"| torchani\", torchani.__version__)\n",
    "\n",
    "# H, C, N, O -> torchani's internal species index\n",
    "IDX = {1: 0, 6: 1, 7: 2, 8: 3}\n",
    "\n",
    "\n",
    "def xnn_graph(Z, pos, cutoff, requires_grad=False):\n",
    "    \"Build a single-structure AtomicGraph from atomic numbers + positions.\"\n",
    "    Z = torch.as_tensor(Z, dtype=torch.long)\n",
    "    pos = torch.as_tensor(pos, dtype=torch.float64)\n",
    "    if requires_grad:\n",
    "        pos = pos.clone().requires_grad_(True)\n",
    "    ei, cs = build_neighbor_list(pos, cutoff)\n",
    "    return AtomicGraph(pos=pos, atomic_numbers=Z, edge_index=ei, cell_shifts=cs,\n",
    "                       batch=torch.zeros(len(Z), dtype=torch.long),\n",
    "                       n_atoms=torch.tensor([len(Z)]))\n",
    "\n",
    "\n",
    "def torchani_aev(Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, ns=4):\n",
    "    \"A torchani AEVComputer built from explicit constants.\"\n",
    "    t = lambda x: torch.tensor(x, dtype=torch.float64)\n",
    "    return torchani.AEVComputer(Rcr, Rca, t(EtaR), t(ShfR), t(EtaA), t(Zeta),\n",
    "                                t(ShfA), t(ShfZ), ns)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c312bd38",
   "metadata": {},
   "source": [
    "### A handful of H/C/N/O molecules\n",
    "\n",
    "Small, distinct organic geometries (water, ammonia, methanol, formamide) with a\n",
    "little random distortion so no symmetry hides a bug. We compare `xnn` and\n",
    "torchani on every one of them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d173ae71",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T06:17:32.398351Z",
     "iopub.status.busy": "2026-07-20T06:17:32.397983Z",
     "iopub.status.idle": "2026-07-20T06:17:32.411710Z",
     "shell.execute_reply": "2026-07-20T06:17:32.410827Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "molecules: water, ammonia, methanol, formamide\n"
     ]
    }
   ],
   "source": [
    "rng = np.random.default_rng(0)\n",
    "\n",
    "_MOLS = {\n",
    "    \"water\":     ([8, 1, 1],\n",
    "                  [[0, 0, 0], [0.76, 0.59, 0], [-0.76, 0.59, 0]]),\n",
    "    \"ammonia\":   ([7, 1, 1, 1],\n",
    "                  [[0, 0, 0.12], [0, 0.94, -0.27], [0.81, -0.47, -0.27],\n",
    "                   [-0.81, -0.47, -0.27]]),\n",
    "    \"methanol\":  ([6, 8, 1, 1, 1, 1],\n",
    "                  [[-0.05, 0.66, 0], [-0.05, -0.75, 0], [1.0, 1.0, 0],\n",
    "                   [-0.56, 1.05, 0.88], [-0.56, 1.05, -0.88], [0.85, -1.09, 0]]),\n",
    "    \"formamide\": ([6, 7, 8, 1, 1, 1],\n",
    "                  [[0, 0.10, 0], [1.27, 0.55, 0], [-0.98, 0.82, 0],\n",
    "                   [-0.20, -0.95, 0], [1.42, 1.55, 0], [2.05, -0.09, 0]]),\n",
    "}\n",
    "MOLS = {name: (np.array(Z, dtype=np.int64),\n",
    "               np.array(pos, dtype=np.float64) + 0.03 * rng.standard_normal((len(Z), 3)))\n",
    "        for name, (Z, pos) in _MOLS.items()}\n",
    "print(\"molecules:\", \", \".join(MOLS))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b779aafa",
   "metadata": {},
   "source": [
    "## 1. Cutoff function (paper eqn 2)\n",
    "\n",
    "$f_C(r) = \\tfrac12\\cos(\\pi r / R_C) + \\tfrac12$ for $r \\le R_C$, else 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "15e5d426",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T06:17:32.413983Z",
     "iopub.status.busy": "2026-07-20T06:17:32.413668Z",
     "iopub.status.idle": "2026-07-20T06:17:33.104126Z",
     "shell.execute_reply": "2026-07-20T06:17:33.102734Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "max|Δ| cutoff = 2.22e-16\n"
     ]
    }
   ],
   "source": [
    "r = torch.linspace(0.05, 5.2, 200, dtype=torch.float64)\n",
    "xn = CosineCutoff(5.2)(r)\n",
    "ta = torchani.aev.cutoff_cosine(r, 5.2) * (r < 5.2)\n",
    "print(f\"max|Δ| cutoff = {(xn - ta).abs().max():.2e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e67ffa94",
   "metadata": {},
   "source": [
    "## 2. AEV: the ANI-1x grid (384 elements)\n",
    "\n",
    "`AEV.ani1x()` uses the exact constants torchani ships for ANI-1x (radial cutoff\n",
    "5.2 Å, 16 radial shifts; angular cutoff 3.5 Å, 4 radial × 8 angular shifts). We\n",
    "build a `torchani.AEVComputer` from the *same* constants and compare the full\n",
    "384-vector per atom, and separately the radial (0:64) and angular (64:384)\n",
    "blocks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ad2250d3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T06:17:33.107272Z",
     "iopub.status.busy": "2026-07-20T06:17:33.106991Z",
     "iopub.status.idle": "2026-07-20T06:17:33.797663Z",
     "shell.execute_reply": "2026-07-20T06:17:33.796958Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "molecule      max|Δ| full    max|Δ| radial   max|Δ| angular\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "water            2.78e-17         2.78e-17         0.00e+00\n",
      "ammonia          1.11e-16         5.55e-17         1.11e-16\n",
      "methanol         2.50e-16         8.33e-17         2.50e-16\n",
      "formamide        2.78e-16         2.78e-17         2.78e-16\n",
      "\n",
      "worst over all molecules: 2.78e-16  (AEV length 384)\n"
     ]
    }
   ],
   "source": [
    "aev_x = AEV.ani1x()\n",
    "tani = torchani_aev(5.2, 3.5, [16.0], _even_shifts(5.2, 16),\n",
    "                    [8.0], [32.0], _even_shifts(3.5, 4), _angle_shifts(8))\n",
    "\n",
    "print(f\"{'molecule':<10} {'max|Δ| full':>14} {'max|Δ| radial':>16} {'max|Δ| angular':>16}\")\n",
    "worst = 0.0\n",
    "for name, (Z, pos) in MOLS.items():\n",
    "    g = xnn_graph(Z, pos, aev_x.cutoff)\n",
    "    x = aev_x(g)\n",
    "    sp = torch.tensor([[IDX[z] for z in Z]])\n",
    "    _, t = tani((sp, torch.as_tensor(pos[None])))\n",
    "    t = t[0]\n",
    "    d_full = (x - t).abs().max().item()\n",
    "    d_rad = (x[:, :64] - t[:, :64]).abs().max().item()\n",
    "    d_ang = (x[:, 64:] - t[:, 64:]).abs().max().item()\n",
    "    worst = max(worst, d_full)\n",
    "    print(f\"{name:<10} {d_full:>14.2e} {d_rad:>16.2e} {d_ang:>16.2e}\")\n",
    "print(f\"\\nworst over all molecules: {worst:.2e}  (AEV length {aev_x.output_dim})\")\n",
    "assert worst < 1e-10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b28a7ab1",
   "metadata": {},
   "source": [
    "## 3. AEV: the original ANI-1 grid (768 elements)\n",
    "\n",
    "`AEV.ani1()` is the paper's parameterisation: radial cutoff 4.6 Å with 32\n",
    "shifts, angular cutoff 3.1 Å with 8 × 8 shifts → a 768-vector for H, C, N, O\n",
    "(paper §3.4). Same check, against a torchani computer built with these\n",
    "constants."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "01aa6dd4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T06:17:33.799598Z",
     "iopub.status.busy": "2026-07-20T06:17:33.799438Z",
     "iopub.status.idle": "2026-07-20T06:17:33.820622Z",
     "shell.execute_reply": "2026-07-20T06:17:33.819863Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "water      max|Δ| = 3.33e-16\n",
      "ammonia    max|Δ| = 8.88e-16\n",
      "methanol   max|Δ| = 8.88e-16\n",
      "formamide  max|Δ| = 6.66e-16\n",
      "\n",
      "worst: 8.88e-16  (AEV length 768)\n"
     ]
    }
   ],
   "source": [
    "aev_1 = AEV.ani1()\n",
    "tani1 = torchani_aev(4.6, 3.1, [16.0], _even_shifts(4.6, 32),\n",
    "                     [8.0], [8.0], _even_shifts(3.1, 8), _angle_shifts(8))\n",
    "\n",
    "worst = 0.0\n",
    "for name, (Z, pos) in MOLS.items():\n",
    "    g = xnn_graph(Z, pos, aev_1.cutoff)\n",
    "    x = aev_1(g)\n",
    "    sp = torch.tensor([[IDX[z] for z in Z]])\n",
    "    _, t = tani1((sp, torch.as_tensor(pos[None])))\n",
    "    d = (x - t[0]).abs().max().item()\n",
    "    worst = max(worst, d)\n",
    "    print(f\"{name:<10} max|Δ| = {d:.2e}\")\n",
    "print(f\"\\nworst: {worst:.2e}  (AEV length {aev_1.output_dim})\")\n",
    "assert worst < 1e-10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b0f56ca",
   "metadata": {},
   "source": [
    "## 4. Element networks → energies and forces (pretrained weights)\n",
    "\n",
    "Now the whole model. We load torchani's **pretrained** ANI-1x (member 0 of its\n",
    "8-model ensemble), transplant its per-element `Linear` weights into an\n",
    "`xnn` `ANI.ani1x()`, and compare total energies and (autograd) forces. The\n",
    "per-element architectures (H `160:128:96`, C `144:112:96`, N/O `128:112:96`),\n",
    "the `CELU(0.1)` activation and the per-element self energies all line up, so the\n",
    "transplant is one-to-one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b2ef26d6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T06:17:33.822157Z",
     "iopub.status.busy": "2026-07-20T06:17:33.822002Z",
     "iopub.status.idle": "2026-07-20T06:17:41.194902Z",
     "shell.execute_reply": "2026-07-20T06:17:41.194112Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/D3/sina/xnn/.venv-ani/lib/python3.13/site-packages/torchani/resources/\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "molecule        E xnn (Ha)       E torchani       |ΔE|    max|ΔF|\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "water          -76.38793938     -76.38793938   9.62e-10   2.03e-08\n",
      "ammonia        -56.52444686     -56.52444686   1.48e-09   3.75e-09\n",
      "methanol      -115.67180357    -115.67180357   7.47e-10   1.45e-08\n",
      "formamide     -169.81418123    -169.81418123   2.51e-10   2.20e-08\n",
      "\n",
      "worst |ΔE| = 1.48e-09 Ha   worst max|ΔF| = 2.20e-08 Ha/Å\n"
     ]
    }
   ],
   "source": [
    "model = torchani.models.ANI1x(periodic_table_index=False).double()\n",
    "member0 = model.neural_networks[0]\n",
    "zsym = {1: \"H\", 6: \"C\", 7: \"N\", 8: \"O\"}\n",
    "\n",
    "\n",
    "def transplant(member, xa):\n",
    "    \"Copy a torchani ANIModel's per-element Linear weights into an xnn ANI.\"\n",
    "    nets = dict(member.named_children())\n",
    "    for z in [1, 6, 7, 8]:\n",
    "        src = [l for l in nets[zsym[z]] if isinstance(l, torch.nn.Linear)]\n",
    "        dst = [l for l in xa.element_nets.nets[str(z)] if isinstance(l, torch.nn.Linear)]\n",
    "        for s, d in zip(src, dst):\n",
    "            d.weight.data = s.weight.data.clone()\n",
    "            d.bias.data = s.bias.data.clone()\n",
    "\n",
    "\n",
    "xa = ANI.ani1x()          # self energies default to torchani's (Hartree)\n",
    "transplant(member0, xa)\n",
    "wrapped = ForceStressOutput(xa)\n",
    "\n",
    "print(f\"{'molecule':<10} {'E xnn (Ha)':>16} {'E torchani':>16} {'|ΔE|':>10} {'max|ΔF|':>10}\")\n",
    "wE = wF = 0.0\n",
    "for name, (Z, pos) in MOLS.items():\n",
    "    g = xnn_graph(Z, pos, xa.cutoff, requires_grad=True)\n",
    "    out = wrapped(g)\n",
    "    E_x = out[\"energy\"].item()\n",
    "    F_x = out[\"forces\"].detach().numpy()\n",
    "\n",
    "    coords = torch.as_tensor(pos[None]).clone().requires_grad_(True)\n",
    "    sp = torch.tensor([[IDX[z] for z in Z]])\n",
    "    e = model.energy_shifter(member0(model.aev_computer((sp, coords)))).energies\n",
    "    F_t = -torch.autograd.grad(e.sum(), coords)[0][0].numpy()\n",
    "    E_t = e.item()\n",
    "\n",
    "    dE, dF = abs(E_x - E_t), np.abs(F_x - F_t).max()\n",
    "    wE, wF = max(wE, dE), max(wF, dF)\n",
    "    print(f\"{name:<10} {E_x:>16.8f} {E_t:>16.8f} {dE:>10.2e} {dF:>10.2e}\")\n",
    "print(f\"\\nworst |ΔE| = {wE:.2e} Ha   worst max|ΔF| = {wF:.2e} Ha/Å\")\n",
    "assert wE < 1e-6 and wF < 1e-6"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd4db76a",
   "metadata": {},
   "source": [
    "## 5. The full 8-model ANI-1x ensemble\n",
    "\n",
    "The released ANI-1x is the **mean** of 8 networks. We transplant all 8 into 8\n",
    "`xnn` models and average, then compare to torchani's built-in ensemble output\n",
    "(which also adds the self energies)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "951fafdd",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T06:17:41.196822Z",
     "iopub.status.busy": "2026-07-20T06:17:41.196567Z",
     "iopub.status.idle": "2026-07-20T06:17:41.541730Z",
     "shell.execute_reply": "2026-07-20T06:17:41.540875Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "molecule        E xnn mean       E torchani       |ΔE|\n",
      "water          -76.38818940     -76.38818940   3.61e-11\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ammonia        -56.52408408     -56.52408408   1.96e-09\n",
      "methanol      -115.67195817    -115.67195817   9.51e-10\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "formamide     -169.81364721    -169.81364721   7.59e-10\n",
      "\n",
      "worst |ΔE| = 1.96e-09 Ha\n"
     ]
    }
   ],
   "source": [
    "members = [ANI.ani1x() for _ in range(len(model.neural_networks))]\n",
    "for xa_k, m_k in zip(members, model.neural_networks):\n",
    "    transplant(m_k, xa_k)\n",
    "\n",
    "wrapped_members = [ForceStressOutput(m) for m in members]\n",
    "\n",
    "print(f\"{'molecule':<10} {'E xnn mean':>16} {'E torchani':>16} {'|ΔE|':>10}\")\n",
    "worst = 0.0\n",
    "for name, (Z, pos) in MOLS.items():\n",
    "    g = xnn_graph(Z, pos, members[0].cutoff)\n",
    "    E_x = float(np.mean([w(g)[\"energy\"].item() for w in wrapped_members]))\n",
    "\n",
    "    coords = torch.as_tensor(pos[None])\n",
    "    sp = torch.tensor([[IDX[z] for z in Z]])\n",
    "    E_t = model((sp, coords)).energies.item()\n",
    "    dE = abs(E_x - E_t)\n",
    "    worst = max(worst, dE)\n",
    "    print(f\"{name:<10} {E_x:>16.8f} {E_t:>16.8f} {dE:>10.2e}\")\n",
    "print(f\"\\nworst |ΔE| = {worst:.2e} Ha\")\n",
    "assert worst < 1e-6"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a10f7386",
   "metadata": {},
   "source": [
    "## 6. ANI-1ccx: the same architecture, coupled-cluster weights\n",
    "\n",
    "The **ANI-1ccx** potential (Smith *et al.*, *Nat. Commun.* **10**, 2903, 2019)\n",
    "is the ANI-1x architecture transfer-learned to CCSD(T)*/CBS coupled-cluster\n",
    "data; in `xnn`, `ANI.ani1ccx()` simply reuses `ANI.ani1x()` with the\n",
    "coupled-cluster self energies. The very same `transplant` therefore moves\n",
    "torchani's pretrained ANI-1ccx across one-to-one: first a single member,\n",
    "then the full 8-network ensemble the potential is released as (paper SI S1.2.4)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6a923200",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T06:17:41.543612Z",
     "iopub.status.busy": "2026-07-20T06:17:41.543441Z",
     "iopub.status.idle": "2026-07-20T06:17:44.993690Z",
     "shell.execute_reply": "2026-07-20T06:17:44.992889Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/D3/sina/xnn/.venv-ani/lib/python3.13/site-packages/torchani/resources/\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "molecule        E xnn (Ha)       E torchani       |ΔE|    max|ΔF|\n",
      "water          -76.38344647     -76.38344647   5.19e-10   1.01e-08\n",
      "ammonia        -56.50648961     -56.50648961   1.04e-09   1.53e-08\n",
      "methanol      -115.61587813    -115.61587813   5.72e-10   2.03e-08\n",
      "formamide     -169.71677096    -169.71677096   1.91e-10   1.87e-08\n",
      "\n",
      "worst |ΔE| = 1.04e-09 Ha   worst max|ΔF| = 2.03e-08 Ha/Å\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "8-model ANI-1ccx ensemble: worst |ΔE| = 1.35e-09 Ha\n"
     ]
    }
   ],
   "source": [
    "model_ccx = torchani.models.ANI1ccx(periodic_table_index=False).double()\n",
    "xccx = ANI.ani1ccx()      # ANI-1x architecture + CCSD(T)*/CBS self energies\n",
    "transplant(model_ccx.neural_networks[0], xccx)\n",
    "wrapped_ccx = ForceStressOutput(xccx)\n",
    "\n",
    "print(f\"{'molecule':<10} {'E xnn (Ha)':>16} {'E torchani':>16} {'|\\u0394E|':>10} {'max|\\u0394F|':>10}\")\n",
    "wE = wF = 0.0\n",
    "for name, (Z, pos) in MOLS.items():\n",
    "    g = xnn_graph(Z, pos, xccx.cutoff, requires_grad=True)\n",
    "    out = wrapped_ccx(g)\n",
    "    E_x, F_x = out[\"energy\"].item(), out[\"forces\"].detach().numpy()\n",
    "\n",
    "    coords = torch.as_tensor(pos[None]).clone().requires_grad_(True)\n",
    "    sp = torch.tensor([[IDX[z] for z in Z]])\n",
    "    e = model_ccx.energy_shifter(\n",
    "        model_ccx.neural_networks[0](model_ccx.aev_computer((sp, coords)))).energies\n",
    "    F_t = -torch.autograd.grad(e.sum(), coords)[0][0].numpy()\n",
    "\n",
    "    dE, dF = abs(E_x - e.item()), np.abs(F_x - F_t).max()\n",
    "    wE, wF = max(wE, dE), max(wF, dF)\n",
    "    print(f\"{name:<10} {E_x:>16.8f} {e.item():>16.8f} {dE:>10.2e} {dF:>10.2e}\")\n",
    "print(f\"\\nworst |\\u0394E| = {wE:.2e} Ha   worst max|\\u0394F| = {wF:.2e} Ha/\\u00c5\")\n",
    "assert wE < 1e-6 and wF < 1e-6\n",
    "\n",
    "# the released ANI-1ccx, like ANI-1x, is the MEAN of 8 networks (paper SI S1.2.4)\n",
    "members_ccx = [ANI.ani1ccx() for _ in range(len(model_ccx.neural_networks))]\n",
    "for xk, mk in zip(members_ccx, model_ccx.neural_networks):\n",
    "    transplant(mk, xk)\n",
    "wrapped_members_ccx = [ForceStressOutput(m) for m in members_ccx]\n",
    "worst = 0.0\n",
    "for name, (Z, pos) in MOLS.items():\n",
    "    g = xnn_graph(Z, pos, xccx.cutoff)\n",
    "    E_x = float(np.mean([w(g)[\"energy\"].item() for w in wrapped_members_ccx]))\n",
    "    coords = torch.as_tensor(pos[None])\n",
    "    sp = torch.tensor([[IDX[z] for z in Z]])\n",
    "    E_t = model_ccx((sp, coords)).energies.item()\n",
    "    worst = max(worst, abs(E_x - E_t))\n",
    "print(f\"\\n8-model ANI-1ccx ensemble: worst |\\u0394E| = {worst:.2e} Ha\")\n",
    "assert worst < 1e-6"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "797fbb38",
   "metadata": {},
   "source": [
    "## 7. ANI-2x: seven elements (adds S, F, Cl)\n",
    "\n",
    "ANI-2x (Devereux et al., *J. Chem. Theory Comput.* **16**, 4192, 2020) extends\n",
    "ANI to seven elements. Its AEV is larger (1008 elements: radial cutoff 5.1 A,\n",
    "angular cutoff 3.5 A, shift grids starting at 0.8 A) and the per-element\n",
    "networks are wider. `ANI.ani2x()` builds it over the torchani element order\n",
    "(H, C, N, O, S, F, Cl); transplanting torchani's pretrained ANI-2x weights\n",
    "reproduces its energies and forces on molecules containing the new elements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c4fc0557",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T06:17:44.995444Z",
     "iopub.status.busy": "2026-07-20T06:17:44.995277Z",
     "iopub.status.idle": "2026-07-20T06:17:53.570669Z",
     "shell.execute_reply": "2026-07-20T06:17:53.569834Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/D3/sina/xnn/.venv-ani/lib/python3.13/site-packages/torchani/resources/\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AEV length: 1008  cutoff: 5.1\n",
      "molecule            E xnn (Ha)       E torchani       |dE|    max|dF|\n",
      "H2S               -399.35810752    -399.35810752   3.17e-10   4.09e-09\n",
      "CH3F              -139.69268288    -139.69268288   2.07e-09   4.11e-09\n",
      "CH3Cl             -500.06826174    -500.06826174   9.95e-10   1.30e-09\n",
      "methanethiol      -438.61078843    -438.61078843   4.19e-10   3.81e-09\n",
      "\n",
      "worst |dE| = 2.07e-09 Ha   worst max|dF| = 4.11e-09 Ha/A\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8-model ANI-2x ensemble: worst |dE| = 8.59e-10 Ha\n"
     ]
    }
   ],
   "source": [
    "model2x = torchani.models.ANI2x(periodic_table_index=False).double()\n",
    "x2x = ANI.ani2x()                                  # 7 elements, 1008-length AEV\n",
    "print(\"AEV length:\", x2x.featurizer.output_dim, \" cutoff:\", x2x.cutoff)\n",
    "\n",
    "# seven-element index and symbol maps (torchani's order for ANI-2x)\n",
    "IDX2 = {z: i for i, z in enumerate(x2x.species)}\n",
    "zsym2 = {1: \"H\", 6: \"C\", 7: \"N\", 8: \"O\", 16: \"S\", 9: \"F\", 17: \"Cl\"}\n",
    "\n",
    "\n",
    "def transplant2x(member, xa):\n",
    "    \"Copy a torchani ANI-2x member's per-element Linear weights into an xnn ANI.\"\n",
    "    nets = dict(member.named_children())\n",
    "    for z in x2x.species:\n",
    "        src = [l for l in nets[zsym2[z]] if isinstance(l, torch.nn.Linear)]\n",
    "        dst = [l for l in xa.element_nets.nets[str(z)] if isinstance(l, torch.nn.Linear)]\n",
    "        for s, d in zip(src, dst):\n",
    "            d.weight.data = s.weight.data.clone()\n",
    "            d.bias.data = s.bias.data.clone()\n",
    "\n",
    "\n",
    "transplant2x(model2x.neural_networks[0], x2x)\n",
    "wrapped2x = ForceStressOutput(x2x)\n",
    "\n",
    "# small molecules that exercise S, F, and Cl\n",
    "rng2 = np.random.default_rng(2)\n",
    "MOLS2X = {\n",
    "    \"H2S\":            ([16, 1, 1],\n",
    "                       [[0, 0, 0], [0.96, 0.94, 0], [-0.96, 0.94, 0]]),\n",
    "    \"CH3F\":           ([6, 9, 1, 1, 1],\n",
    "                       [[0, 0, 0], [0, 0, 1.38], [1.03, 0, -0.36],\n",
    "                        [-0.51, 0.89, -0.36], [-0.51, -0.89, -0.36]]),\n",
    "    \"CH3Cl\":          ([6, 17, 1, 1, 1],\n",
    "                       [[0, 0, 0], [0, 0, 1.78], [1.03, 0, -0.36],\n",
    "                        [-0.51, 0.89, -0.36], [-0.51, -0.89, -0.36]]),\n",
    "    \"methanethiol\":   ([6, 16, 1, 1, 1, 1],\n",
    "                       [[0, 0, 0], [1.42, 0.72, 0], [-0.55, 0.36, 0.89],\n",
    "                        [-0.55, 0.36, -0.89], [-0.10, -1.09, 0], [1.28, 2.00, 0]]),\n",
    "}\n",
    "MOLS2X = {name: (np.array(Z, dtype=np.int64),\n",
    "                 np.array(pos, dtype=np.float64)\n",
    "                 + 0.03 * rng2.standard_normal((len(Z), 3)))\n",
    "          for name, (Z, pos) in MOLS2X.items()}\n",
    "\n",
    "print(f\"{'molecule':<14} {'E xnn (Ha)':>16} {'E torchani':>16} {'|dE|':>10} {'max|dF|':>10}\")\n",
    "wE = wF = 0.0\n",
    "for name, (Z, pos) in MOLS2X.items():\n",
    "    g = xnn_graph(Z, pos, x2x.cutoff, requires_grad=True)\n",
    "    out = wrapped2x(g)\n",
    "    E_x, F_x = out[\"energy\"].item(), out[\"forces\"].detach().numpy()\n",
    "\n",
    "    coords = torch.as_tensor(pos[None]).clone().requires_grad_(True)\n",
    "    sp = torch.tensor([[IDX2[int(z)] for z in Z]])\n",
    "    e = model2x.energy_shifter(\n",
    "        model2x.neural_networks[0](model2x.aev_computer((sp, coords)))).energies\n",
    "    F_t = -torch.autograd.grad(e.sum(), coords)[0][0].numpy()\n",
    "\n",
    "    dE, dF = abs(E_x - e.item()), np.abs(F_x - F_t).max()\n",
    "    wE, wF = max(wE, dE), max(wF, dF)\n",
    "    print(f\"{name:<14} {E_x:>16.8f} {e.item():>16.8f} {dE:>10.2e} {dF:>10.2e}\")\n",
    "print(f\"\\nworst |dE| = {wE:.2e} Ha   worst max|dF| = {wF:.2e} Ha/A\")\n",
    "assert wE < 1e-5 and wF < 1e-5\n",
    "\n",
    "# and the full 8-model ANI-2x ensemble (mean energy)\n",
    "members2x = [ANI.ani2x() for _ in model2x.neural_networks]\n",
    "for xk, mk in zip(members2x, model2x.neural_networks):\n",
    "    transplant2x(mk, xk)\n",
    "wrapped_members2x = [ForceStressOutput(m) for m in members2x]\n",
    "worst = 0.0\n",
    "for name, (Z, pos) in MOLS2X.items():\n",
    "    g = xnn_graph(Z, pos, x2x.cutoff)\n",
    "    E_x = float(np.mean([w(g)[\"energy\"].item() for w in wrapped_members2x]))\n",
    "    coords = torch.as_tensor(pos[None])\n",
    "    sp = torch.tensor([[IDX2[int(z)] for z in Z]])\n",
    "    E_t = model2x((sp, coords)).energies.item()\n",
    "    worst = max(worst, abs(E_x - E_t))\n",
    "print(f\"8-model ANI-2x ensemble: worst |dE| = {worst:.2e} Ha\")\n",
    "assert worst < 1e-5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6021197a",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "Block by block (cutoff, the ANI-1x (384) and ANI-1 (768) AEVs, the per-element\n",
    "networks with pretrained weights, the 8-model ensemble, the ANI-1ccx\n",
    "transplant, and the seven-element ANI-2x transplant), `xnn` reproduces torchani to numerical precision (`float64`\n",
    "round-off), for both energies and forces. The `xnn` `ANI` is a faithful,\n",
    "from-scratch re-implementation of the ANI method; the same `AEV`/`ANI` code\n",
    "trains from scratch on the ANI datasets (see `examples/dnn/ani/`)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xnn-ani",
   "language": "python",
   "name": "xnn-ani"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}