From d5f92ad68dfd0551cb5b522f4d20c99ec8be6626 Mon Sep 17 00:00:00 2001 From: Swaroop Sridhar Date: Wed, 17 Jun 2026 18:23:54 -0500 Subject: [PATCH 01/11] Add per-frame xpts CSV emission to run_hessian_and_build_cache.py (row, col, class columns for X/Omax/Omin points), with a backfill pass that writes CSVs for already-cached frames without re-running the classifier. --- run_hessian_and_build_cache.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/run_hessian_and_build_cache.py b/run_hessian_and_build_cache.py index af97339..865f9f4 100644 --- a/run_hessian_and_build_cache.py +++ b/run_hessian_and_build_cache.py @@ -78,6 +78,21 @@ def discover_all_frames(extract_dir): return sorted(f for f in frame_nums if f > 0) +def _write_xpts_csv(cache_dir, fnum, xpts, optsMax, optsMin): + """Write {fnum}_xpts.csv with row, col, class columns (class in X/Omax/Omin).""" + csv_path = cache_dir / f"{fnum}_xpts.csv" + with open(csv_path, "w") as f: + f.write("row,col,class\n") + for (r, c) in xpts: + f.write(f"{int(r)},{int(c)},X\n") + if optsMax is not None and len(optsMax) > 0: + for (r, c) in optsMax: + f.write(f"{int(r)},{int(c)},Omax\n") + if optsMin is not None and len(optsMin) > 0: + for (r, c) in optsMin: + f.write(f"{int(r)},{int(c)},Omin\n") + + def _process_frame(task): """Worker function: run the Hessian classifier on one frame and write its cache. Returns (frame_num, elapsed).""" param_path, fnum, cache_dir = task @@ -90,6 +105,7 @@ def _process_frame(task): "fileName": fileName, "Bx": bx, "By": by, "Jz": jz} writePgkylDataToCache(cache_dir, fnum, fields) + _write_xpts_csv(cache_dir, fnum, xpts, optsMax, optsMin) elapsed = time.time() - t0 return fnum, elapsed @@ -127,6 +143,21 @@ def main(): cached = len(frames) - len(uncached) print(f"Already cached: {cached}, need to compute: {len(uncached)}") + # Backfill xpts CSVs for any cached frames that don't have one yet. + # No classifier run -- just reads the existing .npy and writes the CSV. + csv_backfilled = 0 + for f in frames: + if cachedPgkylDataExists(cache_dir, f, "psi") and not (cache_dir / f"{f}_xpts.csv").exists(): + xpts = np.load(cache_dir / f"{f}_xpts.npy") + optsMax_path = cache_dir / f"{f}_optsMax.npy" + optsMin_path = cache_dir / f"{f}_optsMin.npy" + optsMax = np.load(optsMax_path) if optsMax_path.exists() else None + optsMin = np.load(optsMin_path) if optsMin_path.exists() else None + _write_xpts_csv(cache_dir, f, xpts, optsMax, optsMin) + csv_backfilled += 1 + if csv_backfilled > 0: + print(f"Backfilled xpts CSV for {csv_backfilled} already-cached frame(s)") + if not uncached: print("All frames already cached!") return From 9b9290ab2a4b122697a61aee45281bd0197fd11b Mon Sep 17 00:00:00 2001 From: Swaroop Sridhar Date: Tue, 23 Jun 2026 20:48:25 -0500 Subject: [PATCH 02/11] Add point_metrics.py for scoring point-list predictions against ground truth at the point level. match_points takes two coordinate arrays and a matching radius, builds a grid index on the gt side, and does greedy 1-to-1 nearest-neighbor matching to figure out tp/fp/fn (with optional confidence ordering so high-confidence preds claim matches first). evaluate_point_predictions runs that across many frames and gives back a global summary plus per-frame metrics. load_xpts_csv and load_xpts_csvs_for_frames are tiny helpers for reading the CSV format the cache builder emits. --- point_metrics.py | 168 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 168 insertions(+) create mode 100644 point_metrics.py diff --git a/point_metrics.py b/point_metrics.py new file mode 100644 index 0000000..fcfd936 --- /dev/null +++ b/point_metrics.py @@ -0,0 +1,168 @@ +"""Point-level precision/recall/F1 for X-point detection on coordinate lists.""" + +from __future__ import annotations + +import csv +from collections import defaultdict +from pathlib import Path +from typing import Iterable, Mapping, Sequence + +import numpy as np + + +# CSV I/O +def load_xpts_csv(csv_path: str | Path, + classes: Sequence[str] = ("X",)) -> np.ndarray: + """Load (row, col) coordinates of the given classes from a per-frame xpts CSV.""" + csv_path = Path(csv_path) + if not csv_path.exists(): + raise FileNotFoundError(f"xpts CSV not found: {csv_path}") + + kept = [] + with open(csv_path) as f: + reader = csv.DictReader(f) + for row in reader: + if row["class"] in classes: + kept.append((int(row["row"]), int(row["col"]))) + + if not kept: + return np.zeros((0, 2), dtype=int) + return np.asarray(kept, dtype=int) + + +# SPATIAL INDEX +class _GridIndex: + """Bucket points into cells so nearest-neighbor lookups touch only the 3x3 cell neighborhood.""" + + def __init__(self, points: np.ndarray, cell_size: float): + self.cell_size = float(cell_size) + self.buckets: dict[tuple[int, int], list[int]] = defaultdict(list) + for i, (r, c) in enumerate(points): + self.buckets[(int(r // cell_size), int(c // cell_size))].append(i) + self.points = points + + def candidates(self, r: float, c: float) -> list[int]: + """Return indices of points in the query cell and its 8 neighbors.""" + cr, cc = int(r // self.cell_size), int(c // self.cell_size) + out: list[int] = [] + for dr in (-1, 0, 1): + for dc in (-1, 0, 1): + bucket = self.buckets.get((cr + dr, cc + dc)) + if bucket: + out.extend(bucket) + return out + + +# PER-FRAME MATCHING +def match_points(pred_pts: np.ndarray, + gt_pts: np.ndarray, + radius: float = 5.0, + pred_confidence: np.ndarray | None = None, + ) -> dict: + """Greedy 1-to-1 nearest-neighbor match within `radius`; returns tp/fp/fn, P/R/F1, and per-pred matched gt index.""" + pred_pts = np.asarray(pred_pts).reshape(-1, 2) + gt_pts = np.asarray(gt_pts).reshape(-1, 2) + P, G = len(pred_pts), len(gt_pts) + + if P == 0 and G == 0: + return {"tp": 0, "fp": 0, "fn": 0, + "precision": 1.0, "recall": 1.0, "f1": 1.0, + "pred_match": np.zeros(0, dtype=int)} + if P == 0: + return {"tp": 0, "fp": 0, "fn": G, + "precision": 0.0, "recall": 0.0, "f1": 0.0, + "pred_match": np.zeros(0, dtype=int)} + if G == 0: + return {"tp": 0, "fp": P, "fn": 0, + "precision": 0.0, "recall": 0.0, "f1": 0.0, + "pred_match": -np.ones(P, dtype=int)} + + order = (np.argsort(-np.asarray(pred_confidence)) + if pred_confidence is not None + else np.arange(P)) + + index = _GridIndex(gt_pts, cell_size=radius) + gt_taken = np.zeros(G, dtype=bool) + pred_match = -np.ones(P, dtype=int) + r2 = radius * radius + + for i in order: + pr, pc = pred_pts[i] + best_j, best_d2 = -1, r2 + 1e-9 + for j in index.candidates(pr, pc): + if gt_taken[j]: + continue + dr = float(pr - gt_pts[j, 0]) + dc = float(pc - gt_pts[j, 1]) + d2 = dr * dr + dc * dc + if d2 <= r2 and d2 < best_d2: + best_d2, best_j = d2, j + if best_j >= 0: + pred_match[i] = best_j + gt_taken[best_j] = True + + tp = int((pred_match >= 0).sum()) + fp = P - tp + fn = G - tp + precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0 + recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0 + f1 = (2 * precision * recall / (precision + recall) + if (precision + recall) > 0 else 0.0) + + return {"tp": tp, "fp": fp, "fn": fn, + "precision": float(precision), "recall": float(recall), "f1": float(f1), + "pred_match": pred_match} + + +# MULTI-FRAME AGGREGATION +def evaluate_point_predictions(pred_by_frame: Mapping[int, np.ndarray], + gt_by_frame: Mapping[int, np.ndarray], + radius: float = 5.0, + pred_confidence_by_frame: Mapping[int, np.ndarray] | None = None, + ) -> dict: + """Aggregate `match_points` across frames into global (summed tp/fp/fn) and per-frame metrics.""" + frame_ids = sorted(set(pred_by_frame) | set(gt_by_frame)) + per_frame = [] + sum_tp = sum_fp = sum_fn = 0 + for fid in frame_ids: + pred = pred_by_frame.get(fid, np.zeros((0, 2), dtype=int)) + gt = gt_by_frame.get(fid, np.zeros((0, 2), dtype=int)) + conf = (pred_confidence_by_frame.get(fid) + if pred_confidence_by_frame is not None else None) + result = match_points(pred, gt, radius=radius, pred_confidence=conf) + sum_tp += result["tp"] + sum_fp += result["fp"] + sum_fn += result["fn"] + per_frame.append({"frame_id": int(fid), + "tp": result["tp"], "fp": result["fp"], "fn": result["fn"], + "precision": result["precision"], + "recall": result["recall"], + "f1": result["f1"]}) + + precision = sum_tp / (sum_tp + sum_fp) if (sum_tp + sum_fp) > 0 else 0.0 + recall = sum_tp / (sum_tp + sum_fn) if (sum_tp + sum_fn) > 0 else 0.0 + f1 = (2 * precision * recall / (precision + recall) + if (precision + recall) > 0 else 0.0) + + return { + "global": {"tp": sum_tp, "fp": sum_fp, "fn": sum_fn, + "precision": float(precision), "recall": float(recall), + "f1": float(f1)}, + "per_frame": per_frame, + } + + +def load_xpts_csvs_for_frames(cache_dir: str | Path, + frame_ids: Iterable[int], + classes: Sequence[str] = ("X",), + ) -> dict[int, np.ndarray]: + """Load {fid}_xpts.csv for each requested frame into a {frame_id: coords} dict.""" + cache_dir = Path(cache_dir) + out: dict[int, np.ndarray] = {} + for fid in frame_ids: + csv_path = cache_dir / f"{fid}_xpts.csv" + if csv_path.exists(): + out[fid] = load_xpts_csv(csv_path, classes=classes) + else: + out[fid] = np.zeros((0, 2), dtype=int) + return out From 251ae044d78969accd3573a886d79b5a00d45f4c Mon Sep 17 00:00:00 2001 From: Swaroop Sridhar Date: Wed, 24 Jun 2026 21:21:45 -0500 Subject: [PATCH 03/11] Add backfill_xpts_csv.py and predict_points.py, the two ways we get point-list CSVs onto disk. backfill_xpts_csv.py walks a cache directory and writes per-frame {N}_xpts.csv files from existing {N}_xpts.npy plus optsMax/optsMin files, with an --output-dir flag for cases where the source cache isn't writable (like the PKPM cache, which is owned by another user). predict_points.py loads a trained checkpoint, runs inference over the cached frames of a dataset, post-processes the predicted heatmap with connected-component NMS to extract one peak per blob above threshold, then writes the predicted points as per-frame CSVs with row/col/confidence columns. Between them they cover both sides of the comparison: ground-truth points from the deterministic classifier, and predicted points from a neural net. --- backfill_xpts_csv.py | 71 +++++++++++++++++++++ predict_points.py | 143 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 214 insertions(+) create mode 100644 backfill_xpts_csv.py create mode 100644 predict_points.py diff --git a/backfill_xpts_csv.py b/backfill_xpts_csv.py new file mode 100644 index 0000000..328de60 --- /dev/null +++ b/backfill_xpts_csv.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 +"""Backfill {N}_xpts.csv files for any cache directory containing {N}_xpts.npy.""" + +import argparse +import re +from pathlib import Path + +import numpy as np + + +def write_xpts_csv(cache_dir, fnum, xpts, optsMax, optsMin): + """Write {fnum}_xpts.csv with row, col, class columns (class in X/Omax/Omin).""" + csv_path = cache_dir / f"{fnum}_xpts.csv" + with open(csv_path, "w") as f: + f.write("row,col,class\n") + for (r, c) in xpts: + f.write(f"{int(r)},{int(c)},X\n") + if optsMax is not None and len(optsMax) > 0: + for (r, c) in optsMax: + f.write(f"{int(r)},{int(c)},Omax\n") + if optsMin is not None and len(optsMin) > 0: + for (r, c) in optsMin: + f.write(f"{int(r)},{int(c)},Omin\n") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("cache_dir", type=Path, + help="Source dir to read {N}_xpts.npy files from") + parser.add_argument("--output-dir", type=Path, default=None, + help="Where to write CSVs (default: same as cache_dir)") + parser.add_argument("--overwrite", action="store_true", + help="Rewrite CSVs that already exist (default: skip)") + args = parser.parse_args() + + if not args.cache_dir.is_dir(): + raise FileNotFoundError(f"Cache directory not found: {args.cache_dir}") + output_dir = args.output_dir if args.output_dir is not None else args.cache_dir + output_dir.mkdir(parents=True, exist_ok=True) + + xpt_files = sorted(args.cache_dir.glob("*_xpts.npy")) + pattern = re.compile(r"^(\d+)_xpts\.npy$") + frames = [] + for f in xpt_files: + m = pattern.match(f.name) + if m: + frames.append(int(m.group(1))) + frames.sort() + print(f"Source dir: {args.cache_dir}") + print(f"Output dir: {output_dir}") + print(f"Found {len(frames)} cached frame(s).") + + n_written = n_skipped = 0 + for fnum in frames: + csv_path = output_dir / f"{fnum}_xpts.csv" + if csv_path.exists() and not args.overwrite: + n_skipped += 1 + continue + xpts = np.load(args.cache_dir / f"{fnum}_xpts.npy") + optsMax_path = args.cache_dir / f"{fnum}_optsMax.npy" + optsMin_path = args.cache_dir / f"{fnum}_optsMin.npy" + optsMax = np.load(optsMax_path) if optsMax_path.exists() else None + optsMin = np.load(optsMin_path) if optsMin_path.exists() else None + write_xpts_csv(output_dir, fnum, xpts, optsMax, optsMin) + n_written += 1 + + print(f"Wrote {n_written} CSV(s), skipped {n_skipped} existing.") + + +if __name__ == "__main__": + main() diff --git a/predict_points.py b/predict_points.py new file mode 100644 index 0000000..dcf601d --- /dev/null +++ b/predict_points.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python3 +"""Run a trained checkpoint on a dataset, extract X-point peaks via NMS, write per-frame predicted CSVs.""" + +import argparse +import os +import sys +import time +from pathlib import Path + +import numpy as np +import torch +from scipy.ndimage import label as cc_label +from torch.amp import autocast + +RC_ROOT = Path(__file__).resolve().parent +sys.path.insert(0, str(RC_ROOT)) + +from XPointMLTest import XPointDataset, UNet, cachedPgkylDataExists + + +DATASET_CONFIG = { + "PKPM": { + "param_path": "/work/nvme/bfim/cwsmith/mlReconnection2025/1024Res_v0/pkpm_2d_turb_p2-params.txt", + "cache_dir": "/work/nvme/bfim/cwsmith/mlReconnection2025/1024Res_v0/cache04082025", + }, + "5M": { + "param_path": "/work/nvme/bfim/ssridhar6/mlReconnection2025/5M/rt_5M_2d_turb_local-params.txt", + "cache_dir": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/5M", + }, + "10M": { + "param_path": "/work/nvme/bfim/ssridhar6/mlReconnection2025/10M/10M/rt_10M_2d_turb_local-params.txt", + "cache_dir": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/10M", + }, +} + + +def extract_peaks(heatmap, threshold=0.3, min_distance=5): + """Extract one (row, col, confidence) peak per connected component above threshold (max-valued pixel in each).""" + above = heatmap > threshold + if not above.any(): + return (np.zeros(0, dtype=int), + np.zeros(0, dtype=int), + np.zeros(0, dtype=float)) + labels, n = cc_label(above) + rows = np.empty(n, dtype=int) + cols = np.empty(n, dtype=int) + confs = np.empty(n, dtype=float) + for k in range(1, n + 1): + mask = labels == k + masked = np.where(mask, heatmap, -np.inf) + r, c = np.unravel_index(int(np.argmax(masked)), heatmap.shape) + rows[k - 1] = r + cols[k - 1] = c + confs[k - 1] = float(heatmap[r, c]) + return rows, cols, confs + + +def write_pred_csv(out_path, rows, cols, confidences): + """Write predicted-points CSV with row, col, confidence columns.""" + with open(out_path, "w") as f: + f.write("row,col,confidence\n") + for r, c, conf in zip(rows, cols, confidences): + f.write(f"{int(r)},{int(c)},{float(conf):.6f}\n") + + +def discover_cached_frames(cache_dir): + """Return sorted list of frame numbers that have a cached _psi.npy in cache_dir.""" + import re + pattern = re.compile(r"^(\d+)_psi\.npy$") + out = [] + for p in Path(cache_dir).glob("*_psi.npy"): + m = pattern.match(p.name) + if m: + out.append(int(m.group(1))) + return sorted(out) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", required=True, type=Path) + parser.add_argument("--datasets", nargs="+", default=["PKPM", "5M", "10M"], + choices=list(DATASET_CONFIG)) + parser.add_argument("--output-root", required=True, type=Path, + help="Base directory; predictions go to //{N}_xpts.csv") + parser.add_argument("--threshold", type=float, default=0.3, + help="Confidence threshold for peak retention (default: 0.3)") + parser.add_argument("--min-distance", type=int, default=5, + help="Min pixel separation between peaks (NMS kernel = 2*N+1)") + parser.add_argument("--base-channels", type=int, default=64) + parser.add_argument("--dropout-rate", type=float, default=0.055) + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"Device: {device}") + if torch.cuda.is_available(): + print(f"GPU: {torch.cuda.get_device_name(0)}") + + model = UNet(input_channels=4, base_channels=args.base_channels, + dropout_rate=args.dropout_rate).to(device) + state_dict = torch.load(str(args.checkpoint), map_location=device, weights_only=False) + model.load_state_dict(state_dict) + model.eval() + print(f"Loaded checkpoint: {args.checkpoint}") + + use_amp = torch.cuda.is_available() + amp_dtype = torch.bfloat16 if (use_amp and torch.cuda.is_bf16_supported()) else torch.float16 + + args.output_root.mkdir(parents=True, exist_ok=True) + print(f"NMS: threshold={args.threshold}, min_distance={args.min_distance} px") + + with torch.no_grad(): + for ds_name in args.datasets: + cfg = DATASET_CONFIG[ds_name] + cache_dir = Path(cfg["cache_dir"]) + out_dir = args.output_root / ds_name + out_dir.mkdir(parents=True, exist_ok=True) + frames = discover_cached_frames(cache_dir) + print(f"\n=== {ds_name}: {len(frames)} cached frames -> {out_dir} ===") + + t0 = time.time() + dataset = XPointDataset(cfg["param_path"], frames, + xptCacheDir=cache_dir, rotateAndReflect=False) + print(f" Loaded dataset in {time.time()-t0:.1f}s") + + n_predicted = 0 + t1 = time.time() + for item in dataset: + fnum = item["fnum"] + all_torch = item["all"].unsqueeze(0).to(device) + with autocast(device_type="cuda", dtype=amp_dtype, enabled=use_amp): + logits = model(all_torch) + probs = torch.sigmoid(logits) + heatmap = probs[0, 0].float().cpu().numpy() + rows, cols, confs = extract_peaks(heatmap, threshold=args.threshold, + min_distance=args.min_distance) + write_pred_csv(out_dir / f"{fnum}_xpts.csv", rows, cols, confs) + n_predicted += len(rows) + print(f" Wrote {len(frames)} CSV(s) in {time.time()-t1:.1f}s " + f"({n_predicted} predicted points total)") + + +if __name__ == "__main__": + main() From 8f4d24b2a8660546e14b221d45dca5b42cf795e7 Mon Sep 17 00:00:00 2001 From: Swaroop Sridhar Date: Mon, 6 Jul 2026 22:57:49 -0500 Subject: [PATCH 04/11] Add score_point_predictions.py, the aggregator that ties the point-list scoring together. Takes a predictions root directory (subdirs per dataset), loads the predicted xpts CSVs with their confidences plus the matching ground-truth CSVs from the per-dataset cache mirror, and runs evaluate_point_predictions from point_metrics on each dataset. Prints a per-dataset F1/precision/recall/TP/FP/FN table for the given checkpoint, and optionally writes the full per-frame breakdown to JSON via --json-out. Takes --radius and --datasets flags so you can tweak the matching radius or pick which subset of datasets to score without editing the script. Together with predict_points.py this closes the loop from checkpoint through NMS extraction through point-level comparison. --- score_point_predictions.py | 99 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 99 insertions(+) create mode 100644 score_point_predictions.py diff --git a/score_point_predictions.py b/score_point_predictions.py new file mode 100644 index 0000000..c011826 --- /dev/null +++ b/score_point_predictions.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +"""Compare predicted-points CSVs against ground-truth-points CSVs at the point level.""" + +import argparse +import json +from pathlib import Path + +from point_metrics import ( + evaluate_point_predictions, + load_xpts_csvs_for_frames, +) + + +GT_CSV_DIRS = { + "PKPM": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/PKPM", + "5M": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/5M", + "10M": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/10M", +} + + +def load_pred_csvs_with_confidence(pred_dir): + """Load predicted xpts CSVs (row, col, confidence) into {frame_id: coords, frame_id: conf} dicts.""" + import re, csv + import numpy as np + pattern = re.compile(r"^(\d+)_xpts\.csv$") + pred_by_frame = {} + conf_by_frame = {} + for path in sorted(Path(pred_dir).glob("*_xpts.csv")): + m = pattern.match(path.name) + if not m: + continue + fid = int(m.group(1)) + rows, confs = [], [] + with open(path) as f: + for row in csv.DictReader(f): + rows.append((int(row["row"]), int(row["col"]))) + confs.append(float(row.get("confidence", 1.0))) + if rows: + pred_by_frame[fid] = np.asarray(rows, dtype=int) + conf_by_frame[fid] = np.asarray(confs, dtype=float) + else: + pred_by_frame[fid] = np.zeros((0, 2), dtype=int) + conf_by_frame[fid] = np.zeros((0,), dtype=float) + return pred_by_frame, conf_by_frame + + +def score_one(pred_dir, gt_dir, radius): + pred_by_frame, conf_by_frame = load_pred_csvs_with_confidence(pred_dir) + frame_ids = sorted(pred_by_frame.keys()) + gt_by_frame = load_xpts_csvs_for_frames(gt_dir, frame_ids, classes=("X",)) + return evaluate_point_predictions(pred_by_frame, gt_by_frame, + radius=radius, + pred_confidence_by_frame=conf_by_frame) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--predictions-root", required=True, type=Path, + help="Directory containing subdirs per dataset (PKPM/, 5M/, 10M/)") + parser.add_argument("--datasets", nargs="+", default=["PKPM", "5M", "10M"], + choices=list(GT_CSV_DIRS)) + parser.add_argument("--radius", type=float, default=5.0, + help="Matching radius in pixels (default: 5.0)") + parser.add_argument("--label", default="model", + help="Label to print in the summary row (e.g. 'PKPM-trained')") + parser.add_argument("--json-out", type=Path, default=None, + help="Optional: write full results as JSON") + args = parser.parse_args() + + print(f"\n{'='*78}") + print(f"POINT-LEVEL F1 (label={args.label}, radius={args.radius:.1f} px)") + print(f"{'='*78}") + print(f"{'Dataset':<8} {'F1':>7} {'Prec':>7} {'Rec':>7} {'TP':>6} {'FP':>6} {'FN':>6} {'Frames':>7}") + print("-" * 78) + + all_results = {} + for ds in args.datasets: + pred_dir = args.predictions_root / ds + if not pred_dir.is_dir(): + print(f"{ds:<8} (no predictions directory at {pred_dir})") + continue + result = score_one(pred_dir, GT_CSV_DIRS[ds], radius=args.radius) + g = result["global"] + n_frames = len(result["per_frame"]) + print(f"{ds:<8} {g['f1']:>7.4f} {g['precision']:>7.4f} {g['recall']:>7.4f} " + f"{g['tp']:>6d} {g['fp']:>6d} {g['fn']:>6d} {n_frames:>7d}") + all_results[ds] = result + + print("=" * 78) + + if args.json_out: + with open(args.json_out, "w") as f: + json.dump({"label": args.label, "radius": args.radius, + "results": all_results}, f, indent=2) + print(f"Wrote JSON: {args.json_out}") + + +if __name__ == "__main__": + main() From 6252c80843e145d7e86e0d08d1fb45fdc8bace31 Mon Sep 17 00:00:00 2001 From: Swaroop Sridhar Date: Fri, 24 Jul 2026 11:04:27 -0500 Subject: [PATCH 05/11] Bring Gaussian heatmap training code (gaussian_xpoints_mask, FocalHeatmapLoss, --gaussianSigma/--targetType/--heatmap* args) onto the point-list branch so a single working tree can both train at varying sigma and score with the point-list metric, for the sigma sweep experiment. --- XPointMLTest.py | 128 ++++++++++++++++++++++++++++++++++++---- test_xpoint_transfer.py | 9 ++- 2 files changed, 123 insertions(+), 14 deletions(-) diff --git a/XPointMLTest.py b/XPointMLTest.py index fb18c5c..f59a27c 100644 --- a/XPointMLTest.py +++ b/XPointMLTest.py @@ -95,9 +95,45 @@ def expand_xpoints_mask(binary_mask, kernel_size=9): # Set the square area to 1 expanded_mask[x_min:x_max, y_min:y_max] = 1 - + return expanded_mask +def gaussian_xpoints_mask(xpts, shape, sigma=3.0): + """ + Render a soft Gaussian heatmap target with one 2D Gaussian centered on each + X-point. Overlapping Gaussians are combined by elementwise max so the peak + of each X-point stays at 1.0 even when X-points are close together. + + Parameters: + xpts : (N, 2) array of (row, col) X-point coordinates + shape : (H, W) output shape + sigma : float, Gaussian std in pixels (default 3.0; the existing 9x9 binary + mask has effective radius ~4 pixels, so sigma=3 gives a slightly + tighter, smoother target) + + Returns: + (H, W) np.ndarray with values in [0, 1]. + """ + H, W = shape + heatmap = np.zeros((H, W), dtype=np.float32) + if len(xpts) == 0: + return heatmap + + # Render each Gaussian into a local window of half-width 3*sigma (covers + # > 99% of the Gaussian's mass) instead of computing over the full frame. + half = int(np.ceil(3.0 * sigma)) + two_sigma_sq = 2.0 * sigma * sigma + for (r, c) in xpts: + r0, r1 = max(0, r - half), min(H, r + half + 1) + c0, c1 = max(0, c - half), min(W, c + half + 1) + if r0 >= r1 or c0 >= c1: + continue + ys, xs = np.ogrid[r0:r1, c0:c1] + g = np.exp(-((ys - r) ** 2 + (xs - c) ** 2) / two_sigma_sq) + # Max-merge so a peak shared by overlapping Gaussians stays at 1.0 + heatmap[r0:r1, c0:c1] = np.maximum(heatmap[r0:r1, c0:c1], g) + return heatmap + def rotate(frameData,deg): if deg not in [90, 180, 270]: print(f"invalid rotation specified... exiting") @@ -223,17 +259,24 @@ class XPointDataset(Dataset): - Returns (psiTensor, maskTensor) as a PyTorch (float) pair. """ def __init__(self, paramFile, fnumList, xptCacheDir=None, - rotateAndReflect=False, verbosity=0): + rotateAndReflect=False, verbosity=0, + target_type='binary', gaussian_sigma=3.0): """ paramFile: Path to parameter file (string). - fnumList: List of frames to iterate. + fnumList: List of frames to iterate. rotateAndReflect: If True, creates static augmented copies (deprecated, use on-the-fly instead) + target_type: 'binary' (current 9x9 dilation, default) or 'gaussian' + (soft heatmap with sigma=gaussian_sigma). + gaussian_sigma: Std (in pixels) for gaussian target rendering. Ignored + unless target_type='gaussian'. """ super().__init__() self.paramFile = paramFile self.fnumList = list(fnumList) # ensure indexable self.xptCacheDir = xptCacheDir self.verbosity = verbosity + self.target_type = target_type + self.gaussian_sigma = float(gaussian_sigma) # We'll store a base 'params' once here, and then customize in __getitem__: self.params = {} @@ -303,11 +346,18 @@ def load(self, fnum): if self.verbosity > 0: print("time (s) to find X and O points: " + str(timer()-t2)) - # Create array of 0s with 1s only at X points - binaryMap = np.zeros(np.shape(fields["psi"])) - binaryMap[fields["xpts"][:, 0], fields["xpts"][:, 1]] = 1 - - binaryMap = expand_xpoints_mask(binaryMap, kernel_size=9) + # Build the per-pixel target map. Binary mode is the original 9x9 + # dilation; gaussian mode places a Gaussian on each X-point and + # produces continuous values in [0, 1] for heatmap regression. + if self.target_type == 'gaussian': + binaryMap = gaussian_xpoints_mask( + fields["xpts"], np.shape(fields["psi"]), + sigma=self.gaussian_sigma, + ) + else: + binaryMap = np.zeros(np.shape(fields["psi"])) + binaryMap[fields["xpts"][:, 0], fields["xpts"][:, 1]] = 1 + binaryMap = expand_xpoints_mask(binaryMap, kernel_size=9) # Normalize input features for better training stability psi_norm = (fields["psi"] - fields["psi"].mean()) / (fields["psi"].std() + 1e-8) @@ -676,6 +726,38 @@ def forward(self, inputs, targets): return (1 - self.dice_weight) * focal_loss + self.dice_weight * dice_loss + +class FocalHeatmapLoss(nn.Module): + """CenterNet-style penalty-reduced focal loss for Gaussian heatmap targets. + + Targets are continuous values in [0, 1] with peak 1.0 at each X-point. + For the peak pixels (target == 1) it applies the standard focal positive + term. For non-peak pixels it applies a penalty-reduced negative term that + discounts the loss near the Gaussian's wings (where the target is close to + 1 but not exactly), so the Gaussian falloff isn't penalized as a hard + background mistake. + + Reference: Law and Deng, CornerNet (2018); Zhou et al., CenterNet (2019). + """ + def __init__(self, alpha=2.0, beta=4.0, eps=1e-6): + super().__init__() + self.alpha = alpha + self.beta = beta + self.eps = eps + + def forward(self, inputs, targets): + # inputs are raw logits; squash to probabilities and clamp for log stability + pred = torch.sigmoid(inputs).clamp(self.eps, 1.0 - self.eps) + pos_mask = targets.eq(1).float() + neg_mask = 1.0 - pos_mask + + pos_loss = -((1 - pred) ** self.alpha) * torch.log(pred) * pos_mask + neg_loss = -((1 - targets) ** self.beta) * (pred ** self.alpha) * torch.log(1 - pred) * neg_mask + + n_pos = pos_mask.sum().clamp(min=1.0) + return (pos_loss.sum() + neg_loss.sum()) / n_pos + + # TRAIN & VALIDATION UTILS def train_one_epoch(model, loader, criterion, optimizer, device, scaler, use_amp, amp_dtype, benchmark=None): model.train() @@ -927,14 +1009,28 @@ def parseCommandLineArgs(): parser.add_argument('--posRatio', type=float, default=0.5, help='target ratio of patches containing X-points (default: 0.5)') parser.add_argument('--lossFunction', type=str, default='dice', - choices=['dice', 'focal_dice'], - help='loss function: dice (default) or focal_dice (combined focal + dice)') + choices=['dice', 'focal_dice', 'heatmap_focal'], + help='loss function: dice (default), focal_dice (focal + dice), ' + 'or heatmap_focal (CenterNet-style penalty-reduced focal ' + 'loss for Gaussian heatmap targets)') parser.add_argument('--focalAlpha', type=float, default=0.75, help='focal loss alpha (class balance weight, default: 0.75)') parser.add_argument('--focalGamma', type=float, default=2.0, help='focal loss gamma (focusing parameter, default: 2.0)') parser.add_argument('--focalDiceWeight', type=float, default=0.5, help='weight of dice component in FocalDiceLoss (default: 0.5)') + parser.add_argument('--heatmapAlpha', type=float, default=2.0, + help='FocalHeatmapLoss alpha exponent on (1-pred) for the ' + 'positive term (default: 2.0, CornerNet value)') + parser.add_argument('--heatmapBeta', type=float, default=4.0, + help='FocalHeatmapLoss beta exponent on (1-target) for the ' + 'penalty-reduced negative term (default: 4.0, CornerNet value)') + parser.add_argument('--targetType', type=str, default='binary', + choices=['binary', 'gaussian'], + help='per-pixel target representation: binary 9x9 mask (default) ' + 'or gaussian heatmap with std=gaussianSigma') + parser.add_argument('--gaussianSigma', type=float, default=3.0, + help='std (in pixels) of the Gaussian for gaussian targets (default: 3.0)') parser.add_argument('--warmupEpochs', type=int, default=0, help='number of linear warmup epochs before cosine decay (default: 0)') parser.add_argument('--swa', action='store_true', @@ -1230,9 +1326,11 @@ def main(): # Set rotateAndReflect=False - we'll use on-the-fly augmentation instead train_dataset = XPointDataset(args.paramFile, train_fnums, - xptCacheDir=args.xptCacheDir, rotateAndReflect=False) + xptCacheDir=args.xptCacheDir, rotateAndReflect=False, + target_type=args.targetType, gaussian_sigma=args.gaussianSigma) val_dataset = XPointDataset(args.paramFile, val_fnums, - xptCacheDir=args.xptCacheDir, rotateAndReflect=False) + xptCacheDir=args.xptCacheDir, rotateAndReflect=False, + target_type=args.targetType, gaussian_sigma=args.gaussianSigma) # Enable augmentation for training, disable for validation train_crop = XPointPatchDataset(train_dataset, patch=64, pos_ratio=args.posRatio, retries=30, @@ -1294,6 +1392,12 @@ def main(): criterion = FocalDiceLoss(alpha=args.focalAlpha, gamma=args.focalGamma, dice_weight=args.focalDiceWeight, smooth=1.0) print(f"Loss function: FocalDiceLoss (alpha={args.focalAlpha}, gamma={args.focalGamma}, dice_weight={args.focalDiceWeight})") + elif args.lossFunction == 'heatmap_focal': + if args.targetType != 'gaussian': + print(f"WARNING: heatmap_focal loss is intended for Gaussian targets, " + f"but --targetType={args.targetType}. Continuing anyway.") + criterion = FocalHeatmapLoss(alpha=args.heatmapAlpha, beta=args.heatmapBeta) + print(f"Loss function: FocalHeatmapLoss (alpha={args.heatmapAlpha}, beta={args.heatmapBeta})") else: criterion = DiceLoss(smooth=1.0) print("Loss function: DiceLoss") diff --git a/test_xpoint_transfer.py b/test_xpoint_transfer.py index 8400c83..1b7103a 100644 --- a/test_xpoint_transfer.py +++ b/test_xpoint_transfer.py @@ -170,14 +170,19 @@ def main(): # the same script can run unchanged at margin=0 (default) or at any # diagnostic margin (e.g. EDGE_MARGIN=10). edge_margin = int(os.environ.get("EDGE_MARGIN", 0)) + threshold = float(os.environ.get("THRESHOLD", 0.5)) suffix_parts = [] if OUTPUT_TAG: suffix_parts.append(OUTPUT_TAG) if edge_margin > 0: suffix_parts.append(f"em{edge_margin}") + if abs(threshold - 0.5) > 1e-9: + suffix_parts.append(f"th{threshold:g}".replace(".", "p")) suffix = ("_" + "_".join(suffix_parts)) if suffix_parts else "" if edge_margin > 0: print(f"\n[edge_margin] Excluding {edge_margin}-pixel border from metrics", flush=True) + if abs(threshold - 0.5) > 1e-9: + print(f"[threshold] Using prediction threshold {threshold} (default 0.5)", flush=True) print(f"\n[config] IN_DOMAIN={IN_DOMAIN}, BEST_MODEL={BEST_MODEL}, OUTPUT_TAG={OUTPUT_TAG or '(none)'}", flush=True) all_results = {} @@ -249,7 +254,7 @@ def main(): t0 = time.time() evaluator = evaluate_model_on_dataset( model, dataset, device, - use_amp=use_amp, amp_dtype=amp_dtype, threshold=0.5, + use_amp=use_amp, amp_dtype=amp_dtype, threshold=threshold, edge_margin=edge_margin, ) elapsed = time.time() - t0 @@ -300,7 +305,7 @@ def main(): t0 = time.time() in_evaluator = evaluate_model_on_dataset( model, in_dataset, device, - use_amp=use_amp, amp_dtype=amp_dtype, threshold=0.5, + use_amp=use_amp, amp_dtype=amp_dtype, threshold=threshold, edge_margin=edge_margin, ) elapsed = time.time() - t0 From 3e2bf6917a95bb5aa0be27b426dcdc58a9e54880 Mon Sep 17 00:00:00 2001 From: Swaroop Sridhar Date: Thu, 20 Aug 2026 15:39:53 -0500 Subject: [PATCH 06/11] move peak extraction into point_metrics and vectorize the connected-component NMS --- point_metrics.py | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/point_metrics.py b/point_metrics.py index fcfd936..7f83056 100644 --- a/point_metrics.py +++ b/point_metrics.py @@ -8,6 +8,41 @@ from typing import Iterable, Mapping, Sequence import numpy as np +from scipy.ndimage import label as cc_label +from scipy.ndimage import maximum as cc_maximum + + +# PEAK EXTRACTION +def extract_peaks(heatmap, threshold=0.3, max_components=20000): + """Extract one (row, col, confidence) peak per connected component above threshold (max-valued pixel in each). + + Ties within a component resolve to the lowest row-major index, matching + np.argmax. Returns empty arrays if the component count exceeds + max_components, which signals a speckled heatmap from an undertrained + model rather than real peaks. + """ + empty = (np.zeros(0, dtype=int), np.zeros(0, dtype=int), np.zeros(0, dtype=float)) + above = heatmap > threshold + if not above.any(): + return empty + + labels, n = cc_label(above) + if max_components is not None and n > max_components: + return empty + + # Per-component maxima, then the first pixel attaining each one. + maxima = np.atleast_1d(cc_maximum(heatmap, labels, index=np.arange(1, n + 1))) + lab_flat = labels.ravel() + at_max = np.flatnonzero((lab_flat > 0) & (heatmap.ravel() == maxima[lab_flat - 1])) + # flatnonzero is ascending, so a stable sort by label keeps the lowest + # flat index first within each component. + first = at_max[np.argsort(lab_flat[at_max], kind="stable")] + _, starts = np.unique(lab_flat[first], return_index=True) + peak_flat = first[starts] + + rows, cols = np.unravel_index(peak_flat, heatmap.shape) + confs = heatmap[rows, cols].astype(float) + return rows, cols, confs # CSV I/O From 85fabba1495e6b306760be3b7be86f91e9a5fd56 Mon Sep 17 00:00:00 2001 From: Swaroop Sridhar Date: Thu, 20 Aug 2026 15:39:53 -0500 Subject: [PATCH 07/11] use the shared extract_peaks in predict_points and thread the second PKPM run through cache build, prediction, and scoring --- predict_points.py | 35 +++++++++------------------------- run_hessian_and_build_cache.py | 6 +++++- score_point_predictions.py | 1 + 3 files changed, 15 insertions(+), 27 deletions(-) diff --git a/predict_points.py b/predict_points.py index dcf601d..07236a6 100644 --- a/predict_points.py +++ b/predict_points.py @@ -9,13 +9,13 @@ import numpy as np import torch -from scipy.ndimage import label as cc_label from torch.amp import autocast RC_ROOT = Path(__file__).resolve().parent sys.path.insert(0, str(RC_ROOT)) from XPointMLTest import XPointDataset, UNet, cachedPgkylDataExists +from point_metrics import extract_peaks DATASET_CONFIG = { @@ -31,30 +31,13 @@ "param_path": "/work/nvme/bfim/ssridhar6/mlReconnection2025/10M/10M/rt_10M_2d_turb_local-params.txt", "cache_dir": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/10M", }, + "PKPMv2": { + "param_path": "/work/nvme/bfim/ssridhar6/mlReconnection2025/1024Res_v2/rt_pkpm_2d_turb_p1-params.txt", + "cache_dir": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/PKPMv2", + }, } -def extract_peaks(heatmap, threshold=0.3, min_distance=5): - """Extract one (row, col, confidence) peak per connected component above threshold (max-valued pixel in each).""" - above = heatmap > threshold - if not above.any(): - return (np.zeros(0, dtype=int), - np.zeros(0, dtype=int), - np.zeros(0, dtype=float)) - labels, n = cc_label(above) - rows = np.empty(n, dtype=int) - cols = np.empty(n, dtype=int) - confs = np.empty(n, dtype=float) - for k in range(1, n + 1): - mask = labels == k - masked = np.where(mask, heatmap, -np.inf) - r, c = np.unravel_index(int(np.argmax(masked)), heatmap.shape) - rows[k - 1] = r - cols[k - 1] = c - confs[k - 1] = float(heatmap[r, c]) - return rows, cols, confs - - def write_pred_csv(out_path, rows, cols, confidences): """Write predicted-points CSV with row, col, confidence columns.""" with open(out_path, "w") as f: @@ -84,8 +67,8 @@ def main(): help="Base directory; predictions go to //{N}_xpts.csv") parser.add_argument("--threshold", type=float, default=0.3, help="Confidence threshold for peak retention (default: 0.3)") - parser.add_argument("--min-distance", type=int, default=5, - help="Min pixel separation between peaks (NMS kernel = 2*N+1)") + parser.add_argument("--max-components", type=int, default=20000, + help="Bail out and return no peaks above this component count (default: 20000)") parser.add_argument("--base-channels", type=int, default=64) parser.add_argument("--dropout-rate", type=float, default=0.055) args = parser.parse_args() @@ -106,7 +89,7 @@ def main(): amp_dtype = torch.bfloat16 if (use_amp and torch.cuda.is_bf16_supported()) else torch.float16 args.output_root.mkdir(parents=True, exist_ok=True) - print(f"NMS: threshold={args.threshold}, min_distance={args.min_distance} px") + print(f"NMS: threshold={args.threshold}, max_components={args.max_components}") with torch.no_grad(): for ds_name in args.datasets: @@ -132,7 +115,7 @@ def main(): probs = torch.sigmoid(logits) heatmap = probs[0, 0].float().cpu().numpy() rows, cols, confs = extract_peaks(heatmap, threshold=args.threshold, - min_distance=args.min_distance) + max_components=args.max_components) write_pred_csv(out_dir / f"{fnum}_xpts.csv", rows, cols, confs) n_predicted += len(rows) print(f" Wrote {len(frames)} CSV(s) in {time.time()-t1:.1f}s " diff --git a/run_hessian_and_build_cache.py b/run_hessian_and_build_cache.py index 865f9f4..204523b 100644 --- a/run_hessian_and_build_cache.py +++ b/run_hessian_and_build_cache.py @@ -56,6 +56,10 @@ def _fixed_GData_init(self, *args, comp=None, **kwargs): "extract_subdir": EXTRACT_DIR / "10M", "param_file": "rt_10M_2d_turb_local-params.txt", }, + "PKPMv2": { + "extract_subdir": EXTRACT_DIR / "1024Res_v2", + "param_file": "rt_pkpm_2d_turb_p1-params.txt", + }, } @@ -112,7 +116,7 @@ def _process_frame(task): def main(): parser = argparse.ArgumentParser(description="Build X-point cache for transfer datasets") - parser.add_argument("--dataset", required=True, choices=["5M", "10M"]) + parser.add_argument("--dataset", required=True, choices=["5M", "10M", "PKPMv2"]) parser.add_argument("--start", type=int, default=None, help="First frame index (inclusive)") parser.add_argument("--end", type=int, default=None, help="Last frame index (inclusive)") parser.add_argument("--workers", type=int, default=1, diff --git a/score_point_predictions.py b/score_point_predictions.py index c011826..09be10b 100644 --- a/score_point_predictions.py +++ b/score_point_predictions.py @@ -15,6 +15,7 @@ "PKPM": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/PKPM", "5M": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/5M", "10M": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/10M", + "PKPMv2": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/PKPMv2", } From 706fd4dedf11f6682334baf2118f95d407fdbefc Mon Sep 17 00:00:00 2001 From: Swaroop Sridhar Date: Thu, 20 Aug 2026 15:39:53 -0500 Subject: [PATCH 08/11] add a gaussian-aware optuna tuner that maximizes point-level F1 --- optuna_tuner_gaussian.py | 675 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 675 insertions(+) create mode 100644 optuna_tuner_gaussian.py diff --git a/optuna_tuner_gaussian.py b/optuna_tuner_gaussian.py new file mode 100644 index 0000000..2e47525 --- /dev/null +++ b/optuna_tuner_gaussian.py @@ -0,0 +1,675 @@ +""" +Optuna hyperparameter tuner for XPointMLTest.py + +Usage: + python optuna_tuner.py \ + --paramFile /path/to/params.txt \ + --xptCacheDir /path/to/cache \ + --n-trials 50 \ + --study-name xpoint-tuning \ + --db sqlite:///optuna_xpoint.db + +The script wraps the existing training pipeline and searches over: + - Learning rate + - Weight decay + - Dropout rate + - Batch size + - Patch size + - Base channels (model capacity) + - Scheduler type + params + - Augmentation positive-patch ratio + +Results persist in a SQLite database so you can: + - Resume after SSH disconnects + - Analyze results with Optuna's built-in visualization + - Run multiple workers in parallel (each with their own process) +""" + +import argparse +import json +import os +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch.optim as optim +from torch.amp import autocast, GradScaler +from torch.utils.data import DataLoader + +try: + import optuna + from optuna.exceptions import TrialPruned +except ImportError: + print("Optuna not installed. Install with:") + print(" pip install optuna --break-system-packages") + print("Optional visualization: pip install plotly --break-system-packages") + sys.exit(1) + +# Import from existing codebase +from XPointMLTest import ( + XPointDataset, + XPointPatchDataset, + UNet, + FocalHeatmapLoss, + train_one_epoch, + validate_one_epoch, + set_seed, +) +from point_metrics import extract_peaks, match_points +from ci_tests import SyntheticXPointDataset + + +def compute_point_f1(model, full_val_dataset, cache_dir, device, + use_amp, amp_dtype, threshold=0.3, radius=5.0): + """Point-level F1 on full val frames: NMS peaks vs ground-truth coords from the cache.""" + model.eval() + tp = fp = fn = 0 + with torch.no_grad(): + for item in full_val_dataset: + fnum = item["fnum"] + all_t = item["all"].unsqueeze(0).to(device) + with autocast(device_type="cuda", dtype=amp_dtype, enabled=use_amp): + probs = torch.sigmoid(model(all_t)) + heatmap = probs[0, 0].float().cpu().numpy() + rows, cols, confs = extract_peaks(heatmap, threshold=threshold) + pred = (np.stack([rows, cols], axis=1) + if len(rows) else np.zeros((0, 2), dtype=int)) + gt = np.load(Path(cache_dir) / f"{fnum}_xpts.npy") + r = match_points(pred, gt, radius=radius, pred_confidence=confs) + tp += r["tp"]; fp += r["fp"]; fn += r["fn"] + model.train() + p = tp / (tp + fp) if (tp + fp) else 0.0 + rec = tp / (tp + fn) if (tp + fn) else 0.0 + return 2 * p * rec / (p + rec) if (p + rec) else 0.0 + + +def objective(trial, args): + """ + Optuna objective function. Trains the model with trial-suggested + hyperparameters and returns the best validation loss. + """ + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # --- Seed: vary per trial for diversity, but keep reproducible --- + seed = args.seed + if seed is not None: + set_seed(seed + trial.number) + + # 1. Suggest ALL hyperparameters in one place + + # Model architecture + base_channels = trial.suggest_categorical("base_channels", [16, 32, 48, 64]) + dropout_rate = trial.suggest_float("dropout_rate", 0.05, 0.5) + + # Optimizer + lr = trial.suggest_float("learning_rate", 1e-5, 5e-3, log=True) + weight_decay = trial.suggest_float("weight_decay", 1e-6, 1e-2, log=True) + + # Data pipeline + batch_size = trial.suggest_categorical("batch_size", [8, 16, 32, 64]) + patch_size = trial.suggest_categorical("patch_size", [48, 64, 96]) + pos_ratio = trial.suggest_float("pos_ratio", 0.3, 0.7) + + # Scheduler + scheduler_name = trial.suggest_categorical("scheduler", ["cosine", "plateau"]) + + # 2. Load data (pre-loaded and cached on args to avoid repeated I/O) + + if args.smoke_test: + train_dataset = SyntheticXPointDataset(nframes=10, shape=(64, 64), nxpoints=3) + val_dataset = SyntheticXPointDataset( + nframes=2, shape=(64, 64), nxpoints=3, seed=123 + ) + else: + train_dataset = args._train_dataset + val_dataset = args._val_dataset + + train_crop = XPointPatchDataset( + train_dataset, + patch=patch_size, + pos_ratio=pos_ratio, + retries=30, + augment=True, + seed=seed, + ) + val_crop = XPointPatchDataset( + val_dataset, + patch=patch_size, + pos_ratio=0.5, + retries=30, + augment=False, + seed=seed, + ) + + train_loader = DataLoader( + train_crop, batch_size=batch_size, shuffle=True, num_workers=0 + ) + val_loader = DataLoader( + val_crop, batch_size=batch_size, shuffle=False, num_workers=0 + ) + + + # 3. Create model, optimizer, scheduler + + model = UNet( + input_channels=4, base_channels=base_channels, dropout_rate=dropout_rate + ).to(device) + + num_epochs = args.epochs + optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay) + + if scheduler_name == "plateau": + plateau_factor = trial.suggest_float("plateau_factor", 0.2, 0.8) + plateau_patience = trial.suggest_int("plateau_patience", 3, 15) + scheduler = optim.lr_scheduler.ReduceLROnPlateau( + optimizer, + mode="min", + factor=plateau_factor, + patience=plateau_patience, + min_lr=1e-6, + ) + else: + scheduler = optim.lr_scheduler.CosineAnnealingLR( + optimizer, T_max=num_epochs, eta_min=1e-6 + ) + + criterion = FocalHeatmapLoss(alpha=args.heatmap_alpha, beta=args.heatmap_beta) + + # --- AMP setup --- + use_amp = args.use_amp and torch.cuda.is_available() + amp_dtype = ( + torch.bfloat16 + if args.amp_dtype == "bfloat16" and torch.cuda.is_bf16_supported() + else torch.float16 + ) + scaler = GradScaler(enabled=(use_amp and amp_dtype == torch.float16)) + + + # 4. Training loop with Optuna pruning + + best_val_loss = float("inf") + best_f1 = 0.0 + patience_counter = 0 + patience = args.patience + epochs_trained = 0 + + for epoch in range(num_epochs): + train_loss = train_one_epoch( + model, + train_loader, + criterion, + optimizer, + device, + scaler, + use_amp, + amp_dtype, + ) + + # Reset validation RNG for deterministic crops each epoch + if seed is not None: + val_crop.reset_rng(seed) + + val_loss = validate_one_epoch( + model, val_loader, criterion, device, use_amp, amp_dtype + ) + + # Non-finite guard: wide sigma / bad LR can drive FocalHeatmapLoss to + # inf/nan. Bail on this config with the worst objective so the sampler + # learns to avoid the region (rather than crashing the whole study). + if not (np.isfinite(train_loss) and np.isfinite(val_loss)): + trial.set_user_attr("non_finite", True) + trial.set_user_attr("epochs_trained", epoch + 1) + return 0.0 + + # LR scheduling + if scheduler_name == "plateau": + scheduler.step(val_loss) + else: + scheduler.step() + + # Track best val loss (used only for early stopping) + if val_loss < best_val_loss: + best_val_loss = val_loss + patience_counter = 0 + else: + patience_counter += 1 + + epochs_trained = epoch + 1 + + # === Point-level F1 (the objective) computed periodically === + # Full-frame inference over ~9 val frames is cheap (<1s), so we can + # report the true metric to the pruner rather than a loss proxy. + if not args.smoke_test and ( + epoch % args.f1_interval == 0 or epoch == num_epochs - 1 + ): + f1 = compute_point_f1( + model, val_dataset, args.xptCacheDir, device, + use_amp, amp_dtype, + threshold=args.nms_threshold, radius=args.match_radius, + ) + best_f1 = max(best_f1, f1) + trial.report(f1, epoch) + if trial.should_prune(): + raise TrialPruned() + + # Early stopping on val loss plateau + if patience_counter >= patience: + break + + # Smoke test has no real cache to score against; fall back to -val_loss. + if args.smoke_test: + best_f1 = -best_val_loss + + trial.set_user_attr("best_val_loss", best_val_loss) + trial.set_user_attr("best_point_f1", best_f1) + trial.set_user_attr("epochs_trained", epochs_trained) + + return best_f1 + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Optuna hyperparameter tuning for X-point classifier" + ) + + # --- Optuna settings --- + parser.add_argument( + "--n-trials", + type=int, + default=50, + help="Number of Optuna trials (default: 50)", + ) + parser.add_argument( + "--study-name", + type=str, + default="xpoint-tuning", + help="Optuna study name (default: xpoint-tuning)", + ) + parser.add_argument( + "--db", + type=str, + default="sqlite:///optuna_xpoint.db", + help="Optuna storage URL (default: sqlite:///optuna_xpoint.db)", + ) + parser.add_argument( + "--timeout", + type=int, + default=None, + help="Stop after this many seconds (default: None, run all trials)", + ) + parser.add_argument( + "--pruner", + type=str, + default="median", + choices=["median", "hyperband", "none"], + help="Pruning strategy (default: median)", + ) + + # --- Data settings (same as XPointMLTest.py) --- + parser.add_argument("--paramFile", type=Path, default=None) + parser.add_argument("--xptCacheDir", type=Path, default=None) + parser.add_argument("--trainFrameFirst", type=int, default=1) + parser.add_argument("--trainFrameLast", type=int, default=140) + parser.add_argument("--validationFrameFirst", type=int, default=141) + parser.add_argument("--validationFrameLast", type=int, default=150) + + # --- Training settings (fixed across all trials) --- + parser.add_argument( + "--epochs", + type=int, + default=300, + help="Max epochs PER TRIAL (default: 300, lower than full training for speed)", + ) + parser.add_argument( + "--patience", + type=int, + default=30, + help="Early stopping patience per trial (default: 30)", + ) + parser.add_argument("--use-amp", action="store_true", help="Enable AMP") + parser.add_argument( + "--amp-dtype", + type=str, + default="bfloat16", + choices=["float16", "bfloat16"], + ) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--require-gpu", action="store_true") + + # --- Gaussian target / heatmap loss (fixed across trials) --- + parser.add_argument("--gaussianSigma", type=float, default=3.0, + help="Gaussian target width in pixels (fixed; default 3.0)") + parser.add_argument("--heatmap-alpha", type=float, default=2.0, + help="FocalHeatmapLoss alpha (default 2.0)") + parser.add_argument("--heatmap-beta", type=float, default=4.0, + help="FocalHeatmapLoss beta (default 4.0)") + + # --- Point-level F1 evaluation (the objective) --- + parser.add_argument("--f1-interval", type=int, default=10, + help="Compute point-level F1 every N epochs for pruning/objective (default 10)") + parser.add_argument("--nms-threshold", type=float, default=0.3, + help="Confidence threshold for peak extraction (default 0.3)") + parser.add_argument("--match-radius", type=float, default=5.0, + help="Point matching radius in pixels (default 5.0)") + + # --- Evaluation --- + parser.add_argument( + "--eval-on-full-frames", + action="store_true", + help="Run full-frame evaluation after each trial (slower but gives F1/IoU)", + ) + + # --- Output --- + parser.add_argument( + "--results-dir", + type=Path, + default="./optuna_results", + help="Directory for result files (default: ./optuna_results)", + ) + + # --- Testing --- + parser.add_argument( + "--smoke-test", + action="store_true", + help="Run 3 trials with synthetic data (no paramFile needed)", + ) + + return parser.parse_args() + + +def print_study_summary(study, results_dir, args): + """Print and save a summary of the Optuna study results.""" + + print("\n" + "=" * 70) + print("OPTUNA STUDY SUMMARY") + print("=" * 70) + + completed = [ + t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE + ] + pruned = [t for t in study.trials if t.state == optuna.trial.TrialState.PRUNED] + failed = [t for t in study.trials if t.state == optuna.trial.TrialState.FAIL] + + print(f"\nStudy name: {study.study_name}") + print(f"Total trials: {len(study.trials)}") + print(f" Completed: {len(completed)}") + print(f" Pruned: {len(pruned)}") + print(f" Failed: {len(failed)}") + + if not completed: + print("\nNo completed trials to summarize.") + return + + best = study.best_trial + print(f"\nBest trial: #{best.number}") + print(f" Best point-level F1: {best.value:.4f}") + print(f" Hyperparameters:") + for key, value in sorted(best.params.items()): + if isinstance(value, float): + print(f" {key:25s} {value:.6g}") + else: + print(f" {key:25s} {value}") + + if best.user_attrs: + print(f" Additional metrics:") + for key, value in sorted(best.user_attrs.items()): + if isinstance(value, float): + print(f" {key:25s} {value:.4f}") + else: + print(f" {key:25s} {value}") + + # --- Top 5 trials (maximize: highest F1 first) --- + sorted_trials = sorted(completed, key=lambda t: t.value, reverse=True) + print(f"\nTop 5 trials:") + print( + f" {'#':>4s} {'PointF1':>10s} {'LR':>10s} {'WD':>10s} " + f"{'Drop':>6s} {'BS':>4s} {'Patch':>5s} {'Ch':>4s} {'Sched':>8s}" + ) + for t in sorted_trials[:5]: + p = t.params + print( + f" {t.number:4d} {t.value:10.4f} " + f"{p.get('learning_rate', 0):10.2e} " + f"{p.get('weight_decay', 0):10.2e} " + f"{p.get('dropout_rate', 0):6.3f} " + f"{p.get('batch_size', 0):4d} " + f"{p.get('patch_size', 0):5d} " + f"{p.get('base_channels', 0):4d} " + f"{p.get('scheduler', 'n/a'):>8s}" + ) + + print("=" * 70) + + # --- Save results to JSON --- + results_dir.mkdir(parents=True, exist_ok=True) + + results = { + "study_name": study.study_name, + "n_trials_total": len(study.trials), + "n_completed": len(completed), + "n_pruned": len(pruned), + "n_failed": len(failed), + "best_trial": { + "number": best.number, + "value": best.value, + "params": best.params, + "user_attrs": best.user_attrs, + }, + "all_completed_trials": [ + { + "number": t.number, + "value": t.value, + "params": t.params, + "user_attrs": t.user_attrs, + } + for t in sorted_trials + ], + } + + results_path = results_dir / "optuna_results.json" + with open(results_path, "w") as f: + json.dump(results, f, indent=2) + print(f"\nResults saved to: {results_path}") + + # --- Generate shell command for retraining with best params --- + p = best.params + cmd_lines = [ + "#!/bin/bash", + f"# Auto-generated from Optuna study: {study.study_name}", + f"# Best trial #{best.number} with point-level F1={best.value:.4f}", + "", + "python -u ${rcRoot}/reconClassifier/XPointMLTest.py \\", + f" --learningRate {p.get('learning_rate', 5e-4):.6g} \\", + f" --weightDecay {p.get('weight_decay', 5e-5):.6g} \\", + f" --dropoutRate {p.get('dropout_rate', 0.15):.4g} \\", + f" --batchSize {p.get('batch_size', 64)} \\", + f" --posRatio {p.get('pos_ratio', 0.5):.4g} \\", + f" --baseChannels {p.get('base_channels', 64)} \\", + " --targetType gaussian \\", + f" --gaussianSigma {args.gaussianSigma} \\", + " --lossFunction heatmap_focal \\", + f" --heatmapAlpha {args.heatmap_alpha} \\", + f" --heatmapBeta {args.heatmap_beta} \\", + " --warmupEpochs 10 \\", + ] + + sched = p.get("scheduler", "cosine") + cmd_lines.append(f" --scheduler {sched} \\") + if sched == "plateau": + cmd_lines.append( + f" --plateau-factor {p.get('plateau_factor', 0.5):.4g} \\" + ) + cmd_lines.append( + f" --plateau-patience {p.get('plateau_patience', 5)} \\" + ) + + cmd_lines.extend( + [ + " --use-amp \\", + " --seed 42 \\", + " --require-gpu \\", + " --fixed-val-crops \\", + " --epochs 1200 \\", + " --patience 200 \\", + " --checkPointFrequency 200 \\", + " --paramFile=${PARAM_FILE} \\", + " --xptCacheDir=${CACHE_DIR}", + ] + ) + + cmd_path = results_dir / "retrain_best_params.sh" + with open(cmd_path, "w") as f: + f.write("\n".join(cmd_lines) + "\n") + os.chmod(cmd_path, 0o755) + + print(f"Retrain command saved to: {cmd_path}") + print(f"\nTo retrain with best hyperparameters:") + print(f" bash {cmd_path}") + print("=" * 70) + + +def main(): + args = parse_args() + + # --- Smoke test overrides --- + if args.smoke_test: + print("=" * 60) + print("SMOKE TEST MODE: 3 trials with synthetic data") + print("=" * 60) + args.n_trials = 3 + args.epochs = 5 + args.patience = 3 + args.eval_on_full_frames = False + + # --- Validate args --- + if not args.smoke_test: + if args.paramFile is None: + print("ERROR: --paramFile is required (or use --smoke-test)") + sys.exit(1) + if not args.paramFile.exists(): + print(f"ERROR: paramFile {args.paramFile} does not exist") + sys.exit(1) + + if args.require_gpu and not torch.cuda.is_available(): + print("ERROR: --require-gpu set but no CUDA device found") + sys.exit(1) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"Using device: {device}") + if torch.cuda.is_available(): + print(f"GPU: {torch.cuda.get_device_name(0)}") + + + # Pre-load datasets ONCE (expensive I/O — shared across all trials) + + if not args.smoke_test: + print("\nLoading datasets (shared across all trials)...") + t0 = time.time() + train_dataset = XPointDataset( + args.paramFile, + range(args.trainFrameFirst, args.trainFrameLast), + xptCacheDir=args.xptCacheDir, + rotateAndReflect=False, + target_type="gaussian", + gaussian_sigma=args.gaussianSigma, + ) + val_dataset = XPointDataset( + args.paramFile, + range(args.validationFrameFirst, args.validationFrameLast), + xptCacheDir=args.xptCacheDir, + rotateAndReflect=False, + target_type="gaussian", + gaussian_sigma=args.gaussianSigma, + ) + print(f"Datasets loaded in {time.time() - t0:.1f}s") + print(f" Training frames: {len(train_dataset)}") + print(f" Validation frames: {len(val_dataset)}") + + # Attach to args for access in objective() + args._train_dataset = train_dataset + args._val_dataset = val_dataset + + + # Create Optuna study + + if args.pruner == "median": + pruner = optuna.pruners.MedianPruner( + n_startup_trials=5, n_warmup_steps=10, interval_steps=5 + ) + elif args.pruner == "hyperband": + pruner = optuna.pruners.HyperbandPruner( + min_resource=10, max_resource=args.epochs, reduction_factor=3 + ) + else: + pruner = optuna.pruners.NopPruner() + + study = optuna.create_study( + study_name=args.study_name, + storage=args.db, + load_if_exists=True, + direction="maximize", # objective is point-level F1 + pruner=pruner, + ) + + n_existing = len(study.trials) + if n_existing > 0: + print(f"\nResuming study '{args.study_name}' with {n_existing} existing trials") + if study.best_trial: + print(f"Current best val_loss: {study.best_trial.value:.6f}") + + + # Run optimization + + print(f"\nStarting Optuna optimization") + print(f" Trials: {args.n_trials}") + print(f" DB: {args.db}") + print(f" Pruner: {args.pruner}") + print(f" Max epochs: {args.epochs} per trial") + print(f" Patience: {args.patience} per trial") + if args.timeout: + print(f" Timeout: {args.timeout}s") + print() + + study.optimize( + lambda trial: objective(trial, args), + n_trials=args.n_trials, + timeout=args.timeout, + show_progress_bar=True, + ) + + + # Results + + print_study_summary(study, args.results_dir, args) + + # --- Try to generate visualization plots --- + try: + from optuna.visualization import ( + plot_param_importances, + plot_optimization_history, + plot_parallel_coordinate, + ) + + fig_dir = args.results_dir / "figures" + fig_dir.mkdir(parents=True, exist_ok=True) + + for name, plot_fn in [ + ("param_importances", plot_param_importances), + ("optimization_history", plot_optimization_history), + ("parallel_coordinate", plot_parallel_coordinate), + ]: + try: + fig = plot_fn(study) + fig.write_html(str(fig_dir / f"{name}.html")) + print(f" Saved: {fig_dir / name}.html") + except Exception as e: + print(f" Skipped {name}: {e}") + + except ImportError: + print("\nNote: pip install plotly for interactive visualizations") + + +if __name__ == "__main__": + main() \ No newline at end of file From deb3d38b6af57dc7625007876feb3f0d786ddeb1 Mon Sep 17 00:00:00 2001 From: Swaroop Sridhar Date: Fri, 4 Sep 2026 12:11:04 -0500 Subject: [PATCH 09/11] derive predict/score dataset paths from RC_EXTRACT_DIR/RC_CACHE_BASE instead of hardcoding them --- predict_points.py | 22 ++++++++++++++-------- score_point_predictions.py | 13 +++++++++---- 2 files changed, 23 insertions(+), 12 deletions(-) diff --git a/predict_points.py b/predict_points.py index 07236a6..e011e99 100644 --- a/predict_points.py +++ b/predict_points.py @@ -18,22 +18,28 @@ from point_metrics import extract_peaks +# Base dirs, overridable via env; same RC_EXTRACT_DIR / RC_CACHE_BASE convention +# as run_hessian_and_build_cache.py so a launcher sets paths once for the whole pipeline. +EXTRACT_DIR = Path(os.environ.get("RC_EXTRACT_DIR", "/work/nvme/bfim/ssridhar6/mlReconnection2025")) +CACHE_BASE = Path(os.environ.get("RC_CACHE_BASE", "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache")) +PKPM_V0_ROOT = Path(os.environ.get("RC_PKPM_V0_ROOT", "/work/nvme/bfim/cwsmith/mlReconnection2025/1024Res_v0")) + DATASET_CONFIG = { "PKPM": { - "param_path": "/work/nvme/bfim/cwsmith/mlReconnection2025/1024Res_v0/pkpm_2d_turb_p2-params.txt", - "cache_dir": "/work/nvme/bfim/cwsmith/mlReconnection2025/1024Res_v0/cache04082025", + "param_path": str(PKPM_V0_ROOT / "pkpm_2d_turb_p2-params.txt"), + "cache_dir": str(PKPM_V0_ROOT / "cache04082025"), }, "5M": { - "param_path": "/work/nvme/bfim/ssridhar6/mlReconnection2025/5M/rt_5M_2d_turb_local-params.txt", - "cache_dir": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/5M", + "param_path": str(EXTRACT_DIR / "5M" / "rt_5M_2d_turb_local-params.txt"), + "cache_dir": str(CACHE_BASE / "5M"), }, "10M": { - "param_path": "/work/nvme/bfim/ssridhar6/mlReconnection2025/10M/10M/rt_10M_2d_turb_local-params.txt", - "cache_dir": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/10M", + "param_path": str(EXTRACT_DIR / "10M" / "10M" / "rt_10M_2d_turb_local-params.txt"), + "cache_dir": str(CACHE_BASE / "10M"), }, "PKPMv2": { - "param_path": "/work/nvme/bfim/ssridhar6/mlReconnection2025/1024Res_v2/rt_pkpm_2d_turb_p1-params.txt", - "cache_dir": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/PKPMv2", + "param_path": str(EXTRACT_DIR / "1024Res_v2" / "rt_pkpm_2d_turb_p1-params.txt"), + "cache_dir": str(CACHE_BASE / "PKPMv2"), }, } diff --git a/score_point_predictions.py b/score_point_predictions.py index 09be10b..b714321 100644 --- a/score_point_predictions.py +++ b/score_point_predictions.py @@ -3,6 +3,7 @@ import argparse import json +import os from pathlib import Path from point_metrics import ( @@ -11,11 +12,15 @@ ) +# Ground-truth CSVs live under the shared cache root; override via RC_CACHE_BASE +# (same convention as run_hessian_and_build_cache.py / predict_points.py). +CACHE_BASE = Path(os.environ.get("RC_CACHE_BASE", "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache")) + GT_CSV_DIRS = { - "PKPM": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/PKPM", - "5M": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/5M", - "10M": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/10M", - "PKPMv2": "/work/nvme/bfim/ssridhar6/mlReconnection2025/cache/PKPMv2", + "PKPM": str(CACHE_BASE / "PKPM"), + "5M": str(CACHE_BASE / "5M"), + "10M": str(CACHE_BASE / "10M"), + "PKPMv2": str(CACHE_BASE / "PKPMv2"), } From b929540a8c0d309359256b3d55a5c35589979ffc Mon Sep 17 00:00:00 2001 From: Swaroop Sridhar Date: Fri, 11 Sep 2026 08:09:17 -0500 Subject: [PATCH 10/11] sweep the peak-extraction threshold inside each optuna trial so every config is scored at its own best cutoff --- optuna_tuner_gaussian.py | 56 +++++++++++++++++++++++++++------------- 1 file changed, 38 insertions(+), 18 deletions(-) diff --git a/optuna_tuner_gaussian.py b/optuna_tuner_gaussian.py index 2e47525..52dc52b 100644 --- a/optuna_tuner_gaussian.py +++ b/optuna_tuner_gaussian.py @@ -62,10 +62,11 @@ def compute_point_f1(model, full_val_dataset, cache_dir, device, - use_amp, amp_dtype, threshold=0.3, radius=5.0): - """Point-level F1 on full val frames: NMS peaks vs ground-truth coords from the cache.""" + use_amp, amp_dtype, thresholds=(0.3,), radius=5.0): + """Point-level F1 swept over a threshold grid; returns (best_f1, best_threshold).""" model.eval() - tp = fp = fn = 0 + thresholds = list(thresholds) + counts = {t: [0, 0, 0] for t in thresholds} # tp, fp, fn per threshold with torch.no_grad(): for item in full_val_dataset: fnum = item["fnum"] @@ -73,16 +74,26 @@ def compute_point_f1(model, full_val_dataset, cache_dir, device, with autocast(device_type="cuda", dtype=amp_dtype, enabled=use_amp): probs = torch.sigmoid(model(all_t)) heatmap = probs[0, 0].float().cpu().numpy() - rows, cols, confs = extract_peaks(heatmap, threshold=threshold) - pred = (np.stack([rows, cols], axis=1) - if len(rows) else np.zeros((0, 2), dtype=int)) gt = np.load(Path(cache_dir) / f"{fnum}_xpts.npy") - r = match_points(pred, gt, radius=radius, pred_confidence=confs) - tp += r["tp"]; fp += r["fp"]; fn += r["fn"] + # The forward pass dominates, so re-extracting peaks per threshold is + # cheap and lets each trial be scored at its own best operating point. + for t in thresholds: + rows, cols, confs = extract_peaks(heatmap, threshold=t) + pred = (np.stack([rows, cols], axis=1) + if len(rows) else np.zeros((0, 2), dtype=int)) + r = match_points(pred, gt, radius=radius, pred_confidence=confs) + c = counts[t] + c[0] += r["tp"]; c[1] += r["fp"]; c[2] += r["fn"] model.train() - p = tp / (tp + fp) if (tp + fp) else 0.0 - rec = tp / (tp + fn) if (tp + fn) else 0.0 - return 2 * p * rec / (p + rec) if (p + rec) else 0.0 + best_f1, best_t = 0.0, thresholds[0] + for t in thresholds: + tp, fp, fn = counts[t] + p = tp / (tp + fp) if (tp + fp) else 0.0 + rec = tp / (tp + fn) if (tp + fn) else 0.0 + f1 = 2 * p * rec / (p + rec) if (p + rec) else 0.0 + if f1 > best_f1: + best_f1, best_t = f1, t + return best_f1, best_t def objective(trial, args): @@ -191,6 +202,7 @@ def objective(trial, args): best_val_loss = float("inf") best_f1 = 0.0 + best_threshold = None patience_counter = 0 patience = args.patience epochs_trained = 0 @@ -244,12 +256,13 @@ def objective(trial, args): if not args.smoke_test and ( epoch % args.f1_interval == 0 or epoch == num_epochs - 1 ): - f1 = compute_point_f1( + f1, thr = compute_point_f1( model, val_dataset, args.xptCacheDir, device, use_amp, amp_dtype, - threshold=args.nms_threshold, radius=args.match_radius, + thresholds=args.threshold_grid, radius=args.match_radius, ) - best_f1 = max(best_f1, f1) + if f1 > best_f1: + best_f1, best_threshold = f1, thr trial.report(f1, epoch) if trial.should_prune(): raise TrialPruned() @@ -264,6 +277,7 @@ def objective(trial, args): trial.set_user_attr("best_val_loss", best_val_loss) trial.set_user_attr("best_point_f1", best_f1) + trial.set_user_attr("best_threshold", best_threshold) trial.set_user_attr("epochs_trained", epochs_trained) return best_f1 @@ -349,8 +363,10 @@ def parse_args(): # --- Point-level F1 evaluation (the objective) --- parser.add_argument("--f1-interval", type=int, default=10, help="Compute point-level F1 every N epochs for pruning/objective (default 10)") - parser.add_argument("--nms-threshold", type=float, default=0.3, - help="Confidence threshold for peak extraction (default 0.3)") + parser.add_argument("--threshold-grid", type=str, + default="0.30,0.35,0.40,0.45,0.50,0.55,0.60,0.65", + help="Comma-separated peak-extraction thresholds swept within each " + "trial; the trial is scored at its best one (default 0.30..0.65)") parser.add_argument("--match-radius", type=float, default=5.0, help="Point matching radius in pixels (default 5.0)") @@ -376,7 +392,9 @@ def parse_args(): help="Run 3 trials with synthetic data (no paramFile needed)", ) - return parser.parse_args() + args = parser.parse_args() + args.threshold_grid = [float(t) for t in args.threshold_grid.split(",") if t.strip()] + return args def print_study_summary(study, results_dir, args): @@ -424,13 +442,15 @@ def print_study_summary(study, results_dir, args): sorted_trials = sorted(completed, key=lambda t: t.value, reverse=True) print(f"\nTop 5 trials:") print( - f" {'#':>4s} {'PointF1':>10s} {'LR':>10s} {'WD':>10s} " + f" {'#':>4s} {'PointF1':>10s} {'Thr':>5s} {'LR':>10s} {'WD':>10s} " f"{'Drop':>6s} {'BS':>4s} {'Patch':>5s} {'Ch':>4s} {'Sched':>8s}" ) for t in sorted_trials[:5]: p = t.params + thr = t.user_attrs.get("best_threshold") print( f" {t.number:4d} {t.value:10.4f} " + f"{(thr if thr is not None else float('nan')):5.2f} " f"{p.get('learning_rate', 0):10.2e} " f"{p.get('weight_decay', 0):10.2e} " f"{p.get('dropout_rate', 0):6.3f} " From 2bf493d04d3f934afefc111d0359d20c40872bee Mon Sep 17 00:00:00 2001 From: Swaroop Sridhar Date: Fri, 11 Sep 2026 08:24:40 -0500 Subject: [PATCH 11/11] document the gaussian target options, point-level evaluation pipeline, and RC_* path variables in the README --- README.md | 140 ++++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 132 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index e5963e7..2be9e0a 100644 --- a/README.md +++ b/README.md @@ -95,11 +95,17 @@ The classifier supports several command line options for training configuration: ### Architecture - `--baseChannels`: Base number of channels in the UNet encoder (default: 64) +### Target Representation +- `--targetType`: Per-pixel training target, `binary` (default; a 9×9 block marking each X-point) or `gaussian` (a 2D Gaussian peak of unit height centered on each X-point, with overlapping peaks combined by elementwise maximum) +- `--gaussianSigma`: Width in pixels of the Gaussian peaks when `--targetType gaussian` (default: 3.0) + ### Loss Function -- `--lossFunction`: Loss function, `dice` (default) or `focal_dice` (combined focal + dice loss for class imbalance) +- `--lossFunction`: Loss function, `dice` (default), `focal_dice` (combined focal + dice loss for class imbalance), or `heatmap_focal` (CornerNet/CenterNet-style penalty-reduced focal loss for Gaussian targets; a warning is printed if it is used with `--targetType binary`) - `--focalAlpha`: Focal loss alpha, class balance weight (default: 0.75) - `--focalGamma`: Focal loss gamma, focusing parameter (default: 2.0) - `--focalDiceWeight`: Weight of the dice component in FocalDiceLoss (default: 0.5) +- `--heatmapAlpha`: `heatmap_focal` focusing exponent, which down-weights pixels the network already predicts well (default: 2.0) +- `--heatmapBeta`: `heatmap_focal` exponent on (1 − target) in the penalty-reduced negative term, so pixels close to an X-point are penalized less for high predictions (default: 4.0) ### Learning Rate Schedule - `--warmupEpochs`: Number of linear warmup epochs before the main scheduler kicks in (default: 0) @@ -165,6 +171,32 @@ python -u ${rcRoot}/reconClassifier/XPointMLTest.py \ --validationFrameLast 120 ``` +To train the Gaussian heatmap model used for point-level detection (the hand-tuned configuration behind the current point-level results): +```bash +python -u ${rcRoot}/reconClassifier/XPointMLTest.py \ +--paramFile=/path/to/params.txt \ +--xptCacheDir=/path/to/cache \ +--targetType gaussian \ +--gaussianSigma 3.0 \ +--lossFunction heatmap_focal \ +--heatmapAlpha 2.0 \ +--heatmapBeta 4.0 \ +--baseChannels 64 \ +--learningRate 0.00224265 \ +--weightDecay 0.00253832 \ +--dropoutRate 0.1784 \ +--batchSize 256 \ +--posRatio 0.4355 \ +--warmupEpochs 10 \ +--scheduler cosine \ +--use-amp \ +--epochs 1200 \ +--patience 200 \ +--fixed-val-crops \ +--seed 42 \ +--require-gpu +``` + ## Hyperparameter Tuning with Optuna The `optuna_tuner.py` script automates hyperparameter search over the knobs above (base channels, dropout, weight decay, learning rate, positive ratio, focal/dice weighting, scheduler choice, SWA start). It uses a Tree-structured Parzen Estimator sampler and a Median Pruner that aborts unpromising runs early based on the validation F1 curve. @@ -180,13 +212,41 @@ python -u ${rcRoot}/reconClassifier/optuna_tuner.py \ The SQLite database is created automatically on first run and reloaded on subsequent runs with the same `--study-name`, so a study can be resumed or extended without re-running completed trials. +### Tuning the Gaussian model for point-level F1 + +`optuna_tuner_gaussian.py` is the corresponding tuner for Gaussian heatmap targets. It searches over the same training knobs (learning rate, weight decay, dropout, batch size, patch size, base channels, positive ratio, scheduler), trains every trial with `heatmap_focal` loss, and **maximizes point-level F1** on the validation frames instead of minimizing validation loss. Every `--f1-interval` epochs it runs full-frame inference on the validation frames, extracts peaks, matches them to the ground-truth X-points, and reports F1 to the pruner. + +The peak-extraction threshold only changes how the network's output is read, so it is not an Optuna search dimension (that would cost a full training run per threshold value). Instead, each F1 evaluation sweeps `--threshold-grid` over the same heatmaps and scores the trial at its best threshold. The winning value is stored as the trial's `best_threshold` user attribute and appears in the summary table. + +``` +python -u ${rcRoot}/reconClassifier/optuna_tuner_gaussian.py \ +--paramFile=/path/to/params.txt \ +--xptCacheDir=/path/to/cache \ +--n-trials 50 \ +--epochs 300 \ +--study-name gauss-tune \ +--db sqlite:///optuna_gauss.db \ +--use-amp \ +--require-gpu +``` + +Options specific to this tuner: +- `--gaussianSigma`: Gaussian target width in pixels, held fixed during the search (default: 3.0) +- `--heatmap-alpha`, `--heatmap-beta`: `heatmap_focal` exponents (defaults: 2.0 and 4.0) +- `--f1-interval`: Epochs between point-level F1 evaluations (default: 10) +- `--threshold-grid`: Comma-separated peak-extraction thresholds swept at each evaluation (default: `0.30,0.35,0.40,0.45,0.50,0.55,0.60,0.65`) +- `--match-radius`: Matching radius in pixels (default: 5.0) +- `--epochs`: Maximum epochs per trial (default: 300) + +Both tuners default to the same `--study-name` and `--db`, but they optimize different objectives (validation loss versus point-level F1), so give the Gaussian tuner its own study to keep the results separate. + ## Cross-regime Transfer Evaluation -The PKPM-trained model can be evaluated zero-shot on additional Gkeyll datasets (currently 5-moment "5M" and 10-moment "10M" fluid simulations). Evaluation runs in two steps: first build the X-point cache for the transfer dataset, then run the evaluator. +The PKPM-trained model can be evaluated zero-shot on additional Gkeyll datasets: the 5-moment ("5M") and 10-moment ("10M") fluid simulations, and a second, independent PKPM run ("PKPMv2") with the same physical and numerical configuration but a different turbulence realization, which is useful as a held-out test set. Evaluation runs in two steps: first build the X-point cache for the dataset, then evaluate a checkpoint on it, either with the point-level pipeline (see [Point-level Evaluation](#point-level-evaluation)) or with the pixel-level `test_xpoint_transfer.py`. -### Building the X-point cache for 5M/10M +### Building the X-point cache for 5M/10M/PKPMv2 -`run_hessian_and_build_cache.py` is the only script that runs the deterministic Hessian X-point classifier; it writes the per-frame results as `.npy` files so the training and evaluation scripts only ever read from cache. Trying to train or evaluate on an uncached frame raises a clear error pointing back to this script. +`run_hessian_and_build_cache.py` is the only script that runs the deterministic Hessian X-point classifier; it writes the per-frame results as `.npy` files so the training and evaluation scripts only ever read from cache. Trying to train or evaluate on an uncached frame raises a clear error pointing back to this script. For each frame it also writes a ground-truth point list, `{N}_xpts.csv`, with `row,col,class` columns (`class` is `X`, `Omax`, or `Omin`), which the point-level scorer reads. `--dataset` accepts `5M`, `10M`, or `PKPMv2`, and `--start`/`--end` are inclusive. ``` python -u ${rcRoot}/reconClassifier/run_hessian_and_build_cache.py \ @@ -195,16 +255,80 @@ python -u ${rcRoot}/reconClassifier/run_hessian_and_build_cache.py \ --workers 30 ``` -The `RC_EXTRACT_DIR` and `RC_CACHE_BASE` environment variables override the default raw-data and cache directories. Pointing `RC_EXTRACT_DIR` at a node-local ramdisk (e.g. `/dev/shm/$USER`) significantly accelerates cache construction on machines where the raw data lives on a slow shared filesystem. +The `RC_EXTRACT_DIR` and `RC_CACHE_BASE` environment variables override the default raw-data and cache directories; `predict_points.py` and `score_point_predictions.py` read the same variables (see [Point-level Evaluation](#point-level-evaluation)). Pointing `RC_EXTRACT_DIR` at a node-local ramdisk (e.g. `/dev/shm/$USER`) significantly accelerates cache construction on machines where the raw data lives on a slow shared filesystem. + +Caches built before the CSV output existed can be given ground-truth point lists without re-running the classifier: + +``` +python -u ${rcRoot}/reconClassifier/backfill_xpts_csv.py /path/to/cache --output-dir /path/to/writable/dir +``` + +`--output-dir` is needed when the cache itself isn't writable (the original PKPM cache, for example, is owned by another user); `--overwrite` rewrites CSVs that already exist. ### Running transfer evaluation -`test_xpoint_transfer.py` loads the best PKPM-trained checkpoint and evaluates it on each transfer dataset, writing per-dataset and combined metrics to `transfer_eval_results/`. The path to the checkpoint is set by the `BEST_MODEL` constant near the top of the script; update it to point at your trained checkpoint before running. Both transfer caches must exist before this script is run. +`test_xpoint_transfer.py` evaluates a checkpoint on the in-domain dataset and each transfer dataset using the pixel-level metrics described under [Model Evaluation Metrics](#model-evaluation-metrics), writing per-dataset and combined results to `transfer_eval_results/`. The caches for every dataset it evaluates must exist before it is run. It is configured through environment variables: +- `BEST_MODEL_PATH`: Checkpoint to evaluate +- `IN_DOMAIN`: Dataset treated as the in-domain reference; the others become zero-shot targets (default: `PKPM`) +- `OUTPUT_TAG`: Suffix added to output filenames so separate runs don't overwrite each other +- `THRESHOLD`: Pixel probability threshold for the pixel-level metrics (default: 0.5) +- `EDGE_MARGIN`: Diagnostic border margin in pixels (default: 0) ``` +BEST_MODEL_PATH=/path/to/checkpoints/best_model.pt \ python -u ${rcRoot}/reconClassifier/test_xpoint_transfer.py ``` +## Point-level Evaluation + +X-points are sparse (tens per 1024² frame), so the main evaluation works on detected points rather than pixels. It has two steps and works with any trained checkpoint: + +1. **Predict.** `predict_points.py` runs the checkpoint on every cached frame of each dataset, reduces the confidence map to one peak per connected region above `--threshold`, and writes `//{N}_xpts.csv` with `row,col,confidence` columns. +2. **Score.** `score_point_predictions.py` matches the predicted points one-to-one against the ground-truth `{N}_xpts.csv` files within `--radius` pixels (greedy, highest confidence first). A matched prediction is a true positive, an unmatched prediction is a false positive, and an unmatched ground-truth X-point is a false negative. It prints F1, precision, recall, TP, FP, and FN for each dataset. + +``` +python -u ${rcRoot}/reconClassifier/predict_points.py \ +--checkpoint /path/to/checkpoints/best_model.pt \ +--datasets PKPM 5M 10M PKPMv2 \ +--output-root ./predictions \ +--base-channels 64 + +python -u ${rcRoot}/reconClassifier/score_point_predictions.py \ +--predictions-root ./predictions \ +--datasets PKPM 5M 10M PKPMv2 \ +--radius 5.0 \ +--json-out ./predictions/score_summary.json +``` + +`predict_points.py` options: +- `--checkpoint`: Trained model weights (e.g. `checkpoints/best_model.pt`) +- `--datasets`: Any of `PKPM`, `5M`, `10M`, `PKPMv2` (default: `PKPM 5M 10M`) +- `--output-root`: Base directory for the per-dataset prediction CSVs +- `--threshold`: Minimum confidence for a peak to count as a detection (default: 0.3) +- `--max-components`: If a frame's thresholded confidence map splits into more regions than this, no peaks are returned for it, since that indicates an undertrained model rather than real detections (default: 20000) +- `--base-channels`: Must match the value the checkpoint was trained with (default: 64) + +`score_point_predictions.py` options: +- `--predictions-root`: Directory with one subdirectory of prediction CSVs per dataset +- `--datasets`: Datasets to score (default: `PKPM 5M 10M`) +- `--radius`: Matching radius in pixels (default: 5.0) +- `--label`: Name shown in the summary table +- `--json-out`: Optional JSON file with the global results and a per-frame breakdown + +**Held-out numbers.** Both scripts process every cached frame, so on a dataset the model was trained on, the printed F1 mixes training and held-out frames and will look better than held-out performance. Report held-out results from the validation or test frames only, for example by aggregating those frames' TP/FP/FN from the per-frame breakdown in `--json-out`. + +**Choosing the threshold.** `--threshold` sets the operating point; it is not learned by the network. Raising it drops low-confidence detections, trading recall for precision, and F1 can be sensitive to it. Choose it on validation frames and report results on frames that were not used to choose it. + +**Data locations.** Dataset paths are built from environment variables, so a launcher script can point the whole pipeline at a different environment by exporting them, just as the training launchers set `--paramFile` and `--xptCacheDir`. The defaults are the shared DeltaAI locations: + +| Variable | Controls | Read by | Default | +|---|---|---|---| +| `RC_EXTRACT_DIR` | Raw 5M / 10M / PKPMv2 data and parameter files | `run_hessian_and_build_cache.py`, `predict_points.py` | `/work/nvme/bfim/ssridhar6/mlReconnection2025` | +| `RC_CACHE_BASE` | 5M / 10M / PKPMv2 caches and all ground-truth CSVs | all three scripts | `/work/nvme/bfim/ssridhar6/mlReconnection2025/cache` | +| `RC_PKPM_V0_ROOT` | Parameter file and cache of the original PKPM run | `predict_points.py` | `/work/nvme/bfim/cwsmith/mlReconnection2025/1024Res_v0` | + +Ground truth for every dataset, including the original PKPM run, is read from `$RC_CACHE_BASE//{N}_xpts.csv`; use `backfill_xpts_csv.py --output-dir` to create these for a cache you can't write to. + ## Resuming Development Work The following commands should be run on `checkers` **every time you create a new shell** to resume work in the existing virtual environment. @@ -217,7 +341,7 @@ source pgkyl/bin/activate ## Model Evaluation Metrics -The model evaluation system measures how well the classifier identifies X-points (magnetic reconnection sites) by treating it as a pixel-level binary classification problem. +During training, `XPointMLTest.py` measures how well the classifier identifies X-points (magnetic reconnection sites) by treating it as a pixel-level binary classification problem, and writes the metrics described here. These pixel metrics are most meaningful with `--targetType binary`; for the Gaussian heatmap model, use the [point-level evaluation](#point-level-evaluation), which is the metric used for reported results. ### Key Metrics @@ -260,4 +384,4 @@ For reconnection studies: - **Precision affects analysis**: False positives corrupt downstream calculations - **IoU indicates localization**: Poor IoU means inaccurate X-point positions -The model uses a 9×9 pixel expansion around X-points to account for localization uncertainty while still requiring accurate region identification. \ No newline at end of file +With `--targetType binary`, the target marks a 9×9 pixel block around each X-point to account for localization uncertainty while still requiring accurate region identification. With `--targetType gaussian`, each X-point is a Gaussian peak, and the localization tolerance is set by the point-level matching radius instead. \ No newline at end of file