hyrax.models.hyrax_ts2vec

TS2Vec: universal self-supervised representation learning for time series.

An implementation of TS2Vec (Yue et al., AAAI 2022, zhihanyue/ts2vec) adapted to the Hyrax model interface. The reference implementation splits its logic between a fit() method that owns the training loop and an encode() method that owns inference; here the per-batch core of fit() lives in HyraxTs2Vec.train_batch() and the pooling core of encode() lives in HyraxTs2Vec.infer_batch(), so Hyrax’s own verbs (train, train_stream, infer, infer_stream) drive the loop.

The method is fully self-supervised: two overlapping random crops of the same series are encoded, and a hierarchical contrastive loss pulls together the representations of the same timestamp across the two views (temporal contrast) while pushing apart different instances at the same timestamp (instance contrast). No labels are required.

Input convention

TSEncoder consumes a single (batch, time, channels) float32 array in which NaN marks a missing or padded timestep, and excludes those positions from the convolution - so ragged light curves need no separate mask tensor once they reach the encoder.

build_event_sequence() produces that array from ordinary padded light-curve columns, and HyraxTs2Vec.prepare_inputs() calls it. No per-survey dataset subclass is needed: LSDBStreamDataset already collates every nested column into <column>_<subcolumn> plus a boolean <column>_<subcolumn>_mask, and the column names, band list, sequence length and normalization are read from [model.HyraxTs2Vec].

Encoding in prepare_inputs rather than in collation is deliberate. Hyrax’s handle_nans runs over every float array produced by collation - warning when data_set.nan_mode is false and overwriting NaN when it is "zero" or "quantile" - and prepare_inputs is the first hook downstream of it, so it is the earliest point at which the NaN sentinel is safe to emit.

Warning

Leave data_set.nan_mode at false for this pipeline. The raw flux column reaches handle_nans with its own NaNs still in it, and "zero" or "quantile" would overwrite a missing measurement at a valid timestep with a number indistinguishable from a real one. With false those NaNs are dropped here instead, by the validity mask. Better still, drop them upstream in LSDB and avoid the per-batch warning.

Attributes

logger
FLUX_CHANNEL
FLUX_ERR_CHANNEL
DELTA_TIME_CHANNEL
N_LEADING_CHANNELS
NORMALIZE_MODES
MAD_TO_SIGMA

Classes

SamePadConv	A 1D convolution padded so that the output keeps the input length.
ConvBlock	Two same-padded convolutions at one dilation, with a residual connection.
DilatedConvEncoder	A stack of ConvBlock whose dilation doubles at every level.
TSEncoder	Project, mask, and dilate-convolve a (batch, time, channels) series.
HierarchicalContrastiveLoss	Sum the instance and temporal contrastive losses over a max-pooling hierarchy.
EncodingSettings	The [model.HyraxTs2Vec] keys that describe the input rather than the architecture.
HyraxTs2Vec	TS2Vec, a self-supervised encoder for multivariate time series.

Functions

generate_binomial_mask(→ torch.Tensor)	Return a (batch_size, length) bool mask keeping each position with probability p.
generate_continuous_mask(→ torch.Tensor)	Return a (batch_size, length) bool mask with n contiguous spans masked out.
instance_contrastive_loss(→ torch.Tensor)	Contrast each instance against the other instances at the same timestamp.
temporal_contrastive_loss(→ torch.Tensor)	Contrast each timestamp against the other timestamps of the same instance.
take_per_row(→ torch.Tensor)	Gather num_elem consecutive timesteps from each row, starting at that row's index.
_select_columns(data, settings)	Pull the light-curve columns and the validity mask out of a collated batch.
_sort_and_trim(columns, valid, length)	Sort each row by time and fit it to exactly length timesteps.
_normalize_flux(flux, flux_err, valid, mode)	Center and rescale each object's flux, applying the same scale to its error.
_band_indices(band, valid, settings)	Map band values to one-hot column offsets, returning -1 for unrecognized ones.
build_event_sequence(→ numpy.ndarray)	Turn padded light-curve columns into the event sequence TSEncoder consumes.

Module Contents

logger

class SamePadConv(in_channels, out_channels, kernel_size, dilation=1, groups=1)

Bases: torch.nn.Module

A 1D convolution padded so that the output keeps the input length.

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

receptive_field

conv

remove = 1

forward(x)

class ConvBlock(in_channels, out_channels, kernel_size, dilation, final=False)

Bases: torch.nn.Module

Two same-padded convolutions at one dilation, with a residual connection.

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

conv1

conv2

projector

forward(x)

class DilatedConvEncoder(in_channels, channels, kernel_size)

Bases: torch.nn.Module

A stack of ConvBlock whose dilation doubles at every level.

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

net

forward(x)

generate_binomial_mask(batch_size: int, length: int, p: float = 0.5) → torch.Tensor

Return a (batch_size, length) bool mask keeping each position with probability p.

generate_continuous_mask(batch_size: int, length: int, n: int | float = 5, mask_length: int | float = 0.1) → torch.Tensor

Return a (batch_size, length) bool mask with n contiguous spans masked out.

n and mask_length may be floats, in which case they are read as fractions of length.

class TSEncoder(input_dims, output_dims, hidden_dims=64, depth=10, mask_mode='binomial')

Bases: torch.nn.Module

Project, mask, and dilate-convolve a (batch, time, channels) series.

NaN timesteps are detected and folded into the mask, so missing observations and padding contribute nothing to the convolution.

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

MASK_MODES = ('binomial', 'continuous', 'all_true', 'all_false', 'mask_last')

input_dims

output_dims

hidden_dims = 64

mask_mode = 'binomial'

input_fc

feature_extractor

repr_dropout

_build_mask(mask, batch_size, length, device)

Resolve a mask mode name into a (batch_size, length) bool tensor.

forward(x, mask=None)

Encode (batch, time, input_dims) into (batch, time, output_dims).

instance_contrastive_loss(z1: torch.Tensor, z2: torch.Tensor) → torch.Tensor

Contrast each instance against the other instances at the same timestamp.

Returns zero for a batch of one, where there are no negatives to contrast against.

temporal_contrastive_loss(z1: torch.Tensor, z2: torch.Tensor) → torch.Tensor

Contrast each timestamp against the other timestamps of the same instance.

Returns zero for a single-timestamp series, where there are no negatives.

class HierarchicalContrastiveLoss(alpha: float = 0.5, temporal_unit: int = 0)

Bases: torch.nn.Module

Sum the instance and temporal contrastive losses over a max-pooling hierarchy.

Each level halves the time axis with a max-pool, so the loss is applied from per-timestamp granularity all the way up to the whole series. This is an nn.Module so that assigning it to self.criterion satisfies the @hyrax_model decorator and suppresses the config-driven criterion.

Parameters:

• alpha (float) – Weight on the instance-contrastive term; 1 - alpha weights the temporal term.

• temporal_unit (int) – Hierarchy level below which the temporal term is skipped.

• state (Initialize internal Module)

• ScriptModule. (shared by both nn.Module and)

alpha = 0.5

temporal_unit = 0

forward(z1, z2)

take_per_row(tensor: torch.Tensor, indices, num_elem: int) → torch.Tensor

Gather num_elem consecutive timesteps from each row, starting at that row’s index.

Parameters:

• tensor (torch.Tensor) – Shape (batch, time, channels).

• indices (array-like) – Per-row start offset, length batch.

• num_elem (int) – Number of consecutive timesteps to take from each row.

Returns:

Shape (batch, num_elem, channels).

Return type:

torch.Tensor

FLUX_CHANNEL = 0

FLUX_ERR_CHANNEL = 1

DELTA_TIME_CHANNEL = 2

N_LEADING_CHANNELS = 3

NORMALIZE_MODES = ('median_mad', 'zscore')

MAD_TO_SIGMA = 1.4826

class EncodingSettings(config: dict)

The [model.HyraxTs2Vec] keys that describe the input rather than the architecture.

Resolved and validated in one place so that build_event_sequence() and HyraxTs2Vec agree on the channel count, and so a bad setting is caught when the model is constructed rather than midway through a run.

time_field, flux_field and flux_err_field are flattened column names, as LSDBStreamDataset exposes nested columns. band_field is None for a single-band survey, where the one-hot block is dropped entirely and bands is empty; otherwise bands gives the one-hot channel order. max_sequence_length is how many observations are kept per object, which fixes the width of the time axis, and normalize is one of NORMALIZE_MODES or False. n_channels is 3 + len(bands), the model’s input width.

Parameters:

config (dict) – The runtime configuration.

time_field = ''

flux_field = ''

flux_err_field = ''

band_field = None

light_curve_fields

bands = []

band_index

max_sequence_length

normalize

n_channels

_select_columns(data: dict, settings: EncodingSettings)

Pull the light-curve columns and the validity mask out of a collated batch.

Returns (time, flux, flux_err, band, valid), each (batch, length). band is None for a single-band catalog. valid starts from the collation mask and then drops non-finite observations, so it is the single source of truth for what is real.

_sort_and_trim(columns, valid, length)

Sort each row by time and fit it to exactly length timesteps.

Invalid entries sort to the tail, so the kept observations are the earliest real ones, matching what a per-object sort-then-truncate would give.

_normalize_flux(flux, flux_err, valid, mode)

Center and rescale each object’s flux, applying the same scale to its error.

Per-object rather than per-batch: raw fluxes span orders of magnitude across bands and objects, and TS2Vec applies no normalization of its own. The shared scale is what keeps the error interpretable next to the flux.

_band_indices(band, valid, settings: EncodingSettings)

Map band values to one-hot column offsets, returning -1 for unrecognized ones.

Accepts integer band codes directly and matches everything else against the configured band names. Mapping goes through np.unique so the dict lookups cost one per distinct band rather than one per observation. Unrecognized bands are reported only where the mask says there is a real observation, so zero padding is not mistaken for a band.

build_event_sequence(data_dict: dict, config: dict) → numpy.ndarray

Turn padded light-curve columns into the event sequence TSEncoder consumes.

Observations are sorted by time and laid out as an ordered sequence of events, one timestep per observation, with channels:

0                      flux, normalized per object
1                      flux error, on the same scale as the flux
2                      log1p of the gap since the previous observation (0 for the first)
3 .. 3 + len(bands)    one-hot band indicator

Uneven sampling is therefore carried as a feature (channel 2) rather than being resampled away, and a light curve that spans several filters stays a single sequence. Cross-matching more surveys to add bands only widens the one-hot block; the encoder sizes its input layer from bands, so no architecture change is needed.

The input columns are whatever the dataset collated under the names configured in [model.HyraxTs2Vec]. LSDBStreamDataset produces them for every nested column without any per-survey subclass: <column>_<subcolumn> padded to the batch’s longest row, plus a boolean <column>_<subcolumn>_mask.

Parameters:

• data_dict (dict) – The collated batch, with a "data" key holding the light-curve columns.

• config (dict) – The runtime configuration.

Returns:

(batch, max_sequence_length, 3 + len(bands)) float32, NaN at every padded or missing timestep.

Return type:

numpy.ndarray

class HyraxTs2Vec(config)

Bases: torch.nn.Module

TS2Vec, a self-supervised encoder for multivariate time series.

Consumes (batch, time, channels) float32 input where NaN marks a missing or padded timestep, and produces one fixed-length representation per object from infer_batch() - ready for reduce_dimensions and the visualize verbs.

channels is the event-sequence encoding built by build_event_sequence(): flux, flux error, log-scaled time gap, then a one-hot band indicator. Adding bands (by cross-matching more surveys) only widens channels; no architecture change is needed.

Notes

This model sets self.optimizer itself (AdamW at [model.HyraxTs2Vec].learning_rate) rather than taking the config-driven optimizer, because the Hyrax default is SGD at lr = 0.01 while TS2Vec is specified on AdamW at 1e-3. That is a supported path, but it does mean Hyrax logs one benign “Both model and config define an optimizer” warning per run.

Differences from the reference implementation, both deliberate:

• No SWA. The reference wraps the encoder in torch.optim.swa_utils.AveragedModel and encodes from the averaged weights. Skipped here: it puts non-gradient parameters into the optimizer, complicates state_dict save/load, and risks DDP unused-parameter errors.

• No max_train_length splitting. The reference chops over-long series into segments stacked along the batch axis before training. Sequence length is instead bounded by the dataset’s max_sequence_length.

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

config

input_dims

output_dims

hidden_dims

depth

mask_mode

temporal_unit

encoding_window

encoder

criterion

optimizer

forward(x, mask=None)

Encode a batch of series into per-timestamp representations.

Parameters:

• x (torch.Tensor) – Shape (batch, time, channels), NaN at missing or padded timesteps.

• mask (str or torch.Tensor, optional) – Overrides the configured masking strategy for this call.

Returns:

Shape (batch, time, output_dims).

Return type:

torch.Tensor

_drop_empty_rows(x)

Drop rows that are NaN everywhere, as the reference’s fit does.

Safe during training because the loss does not need to align with object_id.

_contrastive_loss(x)

Encode two overlapping random crops of x and return their contrastive loss.

train_batch(batch)

Run one self-supervised training step, the inner loop of the reference fit().

Parameters:

batch (torch.Tensor) – Shape (batch, time, channels), NaN at missing or padded timesteps.

Returns:

{"loss": float} for the current batch.

Return type:

dict

validate_batch(batch)

Compute the training loss without updating any weights.

Parameters:

batch (torch.Tensor) – Shape (batch, time, channels), NaN at missing or padded timesteps.

Returns:

{"loss": float} for the current batch.

Return type:

dict

test_batch(batch)

Identical to validate_batch().

Parameters:

batch (torch.Tensor) – Shape (batch, time, channels), NaN at missing or padded timesteps.

Returns:

{"loss": float} for the current batch.

Return type:

dict

infer_batch(batch)

Produce the latent representation of each series, the core of encode().

The model is in eval mode here, so TSEncoder resolves its mask to "all_true" and the output is deterministic.

Series that are entirely missing are not dropped, because the inference output must stay aligned one-to-one with object_id; they encode a fully masked, all-zero input.

Parameters:

batch (torch.Tensor) – Shape (batch, time, channels), NaN at missing or padded timesteps.

Returns:

(batch, output_dims) when encoding_window is "full_series", or (batch, time, output_dims * levels) when it is "multiscale".

Return type:

torch.Tensor

static prepare_inputs(data_dict)

Build the event sequence this model consumes from a collated batch.

The real work is in build_event_sequence(); this is a shim so that the function written out to prepare_inputs.py next to the saved weights stays short and self-contained. That does mean a checkpoint pins the settings of the encoding - they are saved in the run’s runtime_config.toml - rather than its code, which tracks the installed Hyrax.

Parameters:

data_dict (dict) – The collated batch, with a "data" key holding the padded light-curve columns named in [model.HyraxTs2Vec].

Returns:

Shape (batch, max_sequence_length, channels), float32, NaN at padded timesteps. Hyrax converts this to a tensor and moves it to the right device.

Return type:

numpy.ndarray