`rfgen.dataset`¶

Consumer-side abstractions for reading stored records into training pipelines. This module covers two concerns: samplers (rfgen.dataset.samplers) control which sample_id values are yielded and in what order; collators (rfgen.dataset.collators) turn a list of Record objects into batched tensors shaped for a specific ML task.

These abstractions are consumer-side, not part of the producer pipeline. The storage layer stays framework-neutral. The optional rfgen.dataset.collators and rfgen.dataset.RfgenTorchDataset surfaces introduce the PyTorch dependency behind the rfgen[torch] extra; importing sampler-only code does not require PyTorch. from rfgen.dataset import * exports the sampler-only surface. The torch-backed BaseCollator, DetectionCollator, and RfgenTorchDataset symbols remain available through explicit imports or lazy rfgen.dataset attribute access when the torch extra is installed.

Warning

Pre-implementation. This page describes proposed contracts. Class signatures, parameter types, config field names, and behavior are subject to change before code lands. Once implementation exists, content here will be regenerated from docstrings or sourced from running tests.

Class index¶

Class	Kind	Module	Notes
BaseSampler	abc	`rfgen.dataset.samplers`	Yields `sample_id` strings in the desired order
BaseCollator	abc	`rfgen.dataset.collators`	Turns a list of Record objects into batched tensors
SequentialSampler	concrete	`rfgen.dataset.samplers`	Yields in `sample_idx` order; default
RandomSampler	concrete	`rfgen.dataset.samplers`	Deterministic random order; one pass
ClassBalancedSampler	concrete	`rfgen.dataset.samplers`	Equal counts per `class_name`; RadioML benchmark parity
SnrStratifiedSampler	concrete	`rfgen.dataset.samplers`	Sample by SNR bucket
DensityStratifiedSampler	concrete	`rfgen.dataset.samplers`	Sample by emitter count per scene
WeightedSampler	concrete	`rfgen.dataset.samplers`	Custom per-sample weights
DetectionCollator	concrete	`rfgen.dataset.collators`	YOLO-format: `{iq, bboxes, class_ids, emitter_counts}`
RfgenTorchDataset	concrete	`rfgen.dataset`	Optional `IterableDataset` wrapper behind `rfgen[torch]`
SegmentationCollator	concrete	`rfgen.dataset.collators`	`{iq, seg_mask}` with spectrogram segmentation
ClassificationCollator	concrete	`rfgen.dataset.collators`	`{iq, class_id}` single-label per record
AttributeRegressionCollator	concrete	`rfgen.dataset.collators`	PAES-style `{iq, attribute_targets}`
ContrastiveCollator	concrete	`rfgen.dataset.collators`	SSL pairs `{iq_anchor, iq_positive}`
MultiTaskCollator	concrete	`rfgen.dataset.collators`	Multi-head training; dict of sub-collator outputs
BaseAugmentation	abc	`rfgen.dataset.augmentations`	IQ augmentation contract; passed to ContrastiveCollator to produce positive pairs

Samplers¶

`class rfgen.dataset.samplers.BaseSampler`¶

class BaseSampler(ABC):
    @abstractmethod
    def __iter__(self) -> Iterator[str]:
        """Yields rfgen content-hash sample_ids in the desired order."""

    @abstractmethod
    def __len__(self) -> int:
        """Total number of sample_ids this sampler will yield in one pass."""

    def bind(self, source: SampleIdSource) -> BaseSampler:
        """Return a sampler of the same kind bound to a sample-id source."""

    def partition(self, *, worker_id: int, num_workers: int) -> Iterator[str]:
        """Yield this sampler's share of ids for one DataLoader worker."""

class SupportsBindableSampleIdSampler(SupportsSampleIdSampler, Protocol):
    def bind(self, source: SampleIdSource) -> SupportsSampleIdSampler: ...

class SupportsPartitionedSampleIdSampler(SupportsSampleIdSampler, Protocol):
    def partition(self, *, worker_id: int, num_workers: int) -> Iterator[str]: ...

def sampler_has_scalable_partition(
    sampler: SupportsSampleIdSampler,
) -> TypeGuard[SupportsPartitionedSampleIdSampler]:
    """Return whether a sampler exposes a concrete multi-worker partition hook."""

def bind_sampler_to_source(
    sampler: SupportsSampleIdSampler,
    source: SampleIdSource,
) -> SupportsSampleIdSampler:
    """Return the sampler bound to source when its bind hook requires it."""

Warning

Abstract base for all samplers. Yields sample_id strings (content-hash keys as produced by StoreHandle.write) in the implementation-defined order. Consumers call __iter__ to stream sample_id values and look up each one via StoreHandle.read. Source-aware samplers expose bind(source) so a dataset adapter can provide a store handle or sample-id sequence before iteration. Third-party samplers may implement only __iter__ and __len__; RfgenTorchDataset accepts those duck-typed samplers when DataLoader.num_workers <= 1. Multi-worker loading requires either a real partition(worker_id, num_workers) implementation or a source-aware bind(source) path that returns a sampler with that concrete partition hook. The inherited BaseSampler.partition(...) fallback and duck-typed __iter__ fallback are single-worker conveniences. Adapters use the typed SupportsBindableSampleIdSampler and SupportsPartitionedSampleIdSampler capability protocols through bind_sampler_to_source(...) and sampler_has_scalable_partition(...). This avoids having each worker scan or snapshot the full sampler stream independently.

All concrete samplers are deterministic given seed and handle state.

Multi-RX sampling note¶

With RecordAxis.PER_RX, each scene becomes N records (one per RX antenna). Class-balanced sampling over records can over-represent emitters that appear in multi-RX scenes: a 4-RX scene with emitter class A contributes 4 records to the class-A bucket, whereas a 1-RX scene contributes 1. Downstream training pipelines that require strict scene-balance should either use scene_balanced=True in ClassBalancedSampler or implement a custom BaseSampler subclass that de-duplicates by scene_id.

`class rfgen.dataset.samplers.SequentialSampler`¶

class SequentialSampler(BaseSampler):
    def __init__(self, source: SampleIdSource | None = None): ...

Warning

Default sampler. Yields all sample_id values in sorted deterministic order. One complete pass per __iter__ call. If constructed without source, it must be bound with bind(source) before iteration.

Constructor parameters¶

Name	Type	Description
`source`	`SampleIdSource \| None`	Store handle or sample-id sequence to sample from; `None` means the sampler is unbound until `bind(source)` is called

`class rfgen.dataset.samplers.RandomSampler`¶

class RandomSampler(BaseSampler):
    def __init__(self, seed: int, source: SampleIdSource | None = None): ...

Warning

Deterministic random shuffle over sorted sample_id values. Same seed and source state produces the same permutation. One complete pass per __iter__ call. Calling bind(source) scans the source once, stores the sorted sample-id snapshot inside the returned sampler, and permutes that snapshot with numpy.random.default_rng(seed).permutation(...).

Constructor parameters¶

Name	Type	Description
`seed`	int	RNG seed; same seed produces the same permutation
`source`	`SampleIdSource \| None`	Store handle or sample-id sequence to sample from; `None` means the sampler is unbound until `bind(source)` is called

`class rfgen.dataset.samplers.ClassBalancedSampler`¶

class ClassBalancedSampler(BaseSampler):
    def __init__(
        self,
        handle: StoreHandle,
        *,
        seed: int,
        scene_balanced: bool = False,
    ): ...

Warning

Yields equal counts per class_name. Required for RadioML benchmark parity, where each modulation class should appear the same number of times per epoch.

When scene_balanced=True, the sampler de-duplicates by scene_id before balancing, so multi-RX records from the same scene count as one entry. This avoids over-representing emitters in multi-RX scenes. See Multi-RX sampling note above.

Constructor parameters¶

Name	Type	Default	Description
`handle`	StoreHandle	–	Open store session to sample from
`seed`	int	–	RNG seed for within-class shuffling
`scene_balanced`	bool	`False`	When `True`, de-duplicate by `scene_id` before balancing to avoid multi-RX over-representation

`class rfgen.dataset.samplers.SnrStratifiedSampler`¶

class SnrStratifiedSampler(BaseSampler):
    def __init__(
        self,
        handle: StoreHandle,
        *,
        snr_buckets: list[tuple[float, float]],
        seed: int,
    ): ...

Warning

Yields samples grouped by SNR bucket (in dB). Useful for SNR-vs-accuracy evaluation curves.

Constructor parameters¶

Name	Type	Default	Description
`handle`	StoreHandle	–	Open store session
`snr_buckets`	`list[tuple[float, float]]`	–	List of `(low_db, high_db)` half-open intervals. Records whose `metadata.snr_db` falls in bucket `i` are assigned to group `i`
`seed`	int	–	RNG seed for within-bucket shuffling

`class rfgen.dataset.samplers.DensityStratifiedSampler`¶

class DensityStratifiedSampler(BaseSampler):
    def __init__(
        self,
        handle: StoreHandle,
        *,
        density_buckets: list[tuple[int, int]],
        seed: int,
    ): ...

Warning

Yields samples grouped by emitter count per scene (signal density). Useful for studying model performance as a function of scene complexity.

Constructor parameters¶

Name	Type	Default	Description
`handle`	StoreHandle	–	Open store session
`density_buckets`	`list[tuple[int, int]]`	–	List of `(min_emitters, max_emitters)` half-open intervals. Records whose `scene_metadata.num_emitters` falls in bucket `i` are assigned to group `i`
`seed`	int	–	RNG seed for within-bucket shuffling

`class rfgen.dataset.samplers.WeightedSampler`¶

class WeightedSampler(BaseSampler):
    def __init__(
        self,
        handle: StoreHandle,
        *,
        weights: dict[str, float],
        seed: int,
    ): ...

Warning

Custom per-sample weights. weights maps sample_id to a positive float; samples with higher weight are drawn more often. Useful for importance sampling and curriculum learning.

Constructor parameters¶

Name	Type	Default	Description
`handle`	StoreHandle	–	Open store session
`weights`	`dict[str, float]`	–	Map from `sample_id` to sampling weight. Keys not present in the store are silently ignored
`seed`	int	–	RNG seed for the weighted draw

Collators¶

`class rfgen.dataset.collators.BaseCollator`¶

class BaseCollator(ABC):
    @abstractmethod
    def collate(self, records: Sequence[Record]) -> dict[str, torch.Tensor | list[torch.Tensor]]:
        """Turn a list of Records into a batched dict of tensors.

        Output dict keys are task-specific. Values may be tensors or per-record
        tensor lists for variable-length targets.
        """

Warning

Abstract base for all collators. Turns a list of Record objects into a dict[str, torch.Tensor | list[torch.Tensor]] ready for a training loop’s loss function.

Output dict keys are task-specific (e.g. "iq", "bboxes", "class_id"). Tensor values follow the canonical IQ dtype policy: torch.float32 for IQ ((B, 2, N) or (B, num_rx, 2, N) shape), torch.int64 or torch.float32 for labels depending on the task. Detection targets are variable-length per-record lists rather than padded tensors.

Non-torch consumers may write their own collator returning NumPy arrays or JAX arrays. The ABC does not constrain the return type at runtime; the torch.Tensor annotation is a contract for the shipped concrete subclasses.

`class rfgen.dataset.collators.DetectionCollator`¶

class DetectionCollator(BaseCollator):
    def __init__(self): ...

Warning

YOLO-format detection collator. Bounding boxes and class IDs stay as per-record lists because emitter counts vary by record.

Output dict:

Key	Shape	Dtype	Description
`iq`	`(B, 2, N)` or `(B, num_rx, 2, N)`	float32	Batched IQ
`bboxes`	`list[Tensor[num_boxes_i, 4]]`	float32	One tensor per record with normalized YOLO boxes `(cx, cy, w, h)`
`class_ids`	`list[Tensor[num_boxes_i]]`	int64	One tensor per record with class index per bbox
`emitter_counts`	`(B,)`	int64	Number of bbox-bearing emitters per record

`class rfgen.dataset.RfgenTorchDataset`¶

from rfgen.dataset import RfgenTorchDataset

class RfgenTorchDataset(torch.utils.data.IterableDataset):
    def __init__(
        self,
        store: BaseStore,
        *,
        sampler: BaseSampler | SupportsSampleIdSampler | None = None,
        collator: BaseCollator | None = None,
    ): ...

Optional PyTorch bridge behind the rfgen[torch] extra. Each DataLoader worker opens its own StoreHandle from the provided BaseStore, binds the sampler to that handle, reads records by sample_id, and yields either records or collator-produced batches. Import it from rfgen.dataset. Third-party SupportsSampleIdSampler objects only need __iter__() and __len__() for single-worker loading. Multi-worker loading also requires a concrete partition(worker_id, num_workers) hook, or a public bind(source) method that returns such a sampler.

`class rfgen.dataset.collators.SegmentationCollator`¶

class SegmentationCollator(BaseCollator):
    def __init__(self, *, n_fft: int, hop: int): ...

Warning

Time-frequency segmentation collator. Computes an STFT spectrogram and returns the per-bin segmentation mask.

Output dict:

Key	Shape	Dtype	Description
`iq`	`(B, 2, N)`	float32	Batched IQ
`seg_mask`	`(B, F, T)`	int16	Single-label per-bin class mask, where `F = n_fft // 2 + 1` and `T` is the number of STFT frames

For multi-label masks (at higher cost), use a custom BaseSampler subclass or wait for the multi-label segmentation milestone (planned v1+).

Constructor parameters¶

Name	Type	Description
`n_fft`	int	STFT window length in samples
`hop`	int	STFT hop length in samples

`class rfgen.dataset.collators.ClassificationCollator`¶

class ClassificationCollator(BaseCollator):
    def __init__(
        self,
        *,
        label_strategy: Literal["first_emitter", "dominant_emitter"],
    ): ...

Warning

Single-label classification collator. Assigns one class per record.

Output dict:

Key	Shape	Dtype	Description
`iq`	`(B, 2, N)`	float32	Batched IQ
`class_id`	`(B,)`	int64	Class index for the selected emitter

label_strategy controls which emitter’s class label is assigned:

"first_emitter": use component_signals[0].class_name. Deterministic; suitable for single-emitter scenes.
"dominant_emitter": use the emitter with the highest realized SNR. Suitable for multi-emitter scenes where one emitter dominates.

Notes¶

Multi-emitter records contain multiple signals; assigning a single class label is inherently lossy. For multi-emitter classification, use DetectionCollator.

`class rfgen.dataset.collators.AttributeRegressionCollator`¶

class AttributeRegressionCollator(BaseCollator):
    def __init__(self, *, attributes: list[str]): ...

Warning

PAES-style attribute regression collator. Extracts scalar attribute values from per-emitter metadata for regression tasks.

Output dict:

Key	Shape	Dtype	Description
`iq`	`(B, 2, N)`	float32	Batched IQ
`attribute_targets`	`(B, len(attributes))`	float32	Attribute values in `attributes` list order

Constructor parameters¶

Name	Type	Description
`attributes`	`list[str]`	Names of attributes to extract; e.g. `["snr_db", "cfo_hz", "bandwidth_hz"]`. Each name is looked up in `record.metadata.extras`

`class rfgen.dataset.collators.ContrastiveCollator`¶

class ContrastiveCollator(BaseCollator):
    def __init__(self, *, augmentation: BaseAugmentation | None = None): ...

Warning

Contrastive self-supervised learning collator. Produces anchor/positive pairs from one record by applying optional augmentation.

Output dict:

Key	Shape	Dtype	Description
`iq_anchor`	`(B, 2, N)`	float32	Original IQ
`iq_positive`	`(B, 2, N)`	float32	Augmented IQ (or copy if `augmentation=None`)

Constructor parameters¶

Name	Type	Default	Description
`augmentation`	`BaseAugmentation \| None`	`None`	Optional IQ augmentation; when `None`, the positive is a copy of the anchor

`class rfgen.dataset.collators.MultiTaskCollator`¶

class MultiTaskCollator(BaseCollator):
    def __init__(self, *, collators: dict[str, BaseCollator]): ...

Warning

Multi-head training collator. Runs all sub-collators and merges their outputs under namespaced keys.

Output: the union of all sub-collator output dicts, with keys prefixed by the sub-collator name from the collators dict. For example, with collators={"det": DetectionCollator(...), "cls": ClassificationCollator(...)}, the output dict has keys "det/iq", "det/bboxes", "cls/iq", "cls/class_id". Duplicate "iq" keys across sub-collators are deduplicated: only the first occurrence is kept.

Constructor parameters¶

Name	Type	Description
`collators`	`dict[str, BaseCollator]`	Map from name prefix to collator instance

Augmentations¶

`class rfgen.dataset.augmentations.BaseAugmentation`¶

class BaseAugmentation(ABC):
    @abstractmethod
    def augment(self, iq: IQ, *, rng: torch.Generator) -> IQ:
        """Apply a stochastic IQ augmentation and return the result.

        The output IQ MUST have the same shape and dtype as the input.
        All randomness MUST be drawn from `rng` so that augmentations are
        reproducible given the same generator state.
        """

Warning

Abstract base for IQ augmentations. Consumed by ContrastiveCollator to produce positive pairs for contrastive self-supervised learning: the collator calls augment(iq_anchor, rng=...) to produce iq_positive.

Kind. Abstract base class.

Abstract method: `augment`¶

@abstractmethod
def augment(self, iq: IQ, *, rng: torch.Generator) -> IQ

Apply one augmentation pass to iq and return the result.

Parameters¶

Name	Type	Description
`iq`	IQ	Input IQ tensor; shape `(2, N)` float32 (or `(num_rx, 2, N)` for multi-RX records)
`rng`	`torch.Generator`	All randomness for this augmentation drawn from here; same state yields the same output

Returns¶

An IQ tensor with the same shape and dtype as the input. The augmented tensor is a new allocation; the input is not modified in place.

Contract¶

Output shape and dtype MUST equal input shape and dtype.
All randomness MUST be drawn from rng (no random, no torch.rand without generator).
Implementations MUST NOT modify iq in place.

Example¶

import torch
from rfgen.dataset.augmentations import BaseAugmentation
from rfgen.core.types import IQ

class AdditivePhaseDrift(BaseAugmentation):
    """Adds a slow phase drift to simulate oscillator wander."""

    def __init__(self, max_drift_rad: float = 0.1) -> None:
        self.max_drift_rad = max_drift_rad

    def augment(self, iq: IQ, *, rng: torch.Generator) -> IQ:
        drift = torch.rand(1, generator=rng).item() * self.max_drift_rad
        phase = torch.linspace(0, drift, iq.shape[-1])
        cos_phase = torch.cos(phase)
        sin_phase = torch.sin(phase)
        i = iq[0] * cos_phase - iq[1] * sin_phase
        q = iq[0] * sin_phase + iq[1] * cos_phase
        return torch.stack([i, q])

Notes¶

Determinism. Pass the same rng state to reproduce an augmented pair. ContrastiveCollator derives a per-record rng from the collator seed and the record index; augmentations do not manage their own seeds.
Shape invariant. The shape-equality contract is strict; downstream loss functions (NT-Xent, SimCLR) assume anchor and positive have identical shapes.
No I/O. Augmentations operate on in-memory tensors only. They must not read files, call network services, or cache state between calls.

rfgen.dataset¶

Class index¶

Samplers¶

class rfgen.dataset.samplers.BaseSampler¶

Multi-RX sampling note¶

class rfgen.dataset.samplers.SequentialSampler¶

Constructor parameters¶

class rfgen.dataset.samplers.RandomSampler¶

Constructor parameters¶

class rfgen.dataset.samplers.ClassBalancedSampler¶

Constructor parameters¶

class rfgen.dataset.samplers.SnrStratifiedSampler¶

Constructor parameters¶

class rfgen.dataset.samplers.DensityStratifiedSampler¶

Constructor parameters¶

class rfgen.dataset.samplers.WeightedSampler¶

Constructor parameters¶

Collators¶

class rfgen.dataset.collators.BaseCollator¶

class rfgen.dataset.collators.DetectionCollator¶

class rfgen.dataset.RfgenTorchDataset¶

class rfgen.dataset.collators.SegmentationCollator¶

Constructor parameters¶

class rfgen.dataset.collators.ClassificationCollator¶

Notes¶

class rfgen.dataset.collators.AttributeRegressionCollator¶

Constructor parameters¶

class rfgen.dataset.collators.ContrastiveCollator¶

Constructor parameters¶

class rfgen.dataset.collators.MultiTaskCollator¶

Constructor parameters¶

Augmentations¶

class rfgen.dataset.augmentations.BaseAugmentation¶

Abstract method: augment¶

Parameters¶

Returns¶

Contract¶

Example¶

Notes¶

See Also¶

See Also¶

`rfgen.dataset`¶

`class rfgen.dataset.samplers.BaseSampler`¶

`class rfgen.dataset.samplers.SequentialSampler`¶

`class rfgen.dataset.samplers.RandomSampler`¶

`class rfgen.dataset.samplers.ClassBalancedSampler`¶

`class rfgen.dataset.samplers.SnrStratifiedSampler`¶

`class rfgen.dataset.samplers.DensityStratifiedSampler`¶

`class rfgen.dataset.samplers.WeightedSampler`¶

`class rfgen.dataset.collators.BaseCollator`¶

`class rfgen.dataset.collators.DetectionCollator`¶

`class rfgen.dataset.RfgenTorchDataset`¶

`class rfgen.dataset.collators.SegmentationCollator`¶

`class rfgen.dataset.collators.ClassificationCollator`¶

`class rfgen.dataset.collators.AttributeRegressionCollator`¶

`class rfgen.dataset.collators.ContrastiveCollator`¶

`class rfgen.dataset.collators.MultiTaskCollator`¶

`class rfgen.dataset.augmentations.BaseAugmentation`¶

Abstract method: `augment`¶