sage.factory.callbacks

Training callbacks — objects that hook into SageVanillaTraining’s loop.

A callback can transform the per-batch context (e.g. inject non-astrophysical samples) and/or run work at epoch boundaries (e.g. hard-noise mining). Every hook is a no-op by default, so a trainer constructed with no callbacks behaves exactly like plain vanilla training.

Classes

Callback	Base training callback. Subclass and override only the hooks you need.
MaskingCallback	Assemble the 4-class consistency batch with a NonAstrophysicalMasker.
HardMiningCallback	Per-epoch CMA-MAE hard-noise mining, as a training callback.

Module Contents

class Callback

Base training callback. Subclass and override only the hooks you need.

Hooks

on_sample(ctx, trainer)

Called once per batch after the batch is assembled (post signal injection) and before preprocessing. ctx is a plain dict the trainer threads through the iteration — read/write its tensors (ctx['x'], ctx['targets'] and any extras a callback adds, e.g. ctx['per_det_mask']). Mutate ctx in place.

on_epoch_end(nepoch, trainer)

Called once after an epoch’s training iterations complete (e.g. to mine hard noise and push it to the sampler).

trainer is the SageVanillaTraining instance, exposing model, noise_sampler, signal_sampler, processor, etc.

on_sample(ctx, trainer)

on_epoch_end(nepoch, trainer)

class MaskingCallback(masker)

Bases: Callback

Assemble the 4-class consistency batch with a NonAstrophysicalMasker.

Replaces the trainer’s default injection (sets ctx['x']): the extra signal-pool injections are decohered into non-astrophysical (class-0) pairs dropped into noise slots, producing coherent (class 1) + non-astro (class 0) + pure-noise slots, plus a per-detector supervision mask. Sets:

• ctx['x'] — assembled strain (noise + injections)

• ctx['targets'] — full width [pe..., class, tc_det..., mc_det...]

• ctx['per_det_mask']— (B, D) per-detector supervision mask

Requires a signal sampler built with append_per_det_targets=True and an extra_batch pool (so signal_targets carries the per-detector columns and there are > S signals to decohere). With no masker / no extra pool it reduces to the 2-class coherent-signal vs noise regime.

masker

on_sample(ctx, trainer)

class HardMiningCallback(keep_threshold_raw=None, keep_threshold_sigmoided=None, hard_bias_prob=0.5, warmup_epochs=1, mine_iters=200, hard_dataset_dir=None, max_total_samples=30000000, descriptor_dim=8, n_cells=1024, learning_rate=0.1, n_emitters=1, emitter_batch_size=36, n_warmup=2048, mine_seed=None)

Bases: Callback

Per-epoch CMA-MAE hard-noise mining, as a training callback.

After each epoch (past warmup_epochs) it mines hard noise with CMAMAEMiner, accumulates the windows across epochs, and pushes them to the noise sampler via set_hard_dataset (so subsequent batches draw hard windows with probability hard_bias_prob). Requires a noise sampler with set_hard_dataset and a model exposing the opt-in embedding (consistency model) or get_ranking_statistic.

The “how hard is hard enough to keep” bar is set with either keep_threshold_raw (a raw detection logit) or keep_threshold_sigmoided (the same bar as a probability in (0, 1)); the raw value wins if both are given, and if neither is set every mined window is kept.

pyribs / the miner are imported lazily on first use, so merely importing this module — and pure vanilla training — never requires pyribs. The miner is built lazily from the trainer’s graph on the first mining pass.

keep_threshold_raw = None

keep_threshold_sigmoided = None

hard_bias_prob

warmup_epochs = 1

mine_iters = 200

hard_dataset_dir = None

max_total_samples = 30000000

descriptor_dim = 8

n_cells = 1024

learning_rate

n_emitters = 1

emitter_batch_size = 36

n_warmup = 2048

mine_seed = None

on_epoch_end(nepoch, trainer)