BCN001/JWST-transient-finder

DenseNet169 Ensemble — JWST Transient Detection (F150W)

A 4-model DenseNet169 ensemble that classifies 64×64 pixel cutouts from JWST NIRCam F150W difference imaging as transient or non-transient. Trained as part of a self-learning pipeline that generates its own labeled data by injecting synthetic point sources (via the instrument PSF) into real F150W difference images. This is the F150W-specific ensemble; other filters (if trained separately) are hosted under their own subfolder in this repo.

Quick Start (inference_pipeline.py)

Single-file, resumable inference pipeline. Point it at a folder of FITS difference files — it reads each file's FILTER header, downloads the matching ensemble from this repo's <FILTER>/ subfolder (strictly that filter — no falling back to another filter's weights), finds candidate peaks, and scores them.

pip install torch torchvision timm astropy numpy pandas scipy pillow tqdm huggingface_hub

fits_files/
    new_diff_file_1.fits
    new_diff_file_2.fits

python inference_pipeline.py

Each run only processes files not yet processed (tracked via a .done marker per file), so it's safe to re-run after dropping new files into fits_files/. Use --force to reprocess everything.

Outputs go to output/<fits_stem>/, one folder per input FITS file:

• predictions.csv — every scored cutout (peak location, sigma, per-model and ensemble mean/std probability, prediction, confidence)
• positives/*.png — cutouts predicted as transient

Downloaded model weights are cached under models/<FILTER>/. Useful flags: --sigma-threshold, --confidence-threshold, --repo-id, --input-dir, --output-dir (--help for the full list).

Legacy multi-script pipeline

test_pipeline.py orchestrates find_peaks_above_k_sigma_test.py + testing_script.py against a local test_directory/ → ML_results/, using local DenseNet169_Ensemble_Model{1-4}_best.pth files instead of fetching from Hugging Face. Kept for reference/offline use; inference_pipeline.py above is the actively maintained entry point.

Model Details

• Architecture: DenseNet169 (growth rate 32, block config (6, 12, 32, 32)), trained from scratch (no pretraining)
• Input: 64×64, 3-channel, ZScale-normalized (astropy ZScaleInterval) cutouts from FITS difference images
• Output: single sigmoid unit — probability of being a real transient
• Ensemble size: 4 independently trained models; predictions are averaged, with the standard deviation across models used as a disagreement/uncertainty signal
• Framework: PyTorch
• Model files: F150W/DenseNet169_Ensemble_Model{1-4}_best.pth

The pipeline was originally designed to also train a 2-model DeiT (vision transformer) ensemble alongside the DenseNet models. Only the DenseNet169 ensemble was trained for this release (DeiT model count was 0 in the training run), so this card and repo cover the DenseNet ensemble only.

Intended Use

Given a candidate source cutout from a JWST NIRCam difference image (science minus reference epoch), predict whether it is a genuine astrophysical transient (e.g. supernova, variable star) versus a subtraction artifact, cosmic ray, or other spurious peak. Intended as a triage/ranking step ahead of human vetting, not a fully autonomous discovery pipeline.

Out of scope: classifying transient type, non-JWST instruments/pixel scales without re-validation, or images not normalized with ZScale the same way as training.

Training Data

Training data is self-generated from a small set of real JWST NIRCam F150W difference images (no external labeled catalog):

• Positives: synthetic point sources injected at random locations using the instrument PSF, at signal-to-noise ratios spanning SNR 3–10 (psf_injection_script.py).
• Negatives: a mix of quiet/uniform background regions (find_non_peaks_64.py) and real detected peaks that are not injected sources (find_peaks_above_k_sigma_training.py), so the model also learns to reject real (non-transient) astrophysical sources.
• Split: 80% train / 10% validation / 10% test (splitfolders, seed 42) → 14,565 train / 1,820 val / 1,823 test cutouts, roughly balanced between positives and negatives.
• Normalization: ZScale, matching standard astronomical display conventions, scaled to [0, 1].

Difference images used for training were JWST NIRCam F150W filter data. The most recent inference/testing run on new science data was also on F150W difference images; an earlier run additionally evaluated the model on F200W data as a cross-filter check (see Limitations — not a training filter).

Training Procedure

• Loss: binary cross-entropy (BCELoss)
• Optimizer: Adam, lr = 1e-4, weight_decay = 1e-4
• Batch size: 16
• Epochs: up to 10 per model, with best-validation-accuracy checkpointing
• Each of the 4 ensemble members is trained independently (same architecture/data pipeline, different initialization) to provide ensemble diversity/uncertainty estimates
• Hardware: single NVIDIA RTX 4090 Laptop GPU

Evaluation Results

Best validation accuracy per ensemble member (10-epoch training run):

Model	Best Val. Accuracy
DenseNet169 Model 1	100.00%
DenseNet169 Model 2	99.95%
DenseNet169 Model 3	99.89%
DenseNet169 Model 4	99.95%
Ensemble average	99.95%

These figures are on the held-out validation split of the self-generated (PSF-injection) dataset, which is easier than real-world discovery — see Limitations.

How to Use Directly (without inference_pipeline.py)

import torch
from pathlib import Path

# see testing_script.py in this repo for the DenseNet definition
# (densenet169()) and the ZScale preprocessing used to turn a FITS
# cutout into a normalized 64x64x3 tensor.

models = []
for path in sorted(Path(".").glob("DenseNet169_Ensemble_Model*_best.pth")):
    model = densenet169(num_classes=1)
    checkpoint = torch.load(path, map_location="cpu")
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    models.append(model)

with torch.no_grad():
    probs = torch.stack([m(image_tensor) for m in models])
    mean_prob = probs.mean(0)
    std_prob = probs.std(0)

Limitations and Biases

• Positive examples are synthetic PSF injections, not confirmed real transients — the model has not been validated against a catalog of known real transients, and may not generalize to transient morphologies unlike a clean PSF (e.g. blended, saturated, or host-contaminated sources).
• Trained only on JWST NIRCam F150W difference images from a limited set of source fields; performance on other filters (F200W included, despite the earlier cross-filter check), instruments, or pixel scales is unverified and should be re-validated before trusting predictions there.
• The near-100% validation accuracy reflects the relative ease of the self-supervised task (clean injected point sources vs. background/real peaks) and should not be read as real-world discovery precision/recall. Always inspect flagged detections visually before treating them as candidate transients.
• No independent test against human-vetted or spectroscopically confirmed transients has been performed.

Citation / Provenance

Produced by the self-learning-pipeline-transient-detection project. See the parent repository's training pipeline (train_pipeline.py, training_script.py) for full training code and data-generation details.