Runs with fairchem reference

3_cascade/3_ensemble_member_divergence.ipynb

@@ -0,0 +1,205 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "intro-md",
+   "metadata": {},
+   "source": [
+    "# Ensemble member divergence\n",
+    "## Why did one bootstrapped member's training round diverge, and could bootstrap variance alone explain it?"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "imports-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "from cascade.agents.db_orm import TrajectoryDB, DBTrainingFrame"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "db-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "db_url = 'postgresql://ase:pw@localhost:5432/cascade'\n",
+    "db = TrajectoryDB(db_url)\n",
+    "run_id = '2026.08.24-18:12:33-ffdfbd'\n",
+    "training_round = 1  # the round where member 2's loss diverged"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "provenance-md",
+   "metadata": {},
+   "source": [
+    "## Per-member provenance\n",
+    "`Trainer.train_model` draws each member's bootstrap sample with an unseeded `np.random.default_rng()` and only ever holds the resulting indices in memory (`cascade/agents/agents.py`) -- they're never written to the DB or logs. So there is no way to recover which exact frames (with duplicates) went to which of the 4 members for this run. What follows instead: characterize the actual round's training pool, then simulate the real resampling procedure on that real pool to see how much per-member variance it produces on its own."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "pull-frames-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def get_round_physical_metrics(db, run_id, training_round):\n",
+    "    \"\"\"Physical-absurdity metrics per training frame, computed directly from the labeled\n",
+    "    (reference) Atoms -- independent of any model prediction.\"\"\"\n",
+    "    with db.session() as sess:\n",
+    "        rows = sess.query(DBTrainingFrame).filter_by(run_id=run_id, training_round=training_round).all()\n",
+    "        records = []\n",
+    "        for tf in rows:\n",
+    "            atoms_labeled = db._deserialize_atoms(tf.atoms_labeled_blob)\n",
+    "            forces = atoms_labeled.calc.results['forces']\n",
+    "            dists = atoms_labeled.get_all_distances(mic=True)\n",
+    "            np.fill_diagonal(dists, np.inf)\n",
+    "            records.append({\n",
+    "                'traj_id': tf.traj_id,\n",
+    "                'chunk_id': tf.chunk_id,\n",
+    "                'calibration_error': tf.calibration_error,\n",
+    "                'max_ref_force': float(np.linalg.norm(forces, axis=-1).max()),\n",
+    "                'cell_volume': atoms_labeled.get_volume(),\n",
+    "                'min_interatomic_dist': float(dists.min()),\n",
+    "            })\n",
+    "    return pd.DataFrame(records)\n",
+    "\n",
+    "pool_df = get_round_physical_metrics(db, run_id, training_round)\n",
+    "pool_df.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "describe-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pool_df[['calibration_error', 'max_ref_force', 'cell_volume', 'min_interatomic_dist']].describe(\n",
+    "    percentiles=[.05, .25, .5, .75, .9, .95, .99]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "collided-frac-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for threshold in [1.0, 0.5, 0.2]:\n",
+    "    frac = (pool_df['min_interatomic_dist'] < threshold).mean()\n",
+    "    print(f'fraction of round-{training_round} frames with min interatomic distance < {threshold} A: {frac:.1%}')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "hist-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
+    "axes[0].hist(pool_df['min_interatomic_dist'], bins=30)\n",
+    "axes[0].set_xlabel('min interatomic distance ($\\\\mathrm{\\\\AA}$)')\n",
+    "axes[0].set_ylabel('count')\n",
+    "axes[1].hist(pool_df['max_ref_force'], bins=30)\n",
+    "axes[1].set_xlabel('max reference force ($\\\\mathrm{eV/\\\\AA}$)')\n",
+    "axes[1].set_ylabel('count')\n",
+    "fig.suptitle(f'Round {training_round} training pool, n={len(pool_df)}')\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "mc-md",
+   "metadata": {},
+   "source": [
+    "## Monte Carlo: does bootstrap variance alone explain one member drawing a worse hand than the others?\n",
+    "Simulates the actual resampling procedure (`bootstrap_fraction=1.0` -> sample-with-replacement, `n_sample == len(pool)`, see `cascade/agents/config.py` / `Trainer.train_model`) on the real round's `max_ref_force` values, many times, for 4 independent members per trial."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "mc-sim-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "bootstrap_fraction = 1.0  # matches TrainerConfig default\n",
+    "n_members = 4\n",
+    "n_trials = 5000\n",
+    "\n",
+    "vals = pool_df['max_ref_force'].values\n",
+    "n_sample = int(len(vals) * bootstrap_fraction)\n",
+    "worst_idx = int(np.argmax(vals))\n",
+    "\n",
+    "rng = np.random.default_rng(0)  # seeded only so this simulation itself is reproducible -- the real run's draw was not seeded and is not recoverable\n",
+    "draws = rng.integers(0, len(vals), size=(n_trials, n_members, n_sample))\n",
+    "\n",
+    "worst_frame_copies = (draws == worst_idx).sum(axis=2)  # (n_trials, n_members): copies of the single worst frame each member happened to draw\n",
+    "exposure = (vals[draws] ** 2).sum(axis=2)  # (n_trials, n_members): total force^2 \"gradient shock\" exposure per member's draw\n",
+    "\n",
+    "copy_spread = worst_frame_copies.max(axis=1) - worst_frame_copies.min(axis=1)\n",
+    "exposure_ratio = exposure.max(axis=1) / exposure.min(axis=1)\n",
+    "\n",
+    "print(f'worst single frame in the real pool: max_ref_force={vals[worst_idx]:.1f} eV/A')\n",
+    "print(f'P(most-exposed member draws >=2 more copies of it than least-exposed member): {(copy_spread >= 2).mean():.1%}')\n",
+    "print(f'P(most-exposed member draws >=3 more copies of it than least-exposed member): {(copy_spread >= 3).mean():.1%}')\n",
+    "print()\n",
+    "print('per-trial (max member exposure / min member exposure) ratio:')\n",
+    "print(pd.Series(exposure_ratio).describe())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "mc-plot-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
+    "axes[0].hist(worst_frame_copies.flatten(), bins=np.arange(worst_frame_copies.max() + 2) - 0.5)\n",
+    "axes[0].set_xlabel('copies of the single worst frame in one member\\'s bootstrap draw')\n",
+    "axes[0].set_ylabel('count (across all simulated members)')\n",
+    "axes[1].hist(exposure_ratio, bins=30)\n",
+    "axes[1].set_xlabel('per-trial max/min member force$^2$ exposure ratio')\n",
+    "axes[1].set_ylabel('count (across trials)')\n",
+    "fig.suptitle('Bootstrap-variance simulation, 4 members x 5000 trials')\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

3_cascade/benchmark_trainer.sh

@@ -6,19 +6,21 @@ source /home/michael/miniconda3/etc/profile.d/conda.sh
 conda activate cascade
 source "${SLURM_SUBMIT_DIR:-$(dirname "${BASH_SOURCE[0]}")}/env_setup.sh"
 python scripts/benchmark_trainer.py \
-    --run-id reference-2026.08.07-17:51:00-bd260f \
+    --run-id reference-2026.08.21-19:31:10-43722e \
     --max-frames 1000 \
     --base-model small \
     --batch-sizes 2 \
     --replay-batch-size 2 \
     --num-epochs 200 \
     --patience 15 \
     --device cuda:0 \
+    --n-ensemble 4 \
+    --max-workers 1 \
     --replay-dataset ./datasets/mace-mp/sampled_1000.traj \
-    --out-dir benchmark_trainer_out_batch_small_200
+    --out-dir benchmark_trainer_out_batch_small_200_ensemble4
 
 # python scripts/benchmark_trainer.py \
-#     --run-id reference-2026.08.07-17:51:00-bd260f \
+#     --run-id reference-2026.08.21-19:31:10-43722e \
 #     --max-frames 1000 \
 #     --base-model small \
 #     --batch-sizes 2 \

3_cascade/evaluate_calculator.py

@@ -0,0 +1,141 @@
+"""Evaluate a calculator's energy/forces/stress predictions against ground-truth trajectories
+produced by another (reference) calculator, e.g. run_reference_dynamics.py.
+
+Reads frames already written to the DB, and writes the true and predicted values into an extxyz file.
+"""
+import argparse
+import logging
+import os
+
+from ase.calculators.calculator import all_changes
+from ase.io import write
+
+from cascade.agents.db_orm import TrajectoryDB
+from cascade.calculator import get_calc_factory
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--run-id',
+        type=str,
+        required=True,
+        help='Reference run (written by e.g. run_reference_dynamics.py) to evaluate the calculator against'
+    )
+    parser.add_argument(
+        '--traj-id',
+        type=int,
+        nargs='+',
+        default=None,
+        help='Trajectory ids to evaluate (defaults to all trajectories in the run)'
+    )
+    parser.add_argument(
+        '--calc-type',
+        type=str,
+        choices=['mace', 'fairchem'],
+        default='mace',
+        help='Which calculator family to evaluate'
+    )
+    parser.add_argument(
+        '--calc-model',
+        type=str,
+        default='small',
+        help='For --calc-type=mace, a MACE-MP model size (e.g. "small") or path to a MACE '
+             'checkpoint. For --calc-type=fairchem, the path to a FairChem .pt checkpoint.'
+    )
+    parser.add_argument(
+        '--calc-task',
+        type=str,
+        default=None,
+        help='FairChem task name selecting the model head (e.g. "omol", "omat", "oc20", '
+             '"odac", "omc"), ignored for --calc-type=mace. Only needed for --calc-type=fairchem '
+             'if the checkpoint supports more than one task'
+    )
+    parser.add_argument(
+        '--device',
+        type=str,
+        default='cpu',
+        help='Device for the evaluated calculator'
+    )
+    parser.add_argument(
+        '--stride',
+        type=int,
+        default=1,
+        help='Evaluate every Nth frame of each trajectory'
+    )
+    parser.add_argument(
+        '--output',
+        type=str,
+        required=True,
+        help='Path to write paired true/predicted results (extxyz, appended incrementally)'
+    )
+    parser.add_argument(
+        '--db-url',
+        type=str,
+        default=os.environ.get('CASCADE_DB_URL'),
+        help='Database URL, e.g. postgresql://ase:pw@<host>:5432/cascade '
+             '(defaults to the CASCADE_DB_URL env var)'
+    )
+    parser.add_argument(
+        '--log-level',
+        type=str,
+        default='INFO',
+        help='Logging level'
+    )
+    return parser.parse_args()
+
+
+def evaluate_frame(atoms_true, calc, run_id: str, traj_id: int, frame_index: int):
+    """Evaluate calc on a copy of atoms_true, returning an Atoms with both the true (already
+    attached to atoms_true.calc) and predicted energy/forces/stress stashed side by side."""
+    atoms_pred = atoms_true.copy()
+    atoms_pred.calc = calc
+    calc.calculate(atoms_pred, properties=calc.implemented_properties, system_changes=all_changes)
+
+    results_true = atoms_true.calc.results
+    results_pred = atoms_pred.calc.results
+
+    out = atoms_true.copy()
+    out.info['run_id'] = run_id
+    out.info['traj_id'] = traj_id
+    out.info['frame_index'] = frame_index
+    out.info['energy_true'] = results_true['energy']
+    out.info['energy_pred'] = results_pred['energy']
+    out.new_array('forces_true', results_true['forces'])
+    out.new_array('forces_pred', results_pred['forces'])
+    if 'stress' in results_true and 'stress' in results_pred:
+        out.info['stress_true'] = results_true['stress']
+        out.info['stress_pred'] = results_pred['stress']
+    return out
+
+
+def main():
+    args = parse_args()
+    logging.basicConfig(level=args.log_level)
+    logger = logging.getLogger('evaluate_calculator')
+
+    traj_db = TrajectoryDB(args.db_url)
+    traj_ids = args.traj_id
+    if traj_ids is None:
+        traj_ids = [t['traj_id'] for t in traj_db.list_trajectories_in_run(args.run_id)]
+        logger.info(f'No --traj-id given, evaluating all {len(traj_ids)} trajectories in run {args.run_id}')
+
+    calc_factory = get_calc_factory(args.calc_type, args.calc_model, args.device, args.calc_task)
+    calc = calc_factory()
+
+    n_written = 0
+    for traj_id in traj_ids:
+        atoms_list = traj_db.get_trajectory_atoms(args.run_id, traj_id)
+        logger.info(f'Evaluating traj {traj_id}: {len(atoms_list)} frames, stride={args.stride}')
+        for frame_index in range(0, len(atoms_list), args.stride):
+            out = evaluate_frame(atoms_list[frame_index], calc, args.run_id, traj_id, frame_index)
+            write(args.output, out, format='extxyz', append=True)
+            n_written += 1
+            if n_written % 100 == 0:
+                logger.info(f'Wrote {n_written} evaluated frames so far (traj {traj_id}, frame {frame_index})')
+
+    logger.info(f'Done. Wrote {n_written} evaluated frames to {args.output}')
+
+
+if __name__ == '__main__':
+    main()

3_cascade/evaluate_calculator.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+#SBATCH --partition=main
+#SBATCH --gres=gpu:1
+#SBATCH --exclusive
+
+source /home/michael/miniconda3/etc/profile.d/conda.sh
+conda activate cascade
+
+source "${SLURM_SUBMIT_DIR:-$(dirname "${BASH_SOURCE[0]}")}/env_setup.sh"
+
+python evaluate_calculator.py \
+    --run-id reference-2026.08.21-19:31:10-43722e \
+    --calc-type mace \
+    --calc-model small \
+    --device cuda:0 \
+    --output "eval_mace_small_vs_uma_omat.${SLURM_JOB_ID}.xyz" \
+    --log-level DEBUG > "evaluate_calculator.${SLURM_JOB_ID}.log" 2>&1