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 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/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/optuna_tuner_gaussian.py b/optuna_tuner_gaussian.py new file mode 100644 index 0000000..52dc52b --- /dev/null +++ b/optuna_tuner_gaussian.py @@ -0,0 +1,695 @@ +""" +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, thresholds=(0.3,), radius=5.0): + """Point-level F1 swept over a threshold grid; returns (best_f1, best_threshold).""" + model.eval() + 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"] + 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() + gt = np.load(Path(cache_dir) / f"{fnum}_xpts.npy") + # 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() + 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): + """ + 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 + best_threshold = None + 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, thr = compute_point_f1( + model, val_dataset, args.xptCacheDir, device, + use_amp, amp_dtype, + thresholds=args.threshold_grid, radius=args.match_radius, + ) + if f1 > best_f1: + best_f1, best_threshold = f1, thr + 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("best_threshold", best_threshold) + 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("--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)") + + # --- 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)", + ) + + 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): + """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} {'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} " + 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 diff --git a/point_metrics.py b/point_metrics.py new file mode 100644 index 0000000..7f83056 --- /dev/null +++ b/point_metrics.py @@ -0,0 +1,203 @@ +"""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 +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 +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 diff --git a/predict_points.py b/predict_points.py new file mode 100644 index 0000000..e011e99 --- /dev/null +++ b/predict_points.py @@ -0,0 +1,132 @@ +#!/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 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 + + +# 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": str(PKPM_V0_ROOT / "pkpm_2d_turb_p2-params.txt"), + "cache_dir": str(PKPM_V0_ROOT / "cache04082025"), + }, + "5M": { + "param_path": str(EXTRACT_DIR / "5M" / "rt_5M_2d_turb_local-params.txt"), + "cache_dir": str(CACHE_BASE / "5M"), + }, + "10M": { + "param_path": str(EXTRACT_DIR / "10M" / "10M" / "rt_10M_2d_turb_local-params.txt"), + "cache_dir": str(CACHE_BASE / "10M"), + }, + "PKPMv2": { + "param_path": str(EXTRACT_DIR / "1024Res_v2" / "rt_pkpm_2d_turb_p1-params.txt"), + "cache_dir": str(CACHE_BASE / "PKPMv2"), + }, +} + + +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("--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() + + 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}, max_components={args.max_components}") + + 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, + 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 " + f"({n_predicted} predicted points total)") + + +if __name__ == "__main__": + main() diff --git a/run_hessian_and_build_cache.py b/run_hessian_and_build_cache.py index af97339..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", + }, } @@ -78,6 +82,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,13 +109,14 @@ 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 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, @@ -127,6 +147,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 diff --git a/score_point_predictions.py b/score_point_predictions.py new file mode 100644 index 0000000..b714321 --- /dev/null +++ b/score_point_predictions.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +"""Compare predicted-points CSVs against ground-truth-points CSVs at the point level.""" + +import argparse +import json +import os +from pathlib import Path + +from point_metrics import ( + evaluate_point_predictions, + load_xpts_csvs_for_frames, +) + + +# 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": str(CACHE_BASE / "PKPM"), + "5M": str(CACHE_BASE / "5M"), + "10M": str(CACHE_BASE / "10M"), + "PKPMv2": str(CACHE_BASE / "PKPMv2"), +} + + +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() 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