sage.data.noise.real_noise

Classes

HDF5SingleNoiseSampler	Duration-weighted random noise sampler from GW noise datasets.
MemmapSingleNoiseSampler	Duration-weighted random noise sampler from GW noise datasets.
MemmapNoiseSampler	GPU-resident batch sampler for monolithic .bin noise files with

Module Contents

class HDF5SingleNoiseSampler(source)

Duration-weighted random noise sampler from GW noise datasets.

Supports: - single monolithic HDF5 file - directory of monolithic HDF5 files

Sampling: - segment chosen ∝ usable duration - random contiguous slice returned

Parameters:

source (Union[str, pathlib.Path]) –

• path to monolithic HDF5 file

• OR directory containing multiple monolithic files

files: List[h5py.File] = []

segments = []

seg_lengths = []

sample(requested_nsamples)

Draw a random noise slice.

Parameters:

requested_nsamples (int) – Total samples required (already includes corruption padding)

Returns:

np.ndarray

Return type:

numpy.ndarray

close()

Close all open HDF5 file handles.

class MemmapSingleNoiseSampler(source, return_tensor=False, tensor_dtype=torch.float32)

Duration-weighted random noise sampler from GW noise datasets.

Supports: - single monolithic .bin file with sidecar *_segments.json

Sampling: - segment chosen ∝ usable duration - random contiguous slice returned

Parameters:

source (Union[str, pathlib.Path]) –

• path to monolithic .bin file

bin_file

return_tensor = False

tensor_dtype = Ellipsis

dtype

mm

segments

seg_lengths

sample(requested_nsamples, rng=None)

Draw a random noise slice.

Parameters:

requested_nsamples (int) – Total samples required (already includes corruption padding)

Returns:

np.ndarray of shape (requested_nsamples,)

Return type:

numpy.ndarray

close()

Release memmap

class MemmapNoiseSampler(postprocess_fn=None, prefetch=3, seed=None, training=True, hard_dataset_path=None, hard_dataset_dir=None, hard_bias_prob=0.0, num_read_workers=16)

Bases: torch.nn.Module

GPU-resident batch sampler for monolithic .bin noise files with asynchronous prefetching.

Each detector’s noise data is stored as a flat binary file (float32 or float64) accompanied by a sidecar *_segments.json file that records the GPS start/end times, sample indices, and per-segment metadata for every contiguous data segment.

Sampling is duration-weighted: longer segments are proportionally more likely to be drawn so that no useful data is wasted. Within each chosen segment, a uniformly random start offset is selected.

An async daemon thread continuously fills a bounded queue.Queue with pre-processed FD batches so that the GPU is never starved waiting for disk I/O. Callers consume batches through sample_batch(), which performs a thread-safe queue.get().

Warning

_read_batch() and _sample_starts_batch() access a non-thread-safe NumPy RNG and must not be called from the main thread while the prefetch daemon is running. Always use sample_batch() (queue pop) to consume batches safely.

Parameters:

• postprocess_fn (callable or None) – Applied to (batch_td, segment_ids) to convert from TD to FD. If None, a plain torch.fft.rfft is used. The typical choice is RecolourPostprocess.

• prefetch (int) – Maximum number of batches held in the prefetch queue.

• seed (int or None) – Seed for the internal NumPy RNG (for reproducibility).

• training (bool) – If True, reads from data_cfg.training_noise_files; otherwise from data_cfg.validation_noise_files.

• hard_bias_prob (float)

• num_read_workers (int)

GRAPH_READY

False — the prefetch queue pop cannot be traced by torch.compile.

Type:

bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

GRAPH_READY = False

seq_len

device

prefetch = 3

n_detectors

batch_size

postprocess_fn = None

mmaps = []

seg_index = []

segment_probs = []

dtypes = []

noise_target

rng

num_read_workers = 16

queue

set_hard_dataset(dataset, hard_bias_prob=None, epoch=None, save_dir=None)

Replace the hard-noise dataset and optionally persist it to disk.

Thread-safe: can be called from the main thread while the prefetch daemon is running. The change takes effect on the next batch the prefetch thread starts reading.

Versioned saving

When epoch and save_dir are both given, the dataset is saved as {save_dir}/hard_noise_epoch_{epoch:04d}.npz. Files from earlier epochs are left untouched. If two miners both fire at the same epoch (explorer then refiner), the second call overwrites the first for that epoch number — only the refiner’s file survives on disk, though the explorer’s GPS findings remain in the shared bank. On the next training restart, MemmapNoiseSampler(hard_dataset_dir=save_dir) automatically picks up the latest epoch file.

param dataset:

A freshly mined dataset. Pass None to disable biasing.

type dataset:

StartTimeDataset or None

param hard_bias_prob:

New bias probability. None keeps the existing value.

type hard_bias_prob:

float or None

param epoch:

Training epoch at which this dataset was mined. Used in the filename when save_dir is also provided.

type epoch:

int or None

param save_dir:

Directory in which to save the versioned file.

type save_dir:

str or Path or None

Parameters:

• hard_bias_prob (float | None)

• epoch (int | None)

• save_dir (str | pathlib.Path | None)

sample_batch()

Return a GPU batch. Starts async prefetching if first call.

forward()

shutdown()

Stop the prefetch thread and release memmaps.

Idempotent and safe to call explicitly, from atexit, or twice.