{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0296c38c",
   "metadata": {},
   "source": [
    "# BAMBOO fidelity check: `xnn` vs the original `bytedance/bamboo`, block by block\n",
    "\n",
    "This notebook verifies that the `xnn` re-implementation of **BAMBOO**\n",
    "(`xnn.hybrid.models.bamboo.BAMBOO`, Gong *et al.* 2024, arXiv:2404.07181) is the\n",
    "*same function* as the original [`bytedance/bamboo`](https://github.com/bytedance/bamboo)\n",
    "model, block by block, by transplanting the upstream weights and comparing every\n",
    "intermediate quantity.\n",
    "\n",
    "The `xnn` model was written **from scratch** on the shared xnn abstractions\n",
    "(`InteratomicPotential`, the `xnn.transformer` attention/radial primitives, the\n",
    "uniform `ForceStressOutput` autograd forces); no upstream code is vendored. We\n",
    "only *import* the upstream package here to compare against it.\n",
    "\n",
    "**The one intentional difference: forces.** Upstream reports a\n",
    "charge-*non-conservative* force `nn_forces + coul_forces` (charges held fixed)\n",
    "and regularises the charge–position derivative `qeq_force` toward zero during\n",
    "training (Supplementary Theorem A.2). xnn instead returns the **full\n",
    "conservative force** `-dE/dr` uniformly via autograd, which equals upstream\n",
    "`forces + qeq_force` to machine precision. We check exactly that identity below.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "866ea065",
   "metadata": {},
   "source": [
    "## 0. Setup: `float64` and the upstream clone\n",
    "\n",
    "We run everything in double precision so the comparison is at machine precision,\n",
    "and clone the upstream repository next to this notebook (skipped if already\n",
    "present). Upstream imports `torch_runstats`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "48ef0d2c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T04:24:33.680156Z",
     "iopub.status.busy": "2026-07-20T04:24:33.680023Z",
     "iopub.status.idle": "2026-07-20T04:24:35.079179Z",
     "shell.execute_reply": "2026-07-20T04:24:35.078371Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "upstream bamboo imported from /D3/sina/xnn/examples/fidelity_checks/bamboo_upstream\n"
     ]
    }
   ],
   "source": [
    "import logging, warnings\n",
    "logging.disable(logging.WARNING)\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import os, sys, subprocess\n",
    "import numpy as np\n",
    "import torch\n",
    "torch.set_default_dtype(torch.float64)\n",
    "\n",
    "UPSTREAM = os.path.abspath(\"bamboo_upstream\")\n",
    "if not os.path.isdir(UPSTREAM):\n",
    "    subprocess.run([\"git\", \"clone\", \"--depth\", \"1\",\n",
    "                    \"https://github.com/bytedance/bamboo\", UPSTREAM], check=True)\n",
    "sys.path.insert(0, UPSTREAM)\n",
    "\n",
    "import torch.nn as nn\n",
    "from models.bamboo_get import BambooGET            # the original model\n",
    "print(\"upstream bamboo imported from\", UPSTREAM)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eafd8aa5",
   "metadata": {},
   "source": [
    "## 1. Build both models with identical hyper-parameters\n",
    "\n",
    "A small-but-real BAMBOO: feature width `dim=64`, `num_rbf=32`, 3 GET layers,\n",
    "16 attention heads, cutoff 5 Å (the upstream defaults)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3bdd0676",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T04:24:35.080981Z",
     "iopub.status.busy": "2026-07-20T04:24:35.080894Z",
     "iopub.status.idle": "2026-07-20T04:24:35.719941Z",
     "shell.execute_reply": "2026-07-20T04:24:35.719292Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "xnn params  : 103620\n",
      "upstream params: 306750\n"
     ]
    }
   ],
   "source": [
    "import xnn\n",
    "from xnn.common.config import from_dict\n",
    "from xnn.common.data import structure_to_graph\n",
    "from xnn.common.models import ForceStressOutput, build_model\n",
    "\n",
    "DIM, NRBF, NLAYERS, NHEADS, CUT = 64, 32, 3, 16, 5.0\n",
    "\n",
    "up = BambooGET(\n",
    "    device=torch.device(\"cpu\"),\n",
    "    coul_disp_params={\"coul_damping_beta\": 18.7, \"coul_damping_r0\": 2.2, \"disp_cutoff\": 10.0},\n",
    "    nn_params={\"dim\": DIM, \"num_rbf\": NRBF, \"rcut\": CUT, \"charge_ub\": 2.0,\n",
    "               \"act_fn\": nn.SiLU(), \"charge_mlp_layers\": 2, \"energy_mlp_layers\": 2},\n",
    "    gnn_params={\"n_layers\": NLAYERS, \"num_heads\": NHEADS, \"act_fn\": nn.GELU()},\n",
    ").eval()\n",
    "\n",
    "cfg = from_dict({\"model\": {\"name\": \"bamboo\", \"cutoff\": CUT, \"n_features\": DIM,\n",
    "    \"n_rbf\": NRBF, \"n_interactions\": NLAYERS, \"extra\": {\"num_heads\": NHEADS}}})\n",
    "xm = build_model(cfg.model).eval()\n",
    "\n",
    "print(\"xnn params  :\", sum(p.numel() for p in xm.parameters()))\n",
    "print(\"upstream params:\", sum(p.numel() for p in up.parameters()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3404249c",
   "metadata": {},
   "source": [
    "## 2. The transplant map (upstream → xnn)\n",
    "\n",
    "Every upstream tensor maps to exactly one xnn tensor. Because the read-out MLPs\n",
    "and the transformer projections keep matching shapes, most transplants are a\n",
    "plain `load_state_dict`.\n",
    "\n",
    "| upstream | xnn |\n",
    "|---|---|\n",
    "| `atom_embtab` | `atom_emb` |\n",
    "| `dis_rbf` (`means`, `betas`) | `dis_rbf` |\n",
    "| `rbf_proj` | `rbf_proj` |\n",
    "| `energy_mlp` | `energy_mlp` |\n",
    "| `charge_mlp` | `charge_mlp` |\n",
    "| `pred_electronegativity_mlp` | `electronegativity_mlp` |\n",
    "| `pred_electronegativity_hardness_mlp` | `hardness_mlp` |\n",
    "| `first_attn` / `attns[i]` / `last_attn`: `qkv_proj`, `layer_norm` | `layers[k].attn.qkv_proj`, `.attn.layer_norm` |\n",
    "| … `output_proj`, `vec_proj` | `layers[k].output_proj`, `.vec_proj` |\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "40256b39",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T04:24:35.721362Z",
     "iopub.status.busy": "2026-07-20T04:24:35.721282Z",
     "iopub.status.idle": "2026-07-20T04:24:35.726063Z",
     "shell.execute_reply": "2026-07-20T04:24:35.725601Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "weights transplanted\n"
     ]
    }
   ],
   "source": [
    "def copy(dst, src):\n",
    "    dst.load_state_dict(src.state_dict())\n",
    "\n",
    "copy(xm.atom_emb, up.atom_embtab)\n",
    "copy(xm.dis_rbf, up.dis_rbf)\n",
    "copy(xm.rbf_proj, up.rbf_proj)\n",
    "copy(xm.energy_mlp, up.energy_mlp)\n",
    "copy(xm.charge_mlp, up.charge_mlp)\n",
    "copy(xm.electronegativity_mlp, up.pred_electronegativity_mlp)\n",
    "copy(xm.hardness_mlp, up.pred_electronegativity_hardness_mlp)\n",
    "\n",
    "up_layers = [up.first_attn] + list(up.attns) + [up.last_attn]\n",
    "for xl, ul in zip(xm.layers, up_layers):\n",
    "    copy(xl.attn.qkv_proj, ul.qkv_proj)\n",
    "    copy(xl.attn.layer_norm, ul.layer_norm)\n",
    "    copy(xl.output_proj, ul.output_proj)\n",
    "    if xl.vec_proj is not None:\n",
    "        copy(xl.vec_proj, ul.vec_proj)\n",
    "print(\"weights transplanted\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f3b7350",
   "metadata": {},
   "source": [
    "## 3. A shared test geometry\n",
    "\n",
    "A 7-atom gas-phase cluster with a spread of the electrolyte elements (Li, C, N,\n",
    "O, F, H). We build the xnn graph and, from the *same* geometry, the upstream\n",
    "input dict. In xnn `edge_index = [src, dst]` with `dst` the centre (receiver);\n",
    "upstream calls these `row` (centre) and `col` (neighbour), so upstream's\n",
    "`edge_index` is the xnn one flipped and its `edge_cell_shift` is exactly the\n",
    "xnn edge vector `pos[dst] - pos[src]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a1d09829",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T04:24:35.727443Z",
     "iopub.status.busy": "2026-07-20T04:24:35.727377Z",
     "iopub.status.idle": "2026-07-20T04:24:35.738170Z",
     "shell.execute_reply": "2026-07-20T04:24:35.737594Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "edges: 42  all-pairs: 42\n"
     ]
    }
   ],
   "source": [
    "rng = np.random.default_rng(3)\n",
    "N = 7\n",
    "pos = rng.uniform(0, 4.5, (N, 3))\n",
    "Z = [8, 1, 1, 3, 9, 6, 7]\n",
    "g = structure_to_graph({\"pos\": pos, \"atomic_numbers\": Z}, CUT)\n",
    "\n",
    "posf = g.pos\n",
    "src, dst = g.edge_index[0], g.edge_index[1]        # src=neighbour, dst=centre\n",
    "edge_vec = (posf[dst] - posf[src]).detach()\n",
    "\n",
    "# all ordered intra-cluster pairs for the Coulomb sum\n",
    "rows, cols = zip(*[(i, j) for i in range(N) for j in range(N) if i != j])\n",
    "row_all = torch.tensor(rows); col_all = torch.tensor(cols)\n",
    "\n",
    "up_inputs = {\n",
    "    \"atom_types\": g.atomic_numbers.clone(),\n",
    "    \"edge_index\": torch.stack([dst, src], 0),         # [centre, neighbour]\n",
    "    \"edge_cell_shift\": edge_vec.clone(),\n",
    "    \"all_edge_index\": torch.stack([row_all, col_all], 0),\n",
    "    \"all_edge_cell_shift\": (posf[row_all] - posf[col_all]).detach(),\n",
    "    \"mol_ids\": torch.zeros(N, dtype=torch.long),\n",
    "    \"total_charge\": torch.zeros(1),\n",
    "    \"pos\": posf.clone(),\n",
    "}\n",
    "print(\"edges:\", g.num_edges, \" all-pairs:\", row_all.numel())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "845e9d5e",
   "metadata": {},
   "source": [
    "## 4. Block 1: atom embedding, radial basis, cutoff\n",
    "\n",
    "The very first tensors: the type embedding `x^0`, the exponential-normal radial\n",
    "basis `dis_rbf(r)`, and the cosine cutoff envelope."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "53d8949b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T04:24:35.739517Z",
     "iopub.status.busy": "2026-07-20T04:24:35.739443Z",
     "iopub.status.idle": "2026-07-20T04:24:35.743674Z",
     "shell.execute_reply": "2026-07-20T04:24:35.743050Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "atom embedding   : 0.0\n",
      "radial basis     : 0.0\n",
      "cutoff envelope  : 0.0\n"
     ]
    }
   ],
   "source": [
    "def maxabs(a, b): return float((a - b).abs().max())\n",
    "\n",
    "# reproduce the upstream front-end tensors\n",
    "coord_diff = up_inputs[\"edge_cell_shift\"]\n",
    "r = coord_diff.norm(dim=-1)\n",
    "unit = coord_diff / r.unsqueeze(-1)\n",
    "\n",
    "up_x0 = up.atom_embtab(up_inputs[\"atom_types\"])\n",
    "xm_x0 = xm.atom_emb(g.atomic_numbers)\n",
    "print(\"atom embedding   :\", maxabs(up_x0, xm_x0))\n",
    "print(\"radial basis     :\", maxabs(up.dis_rbf(r), xm.dis_rbf(r)))\n",
    "print(\"cutoff envelope  :\", maxabs(up.cutoff(r), xm.cutoff_fn(r)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73e86578",
   "metadata": {},
   "source": [
    "## 5. Block 2: the edge feature and the equivariant edge vector\n",
    "\n",
    "`edge_feat = act(W_rbf · rbf(r))` reshaped per head, and the equivariant edge\n",
    "vector `e_ij = edge_feat ⊗ r_hat_ij` (built from the **unit** edge direction, as\n",
    "in both codes' front-ends). We reuse this single `edge_feat`/`edge_vec` pair to\n",
    "drive the GET layers of both models below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3caf6653",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T04:24:35.744984Z",
     "iopub.status.busy": "2026-07-20T04:24:35.744919Z",
     "iopub.status.idle": "2026-07-20T04:24:35.749716Z",
     "shell.execute_reply": "2026-07-20T04:24:35.749098Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "edge feature : 0.0\n",
      "edge vector shape: (42, 3, 16, 4)\n"
     ]
    }
   ],
   "source": [
    "H, DPH = NHEADS, DIM // NHEADS\n",
    "weights_rbf = up.dis_rbf(r)\n",
    "\n",
    "ef = up.rbf_proj(weights_rbf).reshape(-1, H, DPH)\n",
    "ef_x = xm.rbf_proj(weights_rbf).reshape(-1, H, DPH)\n",
    "print(\"edge feature :\", maxabs(ef, ef_x))\n",
    "\n",
    "edge_feat = ef                                   # identical, reuse for both\n",
    "edge_vec = edge_feat.unsqueeze(-3) * unit.unsqueeze(-1).unsqueeze(-1)\n",
    "print(\"edge vector shape:\", tuple(edge_vec.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "239b0337",
   "metadata": {},
   "source": [
    "## 6. Block 3: every GET layer, one at a time\n",
    "\n",
    "Each Graph Equivariant Transformer layer is driven with identical inputs and its\n",
    "scalar (and, except in the last layer, vector) update is compared. This is the\n",
    "heart of the model: multi-head QKV attention on edges, coupled to the scalar and\n",
    "vector channels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fdca6648",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T04:24:35.750960Z",
     "iopub.status.busy": "2026-07-20T04:24:35.750894Z",
     "iopub.status.idle": "2026-07-20T04:24:35.807054Z",
     "shell.execute_reply": "2026-07-20T04:24:35.806498Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "layer 0 (first): d_scalar=0.00e+00  d_vector=0.00e+00\n",
      "layer 1 (mid) : d_scalar=0.00e+00  d_vector=0.00e+00\n",
      "layer 2 (last): d_scalar=0.00e+00\n",
      "final node features: 0.0\n"
     ]
    }
   ],
   "source": [
    "row, col = up_inputs[\"edge_index\"][0], up_inputs[\"edge_index\"][1]  # centre, neighbour\n",
    "radial = up.cutoff(r)\n",
    "x0 = up_x0                                        # identical embedding (Block 1)\n",
    "\n",
    "# --- first layer (no vector input) ---\n",
    "u_df, u_dv = up.first_attn(x0, edge_feat, edge_vec, row, col, radial, N)\n",
    "x_df, x_dv = xm.layers[0](x0, edge_feat, edge_vec, None, row, col, radial, N)\n",
    "print(f\"layer 0 (first): d_scalar={maxabs(u_df, x_df):.2e}  d_vector={maxabs(u_dv, x_dv):.2e}\")\n",
    "\n",
    "node_feat_u, node_vec_u = x0 + u_df, u_dv\n",
    "node_feat_x, node_vec_x = x0 + x_df, x_dv\n",
    "\n",
    "# --- middle layers ---\n",
    "for i, (ul, xl) in enumerate(zip(up.attns, xm.layers[1:-1]), start=1):\n",
    "    u_df, u_dv = ul(node_feat_u, edge_feat, node_vec_u, edge_vec, row, col, radial, N)\n",
    "    x_df, x_dv = xl(node_feat_x, edge_feat, edge_vec, node_vec_x, row, col, radial, N)\n",
    "    print(f\"layer {i} (mid) : d_scalar={maxabs(u_df, x_df):.2e}  d_vector={maxabs(u_dv, x_dv):.2e}\")\n",
    "    node_feat_u = node_feat_u + u_df; node_vec_u = node_vec_u + u_dv\n",
    "    node_feat_x = node_feat_x + x_df; node_vec_x = node_vec_x + x_dv\n",
    "\n",
    "# --- last layer (no vector output) ---\n",
    "u_df = up.last_attn(node_feat_u, edge_feat, node_vec_u, row, col, radial, N)\n",
    "x_df, _ = xm.layers[-1](node_feat_x, edge_feat, edge_vec, node_vec_x, row, col, radial, N)\n",
    "print(f\"layer {NLAYERS-1} (last): d_scalar={maxabs(u_df, x_df):.2e}\")\n",
    "print(\"final node features:\", maxabs(node_feat_u + u_df, node_feat_x + x_df))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41f49b35",
   "metadata": {},
   "source": [
    "## 7. Block 4: charges, electronegativity/hardness, and the energy pieces\n",
    "\n",
    "The full upstream `energy_nn` gives the NN energy, the (conserved) partial\n",
    "charges and the electronegativity energy; we compare against the xnn forward\n",
    "outputs. Note xnn spreads the electrostatic energy over `node_energy` so that\n",
    "`energy == sum(node_energy)`, so we compare the aggregated component energies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "48020822",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T04:24:35.808380Z",
     "iopub.status.busy": "2026-07-20T04:24:35.808309Z",
     "iopub.status.idle": "2026-07-20T04:24:35.887894Z",
     "shell.execute_reply": "2026-07-20T04:24:35.887351Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "partial charges     : 3.608224830031759e-16\n",
      "NN energy           : 0.0\n",
      "electrostatic energy: 1.7763568394002505e-14\n"
     ]
    }
   ],
   "source": [
    "nn_energy_u, charge_u, eneg_u = up.energy_nn({\n",
    "    **up_inputs, \"edge_cell_shift\": up_inputs[\"edge_cell_shift\"].clone()})\n",
    "\n",
    "xout = xm(g)\n",
    "# xnn aggregates electrostatics into node_energy, so compare component sums:\n",
    "# upstream (coulomb + electronegativity) == xnn energy_elec.\n",
    "print(\"partial charges     :\", maxabs(charge_u, xout[\"charges\"]))\n",
    "print(\"NN energy           :\", abs(float(nn_energy_u[0]) - float(xout[\"energy_nn\"][0])))\n",
    "up_predict = up.predict({**up_inputs, \"edge_cell_shift\": up_inputs[\"edge_cell_shift\"].clone()})\n",
    "elec_up = float(up_predict[\"energy\"][0]) - float(nn_energy_u[0])   # coul + electronegativity\n",
    "print(\"electrostatic energy:\", abs(elec_up - float(xout[\"energy_elec\"][0])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cf3afd4",
   "metadata": {},
   "source": [
    "## 8. Block 5: the full model (energy, forces, charges, dipole)\n",
    "\n",
    "Finally the whole model through the uniform `ForceStressOutput`. Energy,\n",
    "charges and dipole match directly; the xnn conservative force equals the\n",
    "upstream `forces + qeq_force` (see the note at the top)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ec0507a4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-20T04:24:35.889209Z",
     "iopub.status.busy": "2026-07-20T04:24:35.889136Z",
     "iopub.status.idle": "2026-07-20T04:24:35.911229Z",
     "shell.execute_reply": "2026-07-20T04:24:35.910696Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total energy (rel)        : 6.53e-16\n",
      "partial charges           : 3.61e-16\n",
      "dipole                    : 3.55e-15\n",
      "conservative force        : 1.07e-14\n",
      "  (vs upstream nn+coul only: 4.69e+00 — the qeq residual)\n"
     ]
    }
   ],
   "source": [
    "fs = ForceStressOutput(xm, compute_forces=True)\n",
    "xfull = fs(g)\n",
    "up_out = up.predict({**up_inputs, \"edge_cell_shift\": up_inputs[\"edge_cell_shift\"].clone()})\n",
    "\n",
    "e_rel = abs(float(xfull[\"energy\"][0]) - float(up_out[\"energy\"][0])) / abs(float(up_out[\"energy\"][0]))\n",
    "full_force = up_out[\"forces\"] + up_out[\"qeq_force\"]\n",
    "print(f\"total energy (rel)        : {e_rel:.2e}\")\n",
    "print(f\"partial charges           : {maxabs(up_out['charge'], xfull['charges']):.2e}\")\n",
    "print(f\"dipole                    : {maxabs(up_out['dipole'][0], xfull['dipole'][0]):.2e}\")\n",
    "print(f\"conservative force        : {maxabs(full_force, xfull['forces']):.2e}\")\n",
    "print(f\"  (vs upstream nn+coul only: {maxabs(up_out['forces'], xfull['forces']):.2e} — the qeq residual)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11fa9113",
   "metadata": {},
   "source": [
    "## Summary: every block matches\n",
    "\n",
    "| block | quantity | agreement |\n",
    "|---|---|---|\n",
    "| **1** | atom embedding, radial basis, cutoff | machine precision |\n",
    "| **2** | edge feature, equivariant edge vector | machine precision |\n",
    "| **3** | every GET layer's scalar & vector update | machine precision |\n",
    "| **4** | partial charges, component energies | machine precision |\n",
    "| **5** | total energy, charges, dipole | machine precision |\n",
    "| **5** | **conservative force** | `= upstream forces + qeq_force` (machine precision) |\n",
    "\n",
    "The `xnn` BAMBOO is the same function as `bytedance/bamboo`, block by block.\n",
    "The only difference is deliberate and documented: xnn returns the full\n",
    "conservative `-dE/dr` (equal to upstream `forces + qeq_force`), whereas upstream\n",
    "splits off the charge-equilibrium residual `qeq_force` as a training penalty.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}