From 9fa54484d8775cee156438ec0b21ab96c9c7a51b Mon Sep 17 00:00:00 2001 From: Maximilian Dax Date: Thu, 2 Jul 2026 14:11:58 +0200 Subject: [PATCH 1/9] Add characterization tests for the model build path Pin current behavior ahead of the NN build-system refactor (hackathon/ NN_Build_System_Design.md, step 1): build_model_from_kwargs dispatch, autocomplete_model_kwargs dimensional completion (plain + GNPE), end-to-end forward passes for all three posterior model types, GNPE added_context embedding, unconditional flow, SVD initial-weight seeding, old nsf+embedding schema mapping, and save/rebuild round trip. Also documents (without fixing) a latent bug: get_theta_embedding_net reads 'frequencies' from the wrong settings level, so any nonzero value crashes the FMPE/score build. Co-Authored-By: Claude Fable 5 --- tests/core/test_build_model.py | 383 ++++++++++++++++++++++++++++----- 1 file changed, 332 insertions(+), 51 deletions(-) diff --git a/tests/core/test_build_model.py b/tests/core/test_build_model.py index 25a0c060..75c1727d 100644 --- a/tests/core/test_build_model.py +++ b/tests/core/test_build_model.py @@ -1,87 +1,368 @@ +""" +Characterization tests for the model build path. + +These tests pin the *current* behavior of building posterior models from settings +dictionaries — type dispatch (build_model_from_kwargs), dimensional autocompletion +(autocomplete_model_kwargs), end-to-end construction and forward passes for all three +posterior model types, SVD initial-weight seeding, and loading of old-schema +checkpoints — ahead of the NN build-system refactor (see hackathon/ +NN_Build_System_Design.md). If one of these tests breaks, either the refactor changed +observable behavior (fix the refactor) or the behavior change is intended and +documented (update the test alongside the compatibility shim). +""" + +import copy + import numpy as np import pytest +import torch +from dingo.core.posterior_models.base_model import BasePosteriorModel from dingo.core.posterior_models.build_model import ( autocomplete_model_kwargs, build_model_from_kwargs, ) +from dingo.core.posterior_models.flow_matching import FlowMatchingPosteriorModel from dingo.core.posterior_models.normalizing_flow import NormalizingFlowPosteriorModel +from dingo.core.posterior_models.score_matching import ScoreDiffusionPosteriorModel +from dingo.core.utils.backward_compatibility import update_model_config + +# Data dimensions used throughout: 4 inference parameters, strain data of shape +# (num_blocks=2, num_channels=3, num_bins=20), embedding output dimension 8. +NUM_PARAMETERS = 4 +DATA_SHAPE = (2, 3, 20) +EMBEDDING_OUTPUT_DIM = 8 +GNPE_PROXY_DIM = 2 +BATCH_SIZE = 5 + + +def embedding_kwargs(): + return { + "input_dims": list(DATA_SHAPE), + "svd": {"size": 10}, + "V_rb_list": None, + "output_dim": EMBEDDING_OUTPUT_DIM, + "hidden_dims": [32, 16, 8], + "activation": "elu", + "dropout": 0.0, + "batch_norm": True, + "added_context": False, + } + + +def nsf_posterior_kwargs(): + return { + "input_dim": NUM_PARAMETERS, + "context_dim": EMBEDDING_OUTPUT_DIM, + "num_flow_steps": 2, + "base_transform_kwargs": { + "hidden_dim": 16, + "num_transform_blocks": 1, + "activation": "elu", + "dropout_probability": 0.0, + "batch_norm": True, + "num_bins": 4, + "base_transform_type": "rq-coupling", + }, + } -BASE_TRANSFORM_KWARGS = { - "hidden_dim": 8, - "num_transform_blocks": 1, - "activation": "elu", - "dropout_probability": 0.0, - "batch_norm": False, - "num_bins": 4, - "base_transform_type": "rq-coupling", -} +def cflow_posterior_kwargs(): + return { + "input_dim": NUM_PARAMETERS, + "context_dim": EMBEDDING_OUTPUT_DIM, + "activation": "gelu", + "batch_norm": False, + "dropout": 0.0, + "hidden_dims": [16, 16], + "theta_with_glu": False, + "context_with_glu": False, + "time_prior_exponent": 1, + "theta_embedding_kwargs": { + "embedding_net": { + "activation": "gelu", + "hidden_dims": [8], + "output_dim": 8, + "type": "DenseResidualNet", + }, + # NOTE: frequencies > 0 crashes on main: get_theta_embedding_net reads + # `frequencies` from the top level of theta_embedding_kwargs while + # get_dim_positional_embedding reads it from `encoding` — inconsistent. + # The fmpe example only works because it uses frequencies: 0. Pinned here. + "encoding": {"encode_all": False, "frequencies": 0}, + }, + } -def _settings(posterior_model_type="normalizing_flow"): +def model_settings(posterior_model_type): + if posterior_model_type == "normalizing_flow": + posterior_kwargs = nsf_posterior_kwargs() + else: + posterior_kwargs = cflow_posterior_kwargs() + if posterior_model_type == "flow_matching": + posterior_kwargs["sigma_min"] = 0.001 + elif posterior_model_type == "score_matching": + posterior_kwargs["epsilon"] = 1e-3 + posterior_kwargs["beta_min"] = 0.1 + posterior_kwargs["beta_max"] = 20.0 return { "train_settings": { "model": { "posterior_model_type": posterior_model_type, - "posterior_kwargs": { - "input_dim": 3, - "context_dim": None, - "num_flow_steps": 2, - "base_transform_kwargs": BASE_TRANSFORM_KWARGS, - }, + "posterior_kwargs": posterior_kwargs, + "embedding_kwargs": embedding_kwargs(), } } } -def test_build_model_dispatches_to_normalizing_flow(): - model = build_model_from_kwargs(settings=_settings(), device="cpu") - assert isinstance(model, NormalizingFlowPosteriorModel) +def data_sample(with_gnpe_proxies=False): + """A sample in the format produced by the dataloader (wfd[0] after UnpackDict): + [inference_parameters, strain data(, gnpe_proxies)].""" + sample = [ + np.random.rand(NUM_PARAMETERS).astype(np.float32), + np.random.rand(*DATA_SHAPE).astype(np.float32), + ] + if with_gnpe_proxies: + sample.append(np.random.rand(GNPE_PROXY_DIM).astype(np.float32)) + return sample + + +def batch(model_type="normalizing_flow", with_gnpe_proxies=False): + theta = torch.rand(BATCH_SIZE, NUM_PARAMETERS) + context = [torch.rand(BATCH_SIZE, *DATA_SHAPE)] + if with_gnpe_proxies: + context.append(torch.rand(BATCH_SIZE, GNPE_PROXY_DIM)) + return theta, context + + +# ----------------------------------------------------------------------------------- +# build_model_from_kwargs: type dispatch +# ----------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "posterior_model_type, expected_class", + [ + ("normalizing_flow", NormalizingFlowPosteriorModel), + ("flow_matching", FlowMatchingPosteriorModel), + ("score_matching", ScoreDiffusionPosteriorModel), + ], +) +def test_build_model_from_kwargs_dispatch(posterior_model_type, expected_class): + settings = model_settings(posterior_model_type) + pm = build_model_from_kwargs(settings=settings, device="cpu") + assert type(pm) is expected_class + assert isinstance(pm, BasePosteriorModel) + assert pm.metadata is settings + + +def test_build_model_from_kwargs_rejects_unknown_type(): + settings = model_settings("normalizing_flow") + settings["train_settings"]["model"]["posterior_model_type"] = "not_a_model" + with pytest.raises(ValueError): + build_model_from_kwargs(settings=settings, device="cpu") + + +def test_build_model_from_kwargs_requires_exactly_one_source(): + settings = model_settings("normalizing_flow") + with pytest.raises(ValueError): + build_model_from_kwargs(filename=None, settings=None) + with pytest.raises(ValueError): + build_model_from_kwargs(filename="model.pt", settings=settings) + + +# ----------------------------------------------------------------------------------- +# autocomplete_model_kwargs: dimensional glue +# ----------------------------------------------------------------------------------- + + +def test_autocomplete_model_kwargs_without_gnpe(): + model_kwargs = model_settings("normalizing_flow")["train_settings"]["model"] + # Settings as written by a user: dims absent. + del model_kwargs["embedding_kwargs"]["input_dims"] + del model_kwargs["posterior_kwargs"]["input_dim"] + del model_kwargs["posterior_kwargs"]["context_dim"] + + autocomplete_model_kwargs(model_kwargs, data_sample(with_gnpe_proxies=False)) + + assert model_kwargs["embedding_kwargs"]["input_dims"] == list(DATA_SHAPE) + assert model_kwargs["posterior_kwargs"]["input_dim"] == NUM_PARAMETERS + assert model_kwargs["embedding_kwargs"]["added_context"] is False + assert model_kwargs["posterior_kwargs"]["context_dim"] == EMBEDDING_OUTPUT_DIM + + +def test_autocomplete_model_kwargs_with_gnpe(): + model_kwargs = model_settings("normalizing_flow")["train_settings"]["model"] + + autocomplete_model_kwargs(model_kwargs, data_sample(with_gnpe_proxies=True)) + + assert model_kwargs["embedding_kwargs"]["added_context"] is True + assert ( + model_kwargs["posterior_kwargs"]["context_dim"] + == EMBEDDING_OUTPUT_DIM + GNPE_PROXY_DIM + ) + + +# ----------------------------------------------------------------------------------- +# End-to-end: build + forward passes for all three model types +# ----------------------------------------------------------------------------------- -def test_build_model_dispatch_is_case_insensitive(): - model = build_model_from_kwargs( - settings=_settings("Normalizing_Flow"), device="cpu" +@pytest.mark.parametrize( + "posterior_model_type", ["normalizing_flow", "flow_matching", "score_matching"] +) +def test_model_forward_passes(posterior_model_type): + pm = build_model_from_kwargs( + settings=model_settings(posterior_model_type), device="cpu" ) - assert isinstance(model, NormalizingFlowPosteriorModel) + theta, context = batch() + + loss = pm.loss(theta, *context) + assert loss.shape == () + assert torch.isfinite(loss) + + pm.network.eval() + with torch.no_grad(): + samples = pm.sample(*context, num_samples=3) + assert samples.shape == (BATCH_SIZE, 3, NUM_PARAMETERS) + + log_prob = pm.log_prob(theta, *context) + assert log_prob.shape == (BATCH_SIZE,) + assert torch.isfinite(log_prob).all() + + samples, log_prob = pm.sample_and_log_prob(*context, num_samples=3) + assert samples.shape == (BATCH_SIZE, 3, NUM_PARAMETERS) + assert log_prob.shape == (BATCH_SIZE, 3) + + +def test_normalizing_flow_with_gnpe_context(): + """With added_context=True, the embedding merges (data, proxies) via ModuleMerger, + and sampling/density evaluation take two context tensors.""" + settings = model_settings("normalizing_flow") + model = settings["train_settings"]["model"] + model["embedding_kwargs"]["added_context"] = True + model["posterior_kwargs"]["context_dim"] = EMBEDDING_OUTPUT_DIM + GNPE_PROXY_DIM + + pm = build_model_from_kwargs(settings=settings, device="cpu") + theta, context = batch(with_gnpe_proxies=True) + + loss = pm.loss(theta, *context) + assert torch.isfinite(loss) + + pm.network.eval() + with torch.no_grad(): + log_prob = pm.log_prob(theta, *context) + assert log_prob.shape == (BATCH_SIZE,) + + +def test_normalizing_flow_unconditional(): + """Without embedding_kwargs, an unconditional flow is built (the models-as-priors + path used by unconditional_density_estimation).""" + settings = model_settings("normalizing_flow") + model = settings["train_settings"]["model"] + del model["embedding_kwargs"] + # Convention from unconditional_density_estimation.py:74-75. + model["posterior_kwargs"]["context_dim"] = None + + pm = build_model_from_kwargs(settings=settings, device="cpu") + theta = torch.rand(BATCH_SIZE, NUM_PARAMETERS) + loss = pm.loss(theta) + assert torch.isfinite(loss) -def test_build_model_requires_exactly_one_of_filename_or_settings(): - # Neither provided. - with pytest.raises(ValueError, match="filename or a settings"): - build_model_from_kwargs() - # Both provided. - with pytest.raises(ValueError, match="filename or a settings"): - build_model_from_kwargs(filename="x.pt", settings=_settings()) + pm.network.eval() + with torch.no_grad(): + samples = pm.sample(num_samples=3) + assert samples.shape == (3, NUM_PARAMETERS) + log_prob = pm.log_prob(theta) + assert log_prob.shape == (BATCH_SIZE,) -def test_build_model_rejects_unknown_type(): - with pytest.raises(ValueError, match="No valid posterior model type"): - build_model_from_kwargs(settings=_settings("not_a_model"), device="cpu") +# ----------------------------------------------------------------------------------- +# SVD initial-weight seeding +# ----------------------------------------------------------------------------------- -def test_autocomplete_model_kwargs_without_gnpe_proxies(): - model_kwargs = {"embedding_kwargs": {"output_dim": 8}, "posterior_kwargs": {}} - # data_sample = [parameters, GW data] (no gnpe proxies) - autocomplete_model_kwargs( - model_kwargs, data_sample=[np.zeros(4), np.zeros((2, 3, 20))] +def test_initial_weights_seed_svd_projection(): + """initial_weights['V_rb_list'] seeds the LinearProjectionRB layer weights, and the + settings dict is not polluted with the (large) V matrices.""" + num_bins = DATA_SHAPE[2] + n_rb = 10 + V_rb_list = [ + (np.random.rand(num_bins, n_rb) + 1j * np.random.rand(num_bins, n_rb)) + for _ in range(DATA_SHAPE[0]) + ] + settings = model_settings("normalizing_flow") + settings_before = copy.deepcopy(settings) + + pm = build_model_from_kwargs( + settings=settings, initial_weights={"V_rb_list": V_rb_list}, device="cpu" ) - assert model_kwargs["embedding_kwargs"]["input_dims"] == [2, 3, 20] - assert model_kwargs["posterior_kwargs"]["input_dim"] == 4 - assert model_kwargs["embedding_kwargs"]["added_context"] is False - # context_dim == embedding output_dim. - assert model_kwargs["posterior_kwargs"]["context_dim"] == 8 + projection = pm.network.embedding_net[0] + V = V_rb_list[0][:, :n_rb] + layer_weight = projection.layers_rb[0].weight.data + assert torch.allclose( + layer_weight[:n_rb, :num_bins], + torch.from_numpy(V.real.T).float(), + ) + assert torch.allclose( + layer_weight[n_rb:, :num_bins], + torch.from_numpy(V.imag.T).float(), + ) + # The V matrices must not leak into the saved settings. + assert settings == settings_before + + +# ----------------------------------------------------------------------------------- +# Backward compatibility: old settings schema +# ----------------------------------------------------------------------------------- + + +def test_update_model_config_maps_old_schema(): + old = { + "type": "nsf+embedding", + "nsf_kwargs": nsf_posterior_kwargs(), + "embedding_net_kwargs": embedding_kwargs(), + } + update_model_config(old) + assert old["posterior_model_type"] == "normalizing_flow" + assert old["posterior_kwargs"] == nsf_posterior_kwargs() + assert old["embedding_kwargs"] == embedding_kwargs() + assert "type" not in old and "nsf_kwargs" not in old -def test_autocomplete_model_kwargs_with_gnpe_proxies(): - model_kwargs = {"embedding_kwargs": {"output_dim": 8}, "posterior_kwargs": {}} - # data_sample = [parameters, GW data, gnpe_proxies (len 2)] - autocomplete_model_kwargs( - model_kwargs, data_sample=[np.zeros(4), np.zeros((2, 3, 20)), np.zeros(2)] +def test_build_model_from_old_schema_settings(): + settings = model_settings("normalizing_flow") + model = settings["train_settings"]["model"] + settings["train_settings"]["model"] = { + "type": "nsf+embedding", + "nsf_kwargs": model["posterior_kwargs"], + "embedding_net_kwargs": model["embedding_kwargs"], + } + pm = build_model_from_kwargs(settings=settings, device="cpu") + assert type(pm) is NormalizingFlowPosteriorModel + + +# ----------------------------------------------------------------------------------- +# Save / load round trip through build_model_from_kwargs (filename path) +# ----------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "posterior_model_type", ["normalizing_flow", "flow_matching"] +) +def test_save_and_rebuild_from_file(tmp_path, posterior_model_type): + pm = build_model_from_kwargs( + settings=model_settings(posterior_model_type), device="cpu" ) + filename = str(tmp_path / "model.pt") + pm.save_model(filename) - assert model_kwargs["embedding_kwargs"]["added_context"] is True - # context_dim == output_dim + gnpe_proxy_dim == 8 + 2. - assert model_kwargs["posterior_kwargs"]["context_dim"] == 10 + pm_loaded = build_model_from_kwargs( + filename=filename, device="cpu", load_training_info=False + ) + assert type(pm_loaded) is type(pm) + for p0, p1 in zip(pm.network.parameters(), pm_loaded.network.parameters()): + assert torch.equal(p0.data, p1.data) From 8dba80c350d381f1e929946b7a7ea855b31d3a1c Mon Sep 17 00:00:00 2001 From: Maximilian Dax Date: Thu, 2 Jul 2026 14:25:38 +0200 Subject: [PATCH 2/9] Add component registry; rename BasePosteriorModel to NeuralDistribution Step 2 of the NN build-system refactor (hackathon/NN_Build_System_Design.md): - dingo/core/registry.py: generic Registry with four-step name resolution (registered name -> entry points ('dingo.architectures') -> dotted import path -> 'file.py:Class'), plus the NEURAL_DISTRIBUTIONS / EMBEDDING_NETS / CONTEXT_MERGERS instances. Third-party architectures no longer require editing dingo source. - The three model types register themselves; build_model_from_kwargs dispatches via the registry instead of the hardcoded models_dict, so posterior_model_type can also name a plugin (tested end-to-end with a file-path plugin). - Name lookup is now case-sensitive; the historical case-insensitivity for built-in type names moved into update_model_config (the back-compat boundary), so old checkpoints still load. - BasePosteriorModel -> NeuralDistribution ('posterior' was too narrow: the same class models priors/proposals, e.g. the unconditional flow). Lasting alias BasePosteriorModel kept for existing imports and branches. Behavior pinned by tests/core/test_build_model.py is unchanged (all pass). Co-Authored-By: Claude Fable 5 --- dingo/core/posterior_models/__init__.py | 5 +- dingo/core/posterior_models/base_model.py | 23 ++- dingo/core/posterior_models/build_model.py | 35 ++-- dingo/core/posterior_models/cflow_base.py | 4 +- dingo/core/posterior_models/flow_matching.py | 2 + .../core/posterior_models/normalizing_flow.py | 6 +- dingo/core/posterior_models/score_matching.py | 2 + dingo/core/registry.py | 149 ++++++++++++++++++ dingo/core/samplers.py | 12 +- dingo/core/utils/backward_compatibility.py | 10 ++ dingo/gw/training/train_pipeline.py | 18 +-- tests/core/test_build_model.py | 10 ++ tests/core/test_registry.py | 135 ++++++++++++++++ 13 files changed, 366 insertions(+), 45 deletions(-) create mode 100644 dingo/core/registry.py create mode 100644 tests/core/test_registry.py diff --git a/dingo/core/posterior_models/__init__.py b/dingo/core/posterior_models/__init__.py index 148346f9..8281c1fa 100644 --- a/dingo/core/posterior_models/__init__.py +++ b/dingo/core/posterior_models/__init__.py @@ -1,4 +1,7 @@ -from dingo.core.posterior_models.base_model import BasePosteriorModel +from dingo.core.posterior_models.base_model import ( + BasePosteriorModel, + NeuralDistribution, +) from dingo.core.posterior_models.normalizing_flow import NormalizingFlowPosteriorModel from dingo.core.posterior_models.cflow_base import ContinuousFlowPosteriorModel from dingo.core.posterior_models.flow_matching import FlowMatchingPosteriorModel diff --git a/dingo/core/posterior_models/base_model.py b/dingo/core/posterior_models/base_model.py index 68508d19..de6797b4 100755 --- a/dingo/core/posterior_models/base_model.py +++ b/dingo/core/posterior_models/base_model.py @@ -24,14 +24,19 @@ from dingo.core.utils.trainutils import EarlyStopping -class BasePosteriorModel(ABC): +class NeuralDistribution(ABC): """ - Abstract base class for PosteriorModels. This is intended to construct and hold a - neural network for estimating the posterior density, as well as saving / loading, - and training. + Abstract base class for distributions parameterized by a neural network. - Subclasses must implement methods for constructing the specific network, sampling, - density evaluation, and computing the loss during training. + A NeuralDistribution can be conditional or unconditional, and can model any + distribution over parameters — a posterior, a prior, or a proposal. The contract + is: sampling is always available; log_prob is available where the architecture + affords it (e.g., normalizing flows exactly, score matching only via + probability-flow ODE integration). + + This class constructs and holds the network, and provides saving / loading and + training. Subclasses must implement methods for constructing the specific + network, sampling, density evaluation, and computing the loss during training. """ def __init__( @@ -469,6 +474,12 @@ def train( print(f"Finished training epoch {self.epoch}.\n") +# Lasting alias: the class was called BasePosteriorModel before the NN build-system +# refactor; "posterior" was too narrow (a NeuralDistribution can also model e.g. a +# prior). Kept so that existing imports and branches keep working. +BasePosteriorModel = NeuralDistribution + + def train_epoch(pm, dataloader): pm.network.train() loss_info = dingo.core.utils.trainutils.LossInfo( diff --git a/dingo/core/posterior_models/build_model.py b/dingo/core/posterior_models/build_model.py index 2f757103..2a96d5ee 100644 --- a/dingo/core/posterior_models/build_model.py +++ b/dingo/core/posterior_models/build_model.py @@ -1,7 +1,9 @@ -from dingo.core.posterior_models.base_model import BasePosteriorModel -from dingo.core.posterior_models.flow_matching import FlowMatchingPosteriorModel -from dingo.core.posterior_models.normalizing_flow import NormalizingFlowPosteriorModel -from dingo.core.posterior_models.score_matching import ScoreDiffusionPosteriorModel +# Importing the modules registers the built-in model types with NEURAL_DISTRIBUTIONS. +import dingo.core.posterior_models.flow_matching # noqa: F401 +import dingo.core.posterior_models.normalizing_flow # noqa: F401 +import dingo.core.posterior_models.score_matching # noqa: F401 +from dingo.core.posterior_models.base_model import NeuralDistribution +from dingo.core.registry import NEURAL_DISTRIBUTIONS from dingo.core.utils.backward_compatibility import ( torch_load_with_fallback, update_model_config, @@ -11,12 +13,13 @@ def build_model_from_kwargs( filename: str = None, settings: dict = None, **kwargs -) -> BasePosteriorModel: +) -> NeuralDistribution: """ - Returns a PosteriorModel based on a saved network or settings dict. + Returns a NeuralDistribution based on a saved network or settings dict. - The function is careful to choose the appropriate PosteriorModel class (e.g., - for a normalizing flow, flow matching, or score matching). + The model class is resolved from the settings' posterior_model_type via the + NEURAL_DISTRIBUTIONS registry (e.g., normalizing flow, flow matching, or score + matching, or a plugin type; see dingo.core.registry). Parameters ---------- @@ -29,19 +32,13 @@ def build_model_from_kwargs( Returns ------- - PosteriorModel + NeuralDistribution """ if (filename is None) == (settings is None): raise ValueError( "Either a filename or a settings dict must be provided, but not both." ) - models_dict = { - "normalizing_flow": NormalizingFlowPosteriorModel, - "flow_matching": FlowMatchingPosteriorModel, - "score_matching": ScoreDiffusionPosteriorModel, - } - if filename is not None: d, _ = torch_load_with_fallback(filename, preferred_map_location="meta") if "version" in d: @@ -59,10 +56,10 @@ def build_model_from_kwargs( "posterior_model_type" ] - if not posterior_model_type.lower() in models_dict: - raise ValueError("No valid posterior model type specified.") - - model = models_dict[posterior_model_type.lower()] + try: + model = NEURAL_DISTRIBUTIONS.get(posterior_model_type) + except KeyError as e: + raise ValueError(f"No valid posterior model type specified. {e}") from e return model(model_filename=filename, metadata=settings, **kwargs) diff --git a/dingo/core/posterior_models/cflow_base.py b/dingo/core/posterior_models/cflow_base.py index 2f9b09bb..cdc5f251 100644 --- a/dingo/core/posterior_models/cflow_base.py +++ b/dingo/core/posterior_models/cflow_base.py @@ -5,12 +5,12 @@ from torchdiffeq import odeint from glasflow.nflows.utils.torchutils import repeat_rows, split_leading_dim -from .base_model import BasePosteriorModel +from .base_model import NeuralDistribution from dingo.core.nn.cfnets import create_cf -class ContinuousFlowPosteriorModel(BasePosteriorModel): +class ContinuousFlowPosteriorModel(NeuralDistribution): """ Class for posterior models based on continuous normalizing flows (CNFs). diff --git a/dingo/core/posterior_models/flow_matching.py b/dingo/core/posterior_models/flow_matching.py index 3e53a61c..f85d4e91 100644 --- a/dingo/core/posterior_models/flow_matching.py +++ b/dingo/core/posterior_models/flow_matching.py @@ -5,8 +5,10 @@ from torch import nn from .cflow_base import ContinuousFlowPosteriorModel +from dingo.core.registry import NEURAL_DISTRIBUTIONS +@NEURAL_DISTRIBUTIONS.register("flow_matching") class FlowMatchingPosteriorModel(ContinuousFlowPosteriorModel): __doc__ = ( inspect.getdoc(ContinuousFlowPosteriorModel) diff --git a/dingo/core/posterior_models/normalizing_flow.py b/dingo/core/posterior_models/normalizing_flow.py index 8481aeda..ddf1921e 100644 --- a/dingo/core/posterior_models/normalizing_flow.py +++ b/dingo/core/posterior_models/normalizing_flow.py @@ -1,4 +1,5 @@ -from .base_model import BasePosteriorModel +from .base_model import NeuralDistribution +from dingo.core.registry import NEURAL_DISTRIBUTIONS from dingo.core.nn.nsf import ( create_nsf_with_rb_projection_embedding_net, @@ -6,7 +7,8 @@ ) -class NormalizingFlowPosteriorModel(BasePosteriorModel): +@NEURAL_DISTRIBUTIONS.register("normalizing_flow") +class NormalizingFlowPosteriorModel(NeuralDistribution): """ Posterior model based on a (discrete) normalizing flow. diff --git a/dingo/core/posterior_models/score_matching.py b/dingo/core/posterior_models/score_matching.py index 5fbaca74..5031a906 100644 --- a/dingo/core/posterior_models/score_matching.py +++ b/dingo/core/posterior_models/score_matching.py @@ -4,8 +4,10 @@ import torch from .cflow_base import ContinuousFlowPosteriorModel +from dingo.core.registry import NEURAL_DISTRIBUTIONS +@NEURAL_DISTRIBUTIONS.register("score_matching") class ScoreDiffusionPosteriorModel(ContinuousFlowPosteriorModel): __doc__ = ( inspect.getdoc(ContinuousFlowPosteriorModel) diff --git a/dingo/core/registry.py b/dingo/core/registry.py new file mode 100644 index 00000000..3cc86e63 --- /dev/null +++ b/dingo/core/registry.py @@ -0,0 +1,149 @@ +""" +Registries for pluggable dingo components. + +A Registry maps short, stable names (as stored in train settings and model metadata) +to component classes or builder functions. Components inside dingo register +themselves with the ``register`` decorator. Third-party components can be used +without editing dingo source; ``Registry.get`` resolves a name in this order: + +1. a name registered via the decorator (dingo-internal components), +2. an installed entry point in the registry's entry-point group (pip-installed + plugin packages), +3. a dotted import path, e.g. ``"my_package.nets.MyNN"``, +4. a file path with class name, e.g. ``"/path/to/my_net.py:MyNN"`` (un-packaged + experiments). + +Checkpoints should store the short name (form 1/2) where possible: it is stable +under dingo-internal refactors, unlike full import paths. See +hackathon/NN_Build_System_Design.md §4.2. +""" + +from __future__ import annotations + +import importlib +import importlib.metadata +import importlib.util +import sys +from typing import Any, Callable, Dict, List + +ARCHITECTURE_ENTRY_POINT_GROUP = "dingo.architectures" + + +class Registry: + """ + A name -> component mapping with plugin resolution. + + Parameters + ---------- + kind : str + Human-readable name of the component kind (e.g. "neural_distributions"). + Used in error messages. + entry_point_group : str + Entry-point group searched for pip-installed plugins. + """ + + def __init__( + self, kind: str, entry_point_group: str = ARCHITECTURE_ENTRY_POINT_GROUP + ): + self.kind = kind + self.entry_point_group = entry_point_group + self._components: Dict[str, Any] = {} + + def register(self, name: str) -> Callable: + """ + Class/function decorator registering the component under ``name``. + + Raises ValueError if the name is already taken by a different component. + """ + + def decorator(component): + existing = self._components.get(name) + if existing is not None and existing is not component: + raise ValueError( + f"{self.kind}: name '{name}' is already registered " + f"for {existing!r}." + ) + self._components[name] = component + return component + + return decorator + + def get(self, name: str) -> Any: + """ + Resolve ``name`` to a component (see module docstring for the lookup order). + + Raises KeyError if the name cannot be resolved. + """ + if name in self._components: + return self._components[name] + + for resolve in (self._from_entry_points, self._from_dotted_path, + self._from_file_path): + component = resolve(name) + if component is not None: + # Cache so repeated lookups are cheap and resolve consistently. + self._components[name] = component + return component + + raise KeyError( + f"{self.kind}: '{name}' not found. Available names: {self.names()}. " + f"For a plugin, is the providing package installed (entry-point group " + f"'{self.entry_point_group}')? Alternatively use a dotted import path " + f"('my_package.module.MyClass') or a file path " + f"('/path/to/file.py:MyClass')." + ) + + def names(self) -> List[str]: + """Names registered so far (excluding not-yet-loaded entry points).""" + return sorted(self._components) + + def __contains__(self, name: str) -> bool: + return name in self._components + + def _from_entry_points(self, name: str): + for entry_point in importlib.metadata.entry_points( + group=self.entry_point_group + ): + if entry_point.name == name: + return entry_point.load() + return None + + @staticmethod + def _from_dotted_path(name: str): + module_name, _, attribute = name.rpartition(".") + if not module_name: + return None + try: + module = importlib.import_module(module_name) + except ImportError: + return None + return getattr(module, attribute, None) + + @staticmethod + def _from_file_path(name: str): + path, separator, attribute = name.rpartition(":") + if not separator or not path.endswith(".py"): + return None + module_name = f"_dingo_file_plugin_{abs(hash(path))}" + if module_name in sys.modules: + module = sys.modules[module_name] + else: + spec = importlib.util.spec_from_file_location(module_name, path) + if spec is None or spec.loader is None: + return None + module = importlib.util.module_from_spec(spec) + # Insert before exec so that e.g. dataclasses in the file resolve. + sys.modules[module_name] = module + try: + spec.loader.exec_module(module) + except FileNotFoundError: + del sys.modules[module_name] + return None + return getattr(module, attribute, None) + + +# The registries of the NN build system. Components register themselves where they +# are defined; nothing needs to be added here to introduce a new architecture. +NEURAL_DISTRIBUTIONS = Registry("neural_distributions") +EMBEDDING_NETS = Registry("embedding_networks") +CONTEXT_MERGERS = Registry("context_mergers") diff --git a/dingo/core/samplers.py b/dingo/core/samplers.py index efc4fc1f..b01cf228 100644 --- a/dingo/core/samplers.py +++ b/dingo/core/samplers.py @@ -10,7 +10,7 @@ import torch from torchvision.transforms import Compose -from dingo.core.posterior_models import BasePosteriorModel +from dingo.core.posterior_models import NeuralDistribution from dingo.core.result import Result from dingo.core.result import DATA_KEYS as RESULT_DATA_KEYS from dingo.core.utils import torch_detach_to_cpu, IterationTracker @@ -42,7 +42,7 @@ class Sampler(object): Attributes ---------- - model : BasePosteriorModel + model : NeuralDistribution inference_parameters : list samples : DataFrame Samples produced from the model by run_sampler(). @@ -59,12 +59,12 @@ class Sampler(object): def __init__( self, - model: BasePosteriorModel, + model: NeuralDistribution, ): """ Parameters ---------- - model : BasePosteriorModel + model : NeuralDistribution """ self.model = model self.event_metadata = None @@ -363,14 +363,14 @@ class GNPESampler(Sampler): def __init__( self, - model: BasePosteriorModel, + model: NeuralDistribution, init_sampler: Sampler, num_iterations: int = 1, ): """ Parameters ---------- - model : BasePosteriorModel + model : NeuralDistribution init_sampler : Sampler Used for generating initial samples num_iterations : int diff --git a/dingo/core/utils/backward_compatibility.py b/dingo/core/utils/backward_compatibility.py index af48f34e..d32cb95e 100644 --- a/dingo/core/utils/backward_compatibility.py +++ b/dingo/core/utils/backward_compatibility.py @@ -147,3 +147,13 @@ def update_model_config(model_settings: dict): del model_settings["nsf_kwargs"] model_settings["embedding_kwargs"] = model_settings["embedding_net_kwargs"] del model_settings["embedding_net_kwargs"] + + # The model type used to be matched case-insensitively; registry lookup is + # case-sensitive, so lowercase built-in type names from old checkpoints. + posterior_model_type = model_settings.get("posterior_model_type") + if posterior_model_type is not None and posterior_model_type.lower() in ( + "normalizing_flow", + "flow_matching", + "score_matching", + ): + model_settings["posterior_model_type"] = posterior_model_type.lower() diff --git a/dingo/gw/training/train_pipeline.py b/dingo/gw/training/train_pipeline.py index ac1420b0..48470aea 100644 --- a/dingo/gw/training/train_pipeline.py +++ b/dingo/gw/training/train_pipeline.py @@ -28,7 +28,7 @@ ) from dingo.core.utils.trainutils import EarlyStopping from dingo.gw.dataset import WaveformDataset -from dingo.core.posterior_models import BasePosteriorModel +from dingo.core.posterior_models import NeuralDistribution def copy_files_to_local( @@ -82,7 +82,7 @@ def copy_files_to_local( def prepare_training_new( train_settings: dict, train_dir: str, local_settings: dict -) -> Tuple[BasePosteriorModel, WaveformDataset]: +) -> Tuple[NeuralDistribution, WaveformDataset]: """ Based on a settings dictionary, initialize a WaveformDataset and PosteriorModel. @@ -101,7 +101,7 @@ def prepare_training_new( Returns ------- - (BasePosteriorModel, WaveformDataset) + (NeuralDistribution, WaveformDataset) """ data_settings = deepcopy(train_settings["data"]) # Optionally copy files to local and update path @@ -180,7 +180,7 @@ def prepare_training_new( def prepare_training_resume( checkpoint_name: str, local_settings: dict, train_dir: str -) -> Tuple[BasePosteriorModel, WaveformDataset]: +) -> Tuple[NeuralDistribution, WaveformDataset]: """ Loads a PosteriorModel from a checkpoint, as well as the corresponding WaveformDataset, in order to continue training. It initializes the saved optimizer @@ -197,7 +197,7 @@ def prepare_training_resume( Returns ------- - (BasePosteriorModel, WaveformDataset) + (NeuralDistribution, WaveformDataset) """ pm = build_model_from_kwargs( @@ -232,7 +232,7 @@ def prepare_training_resume( def initialize_stage( - pm: BasePosteriorModel, + pm: NeuralDistribution, wfd: WaveformDataset, stage: dict, num_workers: int, @@ -248,7 +248,7 @@ def initialize_stage( Parameters ---------- - pm : BasePosteriorModel + pm : NeuralDistribution wfd : WaveformDataset stage : dict Settings specific to current stage of training @@ -301,7 +301,7 @@ def initialize_stage( def train_stages( - pm: BasePosteriorModel, wfd: WaveformDataset, train_dir: str, local_settings: dict + pm: NeuralDistribution, wfd: WaveformDataset, train_dir: str, local_settings: dict ) -> bool: """ Train the network, iterating through the sequence of stages. Stages can change @@ -309,7 +309,7 @@ def train_stages( Parameters ---------- - pm : BasePosteriorModel + pm : NeuralDistribution wfd : WaveformDataset train_dir : str Directory for saving checkpoints and train history. diff --git a/tests/core/test_build_model.py b/tests/core/test_build_model.py index 75c1727d..8930e603 100644 --- a/tests/core/test_build_model.py +++ b/tests/core/test_build_model.py @@ -333,6 +333,16 @@ def test_update_model_config_maps_old_schema(): assert "type" not in old and "nsf_kwargs" not in old +def test_update_model_config_lowercases_builtin_type_names(): + """Model types used to be matched case-insensitively; the compat shim lowercases + built-in names from old checkpoints so the case-sensitive registry finds them.""" + settings = model_settings("normalizing_flow") + model = settings["train_settings"]["model"] + model["posterior_model_type"] = "Normalizing_Flow" + pm = build_model_from_kwargs(settings=settings, device="cpu") + assert type(pm) is NormalizingFlowPosteriorModel + + def test_build_model_from_old_schema_settings(): settings = model_settings("normalizing_flow") model = settings["train_settings"]["model"] diff --git a/tests/core/test_registry.py b/tests/core/test_registry.py new file mode 100644 index 00000000..1ac8ea2a --- /dev/null +++ b/tests/core/test_registry.py @@ -0,0 +1,135 @@ +"""Tests for dingo.core.registry: name resolution for pluggable components.""" + +import pytest + +from dingo.core.registry import NEURAL_DISTRIBUTIONS, Registry + + +@pytest.fixture() +def registry(): + return Registry("test_components", entry_point_group="dingo.test_components") + + +def test_register_and_get(registry): + @registry.register("my_component") + class MyComponent: + pass + + assert registry.get("my_component") is MyComponent + assert "my_component" in registry + assert registry.names() == ["my_component"] + + +def test_register_duplicate_name_raises(registry): + @registry.register("taken") + class ComponentA: + pass + + with pytest.raises(ValueError, match="already registered"): + + @registry.register("taken") + class ComponentB: + pass + + +def test_register_same_component_twice_is_idempotent(registry): + class MyComponent: + pass + + registry.register("name")(MyComponent) + registry.register("name")(MyComponent) + assert registry.get("name") is MyComponent + + +def test_get_dotted_path(registry): + from dingo.core.nn.enets import DenseResidualNet + + assert registry.get("dingo.core.nn.enets.DenseResidualNet") is DenseResidualNet + + +def test_get_file_path(registry, tmp_path): + plugin = tmp_path / "my_plugin.py" + plugin.write_text( + "class MyNN:\n" + " marker = 'from-file'\n" + ) + component = registry.get(f"{plugin}:MyNN") + assert component.marker == "from-file" + # Second lookup resolves from the cache to the same class object. + assert registry.get(f"{plugin}:MyNN") is component + + +def test_get_unknown_name_raises_keyerror(registry): + with pytest.raises(KeyError, match="test_components.*'nonexistent' not found"): + registry.get("nonexistent") + + +def test_get_missing_file_raises_keyerror(registry, tmp_path): + with pytest.raises(KeyError): + registry.get(f"{tmp_path}/does_not_exist.py:MyNN") + + +def test_get_missing_attribute_raises_keyerror(registry, tmp_path): + plugin = tmp_path / "my_plugin_2.py" + plugin.write_text("class MyNN:\n pass\n") + with pytest.raises(KeyError): + registry.get(f"{plugin}:WrongName") + + +def test_entry_point_resolution(registry, monkeypatch): + class FakeEntryPoint: + name = "installed_component" + + @staticmethod + def load(): + return "the-component" + + def fake_entry_points(group): + assert group == "dingo.test_components" + return [FakeEntryPoint] + + monkeypatch.setattr("importlib.metadata.entry_points", fake_entry_points) + assert registry.get("installed_component") == "the-component" + + +def test_builtin_distributions_are_registered(): + from dingo.core.posterior_models import ( + FlowMatchingPosteriorModel, + NormalizingFlowPosteriorModel, + ScoreDiffusionPosteriorModel, + ) + + assert NEURAL_DISTRIBUTIONS.get("normalizing_flow") is NormalizingFlowPosteriorModel + assert NEURAL_DISTRIBUTIONS.get("flow_matching") is FlowMatchingPosteriorModel + assert NEURAL_DISTRIBUTIONS.get("score_matching") is ScoreDiffusionPosteriorModel + + +def test_build_model_from_kwargs_with_file_path_plugin(tmp_path): + """A NeuralDistribution defined in a user file (never pip-installed, not in the + dingo source) can be selected as posterior_model_type via the file-path form.""" + from tests.core.test_build_model import model_settings + + plugin = tmp_path / "my_distribution.py" + plugin.write_text( + "from dingo.core.posterior_models import NormalizingFlowPosteriorModel\n" + "\n" + "class MyDistribution(NormalizingFlowPosteriorModel):\n" + " pass\n" + ) + settings = model_settings("normalizing_flow") + settings["train_settings"]["model"][ + "posterior_model_type" + ] = f"{plugin}:MyDistribution" + + from dingo.core.posterior_models.build_model import build_model_from_kwargs + + pm = build_model_from_kwargs(settings=settings, device="cpu") + assert type(pm).__name__ == "MyDistribution" + + +def test_backward_compatible_alias(): + """Other branches and downstream code import BasePosteriorModel; the alias must + survive the NeuralDistribution rename.""" + from dingo.core.posterior_models import BasePosteriorModel, NeuralDistribution + + assert BasePosteriorModel is NeuralDistribution From e17b5695e435b30b740988ea0f92cd1eedb93624 Mon Sep 17 00:00:00 2001 From: Maximilian Dax Date: Fri, 3 Jul 2026 09:54:40 +0200 Subject: [PATCH 3/9] Add parameter-contract accessors to NeuralDistribution inference_parameters / context_parameters / standardization as properties: the interface FlowFactor.from_model (sampler-revamp) consumes, so samplers stop dict-spelunking metadata. Reads the model's own train_settings['data'] (matching core/samplers.py semantics, also for unconditional models); works on old checkpoints unchanged. Proposal doc: hackathon/NN_Sampler_Interface.md. Co-Authored-By: Claude Fable 5 --- dingo/core/posterior_models/base_model.py | 30 ++++++++++++++++++ tests/core/test_build_model.py | 38 +++++++++++++++++++++++ 2 files changed, 68 insertions(+) diff --git a/dingo/core/posterior_models/base_model.py b/dingo/core/posterior_models/base_model.py index de6797b4..a46e32d6 100755 --- a/dingo/core/posterior_models/base_model.py +++ b/dingo/core/posterior_models/base_model.py @@ -99,6 +99,36 @@ def initialize_network(self): """ pass + # Parameter contract read by samplers (e.g. FlowFactor.from_model on the + # factorized-sampler branch). These accessors are the supported interface; the + # location inside the metadata dict is an implementation detail. All three read + # the model's *own* train settings — also for unconditional (density-recovery) + # models, whose ``metadata["base"]`` describes the base model's data pipeline, + # not this network's standardization. + + @property + def inference_parameters(self) -> list: + """Names of the parameters this distribution models, in the order of the + network's theta columns.""" + return list(self.metadata["train_settings"]["data"]["inference_parameters"]) + + @property + def context_parameters(self) -> list: + """Names of parameters the network conditions on in addition to the data + (e.g. GNPE proxies, chained-inference conditioning), in the order of the + network's context-parameter columns. Empty for plain NPE and unconditional + models.""" + return list( + self.metadata["train_settings"]["data"].get("context_parameters") or [] + ) + + @property + def standardization(self) -> dict: + """``{"mean": {name: float}, "std": {name: float}}`` for the affine map to + the network's standardized space; covers ``inference_parameters`` and + ``context_parameters``.""" + return self.metadata["train_settings"]["data"]["standardization"] + @abstractmethod def sample(self, *context: torch.Tensor, num_samples: int = 1): """ diff --git a/tests/core/test_build_model.py b/tests/core/test_build_model.py index 8930e603..75a90be9 100644 --- a/tests/core/test_build_model.py +++ b/tests/core/test_build_model.py @@ -279,6 +279,44 @@ def test_normalizing_flow_unconditional(): assert log_prob.shape == (BATCH_SIZE,) +# ----------------------------------------------------------------------------------- +# Parameter-contract accessors (interface consumed by the factorized sampler) +# ----------------------------------------------------------------------------------- + + +def test_parameter_contract_accessors(): + """inference_parameters / context_parameters / standardization read the model's + own train_settings["data"], the interface FlowFactor.from_model consumes.""" + settings = model_settings("normalizing_flow") + parameters = [f"p{i}" for i in range(NUM_PARAMETERS)] + settings["train_settings"]["data"] = { + "inference_parameters": parameters, + "context_parameters": ["ra", "dec"], + "standardization": { + "mean": {p: 0.0 for p in parameters + ["ra", "dec"]}, + "std": {p: 1.0 for p in parameters + ["ra", "dec"]}, + }, + } + pm = build_model_from_kwargs(settings=settings, device="cpu") + assert pm.inference_parameters == parameters + assert pm.context_parameters == ["ra", "dec"] + assert set(pm.standardization["mean"]) == set(parameters + ["ra", "dec"]) + + +def test_parameter_contract_defaults(): + """context_parameters is [] when absent (plain NPE) or None (written as null by + some configs).""" + settings = model_settings("normalizing_flow") + settings["train_settings"]["data"] = { + "inference_parameters": ["chirp_mass"], + "standardization": {"mean": {"chirp_mass": 30.0}, "std": {"chirp_mass": 5.0}}, + } + pm = build_model_from_kwargs(settings=settings, device="cpu") + assert pm.context_parameters == [] + pm.metadata["train_settings"]["data"]["context_parameters"] = None + assert pm.context_parameters == [] + + # ----------------------------------------------------------------------------------- # SVD initial-weight seeding # ----------------------------------------------------------------------------------- From f917def274bc01f5a02e1a978afce0a945127340 Mon Sep 17 00:00:00 2001 From: Maximilian Dax Date: Sun, 12 Jul 2026 18:25:42 +0200 Subject: [PATCH 4/9] Move training and inference to dict batches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Batches are now dicts of named tensors end to end (NN_Build_System_Design §4.4): SelectKeys replaces UnpackDict at the end of the training transform chain, train/test epochs move a dict to device, and the NeuralDistribution methods sample / sample_and_log_prob / log_prob / loss take a single context dict (None for unconditional models) instead of positional *context tensors. FlowWrapper and ContinuousFlow route context[k] for k in context_keys into the embedding network, so tensor ordering is declared once at build time instead of being implied by the loader's key order; a missing key raises a ValueError naming it. Tensors stay positional inside the network modules. autocomplete_model_kwargs reads the named sample, which removes the GNPE-proxy detection via IndexError and the overloaded third data slot. Sampler call sites and SampleDataset build the context dict directly. The sampler-side dict signature is part 2 of the NN <-> sampler interface (hackathon/NN_Sampler_Interface.md §2); the sampler-revamp call-site updates follow once both branches are on hackathon-1. Co-Authored-By: Claude Fable 5 --- .../unconditional_density_estimation.py | 10 +-- dingo/core/nn/cfnets.py | 25 ++++++ dingo/core/nn/nsf.py | 84 +++++++++++-------- dingo/core/posterior_models/base_model.py | 59 +++++++------ dingo/core/posterior_models/build_model.py | 41 ++++----- dingo/core/posterior_models/cflow_base.py | 37 +++++--- dingo/core/posterior_models/flow_matching.py | 10 +-- .../core/posterior_models/normalizing_flow.py | 16 ++-- dingo/core/posterior_models/score_matching.py | 9 +- dingo/core/samplers.py | 24 +++--- dingo/gw/training/train_builders.py | 6 +- dingo/gw/transforms/general_transforms.py | 22 ++++- tests/core/test_build_model.py | 40 +++++---- tests/core/test_nsf.py | 48 ++++++----- .../gw/transforms/test_general_transforms.py | 15 +++- 15 files changed, 274 insertions(+), 172 deletions(-) diff --git a/dingo/core/density/unconditional_density_estimation.py b/dingo/core/density/unconditional_density_estimation.py index 53cd0885..8656052b 100644 --- a/dingo/core/density/unconditional_density_estimation.py +++ b/dingo/core/density/unconditional_density_estimation.py @@ -13,10 +13,8 @@ class SampleDataset(torch.utils.data.Dataset): """ Dataset class for unconditional density estimation. - This is required, since the training method of dingo.core.posterior_models.Base - expects a tuple of (theta, *context) as output of the DataLoader, but here we have - no context, so len(context) = 0. This SampleDataset therefore returns a tuple - (theta, ) instead of just theta. + The training loop of dingo.core.posterior_models expects dict batches keyed by + name; here there is no context, so a sample is just the parameters. """ def __init__(self, data): @@ -27,8 +25,8 @@ def __len__(self): return self.length def __getitem__(self, index): - """Return the data and labels at the given index as a tuple of length 1.""" - return (self.data[index],) + """Return the parameters at the given index as a dict.""" + return {"inference_parameters": self.data[index]} def train_unconditional_density_estimator( diff --git a/dingo/core/nn/cfnets.py b/dingo/core/nn/cfnets.py index 77529f79..a55dc2cf 100644 --- a/dingo/core/nn/cfnets.py +++ b/dingo/core/nn/cfnets.py @@ -33,6 +33,7 @@ def __init__( theta_embedding_net: nn.Module = torch.nn.Identity(), context_with_glu: bool = False, theta_with_glu: bool = False, + context_keys: tuple = ("waveform",), ): """ Parameters @@ -48,6 +49,9 @@ def __init__( theta_with_glu: bool = False Whether to provide theta (and t) as GLU or main input to the continuous_flow_net. + context_keys: tuple = ("waveform",) + Keys of the context dict that the context embedding network consumes, in + the order of its forward arguments. """ super(ContinuousFlow, self).__init__() self.continuous_flow_net = continuous_flow_net @@ -55,11 +59,25 @@ def __init__( self.theta_embedding_net = theta_embedding_net self.theta_with_glu = theta_with_glu self.context_with_glu = context_with_glu + self.context_keys = tuple(context_keys) self._use_cache = None self._cached_context = None self._cached_context_embedding = None + def unpack_context(self, context: dict = None) -> tuple: + """Select the tensors in context_keys (in order) from the context dict. + Tensors stay positional inside the network modules.""" + if context is None: + return () + missing = [k for k in self.context_keys if k not in context] + if missing: + raise ValueError( + f"Context is missing keys {missing}: expected {self.context_keys}, " + f"got {sorted(context)}." + ) + return tuple(context[k] for k in self.context_keys) + @property def use_cache(self): # unless set explicitly, use_cache is True in eval mode and False in train mode @@ -226,12 +244,19 @@ def create_cf( context_features=glu_dim, ) + # With added_context, the embedding merges (waveform, context_parameters). + if embedding_kwargs is not None and embedding_kwargs.get("added_context"): + context_keys = ("waveform", "context_parameters") + else: + context_keys = ("waveform",) + model = ContinuousFlow( continuous_flow_net, context_embedding, theta_embedding, theta_with_glu=posterior_kwargs.get("theta_with_glu", False), context_with_glu=posterior_kwargs.get("context_with_glu", False), + context_keys=context_keys, ) return model diff --git a/dingo/core/nn/nsf.py b/dingo/core/nn/nsf.py index b0727df1..1e844db1 100644 --- a/dingo/core/nn/nsf.py +++ b/dingo/core/nn/nsf.py @@ -189,56 +189,63 @@ def create_transform( class FlowWrapper(nn.Module): """ - This class wraps the neural spline flow. It is required for multiple - reasons. (i) some embedding networks take tuples as input, which is not - supported by the nflows package. (ii) paralellization across multiple - GPUs requires a forward method, but the relevant flow method for training - is log_prob. + This class wraps the neural spline flow, and routes named context tensors into + the embedding network. It is required for multiple reasons. (i) The embedding + network can consume several context tensors (declared in context_keys), which is + not supported by the nflows package. (ii) Parallelization across multiple GPUs + requires a forward method, but the relevant flow method for training is log_prob. """ - def __init__(self, flow: flows.base.Flow, embedding_net: nn.Module = None): + def __init__( + self, + flow: flows.base.Flow, + embedding_net: nn.Module = None, + context_keys: tuple = ("waveform",), + ): """ - :param flow: flows.base.Flow :param embedding_net: nn.Module + :param context_keys: tuple + Keys of the context dict that the embedding network consumes, in the + order of its forward arguments. """ super(FlowWrapper, self).__init__() self.embedding_net = embedding_net self.flow = flow - - def log_prob(self, y, *x): - if len(x) > 0: - if self.embedding_net is not None: - x = self.embedding_net(*x) - return self.flow.log_prob(y, x) - else: - # if there is no context + self.context_keys = tuple(context_keys) + + def _embed_context(self, context: dict): + """Select the tensors in context_keys (in order) and embed them.""" + missing = [k for k in self.context_keys if k not in context] + if missing: + raise ValueError( + f"Context is missing keys {missing}: expected {self.context_keys}, " + f"got {sorted(context)}." + ) + x = [context[k] for k in self.context_keys] + if self.embedding_net is not None: + return self.embedding_net(*x) + if len(x) != 1: + raise ValueError("Multiple context tensors require an embedding network.") + return x[0] + + def log_prob(self, y, context: dict = None): + if context is None: return self.flow.log_prob(y) + return self.flow.log_prob(y, self._embed_context(context)) - def sample(self, *x, num_samples=1): - if len(x) > 0: - if self.embedding_net is not None: - x = self.embedding_net(*x) - return self.flow.sample(num_samples, x) - else: - # if there is no context, omit the context argument + def sample(self, context: dict = None, num_samples: int = 1): + if context is None: return self.flow.sample(num_samples) + return self.flow.sample(num_samples, self._embed_context(context)) - def sample_and_log_prob(self, *x, num_samples=1): - if len(x) > 0: - if self.embedding_net is not None: - x = self.embedding_net(*x) - return self.flow.sample_and_log_prob(num_samples, x) - else: - # if there is no context, omit the context argument + def sample_and_log_prob(self, context: dict = None, num_samples: int = 1): + if context is None: return self.flow.sample_and_log_prob(num_samples) + return self.flow.sample_and_log_prob(num_samples, self._embed_context(context)) - def forward(self, y, *x): - if len(x) > 0: - return self.log_prob(y, *x) - else: - # if there is no context, omit the context argument - return self.log_prob(y) + def forward(self, y, context: dict = None): + return self.log_prob(y, context) def create_nsf_model( @@ -337,7 +344,12 @@ def create_nsf_with_rb_projection_embedding_net( **embedding_kwargs ) flow = create_nsf_model(**posterior_kwargs) - model = FlowWrapper(flow, embedding_net) + # With added_context, the embedding merges (waveform, context_parameters). + if embedding_kwargs.get("added_context"): + context_keys = ("waveform", "context_parameters") + else: + context_keys = ("waveform",) + model = FlowWrapper(flow, embedding_net, context_keys) return model diff --git a/dingo/core/posterior_models/base_model.py b/dingo/core/posterior_models/base_model.py index a46e32d6..788a0ca9 100755 --- a/dingo/core/posterior_models/base_model.py +++ b/dingo/core/posterior_models/base_model.py @@ -130,7 +130,7 @@ def standardization(self) -> dict: return self.metadata["train_settings"]["data"]["standardization"] @abstractmethod - def sample(self, *context: torch.Tensor, num_samples: int = 1): + def sample(self, context: Optional[dict] = None, num_samples: int = 1): """ Sample parameters theta from the posterior model, @@ -138,9 +138,10 @@ def sample(self, *context: torch.Tensor, num_samples: int = 1): Parameters ---------- - context: torch.Tensor - Context information (typically observed data). Should have a batch - dimension (even if size B = 1). + context: dict = None + Named context tensors (keyed like the training batches, e.g. "waveform", + "context_parameters"). Each tensor should have a batch dimension (even if + size B = 1). None for unconditional models. num_samples: int = 1 Number of samples to generate. @@ -152,7 +153,7 @@ def sample(self, *context: torch.Tensor, num_samples: int = 1): pass @abstractmethod - def sample_and_log_prob(self, *context: torch.Tensor, num_samples: int = 1): + def sample_and_log_prob(self, context: Optional[dict] = None, num_samples: int = 1): """ Sample parameters theta from the posterior model, @@ -164,9 +165,10 @@ def sample_and_log_prob(self, *context: torch.Tensor, num_samples: int = 1): Parameters ---------- - context: torch.Tensor - Context information (typically observed data). Should have a batch - dimension (even if size B = 1). + context: dict = None + Named context tensors (keyed like the training batches). Each tensor + should have a batch dimension (even if size B = 1). None for + unconditional models. num_samples: int = 1 Number of samples to generate. @@ -178,7 +180,7 @@ def sample_and_log_prob(self, *context: torch.Tensor, num_samples: int = 1): pass @abstractmethod - def log_prob(self, theta: torch.Tensor, *context: torch.Tensor): + def log_prob(self, theta: torch.Tensor, context: Optional[dict] = None): """ Evaluate the log posterior density, @@ -188,9 +190,11 @@ def log_prob(self, theta: torch.Tensor, *context: torch.Tensor): ---------- theta: torch.Tensor Parameter values at which to evaluate the density. Should have a batch - dimension (even if size B = 1). - context: torch.Tensor - Context information (typically observed data). Must have context.shape[0] = B. + dimension (even if size B = 1). Columns are ordered as + inference_parameters. + context: dict = None + Named context tensors (keyed like the training batches). Each tensor must + have leading dimension B. None for unconditional models. Returns ------- @@ -200,7 +204,7 @@ def log_prob(self, theta: torch.Tensor, *context: torch.Tensor): pass @abstractmethod - def loss(self, theta: torch.Tensor, *context: torch.Tensor): + def loss(self, theta: torch.Tensor, context: Optional[dict] = None): """ Compute the loss for a batch of data. @@ -209,9 +213,10 @@ def loss(self, theta: torch.Tensor, *context: torch.Tensor): theta: torch.Tensor Parameter values at which to evaluate the density. Should have a batch dimension (even if size B = 1). - context: torch.Tensor - Context information (typically observed data). Must have the same leading - (batch) dimension as theta. + context: dict = None + Named context tensors (keyed like the training batches). Each tensor must + have the same leading (batch) dimension as theta. None for unconditional + models. Returns ------- @@ -523,15 +528,16 @@ def train_epoch(pm, dataloader): for batch_idx, data in enumerate(dataloader): loss_info.update_timer() pm.optimizer.zero_grad() - # data to device - data = [d.to(pm.device, non_blocking=True) for d in data] - # compute loss - loss = pm.loss(data[0], *data[1:]) + # Batches are dicts of named tensors; all entries besides + # inference_parameters are context for the network. + data = {k: v.to(pm.device, non_blocking=True) for k, v in data.items()} + theta = data.pop("inference_parameters") + loss = pm.loss(theta, data if data else None) # backward pass and optimizer step loss.backward() pm.optimizer.step() # update loss for history and logging - loss_info.update(loss.detach().item(), len(data[0])) + loss_info.update(loss.detach().item(), len(theta)) loss_info.print_info(batch_idx) return loss_info.get_avg() @@ -550,12 +556,13 @@ def test_epoch(pm, dataloader): for batch_idx, data in enumerate(dataloader): loss_info.update_timer() - # data to device - data = [d.to(pm.device, non_blocking=True) for d in data] - # compute loss - loss = pm.loss(data[0], *data[1:]) + # Batches are dicts of named tensors; all entries besides + # inference_parameters are context for the network. + data = {k: v.to(pm.device, non_blocking=True) for k, v in data.items()} + theta = data.pop("inference_parameters") + loss = pm.loss(theta, data if data else None) # update loss for history and logging - loss_info.update(loss.item(), len(data[0])) + loss_info.update(loss.item(), len(theta)) loss_info.print_info(batch_idx) return loss_info.get_avg() diff --git a/dingo/core/posterior_models/build_model.py b/dingo/core/posterior_models/build_model.py index 2a96d5ee..0d4c6af5 100644 --- a/dingo/core/posterior_models/build_model.py +++ b/dingo/core/posterior_models/build_model.py @@ -64,38 +64,41 @@ def build_model_from_kwargs( return model(model_filename=filename, metadata=settings, **kwargs) -def autocomplete_model_kwargs(model_kwargs: dict, data_sample: list): +def autocomplete_model_kwargs(model_kwargs: dict, data_sample: dict): """ Autocomplete the model kwargs from train_settings and data_sample from the dataloader: - * set input dimension of embedding net to shape of data_sample[1] - * set dimension of parameter space to len(data_sample[0]) - * set added_context flag of embedding net if required for gnpe proxies - * set context dim of posterior model to output dim of embedding net + gnpe proxy dim + * set input dimension of embedding net to the shape of the waveform data + * set dimension of parameter space to the number of inference parameters + * set added_context flag of embedding net if required for context parameters + (e.g., GNPE proxies) + * set context dim of posterior model to output dim of embedding net + dimension + of the context parameters Parameters ---------- model_kwargs: dict Model settings, which are modified in-place. - data_sample: list - Sample from dataloader (e.g., wfd[0]) used for autocomplection. - Should be of format [parameters, GW data, gnpe_proxies], where the - last element is only there is GNPE proxies are required. + data_sample: dict + Sample from dataloader (e.g., wfd[0]) used for autocompletion, with keys + "inference_parameters", "waveform", and (only if the network is conditioned + on additional parameters) "context_parameters". """ # set input dims from ifo_list and domain information - model_kwargs["embedding_kwargs"]["input_dims"] = list(data_sample[1].shape) + model_kwargs["embedding_kwargs"]["input_dims"] = list(data_sample["waveform"].shape) # set dimension of parameter space of posterior model - model_kwargs["posterior_kwargs"]["input_dim"] = len(data_sample[0]) - # set added_context flag of embedding net if GNPE proxies are required - # set context dim of nsf to output dim of embedding net + GNPE proxy dim - try: - gnpe_proxy_dim = len(data_sample[2]) + model_kwargs["posterior_kwargs"]["input_dim"] = len( + data_sample["inference_parameters"] + ) + # set added_context flag of embedding net if context parameters are required + # set context dim of nsf to output dim of embedding net + context parameter dim + if "context_parameters" in data_sample: model_kwargs["embedding_kwargs"]["added_context"] = True - model_kwargs["posterior_kwargs"]["context_dim"] = ( - model_kwargs["embedding_kwargs"]["output_dim"] + gnpe_proxy_dim - ) - except IndexError: + model_kwargs["posterior_kwargs"]["context_dim"] = model_kwargs[ + "embedding_kwargs" + ]["output_dim"] + len(data_sample["context_parameters"]) + else: model_kwargs["embedding_kwargs"]["added_context"] = False model_kwargs["posterior_kwargs"]["context_dim"] = model_kwargs[ "embedding_kwargs" diff --git a/dingo/core/posterior_models/cflow_base.py b/dingo/core/posterior_models/cflow_base.py index cdc5f251..a35f9cd4 100644 --- a/dingo/core/posterior_models/cflow_base.py +++ b/dingo/core/posterior_models/cflow_base.py @@ -129,7 +129,7 @@ def initialize_network(self): model_kwargs["initial_weights"] = self.initial_weights self.network = create_cf(**model_kwargs) - def sample(self, *context: torch.Tensor, num_samples: int = None): + def sample(self, context: dict = None, num_samples: int = 1): """ Sample parameters theta from the posterior model, @@ -139,9 +139,9 @@ def sample(self, *context: torch.Tensor, num_samples: int = None): Parameters ---------- - context: torch.Tensor - Context information (typically observed data). Should have a batch - dimension (even if size B = 1). + context: dict = None + Named context tensors (keyed like the training batches). Each tensor + should have a batch dimension (even if size B = 1). num_samples: int = 1 Number of samples to generate. @@ -150,6 +150,12 @@ def sample(self, *context: torch.Tensor, num_samples: int = None): samples: torch.Tensor Shape (B, num_samples, dim(theta)) """ + context = self.network.unpack_context(context) + if len(context) == 0: + raise ValueError( + "Sampling requires context; unconditional continuous flows are not " + "supported." + ) context_size = context[0].shape[0] theta_0 = self.sample_theta_0(num_samples * context_size) context = [repeat_rows(c, num_reps=num_samples) for c in context] @@ -169,7 +175,7 @@ def sample(self, *context: torch.Tensor, num_samples: int = None): return theta_1 # MD: rename log_prob_batch, extract eps from self.epsilon_ode_integration - def log_prob(self, theta: torch.Tensor, *context: torch.Tensor, hutchinson=False): + def log_prob(self, theta: torch.Tensor, context: dict = None, hutchinson=False): """ Evaluate the log posterior density, @@ -187,8 +193,9 @@ def log_prob(self, theta: torch.Tensor, *context: torch.Tensor, hutchinson=False theta: torch.Tensor Parameter values at which to evaluate the density. Should have a batch dimension (even if size B = 1). - context: torch.Tensor - Context information (typically observed data). Must have context.shape[0] = B. + context: dict = None + Named context tensors (keyed like the training batches). Each tensor must + have leading dimension B. None for unconditional models. hutchinson Returns @@ -196,6 +203,7 @@ def log_prob(self, theta: torch.Tensor, *context: torch.Tensor, hutchinson=False log_prob: torch.Tensor Shape (B,) """ + context = self.network.unpack_context(context) theta_and_div_init = torch.cat( (theta, torch.zeros((theta.shape[0],), device=theta.device).unsqueeze(1)), dim=1, @@ -218,7 +226,7 @@ def log_prob(self, theta: torch.Tensor, *context: torch.Tensor, hutchinson=False log_prior = compute_log_prior(theta_0) return (log_prior - divergence).detach() - def sample_and_log_prob(self, *context: torch.Tensor, num_samples: int = None): + def sample_and_log_prob(self, context: dict = None, num_samples: int = 1): """ Sample parameters theta from the posterior model, @@ -233,9 +241,9 @@ def sample_and_log_prob(self, *context: torch.Tensor, num_samples: int = None): Parameters ---------- - context: torch.Tensor - Context information (typically observed data). Should have a batch - dimension (even if size B = 1). + context: dict = None + Named context tensors (keyed like the training batches). Each tensor + should have a batch dimension (even if size B = 1). num_samples: int = 1 Number of samples to generate. @@ -244,7 +252,12 @@ def sample_and_log_prob(self, *context: torch.Tensor, num_samples: int = None): samples, log_prob: torch.Tensor, torch.Tensor Shapes (B, num_samples, dim(theta)), (B, num_samples) """ - + context = self.network.unpack_context(context) + if len(context) == 0: + raise ValueError( + "Sampling requires context; unconditional continuous flows are not " + "supported." + ) context_size = context[0].shape[0] theta_0 = self.sample_theta_0(num_samples * context_size) context = [repeat_rows(c, num_reps=num_samples) for c in context] diff --git a/dingo/core/posterior_models/flow_matching.py b/dingo/core/posterior_models/flow_matching.py index f85d4e91..58a333d9 100644 --- a/dingo/core/posterior_models/flow_matching.py +++ b/dingo/core/posterior_models/flow_matching.py @@ -61,7 +61,7 @@ def evaluate_vector_field(self, t, theta_t, *context_data): t = t * torch.ones(len(theta_t), device=theta_t.device) return self.network(t, theta_t, *context_data) - def loss(self, theta, *context): + def loss(self, theta, context: dict = None): """ Calculates loss as the mean squared error between the predicted vector field and the vector field for transporting the parameter data to samples from the prior. @@ -71,9 +71,9 @@ def loss(self, theta, *context): theta: torch.Tensor Parameter values at which to evaluate the density. Should have a batch dimension (even if size B = 1). - context: torch.Tensor - Context information (typically observed data). Must have the same leading - (batch) dimension as theta. + context: dict = None + Named context tensors (keyed like the training batches). Each tensor must + have the same leading (batch) dimension as theta. Returns ------- @@ -89,7 +89,7 @@ def loss(self, theta, *context): theta_t = ot_conditional_flow(theta_0, theta_1, t, self.sigma_min) true_vf = theta - (1 - self.sigma_min) * theta_0 - predicted_vf = self.network(t, theta_t, *context) + predicted_vf = self.network(t, theta_t, *self.network.unpack_context(context)) loss = mse(predicted_vf, true_vf) return loss diff --git a/dingo/core/posterior_models/normalizing_flow.py b/dingo/core/posterior_models/normalizing_flow.py index ddf1921e..ed9b407e 100644 --- a/dingo/core/posterior_models/normalizing_flow.py +++ b/dingo/core/posterior_models/normalizing_flow.py @@ -51,14 +51,14 @@ def initialize_network(self): else: self.network = create_nsf_wrapped(**model_kwargs["posterior_kwargs"]) - def log_prob(self, theta, *context): - return self.network(theta, *context) + def log_prob(self, theta, context: dict = None): + return self.network(theta, context) - def sample(self, *context, num_samples: int = 1): - return self.network.sample(*context, num_samples=num_samples) + def sample(self, context: dict = None, num_samples: int = 1): + return self.network.sample(context, num_samples=num_samples) - def sample_and_log_prob(self, *context, num_samples: int = 1): - return self.network.sample_and_log_prob(*context, num_samples=num_samples) + def sample_and_log_prob(self, context: dict = None, num_samples: int = 1): + return self.network.sample_and_log_prob(context, num_samples=num_samples) - def loss(self, theta, *context): - return -self.network(theta, *context).mean() + def loss(self, theta, context: dict = None): + return -self.network(theta, context).mean() diff --git a/dingo/core/posterior_models/score_matching.py b/dingo/core/posterior_models/score_matching.py index 5031a906..baac9593 100644 --- a/dingo/core/posterior_models/score_matching.py +++ b/dingo/core/posterior_models/score_matching.py @@ -50,7 +50,7 @@ def __init__(self, **kwargs): likelihood_weighting ) - def loss(self, theta, *context_data): + def loss(self, theta, context: dict = None): """ Returns the score matching loss for parameters theta conditioned on context. @@ -58,8 +58,9 @@ def loss(self, theta, *context_data): ---------- theta: torch.tensor parameters (e.g., binary-black hole parameters) - *context_data: list[torch.Tensor] - context data (e.g., gravitational-wave data) + context: dict = None + Named context tensors (keyed like the training batches), e.g. + gravitational-wave data. Returns ------- @@ -67,7 +68,7 @@ def loss(self, theta, *context_data): Loss. """ t, theta_t, score = self.get_t_theta_t_score(theta_1=theta) - pred_score = self.network(t, theta_t, *context_data) + pred_score = self.network(t, theta_t, *self.network.unpack_context(context)) weighting = self.likelihood_weighting(t) losses = torch.square(pred_score - score) diff --git a/dingo/core/samplers.py b/dingo/core/samplers.py index b01cf228..90360d7c 100644 --- a/dingo/core/samplers.py +++ b/dingo/core/samplers.py @@ -153,19 +153,18 @@ def _run_sampler( # requested sample. We therefore apply pre-processing only once. x = self.transform_pre(context) # Require a batch dimension for the embedding network. - x = x.unsqueeze(0) - x = [x] + x = {"waveform": x.unsqueeze(0)} else: if context is not None: print("Unconditional model. Ignoring context.") - x = [] + x = None # For a normalizing flow, we get the log_prob for "free" when sampling, # so we always include this. For other architectures, it may make sense to # have a flag for whether to calculate the log_prob. self.model.network.eval() with torch.no_grad(): - y, log_prob = self.model.sample_and_log_prob(*x, num_samples=num_samples) + y, log_prob = self.model.sample_and_log_prob(x, num_samples=num_samples) if not self.unconditional_model: # Squeeze the batch dimension added earlier. @@ -275,14 +274,15 @@ def log_prob(self, samples: pd.DataFrame | dict) -> np.ndarray: # Context is the same for each sample. Expand across batch dimension after # pre-processing. x = self.transform_pre(self.context) - x = x.expand(len(samples), *x.shape) # TODO: Make this more efficient. - x = [x] + x = { + "waveform": x.expand(len(samples), *x.shape) + } # TODO: Make this more efficient. else: - x = [] + x = None self.model.network.eval() with torch.no_grad(): - log_prob = self.model.log_prob(y, *x) + log_prob = self.model.log_prob(y, x) log_prob = log_prob.cpu().numpy() log_prob -= np.sum(np.log(std)) @@ -490,12 +490,10 @@ def _run_sampler( time_sample_start = time.time() self.model.network.eval() with torch.no_grad(): + model_context = {"waveform": x["data"]} if "context_parameters" in x: - y, log_prob = self.model.sample_and_log_prob( - x["data"], x["context_parameters"] - ) - else: - y, log_prob = self.model.sample_and_log_prob(x["data"]) + model_context["context_parameters"] = x["context_parameters"] + y, log_prob = self.model.sample_and_log_prob(model_context) # Squeeze the extra dimension added by sample_and_log_prob(num_samples=1). y = y.squeeze(1) diff --git a/dingo/gw/training/train_builders.py b/dingo/gw/training/train_builders.py index ca62856b..494c798c 100755 --- a/dingo/gw/training/train_builders.py +++ b/dingo/gw/training/train_builders.py @@ -17,7 +17,7 @@ AddWhiteNoiseComplex, SelectStandardizeRepackageParameters, RepackageStrainsAndASDS, - UnpackDict, + SelectKeys, GNPECoalescenceTimes, SampleExtrinsicParameters, GetDetectorTimes, @@ -181,7 +181,7 @@ def set_train_transforms(wfd, data_settings, asd_dataset_path, omit_transforms=N else: selected_keys = ["inference_parameters", "waveform"] - transforms.append(UnpackDict(selected_keys=selected_keys)) + transforms.append(SelectKeys(selected_keys=selected_keys)) # Drop transforms that are not desired. This is useful for generating, e.g., # noise-free data, or for producing data not formatted for input to the network. @@ -253,7 +253,7 @@ def build_svd_for_embedding_network( AddWhiteNoiseComplex, RepackageStrainsAndASDS, SelectStandardizeRepackageParameters, - UnpackDict, + SelectKeys, CropMaskStrainRandom, ], ) diff --git a/dingo/gw/transforms/general_transforms.py b/dingo/gw/transforms/general_transforms.py index 94eeba7d..cef3d8ab 100644 --- a/dingo/gw/transforms/general_transforms.py +++ b/dingo/gw/transforms/general_transforms.py @@ -7,4 +7,24 @@ def __init__(self, selected_keys): self.selected_keys = selected_keys def __call__(self, input_sample): - return [input_sample[k] for k in self.selected_keys] \ No newline at end of file + return [input_sample[k] for k in self.selected_keys] + + +class SelectKeys(object): + """ + Restricts the sample dictionary to selected_keys, to prepare it for final output + of the dataloader. In contrast to UnpackDict, the sample remains a dictionary, so + that batches are keyed by name rather than by position. + """ + + def __init__(self, selected_keys): + self.selected_keys = selected_keys + + def __call__(self, input_sample): + missing = [k for k in self.selected_keys if k not in input_sample] + if missing: + raise KeyError( + f"Sample is missing keys {missing}: expected {self.selected_keys}, " + f"got {sorted(input_sample)}." + ) + return {k: input_sample[k] for k in self.selected_keys} \ No newline at end of file diff --git a/tests/core/test_build_model.py b/tests/core/test_build_model.py index 75a90be9..306303b9 100644 --- a/tests/core/test_build_model.py +++ b/tests/core/test_build_model.py @@ -117,22 +117,24 @@ def model_settings(posterior_model_type): def data_sample(with_gnpe_proxies=False): - """A sample in the format produced by the dataloader (wfd[0] after UnpackDict): - [inference_parameters, strain data(, gnpe_proxies)].""" - sample = [ - np.random.rand(NUM_PARAMETERS).astype(np.float32), - np.random.rand(*DATA_SHAPE).astype(np.float32), - ] + """A sample in the format produced by the dataloader (wfd[0] after SelectKeys): + a dict with inference_parameters, waveform(, context_parameters).""" + sample = { + "inference_parameters": np.random.rand(NUM_PARAMETERS).astype(np.float32), + "waveform": np.random.rand(*DATA_SHAPE).astype(np.float32), + } if with_gnpe_proxies: - sample.append(np.random.rand(GNPE_PROXY_DIM).astype(np.float32)) + sample["context_parameters"] = np.random.rand(GNPE_PROXY_DIM).astype( + np.float32 + ) return sample def batch(model_type="normalizing_flow", with_gnpe_proxies=False): theta = torch.rand(BATCH_SIZE, NUM_PARAMETERS) - context = [torch.rand(BATCH_SIZE, *DATA_SHAPE)] + context = {"waveform": torch.rand(BATCH_SIZE, *DATA_SHAPE)} if with_gnpe_proxies: - context.append(torch.rand(BATCH_SIZE, GNPE_PROXY_DIM)) + context["context_parameters"] = torch.rand(BATCH_SIZE, GNPE_PROXY_DIM) return theta, context @@ -218,27 +220,27 @@ def test_model_forward_passes(posterior_model_type): ) theta, context = batch() - loss = pm.loss(theta, *context) + loss = pm.loss(theta, context) assert loss.shape == () assert torch.isfinite(loss) pm.network.eval() with torch.no_grad(): - samples = pm.sample(*context, num_samples=3) + samples = pm.sample(context, num_samples=3) assert samples.shape == (BATCH_SIZE, 3, NUM_PARAMETERS) - log_prob = pm.log_prob(theta, *context) + log_prob = pm.log_prob(theta, context) assert log_prob.shape == (BATCH_SIZE,) assert torch.isfinite(log_prob).all() - samples, log_prob = pm.sample_and_log_prob(*context, num_samples=3) + samples, log_prob = pm.sample_and_log_prob(context, num_samples=3) assert samples.shape == (BATCH_SIZE, 3, NUM_PARAMETERS) assert log_prob.shape == (BATCH_SIZE, 3) def test_normalizing_flow_with_gnpe_context(): - """With added_context=True, the embedding merges (data, proxies) via ModuleMerger, - and sampling/density evaluation take two context tensors.""" + """With added_context=True, the embedding merges (waveform, context_parameters) + via ModuleMerger, and sampling/density evaluation consume both context entries.""" settings = model_settings("normalizing_flow") model = settings["train_settings"]["model"] model["embedding_kwargs"]["added_context"] = True @@ -247,14 +249,18 @@ def test_normalizing_flow_with_gnpe_context(): pm = build_model_from_kwargs(settings=settings, device="cpu") theta, context = batch(with_gnpe_proxies=True) - loss = pm.loss(theta, *context) + loss = pm.loss(theta, context) assert torch.isfinite(loss) pm.network.eval() with torch.no_grad(): - log_prob = pm.log_prob(theta, *context) + log_prob = pm.log_prob(theta, context) assert log_prob.shape == (BATCH_SIZE,) + # A missing declared context entry must fail loudly. + with pytest.raises(ValueError, match="missing keys"): + pm.log_prob(theta, {"waveform": context["waveform"]}) + def test_normalizing_flow_unconditional(): """Without embedding_kwargs, an unconditional flow is built (the models-as-priors diff --git a/tests/core/test_nsf.py b/tests/core/test_nsf.py index 0e3f73b0..d68c7931 100644 --- a/tests/core/test_nsf.py +++ b/tests/core/test_nsf.py @@ -116,15 +116,21 @@ def data_setup_nsf_small(): (d.batch_size, d.context_dim - d.embedding_net_kwargs["output_dim"]) ) d.y = torch.ones((d.batch_size, d.input_dim)) + d.context = {"waveform": d.x, "context_parameters": d.z} # build d.yy, which depends on input d.zz d.xx = torch.cat((d.x, d.x)) d.yy = torch.cat((d.y, -d.y)) d.zz = torch.cat((d.z, -d.z)) + d.contextcontext = {"waveform": d.xx, "context_parameters": d.zz} return d +# context_keys for embedding networks built with added_context=True. +CONTEXT_KEYS = ("waveform", "context_parameters") + + def test_nsf_number_of_parameters(data_setup_nsf_large): """ Builds a neural spline flow with the hyperparameters from that used in @@ -151,9 +157,9 @@ def test_sample_method_of_nsf(data_setup_nsf_small): embedding_net = d.embedding_net_builder(**d.embedding_net_kwargs) flow = d.nde_builder(**d.nde_kwargs) - model = FlowWrapper(flow, embedding_net) + model = FlowWrapper(flow, embedding_net, CONTEXT_KEYS) - samples = model.sample(d.x, d.z) + samples = model.sample(d.context) # model.sample(num_samples=1) adds an extra dimension that needs to be squeezed. samples = samples.squeeze(1) @@ -163,12 +169,12 @@ def test_sample_method_of_nsf(data_setup_nsf_small): "normalization seems broken." ) + # missing context key with pytest.raises(ValueError): - model.sample(d.z, d.x) - with pytest.raises(ValueError): - model.sample(d.x, d.z, d.z) + model.sample({"waveform": d.x}) + # wrongly-shaped context tensor with pytest.raises(RuntimeError): - model.sample(d.x, d.x) + model.sample({"waveform": d.x, "context_parameters": d.x.flatten(start_dim=1)}) def test_forward_pass_for_log_prob_of_nsf(data_setup_nsf_small): @@ -180,21 +186,21 @@ def test_forward_pass_for_log_prob_of_nsf(data_setup_nsf_small): embedding_net = d.embedding_net_builder(**d.embedding_net_kwargs) flow = d.nde_builder(**d.nde_kwargs) - model = FlowWrapper(flow, embedding_net) + model = FlowWrapper(flow, embedding_net, CONTEXT_KEYS) - loss = -model(d.y, d.x, d.z) + loss = -model(d.y, d.context) assert list(loss.shape) == [d.batch_size], "Unexpected output shape." assert torch.all(loss > 0) and torch.all(loss < 40), ( "Unexpected log prob encountered. Network initialization or " "normalization seems broken." ) + # missing context key with pytest.raises(ValueError): - model(d.y, d.z, d.x) - with pytest.raises(ValueError): - model(d.y, d.x, d.z, d.z) + model(d.y, {"waveform": d.x}) + # wrongly-shaped context tensor with pytest.raises(RuntimeError): - model(d.y, d.x, d.x) + model(d.y, {"waveform": d.x, "context_parameters": d.x.flatten(start_dim=1)}) def test_backward_pass_for_log_prob_of_nsf(data_setup_nsf_small): @@ -207,7 +213,7 @@ def test_backward_pass_for_log_prob_of_nsf(data_setup_nsf_small): embedding_net = d.embedding_net_builder(**d.embedding_net_kwargs) flow = d.nde_builder(**d.nde_kwargs) - model = FlowWrapper(flow, embedding_net) + model = FlowWrapper(flow, embedding_net, CONTEXT_KEYS) optimizer = optim.Adam(model.parameters(), lr=0.003) # Simple train loop. The learned parameters yy are strongly correlated @@ -217,12 +223,12 @@ def test_backward_pass_for_log_prob_of_nsf(data_setup_nsf_small): for idx in range(40): yy = d.yy + 0.02 * torch.rand_like(d.yy) xx = torch.rand_like(d.xx) - loss = -torch.mean(model(yy, xx, d.zz)) + loss = -torch.mean(model(yy, {"waveform": xx, "context_parameters": d.zz})) losses.append(loss.detach().item()) loss.backward() optimizer.step() optimizer.zero_grad() - samples_n = torch.mean(model.sample(d.xx, d.zz, num_samples=100), axis=1) + samples_n = torch.mean(model.sample(d.contextcontext, num_samples=100), axis=1) assert losses[-1] < losses[0], "Loss did not improve in training." assert ( @@ -240,17 +246,19 @@ def test_model_builder_for_nsf_with_rb_embedding_net(data_setup_nsf_small): model = create_nsf_with_rb_projection_embedding_net( d.nde_kwargs, d.embedding_net_kwargs ) + # The builder derives the context routing from added_context. + assert model.context_keys == CONTEXT_KEYS - loss = -model(d.y, d.x, d.z) + loss = -model(d.y, d.context) assert list(loss.shape) == [d.batch_size], "Unexpected output shape." assert torch.all(loss > 0) and torch.all(loss < 40), ( "Unexpected log prob encountered. Network initialization or " "normalization seems broken." ) + # missing context key with pytest.raises(ValueError): - model(d.y, d.z, d.x) - with pytest.raises(ValueError): - model(d.y, d.x, d.z, d.z) + model(d.y, {"waveform": d.x}) + # wrongly-shaped context tensor with pytest.raises(RuntimeError): - model(d.y, d.x, d.x) + model(d.y, {"waveform": d.x, "context_parameters": d.x.flatten(start_dim=1)}) diff --git a/tests/gw/transforms/test_general_transforms.py b/tests/gw/transforms/test_general_transforms.py index 19be7ec3..be8c654f 100644 --- a/tests/gw/transforms/test_general_transforms.py +++ b/tests/gw/transforms/test_general_transforms.py @@ -1,11 +1,22 @@ import pytest import numpy as np -from dingo.gw.transforms import UnpackDict +from dingo.gw.transforms import SelectKeys, UnpackDict def test_UnpackDict(): sample = {'a': 10, 'b': np.random.rand(100), 'c': None} unpack_dict = UnpackDict(['b', 'a']) b, a = unpack_dict(sample) assert id(b) == id(sample['b']) - assert a == sample['a'] \ No newline at end of file + assert a == sample['a'] + + +def test_SelectKeys(): + sample = {'a': 10, 'b': np.random.rand(100), 'c': None} + select_keys = SelectKeys(['b', 'a']) + out = select_keys(sample) + assert list(out) == ['b', 'a'] + assert id(out['b']) == id(sample['b']) + assert out['a'] == sample['a'] + with pytest.raises(KeyError, match="missing keys"): + SelectKeys(['a', 'd'])(sample) \ No newline at end of file From d54a1d332f44ea1e75366f39fab4df066df54d57 Mon Sep 17 00:00:00 2001 From: Maximilian Dax Date: Sun, 12 Jul 2026 19:47:03 +0200 Subject: [PATCH 5/9] Add registry-built embeddings and delete autocomplete_model_kwargs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New model-settings schema (NN_Build_System_Design §4.3/4.5/4.7): a model declares distribution / embedding_net / context_merger, each a registry type plus a flat kwargs dict. Embedding networks follow a small contract (input_keys, output_dim, complete_settings): dense_svd (the classic LinearProjectionRB + DenseResidualNet stack) and the concat context merger are the first registered components. NeuralDistribution.build_embedding_net assembles them generically, so initialize_network loses its embedding_kwargs-presence branches, and any embedding composes with any distribution type. complete_model_settings replaces autocomplete_model_kwargs: each architecture infers its own input dims from a named sample batch, the generic theta/context dims are computed once, and the completed settings are saved in the checkpoint so loading never needs a data sample. Dimensions in user settings and unused context mergers are errors. The V_rb_list deepcopy hack is gone: initial weights are extra constructor kwargs and never touch the saved settings. save_model now enforces that the standardization covers inference and context parameters (NN_Sampler_Interface §1). update_model_config remains the single back-compat boundary: it maps the old posterior_model_type / posterior_kwargs / embedding_kwargs schema (and the older nsf+embedding form) onto the new one, including added_context -> concat merger; old checkpoints build state-dict-identical networks (pinned by test). Deleted along the way: the never-wired embedding_net_builder parameter, create_nsf_wrapped, create_nsf_with_rb_projection_embedding_net. Example configs, pipe defaults, and the architecture docs notebook use the new schema. Co-Authored-By: Claude Fable 5 --- dingo/core/density/__init__.py | 1 + .../unconditional_density_estimation.py | 7 +- dingo/core/nn/cfnets.py | 90 +++--- dingo/core/nn/enets.py | 140 ++++++++- dingo/core/nn/nsf.py | 78 +---- dingo/core/posterior_models/base_model.py | 50 +++ dingo/core/posterior_models/build_model.py | 125 +++++--- dingo/core/posterior_models/cflow_base.py | 19 +- dingo/core/posterior_models/flow_matching.py | 2 +- .../core/posterior_models/normalizing_flow.py | 25 +- dingo/core/posterior_models/score_matching.py | 9 +- dingo/core/registry.py | 7 +- dingo/core/result.py | 2 - dingo/core/transforms.py | 3 +- dingo/core/utils/backward_compatibility.py | 67 +++- dingo/core/utils/condor_utils.py | 52 ++-- dingo/core/utils/gnpeutils.py | 6 +- dingo/core/utils/plotting.py | 4 +- dingo/core/utils/pt_to_hdf5.py | 68 ++-- dingo/core/utils/torchutils.py | 4 +- dingo/gw/training/train_pipeline.py | 23 +- dingo/pipe/default_settings.py | 24 +- dingo/pipe/sampling.py | 6 +- docs/source/network_architecture.ipynb | 11 +- examples/fmpe_model/train_settings.yaml | 81 ++--- examples/gnpe_model/train_settings.yaml | 50 +-- examples/gnpe_model/train_settings_init.yaml | 50 +-- examples/misc/is_settings.yaml | 23 +- examples/npe_model/train_settings.yaml | 50 +-- examples/toy_npe_model/train_settings.yaml | 44 +-- tests/core/test_build_model.py | 294 ++++++++++++------ tests/core/test_nsf.py | 58 ++-- tests/core/test_registry.py | 11 +- .../test_unconditional_density_estimation.py | 16 +- 34 files changed, 924 insertions(+), 576 deletions(-) diff --git a/dingo/core/density/__init__.py b/dingo/core/density/__init__.py index e465636a..af76a40d 100644 --- a/dingo/core/density/__init__.py +++ b/dingo/core/density/__init__.py @@ -3,6 +3,7 @@ This is required for instance to recover the posterior density from GNPE samples, since the density is intractable with GNPE. """ + from .unconditional_density_estimation import train_unconditional_density_estimator from .interpolation import ( interpolated_sample_and_log_prob_multi, diff --git a/dingo/core/density/unconditional_density_estimation.py b/dingo/core/density/unconditional_density_estimation.py index 8656052b..72535920 100644 --- a/dingo/core/density/unconditional_density_estimation.py +++ b/dingo/core/density/unconditional_density_estimation.py @@ -3,6 +3,7 @@ import torch from dingo.core.utils import build_train_and_test_loaders +from dingo.core.utils.backward_compatibility import update_model_config from dingo.core.utils.trainutils import RuntimeLimits import numpy as np import argparse @@ -69,8 +70,10 @@ def train_unconditional_density_estimator( samples_torch = torch.from_numpy((samples - mean) / std).float() # set up density estimation network - settings["model"]["posterior_kwargs"]["input_dim"] = num_params - settings["model"]["posterior_kwargs"]["context_dim"] = None + update_model_config(settings["model"]) # Map old schemas forward. + distribution_kwargs = settings["model"]["distribution"].setdefault("kwargs", {}) + distribution_kwargs["theta_dim"] = num_params + distribution_kwargs["context_dim"] = None # TODO: Allow for other types of density estimators (e.g., flow matching). model = NormalizingFlowPosteriorModel( metadata={"train_settings": settings, "base": copy.deepcopy(result.metadata)}, diff --git a/dingo/core/nn/cfnets.py b/dingo/core/nn/cfnets.py index a55dc2cf..713bcd3d 100644 --- a/dingo/core/nn/cfnets.py +++ b/dingo/core/nn/cfnets.py @@ -1,12 +1,8 @@ -import copy - import numpy as np import torch import torch.nn as nn from dingo.core.utils import torchutils -from dingo.core.nn.enets import create_enet_with_projection_layer_and_dense_resnet - from dingo.core.nn.enets import DenseResidualNet @@ -175,55 +171,47 @@ def forward(self, t, theta, *context): return self.continuous_flow_net(main_input, glu_context) -def create_cf( - posterior_kwargs: dict, embedding_kwargs: dict = None, initial_weights: dict = None -): +def create_cf(kwargs: dict, embedding_net: nn.Module = None): """ - Build a continuous flow based on settings dictionaries. + Build a continuous flow based on completed distribution kwargs. Parameters ---------- - posterior_kwargs: dict - Settings for the flow. This includes the settings for the parameter embedding. - embedding_kwargs: dict - Settings for the context embedding network. - initial_weights: dict - Initial weights for the embedding network (of SVD projection type). + kwargs: dict + Completed kwargs of the distribution settings: theta_dim, context_dim, + hidden_dims, activation, dropout, batch_norm, and optionally + theta_embedding_kwargs, theta_with_glu, context_with_glu. Extra keys + consumed by the training scheme (e.g., sigma_min) are ignored. + embedding_net: nn.Module = None + Embedding network for the context; None for unconditional models. Returns ------- nn.Module Neural network for the continuous flow. """ - theta_dim = posterior_kwargs["input_dim"] - context_dim = posterior_kwargs["context_dim"] - - # get embeddings modules for context - if embedding_kwargs is not None: - context_embedding_kwargs = copy.deepcopy(embedding_kwargs) - if initial_weights is not None: - context_embedding_kwargs["V_rb_list"] = initial_weights["V_rb_list"] - elif "V_rb_list" not in context_embedding_kwargs: - context_embedding_kwargs["V_rb_list"] = None - - context_embedding = create_enet_with_projection_layer_and_dense_resnet( - **context_embedding_kwargs - ) - else: + theta_dim = kwargs["theta_dim"] + context_dim = kwargs["context_dim"] or 0 + + if embedding_net is None: context_embedding = torch.nn.Identity() + context_keys = () + else: + context_embedding = embedding_net + context_keys = embedding_net.input_keys # get embeddings modules for theta (which is actually cat(t, theta)) - if "theta_embedding_kwargs" in posterior_kwargs: + if "theta_embedding_kwargs" in kwargs: theta_embedding = get_theta_embedding_net( - posterior_kwargs["theta_embedding_kwargs"], + kwargs["theta_embedding_kwargs"], input_dim=theta_dim + 1, ) else: theta_embedding = torch.nn.Identity() # get output dimensions of embedded context and theta - theta_with_glu = posterior_kwargs.get("theta_with_glu", False) - context_with_glu = posterior_kwargs.get("context_with_glu", False) + theta_with_glu = kwargs.get("theta_with_glu", False) + context_with_glu = kwargs.get("context_with_glu", False) embedded_theta_dim = theta_embedding(torch.zeros(10, theta_dim + 1)).shape[1] glu_dim = theta_with_glu * embedded_theta_dim + context_with_glu * context_dim @@ -231,31 +219,23 @@ def create_cf( if glu_dim == 0: glu_dim = None - activation_fn = torchutils.get_activation_function_from_string( - posterior_kwargs["activation"] - ) + activation_fn = torchutils.get_activation_function_from_string(kwargs["activation"]) continuous_flow_net = DenseResidualNet( input_dim=input_dim, output_dim=theta_dim, - hidden_dims=posterior_kwargs["hidden_dims"], + hidden_dims=kwargs["hidden_dims"], activation=activation_fn, - dropout=posterior_kwargs["dropout"], - batch_norm=posterior_kwargs["batch_norm"], + dropout=kwargs["dropout"], + batch_norm=kwargs["batch_norm"], context_features=glu_dim, ) - # With added_context, the embedding merges (waveform, context_parameters). - if embedding_kwargs is not None and embedding_kwargs.get("added_context"): - context_keys = ("waveform", "context_parameters") - else: - context_keys = ("waveform",) - model = ContinuousFlow( continuous_flow_net, context_embedding, theta_embedding, - theta_with_glu=posterior_kwargs.get("theta_with_glu", False), - context_with_glu=posterior_kwargs.get("context_with_glu", False), + theta_with_glu=theta_with_glu, + context_with_glu=context_with_glu, context_keys=context_keys, ) return model @@ -296,14 +276,15 @@ def get_dim_positional_embedding(encoding: dict, input_dim: int): return (1 + 2 * encoding["frequencies"]) * input_dim return 2 * encoding["frequencies"] + input_dim + class PositionalEncoding(nn.Module): """ Implements positional encoding as commonly used in transformer architectures. - - Positional encoding introduces a way to inject information about the order of - the input data (e.g., sequence positions) into a neural network that otherwise - lacks a sense of position due to its permutation-invariant nature. This class - computes sinusoidal encodings based on the position of each element in the input + + Positional encoding introduces a way to inject information about the order of + the input data (e.g., sequence positions) into a neural network that otherwise + lacks a sense of position due to its permutation-invariant nature. This class + computes sinusoidal encodings based on the position of each element in the input and concatenates them with the original input features. Attributes @@ -323,7 +304,7 @@ class PositionalEncoding(nn.Module): The number of sinusoidal frequencies to compute. This determines the dimensionality of the positional encoding for each input feature. encode_all : bool, optional (default=True) - If True, the positional encoding is computed for all features in the input. + If True, the positional encoding is computed for all features in the input. Otherwise, it is computed only for the first feature (e.g., the time dimension). base_freq : float, optional (default=2 * np.pi) The base frequency used for sinusoidal encoding. @@ -331,12 +312,13 @@ class PositionalEncoding(nn.Module): Methods ------- forward(t_theta) - Computes the positional encoding for the input tensor `t_theta` and concatenates + Computes the positional encoding for the input tensor `t_theta` and concatenates it with the original input features. - If `encode_all` is True, the positional encoding is computed for all features. - If `encode_all` is False, the positional encoding is applied only to the first feature, such as time, while other features remain unchanged. """ + def __init__(self, nr_frequencies, encode_all=True, base_freq=2 * np.pi): super(PositionalEncoding, self).__init__() frequencies = base_freq * torch.pow( diff --git a/dingo/core/nn/enets.py b/dingo/core/nn/enets.py index 206730f8..24cda50a 100755 --- a/dingo/core/nn/enets.py +++ b/dingo/core/nn/enets.py @@ -1,4 +1,19 @@ -"""Implementation of embedding networks.""" +"""Implementation of embedding networks. + +Embedding networks registered with EMBEDDING_NETS follow a common contract: + +* ``input_keys``: class attribute naming the batch entries the network consumes, + in the order of its forward arguments. +* ``output_dim``: dimension of the embedded context vector. +* ``complete_settings(settings, sample_batch)``: classmethod inferring the + network's own input dimensions from a sample batch; the completed settings + (which must include ``output_dim``) are saved in the checkpoint, so loading + never needs a data sample. + +Context mergers registered with CONTEXT_MERGERS wrap an embedding network to mix +in the (standardized) context parameters; they follow the same contract, plus a +``merged_output_dim`` method used during settings completion. +""" from typing import Tuple, Callable, Union, List import torch @@ -6,6 +21,7 @@ import torch.nn as nn from torch.nn import functional as F from glasflow.nflows.nn.nets.resnet import ResidualBlock +from dingo.core.registry import CONTEXT_MERGERS, EMBEDDING_NETS from dingo.core.utils import torchutils @@ -221,9 +237,11 @@ def __init__( ) self.resize_layers = nn.ModuleList( [ - nn.Linear(self.hidden_dims[n - 1], self.hidden_dims[n]) - if self.hidden_dims[n - 1] != self.hidden_dims[n] - else nn.Identity() + ( + nn.Linear(self.hidden_dims[n - 1], self.hidden_dims[n]) + if self.hidden_dims[n - 1] != self.hidden_dims[n] + else nn.Identity() + ) for n in range(1, self.num_res_blocks) ] + [nn.Linear(self.hidden_dims[-1], self.output_dim)] @@ -237,6 +255,80 @@ def forward(self, x, context=None): return x +@EMBEDDING_NETS.register("dense_svd") +class DenseSVDEmbedding(nn.Sequential): + """ + The classic dingo embedding network: a linear projection onto a reduced (SVD) + basis (LinearProjectionRB), followed by a dense residual network + (DenseResidualNet). See the docstrings of the two modules for details. + + Consumes the "waveform" entry of the batch, of shape + (batch_size, num_blocks, num_channels, num_bins). Subclasses nn.Sequential so + that the state dict lays out as ("0.*", "1.*"), matching old checkpoints. + """ + + input_keys = ("waveform",) + + def __init__( + self, + input_dims: List[int], + output_dim: int, + hidden_dims: Tuple, + svd: dict, + activation: str = "elu", + dropout: float = 0.0, + batch_norm: bool = True, + V_rb_list: Union[Tuple, None] = None, + ): + """ + Parameters + ---------- + input_dims : list + dimensions of input batch, omitting batch dimension, + input_dims = [num_blocks, num_channels, num_bins]. Inferred from a + sample batch by complete_settings; not a user setting. + output_dim : int + output dimension (dimension of the embedded context) + hidden_dims : tuple + dimensions of the hidden layers of the residual network + svd : dict + SVD settings; "size" is the number of reduced-basis elements used for + the projection (further entries are consumed by the training pipeline + when generating the SVD). + activation : str + activation function used in the residual blocks + dropout : float + dropout probability in the residual blocks, for regularization + batch_norm : bool + whether to use batch normalization + V_rb_list : tuple of np.arrays, or None + V matrices of the SVD projection used to initialize the projection + weights. Passed as initial weights at first build; None when loading a + saved model. + """ + projection = LinearProjectionRB(input_dims, svd["size"], V_rb_list) + resnet = DenseResidualNet( + input_dim=projection.output_dim, + output_dim=output_dim, + hidden_dims=hidden_dims, + activation=torchutils.get_activation_function_from_string(activation), + dropout=dropout, + batch_norm=batch_norm, + ) + super().__init__(projection, resnet) + self.output_dim = output_dim + + @classmethod + def complete_settings(cls, settings: dict, sample_batch: dict) -> dict: + """Infer input_dims from the sample batch; return completed settings.""" + if "input_dims" in settings: + raise ValueError( + "'input_dims' is derived from the data and must not be specified " + "in the embedding net settings." + ) + return {**settings, "input_dims": list(sample_batch["waveform"].shape)} + + class ModuleMerger(nn.Module): """ This is a wrapper used to process multiple different kinds of context @@ -277,6 +369,35 @@ def forward(self, *x): return torch.cat(x, axis=1) +@CONTEXT_MERGERS.register("concat") +class ConcatContextMerger(ModuleMerger): + """ + Default context merger: concatenates the embedded data with the (standardized) + context parameters, e.g. GNPE proxies. Wraps the embedding network and an + identity map via ModuleMerger, which keeps the state-dict layout of old + checkpoints ("enets.0.*"). + """ + + def __init__(self, embedding_net: nn.Module, num_context_parameters: int): + """ + Parameters + ---------- + embedding_net : nn.Module + The wrapped embedding network. + num_context_parameters : int + Number of context parameters concatenated to the embedded data. + Inferred from a sample batch during settings completion. + """ + super().__init__((embedding_net, nn.Identity())) + self.input_keys = (*embedding_net.input_keys, "context_parameters") + self.output_dim = embedding_net.output_dim + num_context_parameters + + @staticmethod + def merged_output_dim(embedding_output_dim: int, num_context_parameters: int): + """Output dimension of the merged embedding, for settings completion.""" + return embedding_output_dim + num_context_parameters + + def create_enet_with_projection_layer_and_dense_resnet( input_dims: List[int], # n_rb: int, @@ -347,17 +468,16 @@ def create_enet_with_projection_layer_and_dense_resnet( a tuple with 2 elements, input = (x, z) rather than just the tensor x. :return: nn.Module """ - activation_fn = torchutils.get_activation_function_from_string(activation) - module_1 = LinearProjectionRB(input_dims, svd["size"], V_rb_list) - module_2 = DenseResidualNet( - input_dim=module_1.output_dim, + enet = DenseSVDEmbedding( + input_dims=input_dims, output_dim=output_dim, hidden_dims=hidden_dims, - activation=activation_fn, + svd=svd, + activation=activation, dropout=dropout, batch_norm=batch_norm, + V_rb_list=V_rb_list, ) - enet = nn.Sequential(module_1, module_2) if not added_context: return enet diff --git a/dingo/core/nn/nsf.py b/dingo/core/nn/nsf.py index 1e844db1..876b5868 100644 --- a/dingo/core/nn/nsf.py +++ b/dingo/core/nn/nsf.py @@ -3,16 +3,12 @@ from the uci.py example from https://github.com/bayesiains/nsf. """ -import copy - import torch import torch.nn as nn import glasflow.nflows as nflows # nflows not maintained, so use this maintained fork from glasflow.nflows import distributions, flows, transforms import glasflow.nflows.nn.nets as nflows_nets from dingo.core.utils import torchutils -from dingo.core.nn.enets import create_enet_with_projection_layer_and_dense_resnet -from typing import Union, Callable, Tuple def create_linear_transform(param_dim: int): @@ -253,8 +249,6 @@ def create_nsf_model( context_dim: int, num_flow_steps: int, base_transform_kwargs: dict, - embedding_net_builder: Union[Callable, str] = None, - embedding_kwargs: dict = None, ): """ Build NSF model. This models the posterior distribution p(y|x). @@ -271,87 +265,17 @@ def create_nsf_model( number of sequential transforms :param base_transform_kwargs: dict, hyperparameters for transform steps - :param embedding_net_builder: Callable=None, - build function for embedding network TODO - :param embedding_kwargs: dict=None, - hyperparameters for embedding network :return: Flow the NSF (posterior model) """ - - if embedding_net_builder is not None: - embedding_net = embedding_net_builder(**embedding_kwargs) - else: - embedding_net = None - - # str(embedding_net_builder).split(' ')[1] - distribution = distributions.StandardNormal((input_dim,)) transform = create_transform( num_flow_steps, input_dim, context_dim, base_transform_kwargs ) - flow = flows.Flow(transform, distribution, embedding_net) + flow = flows.Flow(transform, distribution) return flow -def create_nsf_wrapped(**kwargs): - """ - Wraps the NSF model in a FlowWrapper. This is required for parallel - training, and wraps the log_prob method as a forward method. - """ - flow = create_nsf_model(**kwargs) - return FlowWrapper(flow) - - -def create_nsf_with_rb_projection_embedding_net( - posterior_kwargs: dict, - embedding_kwargs: dict, - initial_weights: dict = None, -): - """Builds a neural spline flow with an embedding network that consists of a - reduced basis projection followed by a residual network. Optionally initializes the - embedding network weights. - - Parameters - ---------- - posterior_kwargs : dict - kwargs for neural spline flow - embedding_kwargs : dict - kwargs for emebedding network - initial_weights : dict - Dictionary containing the initial weights for the SVD projection. This should - have one key 'V_rb_list', with value a list of SVD V matrices (one for each - detector). - - Returns - ------- - nn.Module - Neural spline flow model - """ - # We copy the embedding_kwargs to allow an insert of V_rb_list without - # affecting the original embedding_kwargs. This is because we don't want to - # save the embedding_kwargs with the huge V_rb_list included. This is a bit of - # a hack; improve setting of initial weights later. - - embedding_kwargs = copy.deepcopy(embedding_kwargs) - if initial_weights is not None: - embedding_kwargs["V_rb_list"] = initial_weights["V_rb_list"] - elif "V_rb_list" not in embedding_kwargs: - embedding_kwargs["V_rb_list"] = None - - embedding_net = create_enet_with_projection_layer_and_dense_resnet( - **embedding_kwargs - ) - flow = create_nsf_model(**posterior_kwargs) - # With added_context, the embedding merges (waveform, context_parameters). - if embedding_kwargs.get("added_context"): - context_keys = ("waveform", "context_parameters") - else: - context_keys = ("waveform",) - model = FlowWrapper(flow, embedding_net, context_keys) - return model - - if __name__ == "__main__": pass diff --git a/dingo/core/posterior_models/base_model.py b/dingo/core/posterior_models/base_model.py index 788a0ca9..ffab37a1 100755 --- a/dingo/core/posterior_models/base_model.py +++ b/dingo/core/posterior_models/base_model.py @@ -18,6 +18,7 @@ import json from collections import OrderedDict from typing import Optional +from dingo.core.registry import CONTEXT_MERGERS, EMBEDDING_NETS from dingo.core.utils.backward_compatibility import update_model_config from dingo.core.utils.misc import get_version @@ -72,6 +73,7 @@ def __init__( self.metadata = metadata if self.metadata is not None: + update_model_config(self.metadata["train_settings"]["model"]) self.model_kwargs = self.metadata["train_settings"]["model"] # Expect self.optimizer_settings and self.scheduler_settings to be set # separately, and before calling initialize_optimizer_and_scheduler(). @@ -99,6 +101,31 @@ def initialize_network(self): """ pass + def build_embedding_net(self): + """ + Build the embedding network declared in the model settings (resolved via + the EMBEDDING_NETS registry), optionally wrapped with a context merger + (CONTEXT_MERGERS) that mixes in the context parameters. Initial weights + (e.g. SVD projection matrices) are passed as extra constructor kwargs; they + are not part of the saved settings. + + Returns None if the model declares no embedding network (unconditional + models). + """ + embedding_settings = self.model_kwargs.get("embedding_net") + if embedding_settings is None: + return None + kwargs = dict(embedding_settings.get("kwargs", {})) + if self.initial_weights: + kwargs.update(self.initial_weights) + embedding_net = EMBEDDING_NETS.get(embedding_settings["type"])(**kwargs) + merger_settings = self.model_kwargs.get("context_merger") + if merger_settings is not None: + embedding_net = CONTEXT_MERGERS.get(merger_settings["type"])( + embedding_net, **merger_settings.get("kwargs", {}) + ) + return embedding_net + # Parameter contract read by samplers (e.g. FlowFactor.from_model on the # factorized-sampler branch). These accessors are the supported interface; the # location inside the metadata dict is an implementation detail. All three read @@ -272,6 +299,26 @@ def save_model( saved, e.g. optimizer state dict """ + # Enforce the parameter contract promised to samplers: the standardization + # covers all inference and context parameters, for every architecture + # (built-in or plugin). + if self.metadata is not None: + data_settings = self.metadata.get("train_settings", {}).get("data", {}) + if "inference_parameters" in data_settings: + standardization = data_settings.get("standardization") or {} + missing = [ + p + for p in self.inference_parameters + self.context_parameters + if p not in standardization.get("mean", {}) + or p not in standardization.get("std", {}) + ] + if missing: + raise ValueError( + f"Cannot save model: standardization must cover all " + f"inference and context parameters, but is missing " + f"{missing}." + ) + model_dict = { "model_kwargs": self.model_kwargs, "model_state_dict": self.network.state_dict(), @@ -373,6 +420,9 @@ def load_model( self.epoch = d["epoch"] self.metadata = d["metadata"] + # model_kwargs and the metadata's model section are separate dicts in old + # checkpoints; keep both on the current schema. + update_model_config(self.metadata["train_settings"]["model"]) if "context" in d: self.context = d["context"] diff --git a/dingo/core/posterior_models/build_model.py b/dingo/core/posterior_models/build_model.py index 0d4c6af5..60c1a83a 100644 --- a/dingo/core/posterior_models/build_model.py +++ b/dingo/core/posterior_models/build_model.py @@ -1,9 +1,15 @@ +import copy + # Importing the modules registers the built-in model types with NEURAL_DISTRIBUTIONS. import dingo.core.posterior_models.flow_matching # noqa: F401 import dingo.core.posterior_models.normalizing_flow # noqa: F401 import dingo.core.posterior_models.score_matching # noqa: F401 from dingo.core.posterior_models.base_model import NeuralDistribution -from dingo.core.registry import NEURAL_DISTRIBUTIONS +from dingo.core.registry import ( + CONTEXT_MERGERS, + EMBEDDING_NETS, + NEURAL_DISTRIBUTIONS, +) from dingo.core.utils.backward_compatibility import ( torch_load_with_fallback, update_model_config, @@ -17,7 +23,7 @@ def build_model_from_kwargs( """ Returns a NeuralDistribution based on a saved network or settings dict. - The model class is resolved from the settings' posterior_model_type via the + The model class is resolved from the settings' distribution type via the NEURAL_DISTRIBUTIONS registry (e.g., normalizing flow, flow matching, or score matching, or a plugin type; see dingo.core.registry). @@ -47,59 +53,100 @@ def build_model_from_kwargs( # version was introduced in v0.3.3 check_minimum_version("dingo=0.3.2") update_model_config(d["metadata"]["train_settings"]["model"]) # Backward compat - posterior_model_type = d["metadata"]["train_settings"]["model"][ - "posterior_model_type" - ] + model_type = d["metadata"]["train_settings"]["model"]["distribution"]["type"] else: update_model_config(settings["train_settings"]["model"]) # Backward compat - posterior_model_type = settings["train_settings"]["model"][ - "posterior_model_type" - ] + model_type = settings["train_settings"]["model"]["distribution"]["type"] try: - model = NEURAL_DISTRIBUTIONS.get(posterior_model_type) + model = NEURAL_DISTRIBUTIONS.get(model_type) except KeyError as e: raise ValueError(f"No valid posterior model type specified. {e}") from e return model(model_filename=filename, metadata=settings, **kwargs) -def autocomplete_model_kwargs(model_kwargs: dict, data_sample: dict): +def complete_model_settings(model_settings: dict, sample_batch: dict) -> dict: """ - Autocomplete the model kwargs from train_settings and data_sample from the dataloader: + Complete the model settings based on a sample batch from the dataloader. + + Each embedding architecture infers its own input dimensions via its + complete_settings classmethod; the cross-cutting dimensions (theta_dim, + context_dim) are computed here. The completed settings are saved in the + checkpoint, so loading a model never needs a data sample. + + If the batch provides context_parameters and the embedding network does not + consume them natively (i.e., they are not among its input_keys), a context + merger is added ("concat" unless specified otherwise). - * set input dimension of embedding net to the shape of the waveform data - * set dimension of parameter space to the number of inference parameters - * set added_context flag of embedding net if required for context parameters - (e.g., GNPE proxies) - * set context dim of posterior model to output dim of embedding net + dimension - of the context parameters + Dimensions are derived from the data; specifying them in the settings is an + error. Parameters ---------- - model_kwargs: dict - Model settings, which are modified in-place. - data_sample: dict - Sample from dataloader (e.g., wfd[0]) used for autocompletion, with keys - "inference_parameters", "waveform", and (only if the network is conditioned - on additional parameters) "context_parameters". + model_settings: dict + Model section of the train settings. Old schemas are mapped forward + in-place; the completion itself does not modify the input. + sample_batch: dict + Sample from the dataloader (e.g., wfd[0]), with keys + "inference_parameters", "waveform", optionally "context_parameters", and + any architecture-specific entries. + + Returns + ------- + dict + Completed model settings. """ + update_model_config(model_settings) + model_settings = copy.deepcopy(model_settings) - # set input dims from ifo_list and domain information - model_kwargs["embedding_kwargs"]["input_dims"] = list(data_sample["waveform"].shape) - # set dimension of parameter space of posterior model - model_kwargs["posterior_kwargs"]["input_dim"] = len( - data_sample["inference_parameters"] + distribution_kwargs = model_settings["distribution"].setdefault("kwargs", {}) + for key in ("theta_dim", "context_dim"): + if key in distribution_kwargs: + raise ValueError( + f"'{key}' is derived from the data and must not be specified in " + f"the model settings." + ) + distribution_kwargs["theta_dim"] = len(sample_batch["inference_parameters"]) + + embedding_settings = model_settings.get("embedding_net") + if embedding_settings is None: + if "context_merger" in model_settings: + raise ValueError("A context_merger requires an embedding_net.") + distribution_kwargs["context_dim"] = None + return model_settings + + embedding_cls = EMBEDDING_NETS.get(embedding_settings["type"]) + missing = [k for k in embedding_cls.input_keys if k not in sample_batch] + if missing: + raise ValueError( + f"Embedding net '{embedding_settings['type']}' consumes batch entries " + f"{list(embedding_cls.input_keys)}, but the sample batch is missing " + f"{missing} (batch keys: {sorted(sample_batch)})." + ) + embedding_settings["kwargs"] = embedding_cls.complete_settings( + embedding_settings.get("kwargs", {}), sample_batch ) - # set added_context flag of embedding net if context parameters are required - # set context dim of nsf to output dim of embedding net + context parameter dim - if "context_parameters" in data_sample: - model_kwargs["embedding_kwargs"]["added_context"] = True - model_kwargs["posterior_kwargs"]["context_dim"] = model_kwargs[ - "embedding_kwargs" - ]["output_dim"] + len(data_sample["context_parameters"]) + output_dim = embedding_settings["kwargs"]["output_dim"] + + native_context = "context_parameters" in embedding_cls.input_keys + if "context_parameters" in sample_batch and not native_context: + merger_settings = model_settings.setdefault( + "context_merger", {"type": "concat"} + ) + merger_cls = CONTEXT_MERGERS.get(merger_settings["type"]) + merger_settings["kwargs"] = { + **merger_settings.get("kwargs", {}), + "num_context_parameters": len(sample_batch["context_parameters"]), + } + distribution_kwargs["context_dim"] = merger_cls.merged_output_dim( + output_dim, **merger_settings["kwargs"] + ) else: - model_kwargs["embedding_kwargs"]["added_context"] = False - model_kwargs["posterior_kwargs"]["context_dim"] = model_kwargs[ - "embedding_kwargs" - ]["output_dim"] + if "context_merger" in model_settings: + raise ValueError( + "A context_merger is specified, but the data provides no " + "context_parameters (or the embedding net consumes them natively)." + ) + distribution_kwargs["context_dim"] = output_dim + return model_settings diff --git a/dingo/core/posterior_models/cflow_base.py b/dingo/core/posterior_models/cflow_base.py index a35f9cd4..0cbe4de8 100644 --- a/dingo/core/posterior_models/cflow_base.py +++ b/dingo/core/posterior_models/cflow_base.py @@ -49,12 +49,9 @@ class ContinuousFlowPosteriorModel(NeuralDistribution): def __init__(self, **kwargs): super().__init__(**kwargs) self.eps = 0 - self.time_prior_exponent = self.model_kwargs["posterior_kwargs"][ - "time_prior_exponent" - ] - self.theta_dim = self.metadata["train_settings"]["model"]["posterior_kwargs"][ - "input_dim" - ] + distribution_kwargs = self.model_kwargs["distribution"]["kwargs"] + self.time_prior_exponent = distribution_kwargs["time_prior_exponent"] + self.theta_dim = distribution_kwargs["theta_dim"] def sample_t(self, batch_size): t = (1 - self.eps) * torch.rand(batch_size, device=self.device) @@ -122,12 +119,10 @@ def rhs_of_joint_ode(self, t, theta_and_div_t, *context_data, hutchinson=False): return torch.cat((vf, -div_vf), dim=1) def initialize_network(self): - model_kwargs = { - k: v for k, v in self.model_kwargs.items() if k != "posterior_model_type" - } - if self.initial_weights is not None: - model_kwargs["initial_weights"] = self.initial_weights - self.network = create_cf(**model_kwargs) + embedding_net = self.build_embedding_net() + self.network = create_cf( + self.model_kwargs["distribution"]["kwargs"], embedding_net + ) def sample(self, context: dict = None, num_samples: int = 1): """ diff --git a/dingo/core/posterior_models/flow_matching.py b/dingo/core/posterior_models/flow_matching.py index 58a333d9..2302ec7f 100644 --- a/dingo/core/posterior_models/flow_matching.py +++ b/dingo/core/posterior_models/flow_matching.py @@ -36,7 +36,7 @@ class FlowMatchingPosteriorModel(ContinuousFlowPosteriorModel): def __init__(self, **kwargs): super().__init__(**kwargs) self.eps = 0 - self.sigma_min = self.model_kwargs["posterior_kwargs"]["sigma_min"] + self.sigma_min = self.model_kwargs["distribution"]["kwargs"]["sigma_min"] def evaluate_vector_field(self, t, theta_t, *context_data): """ diff --git a/dingo/core/posterior_models/normalizing_flow.py b/dingo/core/posterior_models/normalizing_flow.py index ed9b407e..d04cd56d 100644 --- a/dingo/core/posterior_models/normalizing_flow.py +++ b/dingo/core/posterior_models/normalizing_flow.py @@ -1,10 +1,7 @@ from .base_model import NeuralDistribution from dingo.core.registry import NEURAL_DISTRIBUTIONS -from dingo.core.nn.nsf import ( - create_nsf_with_rb_projection_embedding_net, - create_nsf_wrapped, -) +from dingo.core.nn.nsf import FlowWrapper, create_nsf_model @NEURAL_DISTRIBUTIONS.register("normalizing_flow") @@ -40,16 +37,18 @@ def __init__(self, **kwargs): super().__init__(**kwargs) def initialize_network(self): - model_kwargs = { - k: v for k, v in self.model_kwargs.items() if k != "posterior_model_type" - } - if self.initial_weights is not None: - model_kwargs["initial_weights"] = self.initial_weights - - if self.model_kwargs.get("embedding_kwargs", False): - self.network = create_nsf_with_rb_projection_embedding_net(**model_kwargs) + embedding_net = self.build_embedding_net() + kwargs = self.model_kwargs["distribution"]["kwargs"] + flow = create_nsf_model( + input_dim=kwargs["theta_dim"], + context_dim=kwargs["context_dim"], + num_flow_steps=kwargs["num_flow_steps"], + base_transform_kwargs=kwargs["base_transform_kwargs"], + ) + if embedding_net is None: + self.network = FlowWrapper(flow) else: - self.network = create_nsf_wrapped(**model_kwargs["posterior_kwargs"]) + self.network = FlowWrapper(flow, embedding_net, embedding_net.input_keys) def log_prob(self, theta, context: dict = None): return self.network(theta, context) diff --git a/dingo/core/posterior_models/score_matching.py b/dingo/core/posterior_models/score_matching.py index baac9593..eece3a2d 100644 --- a/dingo/core/posterior_models/score_matching.py +++ b/dingo/core/posterior_models/score_matching.py @@ -38,11 +38,12 @@ class ScoreDiffusionPosteriorModel(ContinuousFlowPosteriorModel): def __init__(self, **kwargs): super().__init__(**kwargs) - self.eps = self.model_kwargs["posterior_kwargs"]["epsilon"] - self.beta_min = self.model_kwargs["posterior_kwargs"]["beta_min"] - self.beta_max = self.model_kwargs["posterior_kwargs"]["beta_max"] + distribution_kwargs = self.model_kwargs["distribution"]["kwargs"] + self.eps = distribution_kwargs["epsilon"] + self.beta_min = distribution_kwargs["beta_min"] + self.beta_max = distribution_kwargs["beta_max"] - likelihood_weighting = self.model_kwargs["posterior_kwargs"].get( + likelihood_weighting = distribution_kwargs.get( "likelihood_weighting", "score-matching" ) if likelihood_weighting: diff --git a/dingo/core/registry.py b/dingo/core/registry.py index 3cc86e63..3cddca3b 100644 --- a/dingo/core/registry.py +++ b/dingo/core/registry.py @@ -77,8 +77,11 @@ def get(self, name: str) -> Any: if name in self._components: return self._components[name] - for resolve in (self._from_entry_points, self._from_dotted_path, - self._from_file_path): + for resolve in ( + self._from_entry_points, + self._from_dotted_path, + self._from_file_path, + ): component = resolve(name) if component is not None: # Cache so repeated lookups are cheap and resolve consistently. diff --git a/dingo/core/result.py b/dingo/core/result.py index 67017271..5b905a9f 100644 --- a/dingo/core/result.py +++ b/dingo/core/result.py @@ -1090,5 +1090,3 @@ def freeze(d): elif isinstance(d, list): return tuple(freeze(value) for value in d) return d - - diff --git a/dingo/core/transforms.py b/dingo/core/transforms.py index 7bb20804..09c10da8 100644 --- a/dingo/core/transforms.py +++ b/dingo/core/transforms.py @@ -6,13 +6,14 @@ def __init__(self, key): def __call__(self, sample): return sample[self.key] + class RenameKey: def __init__(self, old, new): self.old = old self.new = new - def __call__(self, input_sample : dict): + def __call__(self, input_sample: dict): sample = input_sample.copy() sample[self.new] = sample.pop(self.old) return sample diff --git a/dingo/core/utils/backward_compatibility.py b/dingo/core/utils/backward_compatibility.py index d32cb95e..a362ac4d 100644 --- a/dingo/core/utils/backward_compatibility.py +++ b/dingo/core/utils/backward_compatibility.py @@ -133,13 +133,23 @@ def check_minimum_version(version_str: str, raise_exception: bool = False) -> No def update_model_config(model_settings: dict): """ Update the model settings to ensure backwards compatibility with networks - trained using previous versions of Dingo. + trained using previous versions of Dingo. This maps all old schemas forward to + the current one, + + model: + distribution: {type: ..., kwargs: {...}} + embedding_net: {type: ..., kwargs: {...}} # optional + context_merger: {type: ..., kwargs: {...}} # optional + + and is idempotent. It is the single boundary where old settings and checkpoints + are translated; code elsewhere only handles the current schema. Parameters ---------- model_settings: dict - Model settings to be updated. + Model settings to be updated in-place. """ + # Oldest schema: type nsf+embedding. if model_settings.get("type") == "nsf+embedding": model_settings["posterior_model_type"] = "normalizing_flow" del model_settings["type"] @@ -148,12 +158,47 @@ def update_model_config(model_settings: dict): model_settings["embedding_kwargs"] = model_settings["embedding_net_kwargs"] del model_settings["embedding_net_kwargs"] - # The model type used to be matched case-insensitively; registry lookup is - # case-sensitive, so lowercase built-in type names from old checkpoints. - posterior_model_type = model_settings.get("posterior_model_type") - if posterior_model_type is not None and posterior_model_type.lower() in ( - "normalizing_flow", - "flow_matching", - "score_matching", - ): - model_settings["posterior_model_type"] = posterior_model_type.lower() + # Old schema: posterior_model_type + posterior_kwargs + embedding_kwargs. + if "posterior_model_type" in model_settings: + posterior_model_type = model_settings.pop("posterior_model_type") + # The model type used to be matched case-insensitively; registry lookup is + # case-sensitive, so lowercase built-in type names from old checkpoints. + if posterior_model_type.lower() in ( + "normalizing_flow", + "flow_matching", + "score_matching", + ): + posterior_model_type = posterior_model_type.lower() + + distribution_kwargs = dict(model_settings.pop("posterior_kwargs", None) or {}) + if "input_dim" in distribution_kwargs: + distribution_kwargs["theta_dim"] = distribution_kwargs.pop("input_dim") + model_settings["distribution"] = { + "type": posterior_model_type, + "kwargs": distribution_kwargs, + } + + embedding_kwargs = model_settings.pop("embedding_kwargs", None) + if embedding_kwargs: + embedding_kwargs = dict(embedding_kwargs) + added_context = embedding_kwargs.pop("added_context", False) + if embedding_kwargs.get("V_rb_list", "missing") is None: + # Old settings stored a V_rb_list: None placeholder; initial weights + # are no longer part of the settings. + del embedding_kwargs["V_rb_list"] + model_settings["embedding_net"] = { + "type": "dense_svd", + "kwargs": embedding_kwargs, + } + if added_context: + # Old checkpoints merged (waveform, context_parameters) by + # concatenation, with context_dim = output_dim + num proxies. + context_dim = distribution_kwargs.get("context_dim") + output_dim = embedding_kwargs.get("output_dim") + merger_kwargs = {} + if context_dim is not None and output_dim is not None: + merger_kwargs["num_context_parameters"] = context_dim - output_dim + model_settings["context_merger"] = { + "type": "concat", + "kwargs": merger_kwargs, + } diff --git a/dingo/core/utils/condor_utils.py b/dingo/core/utils/condor_utils.py index 1b0e2bbe..0711b36c 100644 --- a/dingo/core/utils/condor_utils.py +++ b/dingo/core/utils/condor_utils.py @@ -11,20 +11,21 @@ def resubmit_condor_job(train_dir, train_settings, epoch): :param epoch: :return: """ - if 'condor_settings' in train_settings: - print('Copying log files') + if "condor_settings" in train_settings: + print("Copying log files") copy_logfiles(train_dir, epoch=epoch) - if epoch >= train_settings['train_settings']['runtime_limits'][ - 'max_epochs_total']: - print('Training complete, job will not be resubmitted') + if ( + epoch + >= train_settings["train_settings"]["runtime_limits"]["max_epochs_total"] + ): + print("Training complete, job will not be resubmitted") else: - print('Training incomplete, resubmitting job.') + print("Training incomplete, resubmitting job.") create_submission_file_and_submit_job(train_dir) -def create_submission_file_and_submit_job(train_dir, - filename='submission_file.sub'): +def create_submission_file_and_submit_job(train_dir, filename="submission_file.sub"): """ TODO: documentation :param train_dir: @@ -32,55 +33,58 @@ def create_submission_file_and_submit_job(train_dir, :return: """ create_submission_file(train_dir, filename) - with open(join(train_dir, 'train_settings.yaml'), 'r') as fp: - bid = yaml.safe_load(fp)['condor_settings']['bid'] - os.system(f'condor_submit_bid {bid} {join(train_dir, filename)}') + with open(join(train_dir, "train_settings.yaml"), "r") as fp: + bid = yaml.safe_load(fp)["condor_settings"]["bid"] + os.system(f"condor_submit_bid {bid} {join(train_dir, filename)}") -def create_submission_file(train_dir, filename='submission_file.sub'): +def create_submission_file(train_dir, filename="submission_file.sub"): """ TODO: documentation :param train_dir: :param filename: :return: """ - with open(join(train_dir, 'train_settings.yaml'), 'r') as fp: - d = yaml.safe_load(fp)['condor_settings'] + with open(join(train_dir, "train_settings.yaml"), "r") as fp: + d = yaml.safe_load(fp)["condor_settings"] lines = [] lines.append(f'executable = {d["python"]}\n') lines.append(f'request_cpus = {d["num_cpus"]}\n') lines.append(f'request_memory = {d["memory_cpus"]}\n') lines.append(f'request_gpus = {d["num_gpus"]}\n') - lines.append(f'requirements = TARGET.CUDAGlobalMemoryMb > ' - f'{d["memory_gpus"]}\n\n') + lines.append( + f"requirements = TARGET.CUDAGlobalMemoryMb > " f'{d["memory_gpus"]}\n\n' + ) lines.append(f'arguments = {d["train_script"]} --train_dir {train_dir}\n') lines.append(f'error = {join(train_dir, "info.err")}\n') lines.append(f'output = {join(train_dir, "info.out")}\n') lines.append(f'log = {join(train_dir, "info.log")}\n') - lines.append('queue') + lines.append("queue") - with open(join(train_dir, filename), 'w') as f: + with open(join(train_dir, filename), "w") as f: for line in lines: f.write(line) def copyfile(src, dst): - os.system('cp -p %s %s' % (src, dst)) + os.system("cp -p %s %s" % (src, dst)) -def copy_logfiles(log_dir, epoch, name='info', suffixes=('.err','.log','.out')): +def copy_logfiles(log_dir, epoch, name="info", suffixes=(".err", ".log", ".out")): for suffix in suffixes: src = join(log_dir, name + suffix) - dest = join(log_dir, name + '_{:03d}'.format(epoch) + suffix) + dest = join(log_dir, name + "_{:03d}".format(epoch) + suffix) try: copyfile(src, dest) except: - print('Could not copy ' + src) + print("Could not copy " + src) -if __name__ == '__main__': - train_dir = '/Users/mdax/Documents/dingo/devel/dingo-devel/tutorials/02_gwpe/train_dir/' +if __name__ == "__main__": + train_dir = ( + "/Users/mdax/Documents/dingo/devel/dingo-devel/tutorials/02_gwpe/train_dir/" + ) create_submission_file(train_dir) # epoch = pm.epoch - 1 diff --git a/dingo/core/utils/gnpeutils.py b/dingo/core/utils/gnpeutils.py index 81077046..bf785798 100644 --- a/dingo/core/utils/gnpeutils.py +++ b/dingo/core/utils/gnpeutils.py @@ -40,9 +40,9 @@ def update(self, new_data): self.data = {k: v.copy()[None, :] for k, v in y.items()} else: self.data = { - k: np.concatenate((v, y[k][None, :]), axis=0) - for k, v in self.data.items() - } + k: np.concatenate((v, y[k][None, :]), axis=0) + for k, v in self.data.items() + } @property def pvalue_min(self): diff --git a/dingo/core/utils/plotting.py b/dingo/core/utils/plotting.py index da856d9b..dded52a1 100644 --- a/dingo/core/utils/plotting.py +++ b/dingo/core/utils/plotting.py @@ -147,9 +147,7 @@ def plot_corner_multi( # every corner call uses identical bins, so the 1D marginal densities are # on a consistent scale regardless of sample size or weight distribution. all_data = pd.concat([s[common_parameters] for s in samples], ignore_index=True) - common_range = [ - (all_data[p].min(), all_data[p].max()) for p in common_parameters - ] + common_range = [(all_data[p].min(), all_data[p].max()) for p in common_parameters] fig = None handles = [] diff --git a/dingo/core/utils/pt_to_hdf5.py b/dingo/core/utils/pt_to_hdf5.py index 06e6aa96..a41963ce 100644 --- a/dingo/core/utils/pt_to_hdf5.py +++ b/dingo/core/utils/pt_to_hdf5.py @@ -10,29 +10,45 @@ def parse_args(): parser = argparse.ArgumentParser( description="Convert the weights of a trained Dingo model from a PyTorch pickle .pt file to HDF5," " for distribution in the LVK's CVMFS.", - epilog="Training history (optimizer_state_dict) is discarded.") - parser.add_argument("-i", "--in_file", type=str, required=True, - help='Input model ".pt" weights file') - parser.add_argument("-o", "--out_file", type=str, required=True, - help='Output model ".hdf5" weights file') - parser.add_argument("-n", "--model_version_number", type=int, required=True, - help="Model version number (integer). " - "Will be included in the output filename and metadata.") + epilog="Training history (optimizer_state_dict) is discarded.", + ) + parser.add_argument( + "-i", + "--in_file", + type=str, + required=True, + help='Input model ".pt" weights file', + ) + parser.add_argument( + "-o", + "--out_file", + type=str, + required=True, + help='Output model ".hdf5" weights file', + ) + parser.add_argument( + "-n", + "--model_version_number", + type=int, + required=True, + help="Model version number (integer). " + "Will be included in the output filename and metadata.", + ) return parser.parse_args() def main(): args = parse_args() - if os.path.splitext(args.in_file)[-1] != '.pt': - raise ValueError('Expected a .pt input file') - if os.path.splitext(args.out_file)[-1] != '.hdf5': - raise ValueError('Expected a .hdf5 output file') + if os.path.splitext(args.in_file)[-1] != ".pt": + raise ValueError("Expected a .pt input file") + if os.path.splitext(args.out_file)[-1] != ".hdf5": + raise ValueError("Expected a .hdf5 output file") # Build output filename with the version number for this network # This is required for use on CVMFS root, ext = os.path.splitext(args.out_file) - out_file_name = f'{root}_v{args.model_version_number}{ext}' - print('Output will be written to', out_file_name) + out_file_name = f"{root}_v{args.model_version_number}{ext}" + print("Output will be written to", out_file_name) # Load data into CPU memory since we'll be saving it using CPU libraries d = torch.load(args.in_file, map_location=torch.device("cpu")) @@ -40,20 +56,19 @@ def main(): # Collect the names of the dicts that can be serialized to JSON; model_state_dict and # optimizer_state_dict contain torch.tensors and cannot be JSON serialized # In addition, we drop dicts related to training information that is not needed at inference time - dicts_to_serialize = ['model_kwargs', 'epoch', 'metadata'] + dicts_to_serialize = ["model_kwargs", "epoch", "metadata"] - - with h5py.File(out_file_name, 'w') as f: + with h5py.File(out_file_name, "w") as f: # Save small nested dicts as json - grp = f.create_group('serialized_dicts') + grp = f.create_group("serialized_dicts") for k in dicts_to_serialize: dict_str = json.dumps(d[k]) grp.create_dataset(k, data=dict_str) # Save the OrderedDict containing the model weights # The keys are ordered alphanumerically as well. - grp_model = f.create_group('model_weights') - for k, v in d['model_state_dict'].items(): + grp_model = f.create_group("model_weights") + for k, v in d["model_state_dict"].items(): if len(v.size()) > 0: grp_model.create_dataset(k, data=v.numpy(), fletcher32=True) else: @@ -66,17 +81,18 @@ def main(): # Metadata for CVMFS LVK distribution # This needs to be exactly the same as the "basename" of the hdf5 file - f.attrs['CANONICAL_FILE_BASENAME'] = os.path.basename(out_file_name) + f.attrs["CANONICAL_FILE_BASENAME"] = os.path.basename(out_file_name) # Add a few metadata entries as attributes - f.attrs['approximant'] = d['metadata']['dataset_settings']['waveform_generator']['approximant'] - f.attrs['epoch'] = d['epoch'] + f.attrs["approximant"] = d["metadata"]["dataset_settings"][ + "waveform_generator" + ]["approximant"] + f.attrs["epoch"] = d["epoch"] # Add the dingo version used for training - f.attrs['version'] = str(d.get('version')) + f.attrs["version"] = str(d.get("version")) # Make it clear to dingo_ls that this is file contains model weights - f.attrs['dataset_type'] = 'trained_model' + f.attrs["dataset_type"] = "trained_model" if __name__ == "__main__": main() - diff --git a/dingo/core/utils/torchutils.py b/dingo/core/utils/torchutils.py index 2d109d7b..5d6c3312 100644 --- a/dingo/core/utils/torchutils.py +++ b/dingo/core/utils/torchutils.py @@ -9,9 +9,9 @@ def fix_random_seeds(_): """Utility function to set random seeds when using multiple workers for DataLoader.""" - np.random.seed(int(torch.initial_seed()) % (2 ** 32 - 1)) + np.random.seed(int(torch.initial_seed()) % (2**32 - 1)) try: - bilby.core.utils.random.seed(int(torch.initial_seed()) % (2 ** 32 - 1)) + bilby.core.utils.random.seed(int(torch.initial_seed()) % (2**32 - 1)) except AttributeError: # In case using an old version of Bilby. pass diff --git a/dingo/gw/training/train_pipeline.py b/dingo/gw/training/train_pipeline.py index 48470aea..9b51e6d6 100644 --- a/dingo/gw/training/train_pipeline.py +++ b/dingo/gw/training/train_pipeline.py @@ -12,9 +12,10 @@ from threadpoolctl import threadpool_limits from dingo.core.posterior_models.build_model import ( - autocomplete_model_kwargs, build_model_from_kwargs, + complete_model_settings, ) +from dingo.core.utils.backward_compatibility import update_model_config from dingo.gw.training.train_builders import ( build_dataset, set_train_transforms, @@ -86,9 +87,8 @@ def prepare_training_new( """ Based on a settings dictionary, initialize a WaveformDataset and PosteriorModel. - For model type 'nsf+embedding' (the only acceptable type at this point) this also - initializes the embedding network projection stage with SVD V matrices based on - clean detector waveforms. + If the embedding network requests an SVD projection stage, this also initializes + the projection weights with SVD V matrices based on clean detector waveforms. Parameters ---------- @@ -117,11 +117,13 @@ def prepare_training_new( ) # No transforms yet initial_weights = {} - # The embedding network is assumed to have an SVD projection layer. If other types - # of embedding networks are added in the future, update this code. + update_model_config(train_settings["model"]) # Map old schemas forward. - if train_settings["model"].get("embedding_kwargs", None): - # First, build the SVD for seeding the embedding network. + # Build the SVD for seeding the embedding network, if it declares one. + # TODO: replace with the architecture-owned initialization hook (build-system + # step 4), so that the trainer no longer knows about the SVD. + embedding_settings = train_settings["model"].get("embedding_net") + if embedding_settings and "svd" in embedding_settings.get("kwargs", {}): print("\nBuilding SVD for initialization of embedding network.") initial_weights["V_rb_list"] = build_svd_for_embedding_network( wfd, @@ -130,7 +132,7 @@ def prepare_training_new( num_workers=local_settings["num_workers"], batch_size=train_settings["training"]["stage_0"]["batch_size"], out_dir=train_dir, - **train_settings["model"]["embedding_kwargs"]["svd"], + **embedding_settings["kwargs"]["svd"], ) # Now set the transforms for training. We need to do this here so that we can (a) @@ -146,8 +148,7 @@ def prepare_training_new( train_settings["training"]["stage_0"]["asd_dataset_path"], ) - # This modifies the model settings in-place. - autocomplete_model_kwargs(train_settings["model"], wfd[0]) + train_settings["model"] = complete_model_settings(train_settings["model"], wfd[0]) full_settings = { "dataset_settings": wfd.settings, "train_settings": train_settings, diff --git a/dingo/pipe/default_settings.py b/dingo/pipe/default_settings.py index 53edc441..870f7696 100644 --- a/dingo/pipe/default_settings.py +++ b/dingo/pipe/default_settings.py @@ -4,17 +4,19 @@ "threshold_std": 5, "nde_settings": { "model": { - "posterior_model_type": "normalizing_flow", - "posterior_kwargs": { - "num_flow_steps": 5, - "base_transform_kwargs": { - "hidden_dim": 256, - "num_transform_blocks": 4, - "activation": "elu", - "dropout_probability": 0.1, - "batch_norm": True, - "num_bins": 8, - "base_transform_type": "rq-coupling", + "distribution": { + "type": "normalizing_flow", + "kwargs": { + "num_flow_steps": 5, + "base_transform_kwargs": { + "hidden_dim": 256, + "num_transform_blocks": 4, + "activation": "elu", + "dropout_probability": 0.1, + "batch_norm": True, + "num_bins": 8, + "base_transform_type": "rq-coupling", + }, }, }, }, diff --git a/dingo/pipe/sampling.py b/dingo/pipe/sampling.py index 9138066b..184832d4 100644 --- a/dingo/pipe/sampling.py +++ b/dingo/pipe/sampling.py @@ -13,6 +13,7 @@ ) from dingo.core.posterior_models.build_model import build_model_from_kwargs +from dingo.core.utils.backward_compatibility import update_model_config from dingo.gw.data.event_dataset import EventDataset from dingo.gw.inference.gw_samplers import GWSampler, GWSamplerGNPE from dingo.gw.inference.inference_utils import prepare_log_prob @@ -164,8 +165,9 @@ def density_recovery_settings(self, settings): # FIXME: If there are proxies other than time, the condition needs to be updated. if len(self.detectors) == 1: model_settings = self._density_recovery_settings["nde_settings"]["model"] - if model_settings["posterior_model_type"] == "normalizing_flow": - base_transform_kwargs = model_settings["posterior_kwargs"][ + update_model_config(model_settings) # User settings may use old schema. + if model_settings["distribution"]["type"] == "normalizing_flow": + base_transform_kwargs = model_settings["distribution"]["kwargs"][ "base_transform_kwargs" ] if base_transform_kwargs["base_transform_type"] == "rq-coupling": diff --git a/docs/source/network_architecture.ipynb b/docs/source/network_architecture.ipynb index 97b2025e..674a19d5 100644 --- a/docs/source/network_architecture.ipynb +++ b/docs/source/network_architecture.ipynb @@ -20,7 +20,8 @@ "metadata": {}, "outputs": [], "source": [ - "from dingo.core.nn.nsf import create_nsf_with_rb_projection_embedding_net" + "from dingo.core.nn.nsf import FlowWrapper, create_nsf_model\n", + "from dingo.core.nn.enets import ConcatContextMerger, DenseSVDEmbedding" ] }, { @@ -126,7 +127,11 @@ "metadata": {}, "outputs": [], "source": [ - "nde = create_nsf_with_rb_projection_embedding_net(posterior_kwargs, embedding_kwargs)" + "embedding_net = ConcatContextMerger(\n", + " DenseSVDEmbedding(**embedding_kwargs), num_context_parameters=1\n", + ")\n", + "flow = create_nsf_model(**posterior_kwargs)\n", + "nde = FlowWrapper(flow, embedding_net, embedding_net.input_keys)" ] }, { @@ -234,4 +239,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} +} \ No newline at end of file diff --git a/examples/fmpe_model/train_settings.yaml b/examples/fmpe_model/train_settings.yaml index 4aef4a51..801522fe 100644 --- a/examples/fmpe_model/train_settings.yaml +++ b/examples/fmpe_model/train_settings.yaml @@ -32,45 +32,48 @@ data: - psi model: - posterior_model_type: flow_matching - posterior_kwargs: - activation: gelu - batch_norm: true - context_with_glu: false - dropout: 0.0 - hidden_dims: [4096, 4096, 4096, - 2048, 2048, 2048, - 1024, 1024, 1024, 1024, 1024, 1024, - 512, 512, 512, 512, 512, 512, 512, 512, - 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, - 128, 128, 128, 128, 128, - 64, 64, 64, - 32, 32, 32, - 16, 16, 16] - sigma_min: 0.001 - theta_embedding_kwargs: - embedding_net: - activation: gelu - hidden_dims: [16, 32, 64, 128, 256] - output_dim: 256 - type: DenseResidualNet - encoding: - encode_all: false - frequencies: 0 - theta_with_glu: true - time_prior_exponent: 1 - type: DenseResidualNet - embedding_kwargs: - activation: gelu - batch_norm: true - dropout: 0.0 - hidden_dims: - - 2048 - output_dim: 2048 - svd: - num_training_samples: 50000 - num_validation_samples: 10000 - size: 150 + distribution: + type: flow_matching + kwargs: + activation: gelu + batch_norm: true + context_with_glu: false + dropout: 0.0 + hidden_dims: [4096, 4096, 4096, + 2048, 2048, 2048, + 1024, 1024, 1024, 1024, 1024, 1024, + 512, 512, 512, 512, 512, 512, 512, 512, + 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, + 128, 128, 128, 128, 128, + 64, 64, 64, + 32, 32, 32, + 16, 16, 16] + sigma_min: 0.001 + theta_embedding_kwargs: + embedding_net: + activation: gelu + hidden_dims: [16, 32, 64, 128, 256] + output_dim: 256 + type: DenseResidualNet + encoding: + encode_all: false + frequencies: 0 + theta_with_glu: true + time_prior_exponent: 1 + type: DenseResidualNet + embedding_net: + type: dense_svd + kwargs: + activation: gelu + batch_norm: true + dropout: 0.0 + hidden_dims: + - 2048 + output_dim: 2048 + svd: + num_training_samples: 50000 + num_validation_samples: 10000 + size: 150 training: evaluate: true diff --git a/examples/gnpe_model/train_settings.yaml b/examples/gnpe_model/train_settings.yaml index 747352d6..697cc1c0 100644 --- a/examples/gnpe_model/train_settings.yaml +++ b/examples/gnpe_model/train_settings.yaml @@ -36,32 +36,34 @@ data: # Model architecture model: - posterior_model_type: normalizing_flow - # kwargs for neural spline flow - posterior_kwargs: - num_flow_steps: 30 - base_transform_kwargs: - hidden_dim: 1024 - num_transform_blocks: 5 + distribution: + type: normalizing_flow + kwargs: + num_flow_steps: 30 + base_transform_kwargs: + hidden_dim: 1024 + num_transform_blocks: 5 + activation: elu + dropout_probability: 0.0 + batch_norm: True + num_bins: 8 + base_transform_type: rq-coupling + # kwargs for embedding net + embedding_net: + type: dense_svd + kwargs: + output_dim: 128 + hidden_dims: [1024, 1024, 1024, 1024, 1024, 1024, + 512, 512, 512, 512, 512, 512, + 256, 256, 256, 256, 256, 256, + 128, 128, 128, 128, 128, 128] activation: elu - dropout_probability: 0.0 + dropout: 0.0 batch_norm: True - num_bins: 8 - base_transform_type: rq-coupling - # kwargs for embedding net - embedding_kwargs: - output_dim: 128 - hidden_dims: [1024, 1024, 1024, 1024, 1024, 1024, - 512, 512, 512, 512, 512, 512, - 256, 256, 256, 256, 256, 256, - 128, 128, 128, 128, 128, 128] - activation: elu - dropout: 0.0 - batch_norm: True - svd: - num_training_samples: 50000 - num_validation_samples: 10000 - size: 200 + svd: + num_training_samples: 50000 + num_validation_samples: 10000 + size: 200 # Training is divided in stages. They each require all settings as indicated below. training: diff --git a/examples/gnpe_model/train_settings_init.yaml b/examples/gnpe_model/train_settings_init.yaml index 7bdf87ee..c4a04401 100644 --- a/examples/gnpe_model/train_settings_init.yaml +++ b/examples/gnpe_model/train_settings_init.yaml @@ -21,32 +21,34 @@ data: # Model architecture model: - posterior_model_type: normalizing_flow - # kwargs for neural spline flow - posterior_kwargs: - num_flow_steps: 15 - base_transform_kwargs: - hidden_dim: 512 - num_transform_blocks: 5 + distribution: + type: normalizing_flow + kwargs: + num_flow_steps: 15 + base_transform_kwargs: + hidden_dim: 512 + num_transform_blocks: 5 + activation: elu + dropout_probability: 0.0 + batch_norm: True + num_bins: 8 + base_transform_type: rq-coupling + # kwargs for embedding net + embedding_net: + type: dense_svd + kwargs: + output_dim: 128 + hidden_dims: [512, 512, 512, 512, 512, 512, + 256, 256, 256, 256, 256, 256, + 128, 128, 128, 128, 128, 128 + ] activation: elu - dropout_probability: 0.0 + dropout: 0.0 batch_norm: True - num_bins: 8 - base_transform_type: rq-coupling - # kwargs for embedding net - embedding_kwargs: - output_dim: 128 - hidden_dims: [512, 512, 512, 512, 512, 512, - 256, 256, 256, 256, 256, 256, - 128, 128, 128, 128, 128, 128 - ] - activation: elu - dropout: 0.0 - batch_norm: True - svd: - num_training_samples: 50000 - num_validation_samples: 10000 - size: 200 + svd: + num_training_samples: 50000 + num_validation_samples: 10000 + size: 200 # Training is divided in stages. They each require all settings as indicated below. training: diff --git a/examples/misc/is_settings.yaml b/examples/misc/is_settings.yaml index 40ead7e5..6b359309 100755 --- a/examples/misc/is_settings.yaml +++ b/examples/misc/is_settings.yaml @@ -17,17 +17,18 @@ nde: parameters: null # need all parameters for likelihood # nde architecture model: - posterior_model_type: normalizing_flow - posterior_kwargs: - num_flow_steps: 10 - base_transform_kwargs: - hidden_dim: 128 - num_transform_blocks: 2 - activation: elu - dropout_probability: 0.1 - batch_norm: true - num_bins: 8 - base_transform_type: rq-coupling + distribution: + type: normalizing_flow + kwargs: + num_flow_steps: 10 + base_transform_kwargs: + hidden_dim: 128 + num_transform_blocks: 2 + activation: elu + dropout_probability: 0.1 + batch_norm: true + num_bins: 8 + base_transform_type: rq-coupling # nde training training: device: cpu diff --git a/examples/npe_model/train_settings.yaml b/examples/npe_model/train_settings.yaml index 1095f6d5..9bc5af2e 100644 --- a/examples/npe_model/train_settings.yaml +++ b/examples/npe_model/train_settings.yaml @@ -35,32 +35,34 @@ data: # Model architecture model: - posterior_model_type: normalizing_flow - # kwargs for neural spline flow - posterior_kwargs: - num_flow_steps: 30 - base_transform_kwargs: - hidden_dim: 1024 - num_transform_blocks: 5 + distribution: + type: normalizing_flow + kwargs: + num_flow_steps: 30 + base_transform_kwargs: + hidden_dim: 1024 + num_transform_blocks: 5 + activation: elu + dropout_probability: 0.0 + batch_norm: True + num_bins: 8 + base_transform_type: rq-coupling + # kwargs for embedding net + embedding_net: + type: dense_svd + kwargs: + output_dim: 128 + hidden_dims: [1024, 1024, 1024, 1024, 1024, 1024, + 512, 512, 512, 512, 512, 512, + 256, 256, 256, 256, 256, 256, + 128, 128, 128, 128, 128, 128] activation: elu - dropout_probability: 0.0 + dropout: 0.0 batch_norm: True - num_bins: 8 - base_transform_type: rq-coupling - # kwargs for embedding net - embedding_kwargs: - output_dim: 128 - hidden_dims: [1024, 1024, 1024, 1024, 1024, 1024, - 512, 512, 512, 512, 512, 512, - 256, 256, 256, 256, 256, 256, - 128, 128, 128, 128, 128, 128] - activation: elu - dropout: 0.0 - batch_norm: True - svd: - num_training_samples: 50000 - num_validation_samples: 10000 - size: 200 + svd: + num_training_samples: 50000 + num_validation_samples: 10000 + size: 200 # Training is divided in stages. They each require all settings as indicated below. training: diff --git a/examples/toy_npe_model/train_settings.yaml b/examples/toy_npe_model/train_settings.yaml index d3e44f08..5eee2aa2 100644 --- a/examples/toy_npe_model/train_settings.yaml +++ b/examples/toy_npe_model/train_settings.yaml @@ -26,29 +26,31 @@ data: # Model architecture model: - posterior_model_type: normalizing_flow - # kwargs for neural spline flow - posterior_kwargs: - num_flow_steps: 5 - base_transform_kwargs: - hidden_dim: 64 - num_transform_blocks: 5 + distribution: + type: normalizing_flow + kwargs: + num_flow_steps: 5 + base_transform_kwargs: + hidden_dim: 64 + num_transform_blocks: 5 + activation: elu + dropout_probability: 0.0 + batch_norm: True + num_bins: 8 + base_transform_type: rq-coupling + # kwargs for embedding net + embedding_net: + type: dense_svd + kwargs: + output_dim: 128 + hidden_dims: [1024, 512, 256, 128] activation: elu - dropout_probability: 0.0 + dropout: 0.0 batch_norm: True - num_bins: 8 - base_transform_type: rq-coupling - # kwargs for embedding net - embedding_kwargs: - output_dim: 128 - hidden_dims: [1024, 512, 256, 128] - activation: elu - dropout: 0.0 - batch_norm: True - svd: - num_training_samples: 1000 - num_validation_samples: 100 - size: 50 + svd: + num_training_samples: 1000 + num_validation_samples: 100 + size: 50 # The first stage (and only) stage of training. training: diff --git a/tests/core/test_build_model.py b/tests/core/test_build_model.py index 306303b9..6b6b81ef 100644 --- a/tests/core/test_build_model.py +++ b/tests/core/test_build_model.py @@ -1,14 +1,12 @@ """ -Characterization tests for the model build path. +Tests for the model build path. -These tests pin the *current* behavior of building posterior models from settings -dictionaries — type dispatch (build_model_from_kwargs), dimensional autocompletion -(autocomplete_model_kwargs), end-to-end construction and forward passes for all three +These tests cover building posterior models from settings dictionaries — type +dispatch (build_model_from_kwargs), per-architecture dimension inference +(complete_model_settings), end-to-end construction and forward passes for all three posterior model types, SVD initial-weight seeding, and loading of old-schema -checkpoints — ahead of the NN build-system refactor (see hackathon/ -NN_Build_System_Design.md). If one of these tests breaks, either the refactor changed -observable behavior (fix the refactor) or the behavior change is intended and -documented (update the test alongside the compatibility shim). +settings/checkpoints through the update_model_config boundary (see +hackathon/NN_Build_System_Design.md). """ import copy @@ -19,8 +17,8 @@ from dingo.core.posterior_models.base_model import BasePosteriorModel from dingo.core.posterior_models.build_model import ( - autocomplete_model_kwargs, build_model_from_kwargs, + complete_model_settings, ) from dingo.core.posterior_models.flow_matching import FlowMatchingPosteriorModel from dingo.core.posterior_models.normalizing_flow import NormalizingFlowPosteriorModel @@ -36,24 +34,23 @@ BATCH_SIZE = 5 -def embedding_kwargs(): - return { - "input_dims": list(DATA_SHAPE), +def embedding_kwargs(completed=True): + """Embedding-net kwargs; completed=False gives user-style settings (no dims).""" + kwargs = { "svd": {"size": 10}, - "V_rb_list": None, "output_dim": EMBEDDING_OUTPUT_DIM, "hidden_dims": [32, 16, 8], "activation": "elu", "dropout": 0.0, "batch_norm": True, - "added_context": False, } + if completed: + kwargs["input_dims"] = list(DATA_SHAPE) + return kwargs -def nsf_posterior_kwargs(): - return { - "input_dim": NUM_PARAMETERS, - "context_dim": EMBEDDING_OUTPUT_DIM, +def nsf_distribution_kwargs(completed=True): + kwargs = { "num_flow_steps": 2, "base_transform_kwargs": { "hidden_dim": 16, @@ -65,12 +62,14 @@ def nsf_posterior_kwargs(): "base_transform_type": "rq-coupling", }, } + if completed: + kwargs["theta_dim"] = NUM_PARAMETERS + kwargs["context_dim"] = EMBEDDING_OUTPUT_DIM + return kwargs -def cflow_posterior_kwargs(): - return { - "input_dim": NUM_PARAMETERS, - "context_dim": EMBEDDING_OUTPUT_DIM, +def cflow_distribution_kwargs(completed=True): + kwargs = { "activation": "gelu", "batch_norm": False, "dropout": 0.0, @@ -92,25 +91,31 @@ def cflow_posterior_kwargs(): "encoding": {"encode_all": False, "frequencies": 0}, }, } + if completed: + kwargs["theta_dim"] = NUM_PARAMETERS + kwargs["context_dim"] = EMBEDDING_OUTPUT_DIM + return kwargs -def model_settings(posterior_model_type): - if posterior_model_type == "normalizing_flow": - posterior_kwargs = nsf_posterior_kwargs() +def model_settings(model_type, completed=True): + if model_type == "normalizing_flow": + distribution_kwargs = nsf_distribution_kwargs(completed) else: - posterior_kwargs = cflow_posterior_kwargs() - if posterior_model_type == "flow_matching": - posterior_kwargs["sigma_min"] = 0.001 - elif posterior_model_type == "score_matching": - posterior_kwargs["epsilon"] = 1e-3 - posterior_kwargs["beta_min"] = 0.1 - posterior_kwargs["beta_max"] = 20.0 + distribution_kwargs = cflow_distribution_kwargs(completed) + if model_type == "flow_matching": + distribution_kwargs["sigma_min"] = 0.001 + elif model_type == "score_matching": + distribution_kwargs["epsilon"] = 1e-3 + distribution_kwargs["beta_min"] = 0.1 + distribution_kwargs["beta_max"] = 20.0 return { "train_settings": { "model": { - "posterior_model_type": posterior_model_type, - "posterior_kwargs": posterior_kwargs, - "embedding_kwargs": embedding_kwargs(), + "distribution": {"type": model_type, "kwargs": distribution_kwargs}, + "embedding_net": { + "type": "dense_svd", + "kwargs": embedding_kwargs(completed), + }, } } } @@ -124,9 +129,7 @@ def data_sample(with_gnpe_proxies=False): "waveform": np.random.rand(*DATA_SHAPE).astype(np.float32), } if with_gnpe_proxies: - sample["context_parameters"] = np.random.rand(GNPE_PROXY_DIM).astype( - np.float32 - ) + sample["context_parameters"] = np.random.rand(GNPE_PROXY_DIM).astype(np.float32) return sample @@ -161,7 +164,7 @@ def test_build_model_from_kwargs_dispatch(posterior_model_type, expected_class): def test_build_model_from_kwargs_rejects_unknown_type(): settings = model_settings("normalizing_flow") - settings["train_settings"]["model"]["posterior_model_type"] = "not_a_model" + settings["train_settings"]["model"]["distribution"]["type"] = "not_a_model" with pytest.raises(ValueError): build_model_from_kwargs(settings=settings, device="cpu") @@ -175,37 +178,72 @@ def test_build_model_from_kwargs_requires_exactly_one_source(): # ----------------------------------------------------------------------------------- -# autocomplete_model_kwargs: dimensional glue +# complete_model_settings: per-architecture dimension inference # ----------------------------------------------------------------------------------- -def test_autocomplete_model_kwargs_without_gnpe(): - model_kwargs = model_settings("normalizing_flow")["train_settings"]["model"] - # Settings as written by a user: dims absent. - del model_kwargs["embedding_kwargs"]["input_dims"] - del model_kwargs["posterior_kwargs"]["input_dim"] - del model_kwargs["posterior_kwargs"]["context_dim"] - - autocomplete_model_kwargs(model_kwargs, data_sample(with_gnpe_proxies=False)) +def test_complete_model_settings_without_context_parameters(): + user = model_settings("normalizing_flow", completed=False)["train_settings"][ + "model" + ] + completed = complete_model_settings(user, data_sample(with_gnpe_proxies=False)) - assert model_kwargs["embedding_kwargs"]["input_dims"] == list(DATA_SHAPE) - assert model_kwargs["posterior_kwargs"]["input_dim"] == NUM_PARAMETERS - assert model_kwargs["embedding_kwargs"]["added_context"] is False - assert model_kwargs["posterior_kwargs"]["context_dim"] == EMBEDDING_OUTPUT_DIM + # The user settings are not modified; the completed ones carry the dims. + assert "input_dims" not in user["embedding_net"]["kwargs"] + assert completed["embedding_net"]["kwargs"]["input_dims"] == list(DATA_SHAPE) + assert completed["distribution"]["kwargs"]["theta_dim"] == NUM_PARAMETERS + assert completed["distribution"]["kwargs"]["context_dim"] == EMBEDDING_OUTPUT_DIM + assert "context_merger" not in completed + # The completed settings build directly. + settings = {"train_settings": {"model": completed}} + pm = build_model_from_kwargs(settings=settings, device="cpu") + assert type(pm) is NormalizingFlowPosteriorModel -def test_autocomplete_model_kwargs_with_gnpe(): - model_kwargs = model_settings("normalizing_flow")["train_settings"]["model"] - autocomplete_model_kwargs(model_kwargs, data_sample(with_gnpe_proxies=True)) +def test_complete_model_settings_with_context_parameters(): + """With context_parameters in the batch, a concat context merger is added and + context_dim grows accordingly.""" + user = model_settings("normalizing_flow", completed=False)["train_settings"][ + "model" + ] + completed = complete_model_settings(user, data_sample(with_gnpe_proxies=True)) - assert model_kwargs["embedding_kwargs"]["added_context"] is True + assert completed["context_merger"] == { + "type": "concat", + "kwargs": {"num_context_parameters": GNPE_PROXY_DIM}, + } assert ( - model_kwargs["posterior_kwargs"]["context_dim"] + completed["distribution"]["kwargs"]["context_dim"] == EMBEDDING_OUTPUT_DIM + GNPE_PROXY_DIM ) +def test_complete_model_settings_rejects_dims_in_user_settings(): + """Dimensions are derived from the data; specifying them is an error.""" + user = model_settings("normalizing_flow", completed=True)["train_settings"]["model"] + with pytest.raises(ValueError, match="derived from the data"): + complete_model_settings(user, data_sample()) + + user = model_settings("normalizing_flow", completed=False)["train_settings"][ + "model" + ] + user["embedding_net"]["kwargs"]["input_dims"] = list(DATA_SHAPE) + with pytest.raises(ValueError, match="derived from the data"): + complete_model_settings(user, data_sample()) + + +def test_complete_model_settings_rejects_unused_context_merger(): + """A context merger without context parameters in the data is an error, not + silently dropped.""" + user = model_settings("normalizing_flow", completed=False)["train_settings"][ + "model" + ] + user["context_merger"] = {"type": "concat"} + with pytest.raises(ValueError, match="context_merger"): + complete_model_settings(user, data_sample(with_gnpe_proxies=False)) + + # ----------------------------------------------------------------------------------- # End-to-end: build + forward passes for all three model types # ----------------------------------------------------------------------------------- @@ -239,12 +277,18 @@ def test_model_forward_passes(posterior_model_type): def test_normalizing_flow_with_gnpe_context(): - """With added_context=True, the embedding merges (waveform, context_parameters) - via ModuleMerger, and sampling/density evaluation consume both context entries.""" + """With a concat context merger, the embedding merges (waveform, + context_parameters), and sampling/density evaluation consume both context + entries.""" settings = model_settings("normalizing_flow") model = settings["train_settings"]["model"] - model["embedding_kwargs"]["added_context"] = True - model["posterior_kwargs"]["context_dim"] = EMBEDDING_OUTPUT_DIM + GNPE_PROXY_DIM + model["context_merger"] = { + "type": "concat", + "kwargs": {"num_context_parameters": GNPE_PROXY_DIM}, + } + model["distribution"]["kwargs"]["context_dim"] = ( + EMBEDDING_OUTPUT_DIM + GNPE_PROXY_DIM + ) pm = build_model_from_kwargs(settings=settings, device="cpu") theta, context = batch(with_gnpe_proxies=True) @@ -263,13 +307,12 @@ def test_normalizing_flow_with_gnpe_context(): def test_normalizing_flow_unconditional(): - """Without embedding_kwargs, an unconditional flow is built (the models-as-priors - path used by unconditional_density_estimation).""" + """Without an embedding_net, an unconditional flow is built (the + models-as-priors path used by unconditional_density_estimation).""" settings = model_settings("normalizing_flow") model = settings["train_settings"]["model"] - del model["embedding_kwargs"] - # Convention from unconditional_density_estimation.py:74-75. - model["posterior_kwargs"]["context_dim"] = None + del model["embedding_net"] + model["distribution"]["kwargs"]["context_dim"] = None pm = build_model_from_kwargs(settings=settings, device="cpu") theta = torch.rand(BATCH_SIZE, NUM_PARAMETERS) @@ -364,39 +407,114 @@ def test_initial_weights_seed_svd_projection(): # ----------------------------------------------------------------------------------- -def test_update_model_config_maps_old_schema(): - old = { +def old_schema_model(model_type="normalizing_flow", added_context=False): + """A completed model config in the old schema, as found in old checkpoints.""" + if model_type == "normalizing_flow": + posterior_kwargs = nsf_distribution_kwargs() + else: + posterior_kwargs = cflow_distribution_kwargs() + posterior_kwargs["sigma_min"] = 0.001 + posterior_kwargs["input_dim"] = posterior_kwargs.pop("theta_dim") + old_embedding_kwargs = embedding_kwargs() + old_embedding_kwargs["V_rb_list"] = None + old_embedding_kwargs["added_context"] = added_context + if added_context: + posterior_kwargs["context_dim"] = EMBEDDING_OUTPUT_DIM + GNPE_PROXY_DIM + return { + "posterior_model_type": model_type, + "posterior_kwargs": posterior_kwargs, + "embedding_kwargs": old_embedding_kwargs, + } + + +def test_update_model_config_maps_old_schemas(): + """All old schemas map forward to the current one, including the oldest + nsf+embedding form; the mapping is idempotent.""" + old = old_schema_model(added_context=True) + oldest = { "type": "nsf+embedding", - "nsf_kwargs": nsf_posterior_kwargs(), - "embedding_net_kwargs": embedding_kwargs(), + "nsf_kwargs": old["posterior_kwargs"], + "embedding_net_kwargs": old["embedding_kwargs"], } - update_model_config(old) - assert old["posterior_model_type"] == "normalizing_flow" - assert old["posterior_kwargs"] == nsf_posterior_kwargs() - assert old["embedding_kwargs"] == embedding_kwargs() - assert "type" not in old and "nsf_kwargs" not in old + for settings in (old, oldest): + update_model_config(settings) + assert settings["distribution"]["type"] == "normalizing_flow" + assert settings["distribution"]["kwargs"]["theta_dim"] == NUM_PARAMETERS + assert settings["embedding_net"]["type"] == "dense_svd" + kwargs = settings["embedding_net"]["kwargs"] + assert "added_context" not in kwargs and "V_rb_list" not in kwargs + # Old concatenated context maps to the concat merger, with the number of + # context parameters recovered from the completed dims. + assert settings["context_merger"] == { + "type": "concat", + "kwargs": {"num_context_parameters": GNPE_PROXY_DIM}, + } + before = copy.deepcopy(settings) + update_model_config(settings) + assert settings == before def test_update_model_config_lowercases_builtin_type_names(): """Model types used to be matched case-insensitively; the compat shim lowercases built-in names from old checkpoints so the case-sensitive registry finds them.""" - settings = model_settings("normalizing_flow") - model = settings["train_settings"]["model"] + model = old_schema_model() model["posterior_model_type"] = "Normalizing_Flow" + settings = {"train_settings": {"model": model}} pm = build_model_from_kwargs(settings=settings, device="cpu") assert type(pm) is NormalizingFlowPosteriorModel -def test_build_model_from_old_schema_settings(): - settings = model_settings("normalizing_flow") - model = settings["train_settings"]["model"] - settings["train_settings"]["model"] = { - "type": "nsf+embedding", - "nsf_kwargs": model["posterior_kwargs"], - "embedding_net_kwargs": model["embedding_kwargs"], +@pytest.mark.parametrize("added_context", [False, True]) +def test_old_schema_builds_state_dict_compatible_network(added_context): + """A network built from old-schema settings has exactly the same state-dict keys + and shapes as one built from the new schema — old checkpoints stay loadable.""" + old_settings = { + "train_settings": {"model": old_schema_model(added_context=added_context)} } - pm = build_model_from_kwargs(settings=settings, device="cpu") - assert type(pm) is NormalizingFlowPosteriorModel + pm_old = build_model_from_kwargs(settings=old_settings, device="cpu") + + new_settings = model_settings("normalizing_flow") + if added_context: + model = new_settings["train_settings"]["model"] + model["context_merger"] = { + "type": "concat", + "kwargs": {"num_context_parameters": GNPE_PROXY_DIM}, + } + model["distribution"]["kwargs"]["context_dim"] = ( + EMBEDDING_OUTPUT_DIM + GNPE_PROXY_DIM + ) + pm_new = build_model_from_kwargs(settings=new_settings, device="cpu") + + state_old = pm_old.network.state_dict() + state_new = pm_new.network.state_dict() + assert list(state_old) == list(state_new) + assert all(state_old[k].shape == state_new[k].shape for k in state_old) + + +@pytest.mark.parametrize("model_type", ["normalizing_flow", "flow_matching"]) +def test_load_old_schema_checkpoint_file(tmp_path, model_type): + """A checkpoint file whose model_kwargs use the old schema loads through the + update_model_config boundary, with identical weights.""" + pm = build_model_from_kwargs(settings=model_settings(model_type), device="cpu") + + old_model = old_schema_model(model_type) + checkpoint = { + "model_kwargs": old_model, + "model_state_dict": pm.network.state_dict(), + "epoch": 3, + "version": "dingo=0.9.9", + "metadata": {"train_settings": {"model": copy.deepcopy(old_model)}}, + } + filename = str(tmp_path / "old_model.pt") + torch.save(checkpoint, filename) + + pm_loaded = build_model_from_kwargs( + filename=filename, device="cpu", load_training_info=False + ) + assert type(pm_loaded) is type(pm) + assert pm_loaded.epoch == 3 + for p0, p1 in zip(pm.network.parameters(), pm_loaded.network.parameters()): + assert torch.equal(p0.data, p1.data) # ----------------------------------------------------------------------------------- @@ -404,9 +522,7 @@ def test_build_model_from_old_schema_settings(): # ----------------------------------------------------------------------------------- -@pytest.mark.parametrize( - "posterior_model_type", ["normalizing_flow", "flow_matching"] -) +@pytest.mark.parametrize("posterior_model_type", ["normalizing_flow", "flow_matching"]) def test_save_and_rebuild_from_file(tmp_path, posterior_model_type): pm = build_model_from_kwargs( settings=model_settings(posterior_model_type), device="cpu" diff --git a/tests/core/test_nsf.py b/tests/core/test_nsf.py index d68c7931..b314b8be 100644 --- a/tests/core/test_nsf.py +++ b/tests/core/test_nsf.py @@ -2,12 +2,12 @@ import types import torch import torch.optim as optim -from dingo.core.nn.nsf import ( - create_nsf_model, - FlowWrapper, - create_nsf_with_rb_projection_embedding_net, +from dingo.core.nn.nsf import create_nsf_model, FlowWrapper +from dingo.core.nn.enets import ( + ConcatContextMerger, + DenseSVDEmbedding, + create_enet_with_projection_layer_and_dense_resnet, ) -from dingo.core.nn.enets import create_enet_with_projection_layer_and_dense_resnet from dingo.core.utils import torchutils @@ -236,18 +236,47 @@ def test_backward_pass_for_log_prob_of_nsf(data_setup_nsf_small): ), "Training may not have worked. Check manually that sampling improves." -def test_model_builder_for_nsf_with_rb_embedding_net(data_setup_nsf_small): +def test_registered_embedding_with_merger(data_setup_nsf_small): """ - Test the builder function create_nsf_with_rb_projection_embedding_net. + Test the registered dense_svd embedding and the concat context merger: contract + attributes, dimension inference, and end-to-end use inside a FlowWrapper. """ d = data_setup_nsf_small + kwargs = { + k: v + for k, v in d.embedding_net_kwargs.items() + if k not in ("added_context", "V_rb_list", "input_dims") + } + num_context_parameters = d.context_dim - kwargs["output_dim"] + + completed = DenseSVDEmbedding.complete_settings( + kwargs, {"waveform": d.x[0], "context_parameters": d.z[0]} + ) + assert completed["input_dims"] == list(d.embedding_net_kwargs["input_dims"]) + + embedding_net = DenseSVDEmbedding(**completed) + assert embedding_net.input_keys == ("waveform",) + assert embedding_net.output_dim == kwargs["output_dim"] - model = create_nsf_with_rb_projection_embedding_net( - d.nde_kwargs, d.embedding_net_kwargs + merged = ConcatContextMerger(embedding_net, num_context_parameters) + assert merged.input_keys == CONTEXT_KEYS + assert merged.output_dim == d.context_dim + assert ( + ConcatContextMerger.merged_output_dim( + embedding_net.output_dim, num_context_parameters + ) + == d.context_dim + ) + + # State-dict layout matches the historic builder (old checkpoints). + legacy = create_enet_with_projection_layer_and_dense_resnet( + **d.embedding_net_kwargs ) - # The builder derives the context routing from added_context. - assert model.context_keys == CONTEXT_KEYS + assert list(merged.state_dict()) == list(legacy.state_dict()) + + flow = d.nde_builder(**d.nde_kwargs) + model = FlowWrapper(flow, merged, merged.input_keys) loss = -model(d.y, d.context) assert list(loss.shape) == [d.batch_size], "Unexpected output shape." @@ -255,10 +284,3 @@ def test_model_builder_for_nsf_with_rb_embedding_net(data_setup_nsf_small): "Unexpected log prob encountered. Network initialization or " "normalization seems broken." ) - - # missing context key - with pytest.raises(ValueError): - model(d.y, {"waveform": d.x}) - # wrongly-shaped context tensor - with pytest.raises(RuntimeError): - model(d.y, {"waveform": d.x, "context_parameters": d.x.flatten(start_dim=1)}) diff --git a/tests/core/test_registry.py b/tests/core/test_registry.py index 1ac8ea2a..0b346b03 100644 --- a/tests/core/test_registry.py +++ b/tests/core/test_registry.py @@ -49,10 +49,7 @@ def test_get_dotted_path(registry): def test_get_file_path(registry, tmp_path): plugin = tmp_path / "my_plugin.py" - plugin.write_text( - "class MyNN:\n" - " marker = 'from-file'\n" - ) + plugin.write_text("class MyNN:\n" " marker = 'from-file'\n") component = registry.get(f"{plugin}:MyNN") assert component.marker == "from-file" # Second lookup resolves from the cache to the same class object. @@ -106,7 +103,7 @@ def test_builtin_distributions_are_registered(): def test_build_model_from_kwargs_with_file_path_plugin(tmp_path): """A NeuralDistribution defined in a user file (never pip-installed, not in the - dingo source) can be selected as posterior_model_type via the file-path form.""" + dingo source) can be selected as the distribution type via the file-path form.""" from tests.core.test_build_model import model_settings plugin = tmp_path / "my_distribution.py" @@ -117,8 +114,8 @@ def test_build_model_from_kwargs_with_file_path_plugin(tmp_path): " pass\n" ) settings = model_settings("normalizing_flow") - settings["train_settings"]["model"][ - "posterior_model_type" + settings["train_settings"]["model"]["distribution"][ + "type" ] = f"{plugin}:MyDistribution" from dingo.core.posterior_models.build_model import build_model_from_kwargs diff --git a/tests/core/test_unconditional_density_estimation.py b/tests/core/test_unconditional_density_estimation.py index 38b6a930..4660106b 100644 --- a/tests/core/test_unconditional_density_estimation.py +++ b/tests/core/test_unconditional_density_estimation.py @@ -15,8 +15,9 @@ def _nde_settings(): """Minimal but valid settings for a fast (1-epoch) unconditional flow. - Uses the ``model.posterior_kwargs`` shape that train_unconditional_density_estimator - reads (matching pipe DENSITY_RECOVERY_SETTINGS), not the flat nde_settings template. + Deliberately uses the old ``model.posterior_kwargs`` schema, so that the + update_model_config shim inside train_unconditional_density_estimator is + exercised (user nde settings files in the wild still use it). """ return { "data": {}, @@ -71,8 +72,8 @@ def test_train_unconditional_density_estimator_basic(result, tmp_path): # The network is configured as an unconditional flow over the chosen parameters. assert settings["data"]["unconditional"] is True - assert settings["model"]["posterior_kwargs"]["input_dim"] == len(PARAMETERS) - assert settings["model"]["posterior_kwargs"]["context_dim"] is None + assert settings["model"]["distribution"]["kwargs"]["theta_dim"] == len(PARAMETERS) + assert settings["model"]["distribution"]["kwargs"]["context_dim"] is None # Standardization is computed from the training samples. expected_mean = result.samples[PARAMETERS].to_numpy().mean(axis=0) @@ -94,7 +95,10 @@ def test_train_unconditional_density_estimator_uses_all_parameters_by_default( # With no "parameters" entry, all sample columns are used. settings = _nde_settings() train_unconditional_density_estimator(result, settings, str(tmp_path)) - assert settings["model"]["posterior_kwargs"]["input_dim"] == result.samples.shape[1] + assert ( + settings["model"]["distribution"]["kwargs"]["theta_dim"] + == result.samples.shape[1] + ) def test_train_unconditional_flow_end_to_end(result, tmp_path): @@ -103,7 +107,7 @@ def test_train_unconditional_flow_end_to_end(result, tmp_path): ) assert isinstance(model, NormalizingFlowPosteriorModel) # Trained over the requested subset only. - assert model.model_kwargs["posterior_kwargs"]["input_dim"] == len(PARAMETERS) + assert model.model_kwargs["distribution"]["kwargs"]["theta_dim"] == len(PARAMETERS) def test_train_unconditional_flow_rejects_too_many_outliers(result): From d2f9315a6a202e61dbe6e383f7f4f71c6237624e Mon Sep 17 00:00:00 2001 From: Maximilian Dax Date: Sun, 12 Jul 2026 20:04:01 +0200 Subject: [PATCH 6/9] Port chained-NPE training-side conditioning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Folds the training side of the chained-NPE branch into the conditioning design (NN_Build_System_Design §4.5): user-declared data.conditioning_parameters feed the same context_parameters list as GNPE proxies (two sources, one declared contract), the ContextMergerMLP lands as the registered "mlp" context merger (mixes context through a learned MLP so context_dim does not grow; dimensions inferred at completion instead of hand-written input_dim), and data.time_alignment=True trains on time-aligned data by skipping the per-detector time shift in ProjectOntoDetectors, with a validation guard on the required conditioning. Tests ported from the chained branch. Its inference side is a factorized-sampler chain (Group A) and is not built here. Co-Authored-By: Claude Fable 5 --- dingo/core/nn/enets.py | 71 +++++++++++++++++++ dingo/gw/training/train_builders.py | 70 ++++++++++++++++-- dingo/gw/transforms/detector_transforms.py | 68 +++++++++++------- tests/core/test_nsf.py | 51 +++++++++++++ tests/gw/test_time_alignment_validation.py | 56 +++++++++++++++ .../gw/transforms/test_detector_projection.py | 46 ++++++++++++ 6 files changed, 331 insertions(+), 31 deletions(-) create mode 100644 tests/gw/test_time_alignment_validation.py diff --git a/dingo/core/nn/enets.py b/dingo/core/nn/enets.py index 24cda50a..1b874fa1 100755 --- a/dingo/core/nn/enets.py +++ b/dingo/core/nn/enets.py @@ -398,6 +398,77 @@ def merged_output_dim(embedding_output_dim: int, num_context_parameters: int): return embedding_output_dim + num_context_parameters +@CONTEXT_MERGERS.register("mlp") +class MLPContextMerger(nn.Module): + """ + Context merger that mixes the embedded data and the (standardized) context + parameters through a learned MLP, in contrast to the concat merger which + simply concatenates them. Ported from the chained-NPE branch + (ContextMergerMLP). + + The data is first embedded, z = embedding_net(x). The context parameters c + are concatenated with z and passed through a DenseResidualNet M, producing + z_new = M(concat(z, c)) with dim(z_new) = output_dim. By default output_dim + equals dim(z), so the conditioning context fed to the downstream flow does + not grow with the number of context parameters. + """ + + def __init__( + self, + embedding_net: nn.Module, + num_context_parameters: int, + hidden_dims: Tuple, + output_dim: int = None, + activation: str = "elu", + dropout: float = 0.0, + batch_norm: bool = True, + ): + """ + Parameters + ---------- + embedding_net : nn.Module + The wrapped embedding network. + num_context_parameters : int + Number of context parameters mixed into the embedded data. Inferred + from a sample batch during settings completion. + hidden_dims : tuple + dimensions of the hidden layers of the merging DenseResidualNet + output_dim : int = None + output dimension of the merged embedding; defaults to the output + dimension of the wrapped embedding network + activation : str + activation function used in the residual blocks + dropout : float + dropout probability in the residual blocks + batch_norm : bool + whether to use batch normalization + """ + super().__init__() + if output_dim is None: + output_dim = embedding_net.output_dim + self.embedding_net = embedding_net + self.context_module = DenseResidualNet( + input_dim=embedding_net.output_dim + num_context_parameters, + output_dim=output_dim, + hidden_dims=tuple(hidden_dims), + activation=torchutils.get_activation_function_from_string(activation), + dropout=dropout, + batch_norm=batch_norm, + ) + self.input_keys = (*embedding_net.input_keys, "context_parameters") + self.output_dim = output_dim + + def forward(self, *x): + *data, context = x + z = self.embedding_net(*data) + return self.context_module(torch.cat([z, context], dim=1)) + + @staticmethod + def merged_output_dim(embedding_output_dim: int, output_dim: int = None, **_unused): + """Output dimension of the merged embedding, for settings completion.""" + return output_dim if output_dim is not None else embedding_output_dim + + def create_enet_with_projection_layer_and_dense_resnet( input_dims: List[int], # n_rb: int, diff --git a/dingo/gw/training/train_builders.py b/dingo/gw/training/train_builders.py index 494c798c..540ceb92 100755 --- a/dingo/gw/training/train_builders.py +++ b/dingo/gw/training/train_builders.py @@ -29,6 +29,47 @@ from dingo.core.utils import * +TIME_ALIGNMENT_REQUIRED_CONDITIONING = ("ra", "dec", "geocent_time") + + +def validate_time_alignment_settings(data_settings: dict) -> None: + """ + Validate that the data_settings dict is consistent with + ``data.time_alignment=True``. + + The aligned model targets the factorisation + q(theta | d) = q(theta_hat | d_aligned, ra, dec, geocent_time) + * q(ra, dec, geocent_time | d). + The transform pipeline this function gates must: + * have {ra, dec, geocent_time} as conditioning parameters (so the network + receives them as inputs), and + * NOT have them as inference targets (the aligned model does not predict + sky/time -- the sky-position model does). + It must also not coexist with ``gnpe_time_shifts``, which manipulates the + same per-detector arrival times stochastically. + """ + required = set(TIME_ALIGNMENT_REQUIRED_CONDITIONING) + conditioning_set = set(data_settings.get("conditioning_parameters", [])) + missing = required - conditioning_set + if missing: + raise ValueError( + f"data.time_alignment=True requires {sorted(required)} to be in " + f"data.conditioning_parameters; missing: {sorted(missing)}." + ) + overlap = required & set(data_settings["inference_parameters"]) + if overlap: + raise ValueError( + f"data.time_alignment=True is incompatible with having " + f"{sorted(overlap)} in data.inference_parameters; these are " + f"conditioning quantities, not inference targets." + ) + if "gnpe_time_shifts" in data_settings: + raise ValueError( + "data.time_alignment=True is incompatible with data.gnpe_time_shifts; " + "both manipulate per-detector arrival times." + ) + + def build_dataset( data_settings: dict, leave_waveforms_on_disk: Optional[bool] = False, @@ -125,9 +166,24 @@ def set_train_transforms(wfd, data_settings, asd_dataset_path, omit_transforms=N ) extra_context_parameters += transforms[-1].context_parameters - # Add the GNPE context to context_parameters the first time the transforms are - # constructed. We do not want to overwrite the ordering of the parameters in - # subsequent runs. + # User-declared conditioning parameters for conditional NPE. These are added + # to the context_parameters tensor alongside any GNPE proxies and share the + # same standardization / repackaging path. They are assumed to be parameters + # already present in either the waveform dataset (intrinsic) or the + # extrinsic prior, so no additional sampling transform is required. + extra_context_parameters += data_settings.get("conditioning_parameters", []) + + # Chained-NPE time alignment: the network sees data with no detector-frame + # time-of-arrival info, and conditions on (ra, dec, geocent_time) to recover + # antenna-pattern dependence. Implemented by skipping the per-detector time + # shift in ProjectOntoDetectors. + time_alignment = data_settings.get("time_alignment", False) + if time_alignment: + validate_time_alignment_settings(data_settings) + + # Add the auto-derived and user-declared context parameters to + # context_parameters the first time the transforms are constructed. We do not + # want to overwrite the ordering of the parameters in subsequent runs. if "context_parameters" not in data_settings: data_settings["context_parameters"] = [] for p in extra_context_parameters: @@ -153,7 +209,11 @@ def set_train_transforms(wfd, data_settings, asd_dataset_path, omit_transforms=N ) data_settings["standardization"] = standardization_dict - transforms.append(ProjectOntoDetectors(ifo_list, domain, ref_time)) + transforms.append( + ProjectOntoDetectors( + ifo_list, domain, ref_time, apply_time_shift=not time_alignment + ) + ) transforms.append(SampleNoiseASD(asd_dataset)) transforms.append(WhitenAndScaleStrain(domain.noise_std)) # We typically add white detector noise. For debugging purposes, this can be turned @@ -277,7 +337,7 @@ def build_svd_for_embedding_network( loader = DataLoader( wfd, batch_size=batch_size, - num_workers= 0, + num_workers=0, worker_init_fn=fix_random_seeds, ) with threadpool_limits(limits=1, user_api="blas"): diff --git a/dingo/gw/transforms/detector_transforms.py b/dingo/gw/transforms/detector_transforms.py index 4daf4de0..da8c45a5 100644 --- a/dingo/gw/transforms/detector_transforms.py +++ b/dingo/gw/transforms/detector_transforms.py @@ -133,12 +133,20 @@ class ProjectOntoDetectors(object): the extrinsic parameters (ra, dec, psi) (3) Time shift the strains in the individual detectors according to the times _time provided in the extrinsic parameters. + + Step (3) can be disabled with ``apply_time_shift=False``. This is used by the + chained-NPE / time-alignment training path, in which the network is meant to + see data with no detector-frame time-of-arrival information. The + ``_time`` parameters are still computed and moved into ``parameters`` + (so they remain available as conditioning quantities or for downstream + transforms); only the strain time-translation is skipped. """ - def __init__(self, ifo_list, domain, ref_time): + def __init__(self, ifo_list, domain, ref_time, apply_time_shift: bool = True): self.ifo_list = ifo_list self.domain = domain self.ref_time = ref_time + self.apply_time_shift = apply_time_shift def __call__(self, input_sample): sample = input_sample.copy() @@ -204,8 +212,11 @@ def __call__(self, input_sample): # (3) time shift the strain. If polarizations are timeshifted by # tc_ref != 0, undo this here by subtracting it from dt. - dt = extrinsic_parameters[f"{ifo.name}_time"] - tc_ref - strains[ifo.name] = self.domain.time_translate_data(strain, dt) + if self.apply_time_shift: + dt = extrinsic_parameters[f"{ifo.name}_time"] - tc_ref + strains[ifo.name] = self.domain.time_translate_data(strain, dt) + else: + strains[ifo.name] = strain # Add extrinsic parameters corresponding to the transformations # applied in the loop above to parameters. These have all been popped off of @@ -280,7 +291,7 @@ class SampleCalibrationParameters(object): calibration envelope, and applies them to generate $N$ observed waveforms $\{h^n_{ obs}(f)\}$. This is intended to be used for marginalizing over the calibration uncertainty when evaluating the likelihood for importance sampling. - + This transform should be followed by ApplyCalibrationToWaveform to apply the sampled calibration curves to the waveform. """ @@ -318,7 +329,8 @@ def __init__( if correction_type is None: correction_type_dict = { - ifo.name: CALIBRATION_CORRECTION_TYPE_LOOKUP[ifo.name] for ifo in self.ifo_list + ifo.name: CALIBRATION_CORRECTION_TYPE_LOOKUP[ifo.name] + for ifo in self.ifo_list } elif correction_type == "data" or correction_type == "template": correction_type_dict = {ifo.name: correction_type for ifo in self.ifo_list} @@ -331,7 +343,7 @@ def __init__( if all([s.endswith(".txt") for s in calibration_envelope.values()]): self.calibration_envelope = calibration_envelope for ifo in self.ifo_list: - # Setting a calibration prior. + # Setting a calibration prior. # Take the calibration envelope and use it to set a spline on # the median and sigma of the amplitude and phase. Then in log # frequency it will setup node points at frequency points, f_i @@ -339,15 +351,15 @@ def __init__( # spaced between f_min and f_max. Then for each node point f_i, # it will create a gaussian prior according to the spline of # the median and sigma found earlier - self.calibration_prior[ - ifo.name - ] = CalibrationPriorDict.from_envelope_file( - self.calibration_envelope[ifo.name], - self.data_domain.f_min, - self.data_domain.f_max, - num_calibration_nodes, - ifo.name, - correction_type=correction_type_dict[ifo.name], + self.calibration_prior[ifo.name] = ( + CalibrationPriorDict.from_envelope_file( + self.calibration_envelope[ifo.name], + self.data_domain.f_min, + self.data_domain.f_max, + num_calibration_nodes, + ifo.name, + correction_type=correction_type_dict[ifo.name], + ) ) else: raise Exception("Calibration envelope must be specified in a .txt file!") @@ -432,8 +444,12 @@ def _ensure_calibration_model(self, ifo, num_calibration_nodes): Ensure the calibration model is set up on the ifo. Creates it if not present or if it has a different number of nodes. """ - if not hasattr(ifo, "calibration_model") or ifo.calibration_model is None or isinstance(ifo.calibration_model, calibration.Recalibrate): - # using https://dcc.ligo.org/LIGO-T2300140 + if ( + not hasattr(ifo, "calibration_model") + or ifo.calibration_model is None + or isinstance(ifo.calibration_model, calibration.Recalibrate) + ): + # using https://dcc.ligo.org/LIGO-T2300140 ifo.calibration_model = calibration.CubicSpline( f"recalib_{ifo.name}_", minimum_frequency=self.data_domain.f_min, @@ -454,9 +470,7 @@ def __call__(self, input_sample): prefix = f"recalib_{ifo.name}_" # Extract calibration parameters for this ifo - calib_params = { - k: v for k, v in extrinsic.items() if k.startswith(prefix) - } + calib_params = {k: v for k, v in extrinsic.items() if k.startswith(prefix)} if not calib_params: continue @@ -486,12 +500,14 @@ def __call__(self, input_sample): # Compute calibration curve for each parameter set for i in range(num_curves): params_i = {k: v[i] for k, v in calib_params.items()} - calibration_draws[ - i, self.data_domain.frequency_mask - ] = ifo.calibration_model.get_calibration_factor( - self.data_domain.sample_frequencies[self.data_domain.frequency_mask], - prefix=prefix, - **params_i, + calibration_draws[i, self.data_domain.frequency_mask] = ( + ifo.calibration_model.get_calibration_factor( + self.data_domain.sample_frequencies[ + self.data_domain.frequency_mask + ], + prefix=prefix, + **params_i, + ) ) # Squeeze out leading dimension if input was scalar diff --git a/tests/core/test_nsf.py b/tests/core/test_nsf.py index b314b8be..fd77310d 100644 --- a/tests/core/test_nsf.py +++ b/tests/core/test_nsf.py @@ -6,6 +6,7 @@ from dingo.core.nn.enets import ( ConcatContextMerger, DenseSVDEmbedding, + MLPContextMerger, create_enet_with_projection_layer_and_dense_resnet, ) from dingo.core.utils import torchutils @@ -284,3 +285,53 @@ def test_registered_embedding_with_merger(data_setup_nsf_small): "Unexpected log prob encountered. Network initialization or " "normalization seems broken." ) + + +def test_mlp_context_merger(data_setup_nsf_small): + """ + Test the mlp context merger (ported ContextMergerMLP): the context parameters + are mixed in through a learned MLP, so the merged output dimension does not + grow with the number of context parameters. + """ + + d = data_setup_nsf_small + kwargs = { + k: v + for k, v in d.embedding_net_kwargs.items() + if k not in ("added_context", "V_rb_list") + } + num_context_parameters = d.z.shape[1] + embedding_net = DenseSVDEmbedding(**kwargs) + + merged = MLPContextMerger( + embedding_net, num_context_parameters, hidden_dims=[16, 16] + ) + assert merged.input_keys == CONTEXT_KEYS + # By default, the merged output dimension equals the embedding output dim. + assert merged.output_dim == embedding_net.output_dim + assert ( + MLPContextMerger.merged_output_dim( + embedding_net.output_dim, + num_context_parameters=num_context_parameters, + hidden_dims=[16, 16], + ) + == embedding_net.output_dim + ) + + out = merged(d.x, d.z) + assert out.shape == (d.batch_size, embedding_net.output_dim) + + # An explicit output_dim overrides the default. + merged_wide = MLPContextMerger( + embedding_net, num_context_parameters, hidden_dims=[16], output_dim=12 + ) + assert merged_wide(d.x, d.z).shape == (d.batch_size, 12) + assert ( + MLPContextMerger.merged_output_dim( + embedding_net.output_dim, + num_context_parameters=num_context_parameters, + hidden_dims=[16], + output_dim=12, + ) + == 12 + ) diff --git a/tests/gw/test_time_alignment_validation.py b/tests/gw/test_time_alignment_validation.py new file mode 100644 index 00000000..c1ecedc6 --- /dev/null +++ b/tests/gw/test_time_alignment_validation.py @@ -0,0 +1,56 @@ +""" +Tests for the validation guard around data.time_alignment=True. The validation +is decoupled from the rest of set_train_transforms so it can be exercised +without building a WaveformDataset. +""" + +import pytest + +from dingo.gw.training.train_builders import validate_time_alignment_settings + + +def _base_settings(): + return { + "inference_parameters": ["mass_1", "mass_2", "luminosity_distance"], + "conditioning_parameters": ["ra", "dec", "geocent_time"], + } + + +def test_validate_time_alignment_accepts_well_formed_settings(): + validate_time_alignment_settings(_base_settings()) + + +@pytest.mark.parametrize("missing", ["ra", "dec", "geocent_time"]) +def test_validate_time_alignment_rejects_missing_conditioning(missing): + s = _base_settings() + s["conditioning_parameters"] = [ + p for p in s["conditioning_parameters"] if p != missing + ] + with pytest.raises(ValueError, match=missing): + validate_time_alignment_settings(s) + + +@pytest.mark.parametrize("inf_param", ["ra", "dec", "geocent_time"]) +def test_validate_time_alignment_rejects_conditioning_in_inference(inf_param): + s = _base_settings() + s["inference_parameters"] = s["inference_parameters"] + [inf_param] + with pytest.raises(ValueError, match=inf_param): + validate_time_alignment_settings(s) + + +def test_validate_time_alignment_rejects_gnpe_combination(): + s = _base_settings() + s["gnpe_time_shifts"] = { + "kernel": "bilby.core.prior.Uniform(0, 1)", + "exact_equiv": True, + } + with pytest.raises(ValueError, match="gnpe_time_shifts"): + validate_time_alignment_settings(s) + + +def test_validate_time_alignment_handles_missing_conditioning_key(): + """If conditioning_parameters is absent entirely (e.g., user forgot to add it), + we should still get a clear error listing all three required names.""" + s = {"inference_parameters": ["mass_1"]} + with pytest.raises(ValueError, match="ra"): + validate_time_alignment_settings(s) diff --git a/tests/gw/transforms/test_detector_projection.py b/tests/gw/transforms/test_detector_projection.py index 267dc014..e71edd72 100644 --- a/tests/gw/transforms/test_detector_projection.py +++ b/tests/gw/transforms/test_detector_projection.py @@ -67,6 +67,52 @@ def test_detector_projection_against_research_code( assert np.max(deviation) / np.max(np.abs(strain)) < 5e-2 +def test_project_onto_detectors_skip_time_shift( + reference_data_research_code, setup_detector_projection +): + """ + With apply_time_shift=False, ProjectOntoDetectors must: + * leave the strain at t=0 (i.e., not apply the per-detector time shift), + * still populate _time in sample['parameters'] for downstream use. + Concretely, applying the inverse time shift to the apply_time_shift=True + output must yield the apply_time_shift=False output. + """ + sample_in, parameters_ref, _ = reference_data_research_code + _, get_detector_times, project_with_shift = setup_detector_projection + ifo_list = project_with_shift.ifo_list + domain = project_with_shift.domain + ref_time = project_with_shift.ref_time + + project_no_shift = ProjectOntoDetectors( + ifo_list, domain, ref_time, apply_time_shift=False + ) + + def _prep(sample): + s = get_detector_times(sample) + # The fixture's _time values are precomputed; copy them in to match + # the existing reference test. + s["extrinsic_parameters"]["H1_time"] = parameters_ref["H1_time"] + s["extrinsic_parameters"]["L1_time"] = parameters_ref["L1_time"] + return s + + out_shift = project_with_shift(_prep(dict(sample_in))) + out_no_shift = project_no_shift(_prep(dict(sample_in))) + + for ifo in ifo_list: + # _time still populated in parameters (independent of the flag). + assert f"{ifo.name}_time" in out_no_shift["parameters"] + + ifo_time = out_no_shift["parameters"][f"{ifo.name}_time"] + # Re-applying the forward shift to the unshifted strain must recover the + # shifted strain (up to FFT round-trip precision). + recovered = domain.time_translate_data( + out_no_shift["waveform"][ifo.name], ifo_time + ) + shifted = out_shift["waveform"][ifo.name] + rel = np.max(np.abs(recovered - shifted)) / np.max(np.abs(shifted)) + assert rel < 1e-5, f"{ifo.name}: round-trip residual {rel:.2e} too large" + + def test_time_delay_from_geocenter(): ifo_list = InterferometerList(["H1", "L1", "V1"]) for ifo in ifo_list: From 03838a1286e3ab89b278b437eb4761a938e6e22c Mon Sep 17 00:00:00 2001 From: Maximilian Dax Date: Sun, 12 Jul 2026 20:22:30 +0200 Subject: [PATCH 7/9] Make weight initialization architecture-owned MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Data-driven weight initialization moves behind two hooks on the embedding contract (NN_Build_System_Design §4.6): init_data_spec() declares the data variation a network wants (noise on/off, network formatting, pinned prior parameters, sample count), and initialize_weights(batches, out_dir) consumes a matching batch iterator. DenseSVDEmbedding implements them, absorbing the SVD math from build_svd_for_embedding_network: it collects per-block strains, builds one SVD basis per block, saves validation diagnostics, drops the zero rows outside its input range, and seeds the projection layer. The trainer no longer knows about the SVD: prepare_training_new walks pm.network.modules(), answers any module's spec via the new initialization_dataloader (a transform-stack variation using the existing omit_transforms mechanism), and lets the module initialize itself — plugin architectures get data-driven initialization without touching dingo. Seeding triggers on svd.num_training_samples, so loading a saved model never re-runs it. Deleted: build_svd_for_embedding_network, the initial_weights threading through NeuralDistribution/build_model_from_kwargs, and the "embedding network is assumed to have an SVD projection layer" special case. SVDBasis moves from dingo/gw/SVD.py to dingo/core/SVD.py — it was always domain-agnostic (numpy/scipy only) and core may not import from gw; the old module remains as a re-export shim. Co-Authored-By: Claude Fable 5 --- dingo/core/SVD.py | 277 ++++++++++++++++++++++ dingo/core/nn/enets.py | 113 ++++++++- dingo/core/posterior_models/base_model.py | 14 +- dingo/gw/SVD.py | 277 +--------------------- dingo/gw/dataset/generate_dataset.py | 2 +- dingo/gw/dataset/utils.py | 2 +- dingo/gw/dataset/waveform_dataset.py | 2 +- dingo/gw/ls_cli.py | 2 +- dingo/gw/training/train_builders.py | 158 +++--------- dingo/gw/training/train_pipeline.py | 59 +++-- tests/core/test_build_model.py | 94 ++++++-- 11 files changed, 545 insertions(+), 455 deletions(-) create mode 100644 dingo/core/SVD.py diff --git a/dingo/core/SVD.py b/dingo/core/SVD.py new file mode 100644 index 00000000..a15dabb3 --- /dev/null +++ b/dingo/core/SVD.py @@ -0,0 +1,277 @@ +import numpy as np +import pandas as pd +import scipy +from sklearn.utils.extmath import randomized_svd +from dingo.core.dataset import DingoDataset + + +class SVDBasis(DingoDataset): + + dataset_type = "svd_basis" + + def __init__( + self, + file_name=None, + dictionary=None, + ): + self.V = None + self.Vh = None + self.s = None + self.n = None + self.mismatches = None + super().__init__( + file_name=file_name, + dictionary=dictionary, + data_keys=["V", "s", "mismatches"], + ) + + def generate_basis(self, training_data: np.ndarray, n: int, method: str = "scipy"): + """Generate the SVD basis from training data and store it. + + The SVD decomposition takes + + training_data = U @ diag(s) @ Vh + + where U and Vh are unitary. + + Parameters + ---------- + training_data: np.ndarray + Array of waveform data on the physical domain + n: int + Number of basis elements to keep. + n=0 keeps all basis elements. + method: str + Select SVD method, 'random' or 'scipy' + """ + if method == "random": + if n == 0: + n = min(training_data.shape) + + # Using LU as a normalizer in the power iteration "normalizer(A @ + # Q)" in randomized_range_finder() called by randomized_svd() can cause + # segfaults. The QR factorization, while slightly slower than LU is more + # numerically stable. These segfaults also disappear when switching off + # multithreading, but we want to keep this on. + # + # The randomized SVD has complexity O(m n k + k^2 (m + n)), + # for a m x n matrix and k is the target rank, here called n + # For small k this is much faster than the standard SVD. + try: + U, s, Vh = randomized_svd( + training_data, n, random_state=0, power_iteration_normalizer="QR" + ) + except ValueError as e: + raise ValueError( + "randomized_svd failed — possibly due to complex-valued input.\n" + "randomized_svd does not support complex arrays in scikit-learn >=1.2.\n" + "To proceed, downgrade scikit-learn to version 1.1.3:\n\n" + " pip install scikit-learn==1.1.3\n\n" + f"Original error: {e}" + ) + + self.Vh = Vh.astype(np.complex128) # TODO: fix types + self.V = self.Vh.T.conj() + self.n = n + self.s = s + elif method == "scipy": + if (n == 0) or (n >= training_data.shape[1]): + # Code below uses scipy's svd tool. Likely slower. + # The deterministic SVD has Complexity O(mn^2). + U, s, Vh = scipy.linalg.svd(training_data, full_matrices=False) + else: + # Use partial SVD if only a subset of basis elements are requested + U, s, Vh = scipy.sparse.linalg.svds(training_data, k=n) + + # Sort singular values in non-increasing order + idx = np.argsort(s)[::-1] + U, s, Vh = U[:, idx], s[idx], Vh[idx, :] + V = Vh.T.conj() + + if (n == 0) or (n > len(V)): + self.V = V + self.Vh = Vh + else: + self.V = V[:, :n] + self.Vh = Vh[:n, :] + + self.n = len(self.Vh) + self.s = s + else: + raise ValueError(f"Unsupported SVD method: {method}.") + + def compute_test_mismatches( + self, + data: np.ndarray, + parameters: pd.DataFrame = None, + increment: int = 50, + verbose: bool = False, + ): + """ + Test SVD basis by computing mismatches of compressed / decompressed data + against original data. Results are saved as a DataFrame. + + Parameters + ---------- + data : np.ndarray + Array of data sets to validate against. + parameters : pd.DataFrame + Optional labels for the data sets. This is useful for checking performance on + particular regions of the parameter space. + increment : int + Specifies SVD truncations for computing mismatches. E.g., increment = 50 + means that the SVD will be truncated at size [50, 100, 150, ..., len(data)]. + verbose : bool + Whether to print summary statistics. + """ + if len(data) != len(parameters): + raise ValueError( + f"Incompatible data: len(data) == {len(data)} and len(" + f"parameters) == {len(parameters)} do not match." + ) + if parameters is not None: + self.mismatches = parameters.copy() + else: + self.mismatches = pd.DataFrame() + + for n in np.append(np.arange(increment, self.n, increment), self.n): + mismatches = np.empty(len(data)) + for i, d in enumerate(data): + compressed = d @ self.V[:, :n] + reconstructed = compressed @ self.Vh[:n] + norm1 = np.sqrt(np.sum(np.abs(d) ** 2)) + norm2 = np.sqrt(np.sum(np.abs(reconstructed) ** 2)) + inner = np.sum(d.conj() * reconstructed).real + mismatches[i] = 1 - inner / (norm1 * norm2) + self.mismatches[f"mismatch n={n}"] = mismatches + + if verbose: + self.print_validation_summary() + + def print_validation_summary(self): + """ + Print a summary of the validation mismatches. + """ + if self.mismatches is not None: + for col in self.mismatches: + if "mismatch" in col: + n = int(col.split(sep="=")[-1]) + mismatches = self.mismatches[col] + print(f"n = {n}") + print(" Mean mismatch = {}".format(np.mean(mismatches))) + print(" Standard deviation = {}".format(np.std(mismatches))) + print(" Max mismatch = {}".format(np.max(mismatches))) + print(" Median mismatch = {}".format(np.median(mismatches))) + print(" Percentiles:") + print(" 99 -> {}".format(np.percentile(mismatches, 99))) + print(" 99.9 -> {}".format(np.percentile(mismatches, 99.9))) + print(" 99.99 -> {}".format(np.percentile(mismatches, 99.99))) + + def decompress(self, coefficients: np.ndarray): + """ + Convert from basis coefficients back to raw data representation. + + Parameters + ---------- + coefficients : np.ndarray + Array of basis coefficients + + Returns + ------- + array of decompressed data + """ + return coefficients @ self.Vh + + def compress(self, data: np.ndarray): + """ + Convert from data (e.g., frequency series) to compressed representation in + terms of basis coefficients. + + Parameters + ---------- + data : np.ndarray + + Returns + ------- + array of basis coefficients + """ + return data @ self.V + + def from_file(self, filename): + """ + Load the SVD basis from a HDF5 file. + + Parameters + ---------- + filename : str + """ + super().from_file(filename) + if self.V is None: + raise KeyError("File does not contain SVD V matrix. No SVD basis to load.") + self.Vh = self.V.T.conj() + self.n = self.V.shape[1] + + def from_dictionary(self, dictionary: dict): + """ + Load the SVD basis from a dictionary. + + Parameters + ---------- + dictionary : dict + The dictionary should contain at least a 'V' key, and optionally an 's' key. + """ + super().from_dictionary(dictionary) + if self.V is None: + raise KeyError("dict does not contain SVD V matrix. No SVD basis to load.") + self.Vh = self.V.T.conj() + self.n = self.V.shape[1] + + # def truncate(self, n: int): + # """ + # Truncate size of SVD. + # + # Parameters + # ---------- + # n : int + # New SVD size. Should be less than current size. + # """ + # if n > self.n or n < 0: + # print(f"Cannot truncate SVD from size n={self.n} to n={n}.") + # else: + # self.V = self.V[:, :n] + # self.Vh = self.Vh[:n, :] + # self.s = self.s[:n] + # self.n = n + + +class ApplySVD(object): + """Transform operator for applying an SVD compression / decompression.""" + + def __init__(self, svd_basis: SVDBasis, inverse: bool = False): + """ + Parameters + ---------- + svd_basis : SVDBasis + inverse : bool + Whether to apply for the forward (compression) or inverse (decompression) + transform. Default: False. + """ + self.svd_basis = svd_basis + self.inverse = inverse + + def __call__(self, waveform: dict): + """ + Parameters + ---------- + waveform : dict + Values should be arrays containing waveforms to be transformed. + + Returns + ------- + dict of the same form as the input, but with transformed waveforms. + """ + if not self.inverse: + func = self.svd_basis.compress + else: + func = self.svd_basis.decompress + return {k: func(v) for k, v in waveform.items()} diff --git a/dingo/core/nn/enets.py b/dingo/core/nn/enets.py index 1b874fa1..b2923428 100755 --- a/dingo/core/nn/enets.py +++ b/dingo/core/nn/enets.py @@ -9,18 +9,30 @@ network's own input dimensions from a sample batch; the completed settings (which must include ``output_dim``) are saved in the checkpoint, so loading never needs a data sample. +* ``init_data_spec()`` (optional): returns a dict describing the data variation + the network wants for data-driven weight initialization (e.g. noise-free, + un-formatted waveforms), or None if no initialization is needed. +* ``initialize_weights(batches, out_dir=None)`` (optional): consumes an iterator + of batches matching the spec and initializes the network weights in-place. + The trainer answers the spec with a matching dataloader and calls this hook; + it does not know about specific architectures. Context mergers registered with CONTEXT_MERGERS wrap an embedding network to mix in the (standardized) context parameters; they follow the same contract, plus a ``merged_output_dim`` method used during settings completion. """ +import os from typing import Tuple, Callable, Union, List + import torch import numpy as np +import pandas as pd import torch.nn as nn from torch.nn import functional as F from glasflow.nflows.nn.nets.resnet import ResidualBlock + +from dingo.core.SVD import SVDBasis from dingo.core.registry import CONTEXT_MERGERS, EMBEDDING_NETS from dingo.core.utils import torchutils @@ -303,8 +315,8 @@ def __init__( whether to use batch normalization V_rb_list : tuple of np.arrays, or None V matrices of the SVD projection used to initialize the projection - weights. Passed as initial weights at first build; None when loading a - saved model. + weights directly. Usually None: the projection is seeded via the + initialize_weights hook instead. """ projection = LinearProjectionRB(input_dims, svd["size"], V_rb_list) resnet = DenseResidualNet( @@ -317,6 +329,7 @@ def __init__( ) super().__init__(projection, resnet) self.output_dim = output_dim + self.svd_settings = svd @classmethod def complete_settings(cls, settings: dict, sample_batch: dict) -> dict: @@ -328,6 +341,102 @@ def complete_settings(cls, settings: dict, sample_batch: dict) -> dict: ) return {**settings, "input_dims": list(sample_batch["waveform"].shape)} + def init_data_spec(self): + """ + Data variation for seeding the SVD projection: clean (noise-free), + un-formatted waveforms at a fixed reference luminosity distance. Returns + None if the svd settings do not request seeding (no + num_training_samples), e.g. when loading a saved model. + """ + if "num_training_samples" not in self.svd_settings: + return None + num_samples = self.svd_settings["num_training_samples"] + self.svd_settings.get( + "num_validation_samples", 0 + ) + return { + "noise": False, + "network_format": False, + "fix_parameters": {"luminosity_distance": 100.0}, + "num_samples": num_samples, + } + + def initialize_weights(self, batches, out_dir=None): + """ + Seed the projection layer with an SVD basis built from clean waveforms. + + Parameters + ---------- + batches : iterable + Batches matching init_data_spec: dicts with "waveform" a + {block: (batch_size, len) complex array} dict (in the GW use case, a + block is a detector) and "parameters", used for validation + diagnostics. Iteration stops once num_samples have been consumed. + out_dir : str = None + If provided, SVD validation diagnostics are computed and saved here. + """ + svd = self.svd_settings + num_training = svd["num_training_samples"] + num_validation = svd.get("num_validation_samples", 0) + num_samples = num_training + num_validation + + waveforms = None + parameters = pd.DataFrame() + collected = 0 + for batch in batches: + strain_data = batch["waveform"] + if waveforms is None: + waveforms = { + block: np.empty( + (num_samples, strains.shape[-1]), dtype=np.complex128 + ) + for block, strains in strain_data.items() + } + batch_size = len(next(iter(strain_data.values()))) + n = min(batch_size, num_samples - collected) + parameters = pd.concat( + [parameters, pd.DataFrame(batch["parameters"]).iloc[:n]], + ignore_index=True, + ) + for block, strains in strain_data.items(): + waveforms[block][collected : collected + n] = strains[:n] + collected += n + if collected == num_samples: + break + if collected < num_samples: + raise IndexError( + f"Requested {num_samples} samples for SVD initialization, but the " + f"dataloader only provided {collected}." + ) + + projection = self[0] + V_rb_list = [] + for block, data in waveforms.items(): + print(f"Generating SVD basis for block {block}.") + basis = SVDBasis() + basis.generate_basis(data[:num_training], svd["size"]) + if out_dir is not None and num_validation > 0: + basis.compute_test_mismatches( + data[num_training:], + parameters=parameters.iloc[num_training:].reset_index(drop=True), + verbose=True, + ) + basis.to_file(os.path.join(out_dir, f"svd_{block}.hdf5")) + # The provided waveforms may be longer than the network input (leading + # entries outside the network's frequency range). These must be zero, + # and the corresponding rows of V are dropped. + V = basis.V + excess = len(V) - projection.num_bins + if not np.allclose(V[:excess], 0): + raise ValueError( + f"Block {block}: SVD basis has non-zero entries outside the " + f"network input range (waveform length {len(V)}, network " + f"num_bins {projection.num_bins})." + ) + V_rb_list.append(V[excess:]) + + projection.test_dimensions(V_rb_list) + projection.init_layers(V_rb_list) + class ModuleMerger(nn.Module): """ diff --git a/dingo/core/posterior_models/base_model.py b/dingo/core/posterior_models/base_model.py index ffab37a1..a2b66b2b 100755 --- a/dingo/core/posterior_models/base_model.py +++ b/dingo/core/posterior_models/base_model.py @@ -44,7 +44,6 @@ def __init__( self, model_filename: str = None, metadata: dict = None, - initial_weights: dict = None, device: str = "cuda", load_training_info: bool = True, ): @@ -57,8 +56,6 @@ def __init__( If given, loads data from the given file. metadata: dict If given, initializes the model from these settings - initial_weights: dict - Initial weights for the model device: str load_training_info: bool """ @@ -69,7 +66,6 @@ def __init__( self.optimizer_kwargs = None self.network_kwargs = None self.scheduler_kwargs = None - self.initial_weights = initial_weights self.metadata = metadata if self.metadata is not None: @@ -105,9 +101,9 @@ def build_embedding_net(self): """ Build the embedding network declared in the model settings (resolved via the EMBEDDING_NETS registry), optionally wrapped with a context merger - (CONTEXT_MERGERS) that mixes in the context parameters. Initial weights - (e.g. SVD projection matrices) are passed as extra constructor kwargs; they - are not part of the saved settings. + (CONTEXT_MERGERS) that mixes in the context parameters. Data-driven weight + initialization (e.g. SVD seeding) happens separately, via the network's + init_data_spec / initialize_weights hooks. Returns None if the model declares no embedding network (unconditional models). @@ -115,9 +111,7 @@ def build_embedding_net(self): embedding_settings = self.model_kwargs.get("embedding_net") if embedding_settings is None: return None - kwargs = dict(embedding_settings.get("kwargs", {})) - if self.initial_weights: - kwargs.update(self.initial_weights) + kwargs = embedding_settings.get("kwargs", {}) embedding_net = EMBEDDING_NETS.get(embedding_settings["type"])(**kwargs) merger_settings = self.model_kwargs.get("context_merger") if merger_settings is not None: diff --git a/dingo/gw/SVD.py b/dingo/gw/SVD.py index e77b7761..df637d77 100644 --- a/dingo/gw/SVD.py +++ b/dingo/gw/SVD.py @@ -1,274 +1,5 @@ -import numpy as np -import pandas as pd -import scipy -from sklearn.utils.extmath import randomized_svd -from dingo.core.dataset import DingoDataset +"""Moved to dingo.core.SVD: the SVD basis is domain-agnostic (used by the core +NN build system for embedding initialization). Kept as a re-export so existing +imports keep working.""" -class SVDBasis(DingoDataset): - - dataset_type = "svd_basis" - - def __init__( - self, - file_name=None, - dictionary=None, - ): - self.V = None - self.Vh = None - self.s = None - self.n = None - self.mismatches = None - super().__init__( - file_name=file_name, - dictionary=dictionary, - data_keys=["V", "s", "mismatches"], - ) - - def generate_basis(self, training_data: np.ndarray, n: int, method: str = "scipy"): - """Generate the SVD basis from training data and store it. - - The SVD decomposition takes - - training_data = U @ diag(s) @ Vh - - where U and Vh are unitary. - - Parameters - ---------- - training_data: np.ndarray - Array of waveform data on the physical domain - n: int - Number of basis elements to keep. - n=0 keeps all basis elements. - method: str - Select SVD method, 'random' or 'scipy' - """ - if method == "random": - if n == 0: - n = min(training_data.shape) - - # Using LU as a normalizer in the power iteration "normalizer(A @ - # Q)" in randomized_range_finder() called by randomized_svd() can cause - # segfaults. The QR factorization, while slightly slower than LU is more - # numerically stable. These segfaults also disappear when switching off - # multithreading, but we want to keep this on. - # - # The randomized SVD has complexity O(m n k + k^2 (m + n)), - # for a m x n matrix and k is the target rank, here called n - # For small k this is much faster than the standard SVD. - try: - U, s, Vh = randomized_svd(training_data, n, random_state=0, - power_iteration_normalizer='QR') - except ValueError as e: - raise ValueError( - "randomized_svd failed — possibly due to complex-valued input.\n" - "randomized_svd does not support complex arrays in scikit-learn >=1.2.\n" - "To proceed, downgrade scikit-learn to version 1.1.3:\n\n" - " pip install scikit-learn==1.1.3\n\n" - f"Original error: {e}" - ) - - self.Vh = Vh.astype(np.complex128) # TODO: fix types - self.V = self.Vh.T.conj() - self.n = n - self.s = s - elif method == "scipy": - if (n == 0) or (n >= training_data.shape[1]): - # Code below uses scipy's svd tool. Likely slower. - # The deterministic SVD has Complexity O(mn^2). - U, s, Vh = scipy.linalg.svd(training_data, full_matrices=False) - else: - # Use partial SVD if only a subset of basis elements are requested - U, s, Vh = scipy.sparse.linalg.svds(training_data, k=n) - - # Sort singular values in non-increasing order - idx = np.argsort(s)[::-1] - U, s, Vh = U[:, idx], s[idx], Vh[idx, :] - V = Vh.T.conj() - - if (n == 0) or (n > len(V)): - self.V = V - self.Vh = Vh - else: - self.V = V[:, :n] - self.Vh = Vh[:n, :] - - self.n = len(self.Vh) - self.s = s - else: - raise ValueError(f"Unsupported SVD method: {method}.") - - def compute_test_mismatches( - self, - data: np.ndarray, - parameters: pd.DataFrame = None, - increment: int = 50, - verbose: bool = False, - ): - """ - Test SVD basis by computing mismatches of compressed / decompressed data - against original data. Results are saved as a DataFrame. - - Parameters - ---------- - data : np.ndarray - Array of data sets to validate against. - parameters : pd.DataFrame - Optional labels for the data sets. This is useful for checking performance on - particular regions of the parameter space. - increment : int - Specifies SVD truncations for computing mismatches. E.g., increment = 50 - means that the SVD will be truncated at size [50, 100, 150, ..., len(data)]. - verbose : bool - Whether to print summary statistics. - """ - if len(data) != len(parameters): - raise ValueError( - f"Incompatible data: len(data) == {len(data)} and len(" - f"parameters) == {len(parameters)} do not match." - ) - if parameters is not None: - self.mismatches = parameters.copy() - else: - self.mismatches = pd.DataFrame() - - for n in np.append(np.arange(increment, self.n, increment), self.n): - mismatches = np.empty(len(data)) - for i, d in enumerate(data): - compressed = d @ self.V[:, :n] - reconstructed = compressed @ self.Vh[:n] - norm1 = np.sqrt(np.sum(np.abs(d) ** 2)) - norm2 = np.sqrt(np.sum(np.abs(reconstructed) ** 2)) - inner = np.sum(d.conj() * reconstructed).real - mismatches[i] = 1 - inner / (norm1 * norm2) - self.mismatches[f"mismatch n={n}"] = mismatches - - if verbose: - self.print_validation_summary() - - def print_validation_summary(self): - """ - Print a summary of the validation mismatches. - """ - if self.mismatches is not None: - for col in self.mismatches: - if "mismatch" in col: - n = int(col.split(sep="=")[-1]) - mismatches = self.mismatches[col] - print(f"n = {n}") - print(" Mean mismatch = {}".format(np.mean(mismatches))) - print(" Standard deviation = {}".format(np.std(mismatches))) - print(" Max mismatch = {}".format(np.max(mismatches))) - print(" Median mismatch = {}".format(np.median(mismatches))) - print(" Percentiles:") - print(" 99 -> {}".format(np.percentile(mismatches, 99))) - print(" 99.9 -> {}".format(np.percentile(mismatches, 99.9))) - print(" 99.99 -> {}".format(np.percentile(mismatches, 99.99))) - - def decompress(self, coefficients: np.ndarray): - """ - Convert from basis coefficients back to raw data representation. - - Parameters - ---------- - coefficients : np.ndarray - Array of basis coefficients - - Returns - ------- - array of decompressed data - """ - return coefficients @ self.Vh - - def compress(self, data: np.ndarray): - """ - Convert from data (e.g., frequency series) to compressed representation in - terms of basis coefficients. - - Parameters - ---------- - data : np.ndarray - - Returns - ------- - array of basis coefficients - """ - return data @ self.V - - def from_file(self, filename): - """ - Load the SVD basis from a HDF5 file. - - Parameters - ---------- - filename : str - """ - super().from_file(filename) - if self.V is None: - raise KeyError("File does not contain SVD V matrix. No SVD basis to load.") - self.Vh = self.V.T.conj() - self.n = self.V.shape[1] - - def from_dictionary(self, dictionary: dict): - """ - Load the SVD basis from a dictionary. - - Parameters - ---------- - dictionary : dict - The dictionary should contain at least a 'V' key, and optionally an 's' key. - """ - super().from_dictionary(dictionary) - if self.V is None: - raise KeyError("dict does not contain SVD V matrix. No SVD basis to load.") - self.Vh = self.V.T.conj() - self.n = self.V.shape[1] - - # def truncate(self, n: int): - # """ - # Truncate size of SVD. - # - # Parameters - # ---------- - # n : int - # New SVD size. Should be less than current size. - # """ - # if n > self.n or n < 0: - # print(f"Cannot truncate SVD from size n={self.n} to n={n}.") - # else: - # self.V = self.V[:, :n] - # self.Vh = self.Vh[:n, :] - # self.s = self.s[:n] - # self.n = n - -class ApplySVD(object): - """Transform operator for applying an SVD compression / decompression.""" - - def __init__(self, svd_basis: SVDBasis, inverse: bool = False): - """ - Parameters - ---------- - svd_basis : SVDBasis - inverse : bool - Whether to apply for the forward (compression) or inverse (decompression) - transform. Default: False. - """ - self.svd_basis = svd_basis - self.inverse = inverse - - def __call__(self, waveform: dict): - """ - Parameters - ---------- - waveform : dict - Values should be arrays containing waveforms to be transformed. - - Returns - ------- - dict of the same form as the input, but with transformed waveforms. - """ - if not self.inverse: - func = self.svd_basis.compress - else: - func = self.svd_basis.decompress - return {k: func(v) for k, v in waveform.items()} +from dingo.core.SVD import SVDBasis, ApplySVD # noqa: F401 diff --git a/dingo/gw/dataset/generate_dataset.py b/dingo/gw/dataset/generate_dataset.py index a52ca986..6e369337 100644 --- a/dingo/gw/dataset/generate_dataset.py +++ b/dingo/gw/dataset/generate_dataset.py @@ -16,7 +16,7 @@ from dingo.gw.dataset.waveform_dataset import WaveformDataset from dingo.gw.domains import build_domain from dingo.gw.prior import build_prior_with_defaults -from dingo.gw.SVD import ApplySVD, SVDBasis +from dingo.core.SVD import ApplySVD, SVDBasis from dingo.gw.transforms import WhitenFixedASD from dingo.gw.waveform_generator import ( NewInterfaceWaveformGenerator, diff --git a/dingo/gw/dataset/utils.py b/dingo/gw/dataset/utils.py index e8911f87..bad3a386 100644 --- a/dingo/gw/dataset/utils.py +++ b/dingo/gw/dataset/utils.py @@ -6,7 +6,7 @@ import yaml from typing import List -from dingo.gw.SVD import SVDBasis +from dingo.core.SVD import SVDBasis from dingo.gw.dataset.generate_dataset import train_svd_basis from dingo.gw.dataset.waveform_dataset import WaveformDataset diff --git a/dingo/gw/dataset/waveform_dataset.py b/dingo/gw/dataset/waveform_dataset.py index af750f5b..86997c40 100644 --- a/dingo/gw/dataset/waveform_dataset.py +++ b/dingo/gw/dataset/waveform_dataset.py @@ -5,7 +5,7 @@ from torchvision.transforms import Compose from dingo.core.dataset import DingoDataset, recursive_hdf5_load -from dingo.gw.SVD import SVDBasis, ApplySVD +from dingo.core.SVD import SVDBasis, ApplySVD from dingo.gw.domains import build_domain from dingo.gw.transforms import WhitenFixedASD diff --git a/dingo/gw/ls_cli.py b/dingo/gw/ls_cli.py index 5c401bee..7ce09d5e 100644 --- a/dingo/gw/ls_cli.py +++ b/dingo/gw/ls_cli.py @@ -12,7 +12,7 @@ from dingo.core.utils.backward_compatibility import torch_load_with_fallback from dingo.gw.dataset import WaveformDataset from dingo.gw.noise.asd_dataset import ASDDataset -from dingo.gw.SVD import SVDBasis +from dingo.core.SVD import SVDBasis def ls(): diff --git a/dingo/gw/training/train_builders.py b/dingo/gw/training/train_builders.py index 540ceb92..e503b599 100755 --- a/dingo/gw/training/train_builders.py +++ b/dingo/gw/training/train_builders.py @@ -1,13 +1,9 @@ from typing import List, Optional import copy -import torch.multiprocessing import torchvision -from threadpoolctl import threadpool_limits from bilby.gw.detector import InterferometerList -from dingo.gw.SVD import SVDBasis - from dingo.gw.dataset.waveform_dataset import WaveformDataset from dingo.gw.domains import build_domain from dingo.gw.transforms import ( @@ -251,151 +247,71 @@ def set_train_transforms(wfd, data_settings, asd_dataset_path, omit_transforms=N wfd.transform = torchvision.transforms.Compose(transforms) -def build_svd_for_embedding_network( +def initialization_dataloader( wfd: WaveformDataset, data_settings: dict, asd_dataset_path: str, - size: int, - num_training_samples: int, - num_validation_samples: int, - num_workers: int = 0, + spec: dict, batch_size: int = 1000, - out_dir: Optional[str] = None, -) -> List: +): """ - Construct SVD matrices V based on clean waveforms in each interferometer. These - will be used to seed the weights of the initial projection part of the embedding - network. + Build a dataloader answering an embedding network's init_data_spec (see the + contract in dingo.core.nn.enets): a transform-stack variation of the training + data, used for data-driven weight initialization (e.g. seeding the SVD + projection layer from clean waveforms). - It first generates a number of training waveforms, and then produces the SVD. + This replaces the waveform dataset's transforms; call set_train_transforms + again afterwards to restore the training configuration. Parameters ---------- wfd : WaveformDataset data_settings : dict + The train data settings; not modified (the spec is applied to a copy). asd_dataset_path : str - Training waveforms will be whitened with respect to these ASDs. - size : int - Number of basis elements to include in the SVD projection. - num_training_samples : int - num_validation_samples : int - num_workers : int + Waveforms are whitened with respect to these ASDs. + spec : dict + The data variation requested by the network: + * "noise": bool -- if False, no noise is added to the waveforms. + * "network_format": bool -- if False, samples are not repackaged / + standardized for network input; they remain dicts with per-detector + complex strains under "waveform". + * "fix_parameters": dict -- prior parameters pinned to a fixed value, + e.g. {"luminosity_distance": 100.0}. + * "num_samples": int -- number of samples the initialization consumes. batch_size : int - out_dir : str - SVD performance diagnostics are saved here. Returns ------- - list of numpy arrays - The V matrices for each interferometer. They are ordered as in data_settings[ - 'detectors']. + torch.utils.data.DataLoader """ - # Building the transforms can alter the data_settings dictionary. We do not want - # the construction of the SVD to impact this, so begin with a fresh copy of this - # dictionary. data_settings = copy.deepcopy(data_settings) - - # This is needed to prevent an occasional error when loading a large dataset into - # memory using a dataloader. This removes a limitation on the number of "open files". - old_sharing_strategy = torch.multiprocessing.get_sharing_strategy() - torch.multiprocessing.set_sharing_strategy("file_system") - - # Fix the luminosity distance to a standard value, just in order to generate the SVD. - data_settings["extrinsic_prior"]["luminosity_distance"] = "100.0" - - # Build the dataset, but with certain transforms omitted. In particular, we want to - # build the SVD based on zero-noise waveforms. They should still be whitened though. - set_train_transforms( - wfd, - data_settings, - asd_dataset_path, - omit_transforms=[ - AddWhiteNoiseComplex, + for name, value in spec.get("fix_parameters", {}).items(): + data_settings["extrinsic_prior"][name] = str(value) + + omit_transforms = [] + if not spec.get("noise", True): + omit_transforms.append(AddWhiteNoiseComplex) + if not spec.get("network_format", True): + omit_transforms += [ RepackageStrainsAndASDS, SelectStandardizeRepackageParameters, SelectKeys, CropMaskStrainRandom, - ], + ] + set_train_transforms( + wfd, data_settings, asd_dataset_path, omit_transforms=omit_transforms or None ) - print("Generating waveforms for embedding network SVD initialization.") - time_start = time.time() - ifos = list(wfd[0]["waveform"].keys()) - waveform_len = len(wfd[0]["waveform"][ifos[0]]) - num_waveforms = num_training_samples + num_validation_samples - if num_waveforms > len(wfd): + num_samples = spec["num_samples"] + if num_samples > len(wfd): raise IndexError( - f"Requested {num_waveforms} samples for generating SVD for embedding " - f"network, but waveform dataset only contains {len(wfd)} samples." + f"Network initialization requests {num_samples} samples, but the " + f"waveform dataset only contains {len(wfd)}." ) - waveforms = { - ifo: np.empty((num_waveforms, waveform_len), dtype=np.complex128) - for ifo in ifos - } - parameters = pd.DataFrame() - - loader = DataLoader( + return DataLoader( wfd, batch_size=batch_size, num_workers=0, worker_init_fn=fix_random_seeds, ) - with threadpool_limits(limits=1, user_api="blas"): - for idx, data in enumerate(loader): - # This is for handling the last batch, which may otherwise push the total - # number of samples above the number requested. - lower = idx * batch_size - n = min(batch_size, num_waveforms - lower) - - parameters = pd.concat( - [parameters, pd.DataFrame(data["parameters"]).iloc[:n]], - ignore_index=True, - ) - strain_data = data["waveform"] - for ifo, strains in strain_data.items(): - waveforms[ifo][lower : lower + n] = strains[:n] - if lower + n == num_waveforms: - break - print(f"...done. This took {time.time() - time_start:.0f} s.") - - # Reset the standard sharing strategy. - torch.multiprocessing.set_sharing_strategy(old_sharing_strategy) - - print("Generating SVD basis for ifo:") - time_start = time.time() - basis_dict = {} - for ifo in ifos: - basis = SVDBasis() - basis.generate_basis(waveforms[ifo][:num_training_samples], size) - basis_dict[ifo] = basis - print(f"...{ifo} done.") - print(f"...this took {time.time() - time_start:.0f} s.") - - if out_dir is not None: - print(f"Testing SVD basis matrices.") - for ifo, basis in basis_dict.items(): - print(f"...{ifo}:") - basis.compute_test_mismatches( - waveforms[ifo][num_training_samples:], - parameters=parameters.iloc[num_training_samples:].reset_index( - drop=True - ), - verbose=True, - ) - basis.to_file(os.path.join(out_dir, f"svd_{ifo}.hdf5")) - print("Done") - - # Return V matrices in standard order. Drop the elements below domain.min_idx, - # since the neural network expects data truncated below these. The dropped elements - # should be 0. - print(f"Truncating SVD matrices below index {wfd.domain.min_idx}.") - print("...V matrix shapes:") - V_rb_list = [] - for ifo in data_settings["detectors"]: - V = basis_dict[ifo].V - assert np.allclose(V[: wfd.domain.min_idx], 0) - V = V[wfd.domain.min_idx :] - print(" " + str(V.shape)) - V_rb_list.append(V) - print("\n") - return V_rb_list diff --git a/dingo/gw/training/train_pipeline.py b/dingo/gw/training/train_pipeline.py index 9b51e6d6..3710eff6 100644 --- a/dingo/gw/training/train_pipeline.py +++ b/dingo/gw/training/train_pipeline.py @@ -2,6 +2,7 @@ import os import numpy as np +import torch.multiprocessing import yaml import argparse import shutil @@ -19,7 +20,7 @@ from dingo.gw.training.train_builders import ( build_dataset, set_train_transforms, - build_svd_for_embedding_network, + initialization_dataloader, ) from dingo.core.utils.trainutils import RuntimeLimits from dingo.core.utils import ( @@ -115,26 +116,9 @@ def prepare_training_new( data_settings=data_settings, leave_waveforms_on_disk=local_settings.get("leave_waveforms_on_disk", True), ) # No transforms yet - initial_weights = {} update_model_config(train_settings["model"]) # Map old schemas forward. - # Build the SVD for seeding the embedding network, if it declares one. - # TODO: replace with the architecture-owned initialization hook (build-system - # step 4), so that the trainer no longer knows about the SVD. - embedding_settings = train_settings["model"].get("embedding_net") - if embedding_settings and "svd" in embedding_settings.get("kwargs", {}): - print("\nBuilding SVD for initialization of embedding network.") - initial_weights["V_rb_list"] = build_svd_for_embedding_network( - wfd, - train_settings["data"], - train_settings["training"]["stage_0"]["asd_dataset_path"], - num_workers=local_settings["num_workers"], - batch_size=train_settings["training"]["stage_0"]["batch_size"], - out_dir=train_dir, - **embedding_settings["kwargs"]["svd"], - ) - # Now set the transforms for training. We need to do this here so that we can (a) # get the data dimensions to configure the network, and (b) save the # parameter standardization dict in the PosteriorModel. In principle, (a) could @@ -142,11 +126,8 @@ def prepare_training_new( # be done outside the transform setup. But for now, this is convenient. The # transforms will be reset later by initialize_stage(). - set_train_transforms( - wfd, - train_settings["data"], - train_settings["training"]["stage_0"]["asd_dataset_path"], - ) + asd_dataset_path = train_settings["training"]["stage_0"]["asd_dataset_path"] + set_train_transforms(wfd, train_settings["data"], asd_dataset_path) train_settings["model"] = complete_model_settings(train_settings["model"], wfd[0]) full_settings = { @@ -160,10 +141,40 @@ def prepare_training_new( pm = build_model_from_kwargs( settings=full_settings, - initial_weights=initial_weights, device=local_settings["device"], ) + # Data-driven weight initialization (e.g. SVD seeding of the projection + # layer): any network module that requests a data spec gets a dataloader + # answering it and initializes itself. The trainer does not know about + # specific architectures. + initialized = False + for module in pm.network.modules(): + spec = getattr(module, "init_data_spec", lambda: None)() + if spec is None: + continue + print(f"\nInitializing weights of {type(module).__name__} from data.") + dataloader = initialization_dataloader( + wfd, + train_settings["data"], + asd_dataset_path, + spec, + batch_size=train_settings["training"]["stage_0"]["batch_size"], + ) + # This is needed to prevent an occasional error when loading a large + # dataset into memory using a dataloader. This removes a limitation on + # the number of "open files". + old_sharing_strategy = torch.multiprocessing.get_sharing_strategy() + torch.multiprocessing.set_sharing_strategy("file_system") + with threadpool_limits(limits=1, user_api="blas"): + module.initialize_weights(dataloader, out_dir=train_dir) + torch.multiprocessing.set_sharing_strategy(old_sharing_strategy) + initialized = True + if initialized: + # initialization_dataloader replaced the transforms; restore the + # training configuration. + set_train_transforms(wfd, train_settings["data"], asd_dataset_path) + if local_settings.get("wandb", False): try: import wandb diff --git a/tests/core/test_build_model.py b/tests/core/test_build_model.py index 6b6b81ef..75225b5e 100644 --- a/tests/core/test_build_model.py +++ b/tests/core/test_build_model.py @@ -367,40 +367,92 @@ def test_parameter_contract_defaults(): # ----------------------------------------------------------------------------------- -# SVD initial-weight seeding +# Data-driven weight initialization (init_data_spec / initialize_weights hooks) # ----------------------------------------------------------------------------------- -def test_initial_weights_seed_svd_projection(): - """initial_weights['V_rb_list'] seeds the LinearProjectionRB layer weights, and the - settings dict is not polluted with the (large) V matrices.""" +def _init_batches(waveform_len, batch_sizes): + """Batches in the format the initialization dataloader provides: per-block + complex strains plus parameters.""" + rng = np.random.default_rng(42) num_bins = DATA_SHAPE[2] + for batch_size in batch_sizes: + waveform = {} + for block in ("H1", "L1"): + strains = rng.normal(size=(batch_size, waveform_len)) + 1j * rng.normal( + size=(batch_size, waveform_len) + ) + # Entries outside the network's input range (leading excess) are zero. + strains[:, : waveform_len - num_bins] = 0.0 + waveform[block] = strains + yield { + "waveform": waveform, + "parameters": {"chirp_mass": rng.uniform(size=batch_size)}, + } + + +def test_svd_initialization_hook(): + """DenseSVDEmbedding requests clean, un-formatted data via init_data_spec and + seeds its projection layer with per-block SVD bases from the provided batches; + rows outside the network's input range are dropped.""" n_rb = 10 - V_rb_list = [ - (np.random.rand(num_bins, n_rb) + 1j * np.random.rand(num_bins, n_rb)) - for _ in range(DATA_SHAPE[0]) - ] + num_bins = DATA_SHAPE[2] + waveform_len = num_bins + 5 # data longer than the network input settings = model_settings("normalizing_flow") + embedding_kwargs = settings["train_settings"]["model"]["embedding_net"]["kwargs"] + embedding_kwargs["svd"] = {"size": n_rb, "num_training_samples": 40} settings_before = copy.deepcopy(settings) - pm = build_model_from_kwargs( - settings=settings, initial_weights={"V_rb_list": V_rb_list}, device="cpu" - ) + pm = build_model_from_kwargs(settings=settings, device="cpu") + embedding = pm.network.embedding_net + + spec = embedding.init_data_spec() + assert spec == { + "noise": False, + "network_format": False, + "fix_parameters": {"luminosity_distance": 100.0}, + "num_samples": 40, + } - projection = pm.network.embedding_net[0] - V = V_rb_list[0][:, :n_rb] - layer_weight = projection.layers_rb[0].weight.data - assert torch.allclose( - layer_weight[:n_rb, :num_bins], - torch.from_numpy(V.real.T).float(), - ) - assert torch.allclose( - layer_weight[n_rb:, :num_bins], - torch.from_numpy(V.imag.T).float(), + # Two batches of 25 provide the 40 samples (iteration stops mid-batch). + embedding.initialize_weights(_init_batches(waveform_len, [25, 25])) + + # The SVD itself is not deterministic (partial SVD with random start vector), + # so check the structure: the weights hold an orthonormal complex basis V of + # the network's input size, in the (real, imag) block layout of + # LinearProjectionRB.init_layers, with zero bias. + for layer in embedding[0].layers_rb: + weight = layer.weight.data + V_real = weight[:n_rb, :num_bins].T + V_imag = weight[n_rb:, :num_bins].T + assert torch.allclose(weight[:n_rb, num_bins : 2 * num_bins].T, -V_imag) + assert torch.allclose(weight[n_rb:, num_bins : 2 * num_bins].T, V_real) + # Third channel (e.g. ASD) initialized to zero. + assert torch.all(weight[:, 2 * num_bins :] == 0) + assert torch.all(layer.bias.data == 0) + V = torch.complex(V_real, V_imag).to(torch.complex128) + gram = V.T.conj() @ V + assert torch.allclose(gram, torch.eye(n_rb, dtype=torch.complex128), atol=1e-5) + # The two blocks received different bases (different data). + assert not torch.allclose( + embedding[0].layers_rb[0].weight, embedding[0].layers_rb[1].weight ) # The V matrices must not leak into the saved settings. assert settings == settings_before + # Too few samples fail loudly. + with pytest.raises(IndexError, match="40"): + embedding.initialize_weights(_init_batches(waveform_len, [25])) + + +def test_init_data_spec_none_without_seeding_request(): + """Without num_training_samples in the svd settings (e.g. when loading a saved + model), no data-driven initialization is requested.""" + pm = build_model_from_kwargs( + settings=model_settings("normalizing_flow"), device="cpu" + ) + assert pm.network.embedding_net.init_data_spec() is None + # ----------------------------------------------------------------------------------- # Backward compatibility: old settings schema From 08d9dbe996a6703ae62967b40591f341de1e073b Mon Sep 17 00:00:00 2001 From: Maximilian Dax Date: Sun, 12 Jul 2026 20:59:16 +0200 Subject: [PATCH 8/9] Merge transformer embedding as a registered architecture Ports the transformer branch (origin/transformer, stage 1 of the T1 merge) onto the registry build system. TransformerEmbedding registers as "transformer": tokenized strain -> Tokenizer (shared DenseResidualNet with GLU position-conditioning) -> TransformerEncoder -> CLS/average pooling -> optional final net. complete_settings infers the tokenizer input dims and num_blocks from a named sample batch (position is its own batch entry now, instead of overloading the GNPE-proxy slot). Because the seam is generic, transformer x flow-matching works with zero extra wiring (pinned by test), and context parameters compose via the concat merger - both impossible on the original branch. StrainTokenization lands as a data transform, wired into set_train_transforms and GWSampler via the tokenization data settings; core samplers accept dict outputs from transform_pre. DenseResidualNet is unified rather than duplicated: the branch's implementation (per-block GLU context, optional layer_norm, arbitrary leading batch dims) moves to core/nn/resnet.py and replaces the glasflow-block-based version in enets.py - MyResidualBlock is a strict superset of glasflow's ResidualBlock with identical parameter names and forward pass (pinned by test), so old checkpoints load unchanged. Only the rq-coupling conditioner is genuinely a different architecture; per the team decision it is selected by an explicit conditioner_type setting (glasflow_residual default | dense_residual with optional layer_norm), not a layer_norm flag. layer_norm threads through the cfnets builders, and ContinuousFlow no longer shares a mutable-default Identity module. Not ported (deliberate): the advertised-but-dead augmentation config keys (drop_detectors, drop_frequency_range - stage 2, from dingo-t1), the allow_tf32 pop-and-discard, the embedding_type branches (registry handles dispatch), and create_nsf_with_transformer_embedding_net. The ConcatContextMerger now routes multi-input embeddings. Registration of built-in architectures moves to dingo/core/nn/__init__.py. Adds examples/transformer_model and the branch's transformer / resnet / tokenization / sampler-transform tests, adapted to the registry API. Co-Authored-By: Claude Fable 5 --- dingo/core/nn/__init__.py | 9 + dingo/core/nn/cfnets.py | 32 +- dingo/core/nn/enets.py | 94 +--- dingo/core/nn/nsf.py | 53 ++- dingo/core/nn/resnet.py | 207 +++++++++ dingo/core/nn/transformer.py | 372 +++++++++++++++ dingo/core/samplers.py | 14 +- dingo/gw/inference/gw_samplers.py | 43 +- dingo/gw/training/train_builders.py | 20 +- dingo/gw/transforms/__init__.py | 4 +- .../gw/transforms/tokenization_transforms.py | 232 ++++++++++ .../transformer_model/train_settings.yaml | 102 ++++ tests/core/test_build_model.py | 127 +++++ tests/core/test_nsf.py | 49 ++ tests/core/test_resnet.py | 125 +++++ tests/core/test_transformer.py | 436 ++++++++++++++++++ tests/gw/inference/__init__.py | 0 .../inference/test_gw_sampler_transforms.py | 201 ++++++++ .../test_tokenization_transforms.py | 321 +++++++++++++ 19 files changed, 2323 insertions(+), 118 deletions(-) create mode 100644 dingo/core/nn/resnet.py create mode 100644 dingo/core/nn/transformer.py create mode 100644 dingo/gw/transforms/tokenization_transforms.py create mode 100644 examples/transformer_model/train_settings.yaml create mode 100644 tests/core/test_resnet.py create mode 100644 tests/core/test_transformer.py create mode 100644 tests/gw/inference/__init__.py create mode 100644 tests/gw/inference/test_gw_sampler_transforms.py create mode 100644 tests/gw/transforms/test_tokenization_transforms.py diff --git a/dingo/core/nn/__init__.py b/dingo/core/nn/__init__.py index e69de29b..ed46a48c 100644 --- a/dingo/core/nn/__init__.py +++ b/dingo/core/nn/__init__.py @@ -0,0 +1,9 @@ +"""Neural network building blocks and registered architectures. + +Importing the architecture modules registers the built-in embedding networks and +context mergers with the registries (dingo.core.registry); importing any +dingo.core.nn submodule triggers this. +""" + +import dingo.core.nn.enets # noqa: F401 +import dingo.core.nn.transformer # noqa: F401 diff --git a/dingo/core/nn/cfnets.py b/dingo/core/nn/cfnets.py index 713bcd3d..c58cf4ef 100644 --- a/dingo/core/nn/cfnets.py +++ b/dingo/core/nn/cfnets.py @@ -1,9 +1,11 @@ +from typing import Optional + import numpy as np import torch import torch.nn as nn from dingo.core.utils import torchutils -from dingo.core.nn.enets import DenseResidualNet +from dingo.core.nn.resnet import DenseResidualNet class ContinuousFlow(nn.Module): @@ -25,8 +27,8 @@ class ContinuousFlow(nn.Module): def __init__( self, continuous_flow_net: nn.Module, - context_embedding_net: nn.Module = torch.nn.Identity(), - theta_embedding_net: nn.Module = torch.nn.Identity(), + context_embedding_net: Optional[nn.Module] = None, + theta_embedding_net: Optional[nn.Module] = None, context_with_glu: bool = False, theta_with_glu: bool = False, context_keys: tuple = ("waveform",), @@ -36,10 +38,12 @@ def __init__( ---------- continuous_flow_net: nn.Module Main network for the continuous flow. - context_embedding_net: nn.Module = torch.nn.Identity() + context_embedding_net: Optional[nn.Module] Embedding network for the context information (e.g., observed data). - theta_embedding_net: nn.Module = torch.nn.Identity() - Embedding network for the parameters. + If None, defaults to nn.Identity(). + theta_embedding_net: Optional[nn.Module] + Embedding network for the parameters. If None, defaults to + nn.Identity(). context_with_glu: bool = False Whether to provide context as GLU or main input to the continuous_flow_net. theta_with_glu: bool = False @@ -51,8 +55,18 @@ def __init__( """ super(ContinuousFlow, self).__init__() self.continuous_flow_net = continuous_flow_net - self.context_embedding_net = context_embedding_net - self.theta_embedding_net = theta_embedding_net + # Default to a fresh nn.Identity() per instance rather than a mutable + # default argument, which would otherwise share a single Identity module + # (and its registration as a submodule) across every ContinuousFlow that + # omits this argument. + self.context_embedding_net = ( + context_embedding_net + if context_embedding_net is not None + else nn.Identity() + ) + self.theta_embedding_net = ( + theta_embedding_net if theta_embedding_net is not None else nn.Identity() + ) self.theta_with_glu = theta_with_glu self.context_with_glu = context_with_glu self.context_keys = tuple(context_keys) @@ -227,6 +241,7 @@ def create_cf(kwargs: dict, embedding_net: nn.Module = None): activation=activation_fn, dropout=kwargs["dropout"], batch_norm=kwargs["batch_norm"], + layer_norm=kwargs.get("layer_norm", False), context_features=glu_dim, ) @@ -262,6 +277,7 @@ def get_theta_embedding_net(embedding_kwargs: dict, input_dim): output_dim=embedding_kwargs["embedding_net"]["output_dim"], hidden_dims=embedding_kwargs["embedding_net"]["hidden_dims"], activation=activation_fn, + layer_norm=embedding_kwargs["embedding_net"].get("layer_norm", False), dropout=embedding_kwargs["embedding_net"].get("dropout", 0.0), batch_norm=embedding_kwargs["embedding_net"].get("batch_norm", True), ) diff --git a/dingo/core/nn/enets.py b/dingo/core/nn/enets.py index b2923428..56d432d9 100755 --- a/dingo/core/nn/enets.py +++ b/dingo/core/nn/enets.py @@ -23,16 +23,15 @@ """ import os -from typing import Tuple, Callable, Union, List +from typing import Tuple, Union, List import torch import numpy as np import pandas as pd import torch.nn as nn -from torch.nn import functional as F -from glasflow.nflows.nn.nets.resnet import ResidualBlock from dingo.core.SVD import SVDBasis +from dingo.core.nn.resnet import DenseResidualNet from dingo.core.registry import CONTEXT_MERGERS, EMBEDDING_NETS from dingo.core.utils import torchutils @@ -185,88 +184,6 @@ def forward(self, x, **_): return x -class DenseResidualNet(nn.Module): - """ - A nn.Module consisting of a sequence of dense residual blocks. This is - used to embed high dimensional input to a compressed output. Linear - resizing layers are used for resizing the input and output to match the - first and last hidden dimension, respectively. - - Module specs - -------- - input dimension: (batch_size, input_dim) - output dimension: (batch_size, output_dim) - """ - - def __init__( - self, - input_dim: int, - output_dim: int, - hidden_dims: Tuple, - activation: Callable = F.elu, - dropout: float = 0.0, - batch_norm: bool = True, - context_features: int = None, - ): - """ - Parameters - ---------- - input_dim : int - dimension of the input to this module - output_dim : int - output dimension of this module - hidden_dims : tuple - tuple with dimensions of hidden layers of this module - activation: callable - activation function used in residual blocks - dropout: float - dropout probability for residual blocks used for reqularization - batch_norm: bool - flag that specifies whether to use batch normalization - context_features: int - Number of additional context features, which are provided to the residual - blocks via gated linear units. If None, no additional context expected. - """ - - super(DenseResidualNet, self).__init__() - self.input_dim = input_dim - self.output_dim = output_dim - self.hidden_dims = hidden_dims - self.num_res_blocks = len(self.hidden_dims) - - self.initial_layer = nn.Linear(self.input_dim, hidden_dims[0]) - self.blocks = nn.ModuleList( - [ - ResidualBlock( - features=self.hidden_dims[n], - context_features=context_features, - activation=activation, - dropout_probability=dropout, - use_batch_norm=batch_norm, - ) - for n in range(self.num_res_blocks) - ] - ) - self.resize_layers = nn.ModuleList( - [ - ( - nn.Linear(self.hidden_dims[n - 1], self.hidden_dims[n]) - if self.hidden_dims[n - 1] != self.hidden_dims[n] - else nn.Identity() - ) - for n in range(1, self.num_res_blocks) - ] - + [nn.Linear(self.hidden_dims[-1], self.output_dim)] - ) - - def forward(self, x, context=None): - x = self.initial_layer(x) - for block, resize_layer in zip(self.blocks, self.resize_layers): - x = block(x, context=context) - x = resize_layer(x) - return x - - @EMBEDDING_NETS.register("dense_svd") class DenseSVDEmbedding(nn.Sequential): """ @@ -501,6 +418,13 @@ def __init__(self, embedding_net: nn.Module, num_context_parameters: int): self.input_keys = (*embedding_net.input_keys, "context_parameters") self.output_dim = embedding_net.output_dim + num_context_parameters + def forward(self, *x): + # Unlike ModuleMerger, the wrapped embedding may consume several inputs + # (e.g. the transformer: waveform, position, mask); the context parameters + # are always the last one. + *data, context = x + return torch.cat([self.enets[0](*data), self.enets[1](context)], dim=1) + @staticmethod def merged_output_dim(embedding_output_dim: int, num_context_parameters: int): """Output dimension of the merged embedding, for settings completion.""" diff --git a/dingo/core/nn/nsf.py b/dingo/core/nn/nsf.py index 876b5868..70e82c7f 100644 --- a/dingo/core/nn/nsf.py +++ b/dingo/core/nn/nsf.py @@ -8,6 +8,7 @@ import glasflow.nflows as nflows # nflows not maintained, so use this maintained fork from glasflow.nflows import distributions, flows, transforms import glasflow.nflows.nn.nets as nflows_nets +from dingo.core.nn.resnet import DenseResidualNet from dingo.core.utils import torchutils @@ -38,6 +39,8 @@ def create_base_transform( activation: str = "relu", dropout_probability: float = 0.0, batch_norm: bool = False, + layer_norm: bool = False, + conditioner_type: str = "glasflow_residual", num_bins: int = 8, tail_bound: float = 1.0, apply_unconditional_transform: bool = False, @@ -78,6 +81,16 @@ def create_base_transform( dropout probability for regularization :param batch_norm: bool = False whether to use batch normalization + :param layer_norm: bool = False + whether to use layer normalization in the conditioner network + (conditioner_type "dense_residual" only) + :param conditioner_type: str = "glasflow_residual" + conditioner network of the rq-coupling transform. "glasflow_residual" + (glasflow's ResidualNet, context concatenated to the input once) or + "dense_residual" (dingo's DenseResidualNet, context injected into every + residual block via a gated linear unit, optional layer_norm). The two are + architecturally different — checkpoints are not interchangeable — so this + is an explicit type, not a flag. :param num_bins: int = 8 number of bins for the spline :param tail_bound: float = 1. @@ -91,6 +104,11 @@ def create_base_transform( """ activation_fn = torchutils.get_activation_function_from_string(activation) + if layer_norm and conditioner_type != "dense_residual": + raise ValueError( + "layer_norm requires conditioner_type 'dense_residual' (glasflow's " + "ResidualNet only supports batch norm)." + ) if base_transform_type == "rq-coupling": if param_dim == 1: @@ -99,9 +117,9 @@ def create_base_transform( mask = nflows.utils.create_alternating_binary_mask( param_dim, even=(i % 2 == 0) ) - return transforms.PiecewiseRationalQuadraticCouplingTransform( - mask=mask, - transform_net_create_fn=( + + if conditioner_type == "glasflow_residual": + transform_net_create_fn = ( lambda in_features, out_features: nflows_nets.ResidualNet( in_features=in_features, out_features=out_features, @@ -112,7 +130,29 @@ def create_base_transform( dropout_probability=dropout_probability, use_batch_norm=batch_norm, ) - ), + ) + elif conditioner_type == "dense_residual": + transform_net_create_fn = ( + lambda in_features, out_features: DenseResidualNet( + input_dim=in_features, + output_dim=out_features, + hidden_dims=(hidden_dim,) * num_transform_blocks, + activation=activation_fn, + context_features=context_dim, + dropout=dropout_probability, + batch_norm=batch_norm, + layer_norm=layer_norm, + ) + ) + else: + raise ValueError( + f"Unknown conditioner_type '{conditioner_type}'; expected " + f"'glasflow_residual' or 'dense_residual'." + ) + + return transforms.PiecewiseRationalQuadraticCouplingTransform( + mask=mask, + transform_net_create_fn=transform_net_create_fn, num_bins=num_bins, tails="linear", tail_bound=tail_bound, @@ -120,6 +160,11 @@ def create_base_transform( ) elif base_transform_type == "rq-autoregressive": + if conditioner_type != "glasflow_residual": + raise ValueError( + "rq-autoregressive only supports conditioner_type " + "'glasflow_residual'." + ) return transforms.MaskedPiecewiseRationalQuadraticAutoregressiveTransform( features=param_dim, hidden_features=hidden_dim, diff --git a/dingo/core/nn/resnet.py b/dingo/core/nn/resnet.py new file mode 100644 index 00000000..2856c9e5 --- /dev/null +++ b/dingo/core/nn/resnet.py @@ -0,0 +1,207 @@ +"""Dense residual network supporting layer normalization and multi-dimensional +(e.g., token-batched) inputs, used by the transformer tokenizer.""" + +from typing import Callable, Optional, Tuple + +import torch +from torch import nn, Tensor +from torch.nn import functional as F, init + + +class MyResidualBlock(nn.Module): + """ + A general-purpose residual block, supporting batch norm or layer norm. + + Context features are injected via a gated linear unit. The GLU is applied along + the last dimension, so this supports both [batch, features] and + [batch, tokens, features] inputs. + """ + + def __init__( + self, + features: int, + context_features: Optional[int] = None, + activation: Callable = F.relu, + dropout_probability: float = 0.0, + use_batch_norm: bool = False, + use_layer_norm: bool = False, + zero_initialization: bool = True, + ): + """ + Parameters + ---------- + features : int + dimension of the residual block input and output + context_features : Optional[int] + number of context features injected via a gated linear unit; if None, + no context is expected + activation : Callable + activation function used between linear layers + dropout_probability : float + dropout probability applied for regularization + use_batch_norm : bool + whether to use batch normalization + use_layer_norm : bool + whether to use layer normalization + zero_initialization : bool + whether to initialize the final linear layer with small weights + """ + super().__init__() + self.activation = activation + + if use_batch_norm and use_layer_norm: + raise ValueError( + "Residual block should not use both batch norm and layer norm." + ) + self.use_batch_norm = use_batch_norm + self.use_layer_norm = use_layer_norm + if use_batch_norm: + self.batch_norm_layers = nn.ModuleList( + [nn.BatchNorm1d(features, eps=1e-3) for _ in range(2)] + ) + if use_layer_norm: + self.layer_norm_layers = nn.ModuleList( + [nn.LayerNorm(features) for _ in range(2)] + ) + if context_features is not None: + self.context_layer = nn.Linear(context_features, features) + self.linear_layers = nn.ModuleList( + [nn.Linear(features, features) for _ in range(2)] + ) + self.dropout = nn.Dropout(p=dropout_probability) + if zero_initialization: + init.uniform_(self.linear_layers[-1].weight, -1e-3, 1e-3) + init.uniform_(self.linear_layers[-1].bias, -1e-3, 1e-3) + + def forward(self, inputs: Tensor, context: Optional[Tensor] = None) -> Tensor: + temps = inputs + if self.use_batch_norm: + temps = self.batch_norm_layers[0](temps) + if self.use_layer_norm: + temps = self.layer_norm_layers[0](temps) + temps = self.activation(temps) + temps = self.linear_layers[0](temps) + if self.use_batch_norm: + temps = self.batch_norm_layers[1](temps) + if self.use_layer_norm: + temps = self.layer_norm_layers[1](temps) + temps = self.activation(temps) + temps = self.dropout(temps) + temps = self.linear_layers[1](temps) + if context is not None: + temps = F.glu( + torch.cat((temps, self.context_layer(context)), dim=-1), dim=-1 + ) + return inputs + temps + + +class DenseResidualNet(nn.Module): + """ + A nn.Module consisting of a sequence of dense residual blocks. This is + used to embed high dimensional input to a compressed output. Linear + resizing layers are used for resizing the input and output to match the + first and last hidden dimension, respectively. + + MyResidualBlock is a superset of glasflow.nflows.nn.nets.ResidualBlock: with + layer_norm=False it has identical parameters (same names) and an identical + forward pass, so checkpoints of networks previously built from glasflow blocks + load unchanged (pinned by test). On top, it supports layer normalization and + inputs with an arbitrary number of leading batch dimensions (e.g. + [batch, tokens, features]), as needed by the transformer tokenizer. + + In contrast, glasflow's ResidualNet (the full network, used as the coupling + conditioner in the normalizing flow) concatenates the context vector with the + input once at the start, whereas this network injects context into every + residual block via a gated linear unit. The two are architecturally different + and their checkpoints are not interchangeable — which is why the coupling + conditioner type is an explicit setting (see create_base_transform), not a + flag. + + Module specs + -------- + input dimension: (..., input_dim) + output dimension: (..., output_dim) + """ + + def __init__( + self, + input_dim: int, + output_dim: int, + hidden_dims: Tuple, + activation: Callable = F.elu, + context_features: Optional[int] = None, + dropout: float = 0.0, + batch_norm: bool = False, + layer_norm: bool = False, + ): + """ + Parameters + ---------- + input_dim : int + dimension of the input to this module + output_dim : int + output dimension of this module + hidden_dims : tuple + tuple with dimensions of hidden layers of this module + activation : Callable + activation function used in residual blocks + context_features : Optional[int] + number of additional context features, which are provided to the + residual blocks via gated linear units; if None, no context expected + dropout : float + dropout probability for residual blocks, used for regularization + batch_norm : bool + whether to use batch normalization + layer_norm : bool + whether to use layer normalization + """ + super().__init__() + self.input_dim = input_dim + self.output_dim = output_dim + self.hidden_dims = hidden_dims + self.num_res_blocks = len(self.hidden_dims) + + self.initial_layer = nn.Linear(self.input_dim, hidden_dims[0]) + self.blocks = nn.ModuleList( + [ + MyResidualBlock( + features=self.hidden_dims[n], + context_features=context_features, + activation=activation, + dropout_probability=dropout, + use_batch_norm=batch_norm, + use_layer_norm=layer_norm, + ) + for n in range(self.num_res_blocks) + ] + ) + self.resize_layers = nn.ModuleList( + [ + ( + nn.Linear(self.hidden_dims[n - 1], self.hidden_dims[n]) + if self.hidden_dims[n - 1] != self.hidden_dims[n] + else nn.Identity() + ) + for n in range(1, self.num_res_blocks) + ] + + [nn.Linear(self.hidden_dims[-1], self.output_dim)] + ) + + def forward(self, x: Tensor, context: Optional[Tensor] = None) -> Tensor: + x = self.initial_layer(x) + for block, resize_layer in zip(self.blocks, self.resize_layers): + x = block(x, context=context) + x = resize_layer(x) + return x + + +class LinearLayer(nn.Module): + """A single linear layer followed by an activation function.""" + + def __init__(self, input_dim: int, output_dim: int, activation: Callable): + super().__init__() + self.linear = nn.Linear(input_dim, output_dim) + self.activation = activation + + def forward(self, x: Tensor) -> Tensor: + return self.activation(self.linear(x)) diff --git a/dingo/core/nn/transformer.py b/dingo/core/nn/transformer.py new file mode 100644 index 00000000..09f01b6a --- /dev/null +++ b/dingo/core/nn/transformer.py @@ -0,0 +1,372 @@ +"""Transformer embedding network for tokenized strain data (DINGO-T1).""" + +from typing import Callable, List, Optional + +import torch +from torch import nn, Tensor +from torch.nn import TransformerEncoder, TransformerEncoderLayer + +from dingo.core.nn.resnet import DenseResidualNet, LinearLayer +from dingo.core.registry import EMBEDDING_NETS +from dingo.core.utils import torchutils + + +class Tokenizer(nn.Module): + """ + Maps each token's raw features to a d_model-dimensional embedding via a shared + DenseResidualNet, conditioned on the token's position (f_min, f_max, detector). + + Methods + ------- + forward: + Obtain the token embedding for a Tensor of shape + [..., num_tokens, num_features], conditioned on position. + """ + + def __init__( + self, + input_dims: List[int], + hidden_dims: List[int], + output_dim: int, + activation: Callable, + num_blocks: int, + dropout: float = 0.0, + batch_norm: bool = False, + layer_norm: bool = False, + ): + """ + Parameters + ---------- + input_dims : List[int] + [num_tokens, num_features], i.e., the shape of the tokenized waveform, + omitting batch dimensions. Only num_features (the last entry) is used. + hidden_dims : List[int] + dimensions of hidden layers for the underlying DenseResidualNet + output_dim : int + output dimension of the token embedding (typically d_model) + activation : Callable + activation function for the DenseResidualNet + num_blocks : int + number of blocks (detectors, in the GW use case); determines the size of + the one-hot detector encoding used as part of the conditioning context + dropout : float + dropout rate for the DenseResidualNet + batch_norm : bool + whether to use batch normalization in the DenseResidualNet + layer_norm : bool + whether to use layer normalization in the DenseResidualNet + """ + super().__init__() + if len(input_dims) != 2: + raise ValueError( + f"Invalid shape in Tokenizer, expected len(input_dims) == 2, got " + f"{input_dims}." + ) + self.num_features = input_dims[-1] + self.num_blocks = num_blocks + self.tokenizer_net = DenseResidualNet( + input_dim=self.num_features, + output_dim=output_dim, + hidden_dims=tuple(hidden_dims), + activation=activation, + context_features=2 + num_blocks, + dropout=dropout, + batch_norm=batch_norm, + layer_norm=layer_norm, + ) + + def forward(self, x: Tensor, position: Tensor) -> Tensor: + """ + Parameters + ---------- + x : Tensor + shape [..., num_tokens, num_features] + position : Tensor + shape [..., num_tokens, 3], last dim = [f_min, f_max, detector_index] + + Returns + ------- + Tensor + shape [..., num_tokens, output_dim] + """ + if x.shape[-1] != self.num_features: + raise ValueError( + f"Invalid shape for token embedding layer. " + f"Expected last dimension to be {self.num_features}, got " + f"{x.shape[-1]}." + ) + detector_per_token = position[..., 2] + detector_one_hot = torch.eye(self.num_blocks, device=position.device)[ + detector_per_token.long() + ] + context = torch.cat((position[..., :2], detector_one_hot), dim=-1) + return self.tokenizer_net(x=x, context=context) + + +class TransformerModel(nn.Module): + """ + Transformer encoder used as an embedding network for the normalizing flow. Each + token is embedded via a conditional Tokenizer (conditioned on position), then + processed by a standard TransformerEncoder. The resulting sequence of token + embeddings is pooled (CLS token or average) into a single vector, optionally + followed by a final network. + """ + + def __init__( + self, + tokenizer: Tokenizer, + d_model: int, + dim_feedforward: int, + nhead: int, + num_layers: int, + dropout: float = 0.1, + norm_first: bool = False, + pooling: str = "cls", + final_net: Optional[nn.Module] = None, + ): + """ + Parameters + ---------- + tokenizer : Tokenizer + Maps raw per-token features (conditioned on position) to d_model-dim + token embeddings. + d_model : int + embedding size of the transformer + dim_feedforward : int + number of hidden dimensions in the feedforward networks of the + transformer encoder layers + nhead : int + number of transformer attention heads + num_layers : int + number of transformer encoder layers + dropout : float + dropout probability in the transformer encoder layers + norm_first : bool + if True, layer normalization is applied before the attention and + feedforward operations in each encoder layer, otherwise after + pooling : str + one of ["average", "cls"]; how to pool the sequence of token embeddings + into a single vector + final_net : Optional[nn.Module] + network applied to the pooled output, e.g., to project it to the + context dimension expected by the normalizing flow. If None, the pooled + output is returned directly. + """ + super().__init__() + if pooling not in ("average", "cls"): + raise ValueError( + f"Invalid pooling operation {pooling}, expected one of " + f"['average', 'cls']." + ) + + self.tokenizer = tokenizer + self.d_model = d_model + self.pooling = pooling + self.final_net = final_net + + encoder_layer = TransformerEncoderLayer( + d_model=d_model, + nhead=nhead, + dim_feedforward=dim_feedforward, + dropout=dropout, + batch_first=True, + norm_first=norm_first, + ) + self.transformer_encoder = TransformerEncoder( + encoder_layer=encoder_layer, num_layers=num_layers + ) + + if self.pooling == "cls": + self.class_token = nn.Parameter(torch.randn((1, 1, d_model))) + + self.init_weights() + + def init_weights(self) -> None: + """ + Initialize parameters of the transformer encoder explicitly, due to + https://github.com/pytorch/pytorch/issues/72253. Parameters are initialized + with xavier uniform. + """ + for p in self.transformer_encoder.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + + def forward( + self, + x: Tensor, + position: Tensor, + src_key_padding_mask: Optional[Tensor] = None, + ) -> Tensor: + """ + Parameters + ---------- + x : Tensor + shape [batch_size, num_tokens, num_features] + position : Tensor + shape [batch_size, num_tokens, 3], last dim = [f_min, f_max, detector] + src_key_padding_mask : Optional[Tensor] + shape [batch_size, num_tokens]; PyTorch transformer convention, True = + masked out (not allowed to attend) + + Returns + ------- + Tensor + shape [batch_size, output_dim of final_net if final_net else d_model] + """ + x = self.tokenizer(x=x, position=position) + + if self.pooling == "cls": + batch_size = x.shape[0] + x = torch.cat((self.class_token.expand(batch_size, -1, -1), x), dim=1) + if src_key_padding_mask is not None: + mask_cls_token = torch.zeros( + (batch_size, 1), + dtype=torch.bool, + device=src_key_padding_mask.device, + ) + src_key_padding_mask = torch.cat( + (mask_cls_token, src_key_padding_mask), dim=1 + ) + + x = self.transformer_encoder(src=x, src_key_padding_mask=src_key_padding_mask) + + if self.pooling == "average": + if src_key_padding_mask is not None: + denominator = torch.sum(~src_key_padding_mask, dim=-1, keepdim=True) + x = ( + torch.sum(x * (~src_key_padding_mask).unsqueeze(-1), dim=-2) + / denominator + ) + else: + x = torch.mean(x, dim=-2) + else: # pooling == "cls" + x = x[..., 0, :] + + if self.final_net is not None: + x = self.final_net(x) + + return x + + +@EMBEDDING_NETS.register("transformer") +class TransformerEmbedding(TransformerModel): + """ + TransformerModel as a registered embedding network (see the contract in + dingo.core.nn.enets): consumes the tokenized batch entries produced by + StrainTokenization, and builds tokenizer / final net from settings dicts. + """ + + input_keys = ("waveform", "position", "drop_token_mask") + + def __init__( + self, + tokenizer_kwargs: dict, + transformer_kwargs: dict, + output_dim: int, + pooling: str = "cls", + final_net_kwargs: Optional[dict] = None, + ): + """ + Parameters + ---------- + tokenizer_kwargs : dict + Settings for the Tokenizer: hidden_dims, activation (str), and + optionally dropout, batch_norm, layer_norm. input_dims and num_blocks + are inferred from a sample batch by complete_settings; the tokenizer + output_dim is transformer_kwargs["d_model"]. + transformer_kwargs : dict + Settings for the transformer encoder: d_model, dim_feedforward, nhead, + num_layers, and optionally dropout, norm_first. + output_dim : int + Dimension of the embedded context: the output_dim of final_net_kwargs + if given, else d_model. Inferred by complete_settings; not a user + setting. + pooling : str + one of ["average", "cls"] + final_net_kwargs : Optional[dict] + Settings for the network applied after pooling. Must contain + output_dim and activation (str). With hidden_dims, a DenseResidualNet + is built (dropout, batch_norm, layer_norm are then read as well); + otherwise a LinearLayer. If None, the pooled d_model-dim vector is + returned directly. + """ + tokenizer_kwargs = dict(tokenizer_kwargs) + tokenizer_kwargs["activation"] = torchutils.get_activation_function_from_string( + tokenizer_kwargs["activation"] + ) + tokenizer = Tokenizer( + output_dim=transformer_kwargs["d_model"], + **tokenizer_kwargs, + ) + + final_net = None + if final_net_kwargs is not None: + final_net_kwargs = dict(final_net_kwargs) + final_net_output_dim = final_net_kwargs.pop("output_dim") + final_net_kwargs["activation"] = ( + torchutils.get_activation_function_from_string( + final_net_kwargs["activation"] + ) + ) + if "hidden_dims" in final_net_kwargs: + final_net_kwargs["hidden_dims"] = tuple(final_net_kwargs["hidden_dims"]) + final_net = DenseResidualNet( + input_dim=transformer_kwargs["d_model"], + output_dim=final_net_output_dim, + **final_net_kwargs, + ) + else: + final_net = LinearLayer( + input_dim=transformer_kwargs["d_model"], + output_dim=final_net_output_dim, + **final_net_kwargs, + ) + else: + final_net_output_dim = transformer_kwargs["d_model"] + if output_dim != final_net_output_dim: + raise ValueError( + f"Inconsistent settings: output_dim is {output_dim}, but the " + f"network produces {final_net_output_dim} " + f"(final_net output_dim, or d_model without a final net)." + ) + + super().__init__( + tokenizer=tokenizer, + pooling=pooling, + final_net=final_net, + **transformer_kwargs, + ) + self.output_dim = output_dim + + @classmethod + def complete_settings(cls, settings: dict, sample_batch: dict) -> dict: + """Infer the tokenizer input dims and number of blocks (detectors) plus the + embedding output_dim from a sample batch; return completed settings.""" + tokenizer_kwargs = dict(settings["tokenizer_kwargs"]) + for key in ("input_dims", "num_blocks"): + if key in tokenizer_kwargs: + raise ValueError( + f"'{key}' is derived from the data and must not be specified " + f"in the tokenizer settings." + ) + if "output_dim" in settings: + raise ValueError( + "'output_dim' is derived from the network settings and must not " + "be specified." + ) + # sample_batch["waveform"] is the tokenized strain, [num_tokens, + # num_features]; column 2 of position holds integer detector indices + # 0..num_blocks-1. + tokenizer_kwargs["input_dims"] = list(sample_batch["waveform"].shape) + tokenizer_kwargs["num_blocks"] = int(sample_batch["position"][:, 2].max()) + 1 + + final_net_kwargs = settings.get("final_net_kwargs") + if final_net_kwargs is not None: + output_dim = final_net_kwargs["output_dim"] + else: + output_dim = settings["transformer_kwargs"]["d_model"] + return { + **settings, + "tokenizer_kwargs": tokenizer_kwargs, + "output_dim": output_dim, + } diff --git a/dingo/core/samplers.py b/dingo/core/samplers.py index 90360d7c..0479fb98 100644 --- a/dingo/core/samplers.py +++ b/dingo/core/samplers.py @@ -151,9 +151,14 @@ def _run_sampler( # transforms_pre are expected to transform the data in the same way for each # requested sample. We therefore apply pre-processing only once. + # transform_pre yields either the prepared strain tensor, or a dict of + # named tensors (e.g. tokenized models: waveform, position, + # drop_token_mask). x = self.transform_pre(context) + if not isinstance(x, dict): + x = {"waveform": x} # Require a batch dimension for the embedding network. - x = {"waveform": x.unsqueeze(0)} + x = {k: v.unsqueeze(0) for k, v in x.items()} else: if context is not None: print("Unconditional model. Ignoring context.") @@ -274,9 +279,10 @@ def log_prob(self, samples: pd.DataFrame | dict) -> np.ndarray: # Context is the same for each sample. Expand across batch dimension after # pre-processing. x = self.transform_pre(self.context) - x = { - "waveform": x.expand(len(samples), *x.shape) - } # TODO: Make this more efficient. + if not isinstance(x, dict): + x = {"waveform": x} + # TODO: Make this more efficient. + x = {k: v.expand(len(samples), *v.shape) for k, v in x.items()} else: x = None diff --git a/dingo/gw/inference/gw_samplers.py b/dingo/gw/inference/gw_samplers.py index 2a1b63d8..d5e38aaf 100644 --- a/dingo/gw/inference/gw_samplers.py +++ b/dingo/gw/inference/gw_samplers.py @@ -32,14 +32,15 @@ GetDetectorTimes, DecimateWaveformsAndASDS, MaskDataForFrequencyRangeUpdate, + SelectKeys, + StrainTokenization, ) class SamplerProtocol(Protocol): base_model_metadata: dict - def _initialize_transforms(self) -> None: - ... + def _initialize_transforms(self) -> None: ... class _GWMixinProtocol(SamplerProtocol): @@ -173,7 +174,6 @@ def _build_domain(self: Sampler): if "domain_update" in data_settings: self.domain.update(data_settings["domain_update"]) - def _correct_reference_time( self: Sampler, samples: Union[dict, pd.DataFrame], inverse: bool = False ): @@ -309,19 +309,36 @@ def _initialize_transforms(self): ) ) # * repackage strains and asds from dicts to an array - # * convert array to torch tensor on the correct device - # * extract only strain/waveform from the sample - transform_pre += [ - # Use base metadata so that unconditional samplers still know how to - # transform data, since this transform is used by the GNPE sampler as - # well. + # * optionally tokenize the strain (tokenized models, e.g. transformer) + # * convert array(s) to torch tensor(s) on the correct device + # * extract the network inputs from the sample + # Use base metadata so that unconditional samplers still know how to + # transform data, since this transform is used by the GNPE sampler as well. + transform_pre.append( RepackageStrainsAndASDS( ifos=self.detectors, first_index=self.domain.min_idx, - ), - ToTorch(device=self.model.device), - GetItem("waveform"), - ] + ) + ) + tokenization = self.metadata["train_settings"]["data"].get("tokenization") + if tokenization: + # StrainTokenization operates on numpy arrays, so it precedes ToTorch. + transform_pre.append( + StrainTokenization( + domain=self.domain, + token_size=tokenization.get("token_size"), + num_tokens_per_block=tokenization.get("num_tokens_per_block"), + drop_last_token=tokenization.get("drop_last_token", False), + ) + ) + transform_pre.append(ToTorch(device=self.model.device)) + if tokenization: + # Dict of named network inputs; the model routes them by key. + transform_pre.append( + SelectKeys(["waveform", "position", "drop_token_mask"]) + ) + else: + transform_pre.append(GetItem("waveform")) self.transform_pre = Compose(transform_pre) # postprocessing transforms: diff --git a/dingo/gw/training/train_builders.py b/dingo/gw/training/train_builders.py index e503b599..33190208 100755 --- a/dingo/gw/training/train_builders.py +++ b/dingo/gw/training/train_builders.py @@ -18,6 +18,7 @@ SampleExtrinsicParameters, GetDetectorTimes, CropMaskStrainRandom, + StrainTokenization, ) from dingo.gw.noise.asd_dataset import ASDDataset from dingo.gw.prior import default_inference_parameters @@ -232,10 +233,22 @@ def set_train_transforms(wfd, data_settings, asd_dataset_path, omit_transforms=N transforms.append( CropMaskStrainRandom(domain, **data_settings["random_strain_cropping"]) ) + if "tokenization" in data_settings: + tokenization = data_settings["tokenization"] + transforms.append( + StrainTokenization( + domain=domain, + token_size=tokenization.get("token_size"), + num_tokens_per_block=tokenization.get("num_tokens_per_block"), + drop_last_token=tokenization.get("drop_last_token", False), + ) + ) + + selected_keys = ["inference_parameters", "waveform"] + if "tokenization" in data_settings: + selected_keys += ["position", "drop_token_mask"] if data_settings["context_parameters"]: - selected_keys = ["inference_parameters", "waveform", "context_parameters"] - else: - selected_keys = ["inference_parameters", "waveform"] + selected_keys += ["context_parameters"] transforms.append(SelectKeys(selected_keys=selected_keys)) @@ -298,6 +311,7 @@ def initialization_dataloader( SelectStandardizeRepackageParameters, SelectKeys, CropMaskStrainRandom, + StrainTokenization, ] set_train_transforms( wfd, data_settings, asd_dataset_path, omit_transforms=omit_transforms or None diff --git a/dingo/gw/transforms/__init__.py b/dingo/gw/transforms/__init__.py index 75e3785e..e601c190 100644 --- a/dingo/gw/transforms/__init__.py +++ b/dingo/gw/transforms/__init__.py @@ -5,4 +5,6 @@ from .gnpe_transforms import * from .inference_transforms import * from .utils import * -from .waveform_transforms import * \ No newline at end of file +from .waveform_transforms import * +from .tokenization_transforms import StrainTokenization, DETECTOR_DICT + diff --git a/dingo/gw/transforms/tokenization_transforms.py b/dingo/gw/transforms/tokenization_transforms.py new file mode 100644 index 00000000..f683d850 --- /dev/null +++ b/dingo/gw/transforms/tokenization_transforms.py @@ -0,0 +1,232 @@ +from typing import Optional + +import numpy as np + +from dingo.gw.domains import UniformFrequencyDomain, MultibandedFrequencyDomain + +DETECTOR_DICT = {"H1": 0, "L1": 1, "V1": 2} + + +class StrainTokenization: + """ + Divide strain frequency bins into fixed-size tokens and attach per-token position + information (f_min, f_max, detector index). + + The input waveform is expected to have shape + [..., num_detectors, num_channels, num_bins] + where num_channels >= 1 (e.g. real, imaginary, ASD). + + The output contains: + - 'waveform': [..., num_detectors * num_tokens_per_detector, + num_channels * num_bins_per_token] + - 'position': [..., num_tokens, 3] + last dim = [f_min, f_max, detector_index] + - 'drop_token_mask': [..., num_tokens] bool, False = keep token + (PyTorch transformer convention: True = masked out). + """ + + def __init__( + self, + domain: UniformFrequencyDomain | MultibandedFrequencyDomain, + num_tokens_per_block: Optional[int] = None, + token_size: Optional[int] = None, + drop_last_token: bool = False, + print_output: bool = True, + ): + """ + Parameters + ---------- + domain: + Domain carrying f_min, f_max, delta_f, sample_frequencies. + num_tokens_per_block: + Number of tokens per detector. Mutually exclusive with token_size. + token_size: + Number of frequency bins per token. Mutually exclusive with + num_tokens_per_block. + drop_last_token: + If True and the bins do not divide evenly, drop the trailing incomplete + token. If False, pad it with zeros. + print_output: + Write a summary to stdout on construction. + """ + if (num_tokens_per_block is None) == (token_size is None): + raise ValueError( + "Specify exactly one of num_tokens_per_block or token_size." + ) + + num_f = domain.frequency_mask_length + + if token_size is not None: + self.num_bins_per_token = token_size + n_full = num_f // token_size + remainder = num_f % token_size + num_tokens_per_block = ( + n_full + if (drop_last_token and remainder) + else (n_full if remainder == 0 else n_full + 1) + ) + else: + remainder = num_f % num_tokens_per_block + # Ceiling ensures the given number of tokens covers the full frequency range. + self.num_bins_per_token = int(np.ceil(num_f / num_tokens_per_block)) + if drop_last_token and remainder: + num_tokens_per_block -= 1 + + self.drop_last_token = drop_last_token + self.num_tokens_per_detector = num_tokens_per_block + + # f_min / f_max for every token (same for all detectors) + freqs = domain.sample_frequencies + start = domain.min_idx + self.f_min_per_token = freqs[start :: self.num_bins_per_token][ + :num_tokens_per_block + ] + self.f_max_per_token = freqs[ + start + self.num_bins_per_token - 1 :: self.num_bins_per_token + ][:num_tokens_per_block] + + # Number of zero-padding bins needed in the last token + self.num_padded_f_bins = 0 + if ( + len(self.f_min_per_token) > len(self.f_max_per_token) + and not drop_last_token + ): + # Last token is incomplete: extrapolate f_max + if isinstance(domain, MultibandedFrequencyDomain): + last_delta_f = domain.delta_f[-1] + else: + last_delta_f = domain.delta_f + f_max_pad = ( + self.f_max_per_token[-1] + self.num_bins_per_token * last_delta_f + ) + self.f_max_per_token = np.append(self.f_max_per_token, f_max_pad) + self.num_padded_f_bins = int((f_max_pad - freqs[-1]) / last_delta_f) + + if not ( + num_tokens_per_block + == len(self.f_min_per_token) + == len(self.f_max_per_token) + ): + raise ValueError( + "f_min_per_token and f_max_per_token lengths do not match num_tokens_per_block." + ) + + if isinstance(domain, MultibandedFrequencyDomain): + _check_mfd_node_compatibility( + f_mins=self.f_min_per_token, + f_maxs=self.f_max_per_token, + mfd_nodes=domain.nodes, + drop_last_token=drop_last_token, + ) + + if print_output: + print( + f"StrainTokenization:\n" + f" token_size: {self.num_bins_per_token} bins\n" + f" tokens per detector: {self.num_tokens_per_detector}\n" + f" drop last token: {self.drop_last_token}\n" + f" first token width: {self.f_min_per_token[1] - self.f_min_per_token[0]:.3f} Hz\n" + f" last token width: {self.f_min_per_token[-1] - self.f_min_per_token[-2]:.3f} Hz" + ) + if self.num_padded_f_bins > 0: + print(f" zero-padded bins in last token: {self.num_padded_f_bins}") + + def __call__(self, input_sample: dict) -> dict: + """ + Parameters + ---------- + input_sample: + Must contain: + - 'waveform': array of shape [..., num_detectors, num_channels, num_bins] + - 'asds': dict {detector_name: asd_array} used to read detector order + + Returns + ------- + dict with keys 'waveform', 'position', 'drop_token_mask' (see class docstring). + """ + sample = input_sample.copy() + strain = sample["waveform"] + *batch_dims, num_blocks, num_channels, _ = strain.shape + + # (0) Cut or zero-pad the frequency axis to a multiple of num_bins_per_token + target_bins = self.num_tokens_per_detector * self.num_bins_per_token + if self.num_padded_f_bins == 0: + strain = strain[..., :target_bins] + else: + pad = [(0, 0)] * (strain.ndim - 1) + [(0, self.num_padded_f_bins)] + strain = np.pad(strain, pad, mode="constant") + + # (1) Split frequency axis into tokens: + # [..., D, C, F] → [..., D, C, T, P] + strain = strain.reshape( + *batch_dims, + num_blocks, + num_channels, + self.num_tokens_per_detector, + self.num_bins_per_token, + ) + + # (2) Move channels before tokens: + # [..., D, C, T, P] → [..., D, T, C, P] + strain = np.moveaxis(strain, source=-2, destination=-3) + + # (3) Flatten block + token, and channel + bin into the final two axes: + # [..., D, T, C, P] → [..., D*T, C*P] + sample["waveform"] = strain.reshape( + *batch_dims, + num_blocks * self.num_tokens_per_detector, + num_channels * self.num_bins_per_token, + ) + + # Position: [f_min, f_max, detector_index] per token + num_tokens = num_blocks * self.num_tokens_per_detector + token_f_min = np.tile(self.f_min_per_token, num_blocks) + token_f_max = np.tile(self.f_max_per_token, num_blocks) + detector_indices = np.array( + [DETECTOR_DICT[k] for k in input_sample["asds"]], dtype=strain.dtype + ) + token_detector = np.repeat(detector_indices, self.num_tokens_per_detector) + token_position = np.stack([token_f_min, token_f_max, token_detector], axis=-1) + + if batch_dims: + token_position = np.broadcast_to( + token_position, (*batch_dims, num_tokens, 3) + ).copy() + + sample["position"] = token_position + sample["drop_token_mask"] = np.zeros((*batch_dims, num_tokens), dtype=bool) + + return sample + + +def _check_mfd_node_compatibility( + f_mins: np.ndarray, + f_maxs: np.ndarray, + mfd_nodes: np.ndarray, + drop_last_token: bool, +) -> None: + """ + Verify that every MFD node falls in a gap between consecutive tokens, not inside + a token. This is required so that all bins within a token share the same delta_f. + + Each node must lie in (f_max[i-1], f_min[i]) for some i. + """ + left_bounds = np.concatenate([[0], f_maxs[:-1]]) + right_bounds = f_mins + intervals = np.stack([left_bounds, right_bounds], axis=1) + + covered = np.any( + (mfd_nodes[:, None] >= intervals[:, 0]) + & (mfd_nodes[:, None] <= intervals[:, 1]), + axis=1, + ) + + # The last node may lie beyond the last token's f_max when not dropping the last token + if not covered[-1] and (mfd_nodes[~covered][0] > f_maxs[-1] or not drop_last_token): + covered[-1] = True + + if not np.all(covered): + raise ValueError( + f"MFD nodes {mfd_nodes[~covered]} fall within a token rather than " + f"between tokens. Adjust token_size or MFD nodes." + ) diff --git a/examples/transformer_model/train_settings.yaml b/examples/transformer_model/train_settings.yaml new file mode 100644 index 00000000..98ea5bd6 --- /dev/null +++ b/examples/transformer_model/train_settings.yaml @@ -0,0 +1,102 @@ +data: + waveform_dataset_path: training_data/waveform_dataset.hdf5 + train_fraction: 0.95 + window: + type: tukey + f_s: 4096 + T: 8.0 + roll_off: 0.4 + detectors: + - H1 + - L1 + - V1 + extrinsic_prior: + dec: default + ra: default + geocent_time: bilby.core.prior.Uniform(minimum=-0.1, maximum=0.1) + psi: default + luminosity_distance: bilby.core.prior.Uniform(minimum=100.0, maximum=6000.0) # Mpc + ref_time: 1126259462.391 + inference_parameters: + - chirp_mass + - mass_ratio + - a_1 + - a_2 + - tilt_1 + - tilt_2 + - phi_12 + - phi_jl + - theta_jn + - luminosity_distance + - geocent_time + - ra + - dec + - psi + # Tokenized strain representation: the frequency axis is split into fixed-size + # tokens with per-token position information (f_min, f_max, detector). + tokenization: + token_size: 16 + +model: + distribution: + type: normalizing_flow + kwargs: + num_flow_steps: 30 + base_transform_kwargs: + hidden_dim: 512 + num_transform_blocks: 5 + activation: elu + batch_norm: False + # The dense_residual conditioner injects context into every residual + # block via GLUs and supports layer_norm; it is a different architecture + # from the default glasflow_residual conditioner. + conditioner_type: dense_residual + layer_norm: True + dropout_probability: 0.0 + num_bins: 8 + base_transform_type: rq-coupling + embedding_net: + type: transformer + kwargs: + tokenizer_kwargs: + hidden_dims: [512] + activation: elu + batch_norm: False + layer_norm: True + transformer_kwargs: + d_model: 1024 + dim_feedforward: 2048 + nhead: 16 + dropout: 0.0 + num_layers: 8 + norm_first: True + pooling: cls + final_net_kwargs: + activation: elu + output_dim: 128 + +training: + stage_0: + epochs: 300 + asd_dataset_path: training_data/asd_dataset/asds_O3.hdf5 + optimizer: + type: adamw + lr: 0.0001 + scheduler: + type: reduce_on_plateau + mode: min + factor: 0.5 + patience: 10 + early_stopping: + patience: 30 # Has to be larger than the ReduceLROnPlateau patience. + delta: 0.0 + metric: validation # one of ['training', 'validation'] + batch_size: 4096 + +local: + device: cuda + num_workers: 8 + runtime_limits: + max_time_per_run: 1_000_000 + max_epochs_per_run: 500 + checkpoint_epochs: 25 diff --git a/tests/core/test_build_model.py b/tests/core/test_build_model.py index 75225b5e..f521634b 100644 --- a/tests/core/test_build_model.py +++ b/tests/core/test_build_model.py @@ -588,3 +588,130 @@ def test_save_and_rebuild_from_file(tmp_path, posterior_model_type): assert type(pm_loaded) is type(pm) for p0, p1 in zip(pm.network.parameters(), pm_loaded.network.parameters()): assert torch.equal(p0.data, p1.data) + + +# ----------------------------------------------------------------------------------- +# Transformer embedding through the generic build path +# ----------------------------------------------------------------------------------- + +NUM_TOKENS = 6 +NUM_FEATURES = 12 +NUM_BLOCKS = 2 + + +def transformer_embedding_settings(): + return { + "type": "transformer", + "kwargs": { + "tokenizer_kwargs": { + "hidden_dims": [16], + "activation": "elu", + "batch_norm": False, + "layer_norm": True, + }, + "transformer_kwargs": { + "d_model": 16, + "dim_feedforward": 32, + "nhead": 4, + "dropout": 0.0, + "num_layers": 1, + "norm_first": True, + }, + "pooling": "cls", + "final_net_kwargs": { + "activation": "elu", + "output_dim": EMBEDDING_OUTPUT_DIM, + }, + }, + } + + +def tokenized_data_sample(with_context_parameters=False): + position = np.stack( + [ + np.linspace(20.0, 100.0, NUM_TOKENS), + np.linspace(30.0, 110.0, NUM_TOKENS), + np.repeat(np.arange(NUM_BLOCKS), NUM_TOKENS // NUM_BLOCKS), + ], + axis=-1, + ).astype(np.float32) + sample = { + "inference_parameters": np.random.rand(NUM_PARAMETERS).astype(np.float32), + "waveform": np.random.rand(NUM_TOKENS, NUM_FEATURES).astype(np.float32), + "position": position, + "drop_token_mask": np.zeros(NUM_TOKENS, dtype=bool), + } + if with_context_parameters: + sample["context_parameters"] = np.random.rand(GNPE_PROXY_DIM).astype(np.float32) + return sample + + +def tokenized_batch(with_context_parameters=False): + sample = tokenized_data_sample(with_context_parameters) + context = { + k: torch.from_numpy(np.stack([v] * BATCH_SIZE)) + for k, v in sample.items() + if k != "inference_parameters" + } + theta = torch.rand(BATCH_SIZE, NUM_PARAMETERS) + return theta, context + + +@pytest.mark.parametrize("model_type", ["normalizing_flow", "flow_matching"]) +def test_transformer_embedding_composes_with_any_distribution(model_type): + """The transformer works with any registered distribution type through the + generic build path — including flow matching, which the original transformer + branch never wired up.""" + settings = model_settings(model_type, completed=False) + model = settings["train_settings"]["model"] + model["embedding_net"] = transformer_embedding_settings() + + model = complete_model_settings(model, tokenized_data_sample()) + assert model["embedding_net"]["kwargs"]["tokenizer_kwargs"]["num_blocks"] == ( + NUM_BLOCKS + ) + assert model["distribution"]["kwargs"]["context_dim"] == EMBEDDING_OUTPUT_DIM + + pm = build_model_from_kwargs( + settings={"train_settings": {"model": model}}, device="cpu" + ) + theta, context = tokenized_batch() + + loss = pm.loss(theta, context) + assert torch.isfinite(loss) + + pm.network.eval() + with torch.no_grad(): + samples = pm.sample(context, num_samples=2) + assert samples.shape == (BATCH_SIZE, 2, NUM_PARAMETERS) + + +def test_transformer_embedding_with_context_parameters(): + """Context parameters compose with the transformer via the generic concat + merger — on the original branch this was impossible (the batch slot for + proxies was occupied by the position tensor).""" + settings = model_settings("normalizing_flow", completed=False) + model = settings["train_settings"]["model"] + model["embedding_net"] = transformer_embedding_settings() + + model = complete_model_settings( + model, tokenized_data_sample(with_context_parameters=True) + ) + assert model["context_merger"]["kwargs"]["num_context_parameters"] == ( + GNPE_PROXY_DIM + ) + assert model["distribution"]["kwargs"]["context_dim"] == ( + EMBEDDING_OUTPUT_DIM + GNPE_PROXY_DIM + ) + + pm = build_model_from_kwargs( + settings={"train_settings": {"model": model}}, device="cpu" + ) + assert pm.network.context_keys == ( + "waveform", + "position", + "drop_token_mask", + "context_parameters", + ) + theta, context = tokenized_batch(with_context_parameters=True) + assert torch.isfinite(pm.loss(theta, context)) diff --git a/tests/core/test_nsf.py b/tests/core/test_nsf.py index fd77310d..df87fcb6 100644 --- a/tests/core/test_nsf.py +++ b/tests/core/test_nsf.py @@ -335,3 +335,52 @@ def test_mlp_context_merger(data_setup_nsf_small): ) == 12 ) + + +def test_dense_residual_conditioner(data_setup_nsf_small): + """The rq-coupling conditioner network is an explicit type: dense_residual + (GLU context, optional layer_norm) builds a working flow; layer_norm without + it, or unknown types, are errors.""" + from dingo.core.nn.resnet import DenseResidualNet as DingoDenseResidualNet + + d = data_setup_nsf_small + kwargs = dict(d.nde_kwargs) + kwargs["base_transform_kwargs"] = { + **d.base_transform_kwargs, + "batch_norm": False, + "conditioner_type": "dense_residual", + "layer_norm": True, + } + flow = create_nsf_model(**kwargs) + # The conditioner networks inside the coupling transforms are dingo's + # DenseResidualNet, one per flow step. + conditioners = [m for m in flow.modules() if isinstance(m, DingoDenseResidualNet)] + assert len(conditioners) == d.num_flow_steps + + context_vector = torch.rand(d.batch_size, d.context_dim) + log_prob = flow.log_prob(d.y, context_vector) + assert log_prob.shape == (d.batch_size,) + assert torch.isfinite(log_prob).all() + + # layer_norm requires the dense_residual conditioner. + with pytest.raises(ValueError, match="layer_norm"): + create_nsf_model( + **{ + **d.nde_kwargs, + "base_transform_kwargs": { + **d.base_transform_kwargs, + "layer_norm": True, + }, + } + ) + # Unknown conditioner types are an error. + with pytest.raises(ValueError, match="conditioner_type"): + create_nsf_model( + **{ + **d.nde_kwargs, + "base_transform_kwargs": { + **d.base_transform_kwargs, + "conditioner_type": "foo", + }, + } + ) diff --git a/tests/core/test_resnet.py b/tests/core/test_resnet.py new file mode 100644 index 00000000..5051b55b --- /dev/null +++ b/tests/core/test_resnet.py @@ -0,0 +1,125 @@ +import pytest +import torch +from torch.nn import functional as F + +from dingo.core.nn.resnet import DenseResidualNet, LinearLayer, MyResidualBlock +from testutils_enets import check_model_forward_pass, check_model_backward_pass + + +def test_forward_pass_of_LinearLayer(): + batch_size, input_dim, output_dim = 10, 16, 4 + layer = LinearLayer(input_dim=input_dim, output_dim=output_dim, activation=F.elu) + check_model_forward_pass(layer, [output_dim], [input_dim], batch_size) + + +def test_backward_pass_of_LinearLayer(): + batch_size, input_dim, output_dim = 10, 16, 4 + layer = LinearLayer(input_dim=input_dim, output_dim=output_dim, activation=F.elu) + check_model_backward_pass(layer, [input_dim], batch_size) + + +def test_forward_pass_of_DenseResidualNet(): + """Forward pass with plain 2D [batch, features] input.""" + batch_size = 100 + input_dim, output_dim, hidden_dims = 120, 8, (128, 64, 32, 64, 16, 16) + enet = DenseResidualNet(input_dim, output_dim, hidden_dims) + check_model_forward_pass(enet, [output_dim], [input_dim], batch_size) + + +def test_backward_pass_of_DenseResidualNet(): + """Backward pass / optimizer step with plain 2D [batch, features] input.""" + batch_size = 100 + input_dim, output_dim, hidden_dims = 120, 8, (128, 64, 32, 64, 16, 16) + enet = DenseResidualNet(input_dim, output_dim, hidden_dims) + check_model_backward_pass(enet, [input_dim], batch_size) + + +def test_forward_pass_with_3d_input(): + """Forward pass with token-batched [batch, tokens, features] input, as used by + the transformer tokenizer. batch_norm is disabled since nn.BatchNorm1d treats + dim 1 as the channel axis, which for 3D input is the token axis, not features; + only layer_norm supports 3D input.""" + batch_size, num_tokens = 100, 7 + input_dim, output_dim, hidden_dims = 120, 8, (64, 32, 64) + enet = DenseResidualNet( + input_dim, output_dim, hidden_dims, batch_norm=False, layer_norm=True + ) + x = torch.rand(batch_size, num_tokens, input_dim) + y = enet(x) + assert y.shape == (batch_size, num_tokens, output_dim) + + +def test_layer_norm_and_batch_norm_are_mutually_exclusive(): + with pytest.raises(ValueError): + MyResidualBlock(features=16, use_batch_norm=True, use_layer_norm=True) + + +def test_layer_norm_runs_and_normalizes(): + """With layer_norm enabled, the normalized pre-activation within the block should + have ~zero mean and ~unit variance across the feature dimension.""" + block = MyResidualBlock(features=32, use_batch_norm=False, use_layer_norm=True) + x = torch.rand(50, 32) * 100 + 1000 # large offset/scale to make norm effect clear + normalized = block.layer_norm_layers[0](x) + assert torch.allclose(normalized.mean(dim=-1), torch.zeros(50), atol=1e-5) + assert torch.allclose( + normalized.std(dim=-1, unbiased=False), torch.ones(50), atol=1e-3 + ) + + +def test_context_glu_does_not_mix_across_tokens(): + """Regression test for the dim=1 -> dim=-1 GLU fix. + + With 3D [batch, tokens, features] input, each token's output must depend only on + its own context, not on other tokens' context. The original glasflow ResidualBlock + used dim=1 for the GLU, which for 3D input operates over the token axis instead of + the feature axis, leaking context across tokens. + """ + torch.manual_seed(0) + features, context_features, num_tokens, batch_size = 16, 4, 3, 5 + block = MyResidualBlock(features=features, context_features=context_features) + block.eval() + + x = torch.rand(batch_size, num_tokens, features) + context = torch.rand(batch_size, num_tokens, context_features) + + out_reference = block(x, context=context) + + # Change only the context of token 1; token 0 and token 2 outputs must be + # unaffected if the GLU correctly operates per-token along the feature axis. + context_modified = context.clone() + context_modified[:, 1, :] = torch.rand(batch_size, context_features) + out_modified = block(x, context=context_modified) + + assert torch.allclose(out_reference[:, 0, :], out_modified[:, 0, :]) + assert torch.allclose(out_reference[:, 2, :], out_modified[:, 2, :]) + assert not torch.allclose(out_reference[:, 1, :], out_modified[:, 1, :]) + + +def test_residual_block_equivalent_to_glasflow(): + """With layer_norm=False, MyResidualBlock is a drop-in replacement for + glasflow's ResidualBlock: identical parameter names (so old checkpoints load) + and an identical forward pass.""" + from glasflow.nflows.nn.nets.resnet import ResidualBlock + + features, context_features, batch_size = 12, 3, 7 + glasflow_block = ResidualBlock( + features=features, + context_features=context_features, + activation=F.elu, + use_batch_norm=True, + ) + block = MyResidualBlock( + features=features, + context_features=context_features, + activation=F.elu, + use_batch_norm=True, + ) + state_dict = glasflow_block.state_dict() + assert list(block.state_dict()) == list(state_dict) + block.load_state_dict(state_dict) + + block.eval() + glasflow_block.eval() + x = torch.rand(batch_size, features) + context = torch.rand(batch_size, context_features) + assert torch.allclose(block(x, context), glasflow_block(x, context)) diff --git a/tests/core/test_transformer.py b/tests/core/test_transformer.py new file mode 100644 index 00000000..e0d09801 --- /dev/null +++ b/tests/core/test_transformer.py @@ -0,0 +1,436 @@ +import copy + +import pytest +import torch +from torch.nn import functional as F + +from dingo.core.nn.resnet import DenseResidualNet, LinearLayer +from dingo.core.nn.transformer import Tokenizer, TransformerEmbedding, TransformerModel + + +def make_embedding( + tokenizer_kwargs, transformer_kwargs, pooling="cls", final_net_kwargs=None +): + """Build a TransformerEmbedding, deriving output_dim as complete_settings does.""" + if final_net_kwargs is not None: + output_dim = final_net_kwargs["output_dim"] + else: + output_dim = transformer_kwargs["d_model"] + return TransformerEmbedding( + tokenizer_kwargs=tokenizer_kwargs, + transformer_kwargs=transformer_kwargs, + output_dim=output_dim, + pooling=pooling, + final_net_kwargs=final_net_kwargs, + ) + + +NUM_TOKENS = 6 +NUM_FEATURES = 12 +NUM_BLOCKS = 2 +OUTPUT_DIM = 8 + + +def make_tokenizer(num_blocks=NUM_BLOCKS, layer_norm=False, batch_norm=False): + return Tokenizer( + input_dims=[NUM_TOKENS, NUM_FEATURES], + hidden_dims=[16, 16], + output_dim=OUTPUT_DIM, + activation=F.elu, + num_blocks=num_blocks, + layer_norm=layer_norm, + batch_norm=batch_norm, + ) + + +def make_position(batch_size, num_tokens=NUM_TOKENS, num_blocks=NUM_BLOCKS): + """position[..., 0]=f_min, position[..., 1]=f_max, position[..., 2]=detector idx.""" + f_min = torch.rand(batch_size, num_tokens) + f_max = f_min + torch.rand(batch_size, num_tokens) + detector = torch.randint(0, num_blocks, (batch_size, num_tokens)).float() + return torch.stack([f_min, f_max, detector], dim=-1) + + +def test_output_shape_batched(): + tokenizer = make_tokenizer() + x = torch.rand(10, NUM_TOKENS, NUM_FEATURES) + position = make_position(batch_size=10) + out = tokenizer(x, position) + assert out.shape == (10, NUM_TOKENS, OUTPUT_DIM) + + +def test_output_shape_unbatched(): + """No leading batch dimension, exercising DenseResidualNet's support for an + arbitrary number of leading dims.""" + tokenizer = make_tokenizer() + x = torch.rand(NUM_TOKENS, NUM_FEATURES) + f_min = torch.rand(NUM_TOKENS) + f_max = f_min + torch.rand(NUM_TOKENS) + detector = torch.randint(0, NUM_BLOCKS, (NUM_TOKENS,)).float() + position = torch.stack([f_min, f_max, detector], dim=-1) + out = tokenizer(x, position) + assert out.shape == (NUM_TOKENS, OUTPUT_DIM) + + +def test_invalid_input_dims_raises(): + with pytest.raises(ValueError): + Tokenizer( + input_dims=[NUM_TOKENS, NUM_FEATURES, 1], + hidden_dims=[16], + output_dim=OUTPUT_DIM, + activation=F.elu, + num_blocks=NUM_BLOCKS, + ) + + +def test_wrong_feature_dim_raises(): + tokenizer = make_tokenizer() + x = torch.rand(10, NUM_TOKENS, NUM_FEATURES + 1) + position = make_position(batch_size=10) + with pytest.raises(ValueError): + tokenizer(x, position) + + +def test_position_affects_output(): + """Changing f_min/f_max (with x and detector held fixed) must change the output, + since the tokenizer is conditioned on position.""" + tokenizer = make_tokenizer() + tokenizer.eval() + x = torch.rand(5, NUM_TOKENS, NUM_FEATURES) + position = make_position(batch_size=5) + + out_reference = tokenizer(x, position) + + position_modified = position.clone() + position_modified[..., 0] += 1.0 # shift f_min + out_modified = tokenizer(x, position_modified) + + assert not torch.allclose(out_reference, out_modified) + + +def test_detector_one_hot_distinguishes_tokens(): + """Two tokens with identical features and f_min/f_max but different detector + indices must produce different embeddings.""" + tokenizer = make_tokenizer(num_blocks=2) + tokenizer.eval() + x = torch.rand(1, 2, NUM_FEATURES).expand(1, 2, NUM_FEATURES).clone() + x[:, 1, :] = x[:, 0, :] # identical features for both tokens + + f_min = torch.tensor([[0.5, 0.5]]) + f_max = torch.tensor([[0.8, 0.8]]) + detector = torch.tensor([[0.0, 1.0]]) # only detector index differs + position = torch.stack([f_min, f_max, detector], dim=-1) + + out = tokenizer(x, position) + assert not torch.allclose(out[:, 0, :], out[:, 1, :]) + + +def test_position_does_not_mix_across_tokens(): + """Regression test mirroring the MyResidualBlock GLU fix: changing one token's + position must not affect another token's output.""" + tokenizer = make_tokenizer() + tokenizer.eval() + x = torch.rand(4, NUM_TOKENS, NUM_FEATURES) + position = make_position(batch_size=4) + + out_reference = tokenizer(x, position) + + position_modified = position.clone() + position_modified[:, 1, :] = make_position(batch_size=4)[:, 1, :] + out_modified = tokenizer(x, position_modified) + + assert torch.allclose(out_reference[:, 0, :], out_modified[:, 0, :]) + assert torch.allclose(out_reference[:, 2:, :], out_modified[:, 2:, :]) + assert not torch.allclose(out_reference[:, 1, :], out_modified[:, 1, :]) + + +def test_backward_pass(): + tokenizer = make_tokenizer(layer_norm=True) + x = torch.rand(8, NUM_TOKENS, NUM_FEATURES) + position = make_position(batch_size=8) + target = torch.rand(8, NUM_TOKENS, OUTPUT_DIM) + loss_fn = torch.nn.L1Loss() + + out_0 = tokenizer(x, position) + loss_before = loss_fn(out_0, target) + optimizer = torch.optim.Adam(tokenizer.parameters(), lr=0.001) + loss_before.backward() + optimizer.step() + + out_1 = tokenizer(x, position) + loss_after = loss_fn(out_1, target) + assert loss_after < loss_before + + +# --------------------------------------------------------------------------- +# TransformerEmbedding +# --------------------------------------------------------------------------- + +D_MODEL = 16 + + +def make_enet_kwargs(): + tokenizer_kwargs = { + "input_dims": [NUM_TOKENS, NUM_FEATURES], + "num_blocks": NUM_BLOCKS, + "hidden_dims": [16], + "activation": "elu", + "batch_norm": False, + "layer_norm": True, + } + transformer_kwargs = { + "d_model": D_MODEL, + "dim_feedforward": 32, + "nhead": 4, + "dropout": 0.0, + "num_layers": 2, + "norm_first": True, + } + return tokenizer_kwargs, transformer_kwargs + + +def make_enet_inputs(batch_size): + x = torch.rand(batch_size, NUM_TOKENS, NUM_FEATURES) + position = make_position(batch_size=batch_size) + return x, position + + +def test_create_transformer_enet_default_pooling_is_cls(): + tokenizer_kwargs, transformer_kwargs = make_enet_kwargs() + model = make_embedding( + tokenizer_kwargs=tokenizer_kwargs, transformer_kwargs=transformer_kwargs + ) + assert model.pooling == "cls" + assert hasattr(model, "class_token") + + +@pytest.mark.parametrize("pooling", ["cls", "average"]) +def test_create_transformer_enet_without_final_net(pooling): + """If final_net_kwargs is None, the pooled d_model-dim vector is returned as is.""" + tokenizer_kwargs, transformer_kwargs = make_enet_kwargs() + model = make_embedding( + tokenizer_kwargs=tokenizer_kwargs, + transformer_kwargs=transformer_kwargs, + pooling=pooling, + ) + assert model.final_net is None + + x, position = make_enet_inputs(batch_size=5) + out = model(x=x, position=position) + assert out.shape == (5, D_MODEL) + + +def test_create_transformer_enet_with_linear_final_net(): + """final_net_kwargs without hidden_dims builds a LinearLayer.""" + tokenizer_kwargs, transformer_kwargs = make_enet_kwargs() + final_net_kwargs = {"activation": "elu", "output_dim": 5} + model = make_embedding( + tokenizer_kwargs=tokenizer_kwargs, + transformer_kwargs=transformer_kwargs, + final_net_kwargs=final_net_kwargs, + ) + assert isinstance(model.final_net, LinearLayer) + + x, position = make_enet_inputs(batch_size=5) + out = model(x=x, position=position) + assert out.shape == (5, 5) + + +def test_create_transformer_enet_with_dense_residual_final_net(): + """final_net_kwargs with hidden_dims builds a DenseResidualNet.""" + tokenizer_kwargs, transformer_kwargs = make_enet_kwargs() + final_net_kwargs = { + "activation": "elu", + "output_dim": 5, + "hidden_dims": [8, 8], + "layer_norm": True, + } + model = make_embedding( + tokenizer_kwargs=tokenizer_kwargs, + transformer_kwargs=transformer_kwargs, + final_net_kwargs=final_net_kwargs, + ) + assert isinstance(model.final_net, DenseResidualNet) + + x, position = make_enet_inputs(batch_size=5) + out = model(x=x, position=position) + assert out.shape == (5, 5) + + +def test_create_transformer_enet_does_not_mutate_input_kwargs(): + """Settings dicts (e.g., loaded once from yaml and reused across training stages) + must not be mutated by repeated calls.""" + tokenizer_kwargs, transformer_kwargs = make_enet_kwargs() + final_net_kwargs = {"activation": "elu", "output_dim": 5} + tokenizer_kwargs_ref = copy.deepcopy(tokenizer_kwargs) + final_net_kwargs_ref = copy.deepcopy(final_net_kwargs) + + make_embedding( + tokenizer_kwargs=tokenizer_kwargs, + transformer_kwargs=transformer_kwargs, + final_net_kwargs=final_net_kwargs, + ) + # second call with the same (unmutated) dicts must not raise + make_embedding( + tokenizer_kwargs=tokenizer_kwargs, + transformer_kwargs=transformer_kwargs, + final_net_kwargs=final_net_kwargs, + ) + + assert tokenizer_kwargs == tokenizer_kwargs_ref + assert final_net_kwargs == final_net_kwargs_ref + + +def test_invalid_pooling_raises(): + tokenizer_kwargs, transformer_kwargs = make_enet_kwargs() + with pytest.raises(ValueError, match="pooling"): + make_embedding( + tokenizer_kwargs=tokenizer_kwargs, + transformer_kwargs=transformer_kwargs, + pooling="max", + ) + + +# --------------------------------------------------------------------------- +# TransformerModel — src_key_padding_mask (drop-token masking) +# --------------------------------------------------------------------------- + + +def make_full_enet(pooling="cls"): + tokenizer_kwargs, transformer_kwargs = make_enet_kwargs() + final_net_kwargs = {"activation": "elu", "output_dim": 5} + return make_embedding( + tokenizer_kwargs=tokenizer_kwargs, + transformer_kwargs=transformer_kwargs, + pooling=pooling, + final_net_kwargs=final_net_kwargs, + ) + + +def test_padding_mask_does_not_change_output_shape(): + """Output shape must be identical whether or not a padding mask is supplied.""" + model = make_full_enet(pooling="cls") + model.eval() + x, position = make_enet_inputs(batch_size=4) + mask = torch.zeros(4, NUM_TOKENS, dtype=torch.bool) + mask[:, -2:] = True # mask last two tokens + + out_no_mask = model(x=x, position=position) + out_masked = model(x=x, position=position, src_key_padding_mask=mask) + assert out_no_mask.shape == out_masked.shape + + +def test_padding_mask_changes_output(): + """Masking some tokens must change the CLS-pooled output.""" + model = make_full_enet(pooling="cls") + model.eval() + x, position = make_enet_inputs(batch_size=4) + mask = torch.zeros(4, NUM_TOKENS, dtype=torch.bool) + mask[:, -2:] = True + + out_no_mask = model(x=x, position=position) + out_masked = model(x=x, position=position, src_key_padding_mask=mask) + assert not torch.allclose(out_no_mask, out_masked) + + +def test_average_pooling_ignores_fully_masked_token(): + """For average pooling, a fully masked token should not affect the result.""" + model = make_full_enet(pooling="average") + model.eval() + x, position = make_enet_inputs(batch_size=2) + + # mask with last token dropped + mask_drop = torch.zeros(2, NUM_TOKENS, dtype=torch.bool) + mask_drop[:, -1] = True + + out_drop = model(x=x, position=position, src_key_padding_mask=mask_drop) + + # Replace the last token's features with noise — output should be unchanged + x_noisy = x.clone() + x_noisy[:, -1, :] = torch.rand_like(x_noisy[:, -1, :]) * 1e3 + out_noisy = model(x=x_noisy, position=position, src_key_padding_mask=mask_drop) + + assert torch.allclose(out_drop, out_noisy, atol=1e-5) + + +def test_create_transformer_enet_backward_pass(): + tokenizer_kwargs, transformer_kwargs = make_enet_kwargs() + final_net_kwargs = {"activation": "elu", "output_dim": 5} + model = make_embedding( + tokenizer_kwargs=tokenizer_kwargs, + transformer_kwargs=transformer_kwargs, + final_net_kwargs=final_net_kwargs, + ) + + x, position = make_enet_inputs(batch_size=8) + target = torch.rand(8, 5) + loss_fn = torch.nn.L1Loss() + + out_0 = model(x=x, position=position) + loss_before = loss_fn(out_0, target) + optimizer = torch.optim.Adam(model.parameters(), lr=0.001) + loss_before.backward() + optimizer.step() + + out_1 = model(x=x, position=position) + loss_after = loss_fn(out_1, target) + assert loss_after < loss_before + + +# --------------------------------------------------------------------------- +# Registered-embedding contract +# --------------------------------------------------------------------------- + + +def test_transformer_embedding_contract(): + """TransformerEmbedding follows the embedding contract: registered name, + declared input_keys, output_dim, and complete_settings inferring + tokenizer dims / num_blocks / output_dim from a sample batch.""" + from dingo.core.registry import EMBEDDING_NETS + + assert EMBEDDING_NETS.get("transformer") is TransformerEmbedding + assert TransformerEmbedding.input_keys == ( + "waveform", + "position", + "drop_token_mask", + ) + + tokenizer_kwargs, transformer_kwargs = make_enet_kwargs() + user_tokenizer_kwargs = { + k: v + for k, v in tokenizer_kwargs.items() + if k not in ("input_dims", "num_blocks") + } + settings = { + "tokenizer_kwargs": user_tokenizer_kwargs, + "transformer_kwargs": transformer_kwargs, + "pooling": "cls", + "final_net_kwargs": {"activation": "elu", "output_dim": 5}, + } + sample_batch = { + "waveform": torch.rand(NUM_TOKENS, NUM_FEATURES), + "position": make_position(batch_size=1)[0], + "drop_token_mask": torch.zeros(NUM_TOKENS, dtype=torch.bool), + } + completed = TransformerEmbedding.complete_settings(settings, sample_batch) + assert completed["tokenizer_kwargs"]["input_dims"] == [NUM_TOKENS, NUM_FEATURES] + assert completed["tokenizer_kwargs"]["num_blocks"] == NUM_BLOCKS + assert completed["output_dim"] == 5 + # User settings are not modified. + assert "input_dims" not in settings["tokenizer_kwargs"] + + model = TransformerEmbedding(**completed) + assert model.output_dim == 5 + x, position = make_enet_inputs(batch_size=3) + mask = torch.zeros(3, NUM_TOKENS, dtype=torch.bool) + out = model(x, position, mask) + assert out.shape == (3, 5) + + # Dims in user settings are an error. + with pytest.raises(ValueError, match="derived from the data"): + TransformerEmbedding.complete_settings( + {**settings, "tokenizer_kwargs": tokenizer_kwargs}, sample_batch + ) + # Inconsistent output_dim is an error. + with pytest.raises(ValueError, match="Inconsistent"): + TransformerEmbedding(**{**completed, "output_dim": 7}) diff --git a/tests/gw/inference/__init__.py b/tests/gw/inference/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/gw/inference/test_gw_sampler_transforms.py b/tests/gw/inference/test_gw_sampler_transforms.py new file mode 100644 index 00000000..9018bf32 --- /dev/null +++ b/tests/gw/inference/test_gw_sampler_transforms.py @@ -0,0 +1,201 @@ +"""Tests for GWSampler._initialize_transforms with and without tokenization.""" + +import numpy as np +import torch +from unittest.mock import MagicMock + +from dingo.core.transforms import GetItem +from dingo.gw.domains import UniformFrequencyDomain +from dingo.gw.inference.gw_samplers import GWSampler +from dingo.gw.transforms import ( + StrainTokenization, + SelectKeys, + ToTorch, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +DETECTORS = ["H1", "L1"] +STANDARDIZATION = { + "mean": {"chirp_mass": 20.0}, + "std": {"chirp_mass": 5.0}, +} +INFERENCE_PARAMS = ["chirp_mass"] + + +def _make_domain(f_min=20.0, f_max=128.0, delta_f=0.25): + return UniformFrequencyDomain(f_min=f_min, f_max=f_max, delta_f=delta_f) + + +def _make_sampler_stub(domain, tokenization_settings=None): + """Return a GWSampler with the minimum attributes set to call _initialize_transforms. + + Uses object.__setattr__ to bypass Sampler.__init__, so no real model or dataset + is needed. + """ + data_settings = { + "detectors": DETECTORS, + "standardization": STANDARDIZATION, + "ref_time": 1126259462.391, + } + if tokenization_settings is not None: + data_settings["tokenization"] = tokenization_settings + + mock_model = MagicMock() + mock_model.device = torch.device("cpu") + + sampler = object.__new__(GWSampler) + sampler.domain = domain + sampler.model = mock_model + metadata = { + "train_settings": {"data": data_settings}, + "dataset_settings": {"intrinsic_prior": {}}, + } + sampler.metadata = metadata + # GWSamplerMixin.detectors reads from base_model_metadata (== metadata for non-GNPE). + sampler.base_model_metadata = metadata + sampler.inference_parameters = INFERENCE_PARAMS + sampler._minimum_frequency = None + sampler._maximum_frequency = None + return sampler + + +def _make_context(domain, rng=None): + """Build a minimal {'waveform': ..., 'asds': ...} dict for *domain*.""" + if rng is None: + rng = np.random.default_rng(0) + n = len(domain.sample_frequencies) + return { + "waveform": { + d: (rng.standard_normal(n) + 1j * rng.standard_normal(n)).astype( + np.complex64 + ) + for d in DETECTORS + }, + "asds": {d: rng.uniform(1e-24, 1e-23, n).astype(np.float32) for d in DETECTORS}, + } + + +# --------------------------------------------------------------------------- +# _initialize_transforms — resnet (no tokenization) +# --------------------------------------------------------------------------- + + +def test_resnet_path_uses_get_item(): + domain = _make_domain() + sampler = _make_sampler_stub(domain, tokenization_settings=None) + sampler._initialize_transforms() + + transforms = sampler.transform_pre.transforms + assert isinstance(transforms[-1], GetItem) + assert transforms[-1].key == "waveform" + + +def test_resnet_path_has_no_strain_tokenization(): + domain = _make_domain() + sampler = _make_sampler_stub(domain, tokenization_settings=None) + sampler._initialize_transforms() + + types = [type(t) for t in sampler.transform_pre.transforms] + assert StrainTokenization not in types + assert SelectKeys not in types + + +def test_resnet_path_output_is_tensor(): + domain = _make_domain() + sampler = _make_sampler_stub(domain) + sampler._initialize_transforms() + + context = _make_context(domain) + x = sampler.transform_pre(context) + assert isinstance(x, torch.Tensor) + + +# --------------------------------------------------------------------------- +# _initialize_transforms — transformer (with tokenization) +# --------------------------------------------------------------------------- + +TOK_SETTINGS = { + "token_size": 16, + "num_tokens_per_block": None, + "drop_last_token": False, +} + + +def test_transformer_path_has_strain_tokenization(): + domain = _make_domain() + sampler = _make_sampler_stub(domain, tokenization_settings=TOK_SETTINGS) + sampler._initialize_transforms() + + types = [type(t) for t in sampler.transform_pre.transforms] + assert StrainTokenization in types + + +def test_transformer_path_has_select_keys(): + domain = _make_domain() + sampler = _make_sampler_stub(domain, tokenization_settings=TOK_SETTINGS) + sampler._initialize_transforms() + + types = [type(t) for t in sampler.transform_pre.transforms] + assert SelectKeys in types + + +def test_transformer_path_no_get_item(): + domain = _make_domain() + sampler = _make_sampler_stub(domain, tokenization_settings=TOK_SETTINGS) + sampler._initialize_transforms() + + types = [type(t) for t in sampler.transform_pre.transforms] + assert GetItem not in types + + +def test_transformer_tokenization_precedes_to_torch(): + """StrainTokenization must come before ToTorch in the chain.""" + domain = _make_domain() + sampler = _make_sampler_stub(domain, tokenization_settings=TOK_SETTINGS) + sampler._initialize_transforms() + + transforms = sampler.transform_pre.transforms + indices = {type(t): i for i, t in enumerate(transforms)} + assert indices[StrainTokenization] < indices[ToTorch] + + +def test_transformer_select_keys_follows_to_torch(): + """SelectKeys must come after ToTorch in the chain.""" + domain = _make_domain() + sampler = _make_sampler_stub(domain, tokenization_settings=TOK_SETTINGS) + sampler._initialize_transforms() + + transforms = sampler.transform_pre.transforms + indices = {type(t): i for i, t in enumerate(transforms)} + assert indices[SelectKeys] > indices[ToTorch] + + +def test_transformer_path_output_is_dict_of_three_tensors(): + domain = _make_domain() + sampler = _make_sampler_stub(domain, tokenization_settings=TOK_SETTINGS) + sampler._initialize_transforms() + + context = _make_context(domain) + x = sampler.transform_pre(context) + assert isinstance(x, dict) + assert list(x) == ["waveform", "position", "drop_token_mask"] + assert all(isinstance(v, torch.Tensor) for v in x.values()) + assert x["drop_token_mask"].dtype == torch.bool + + +def test_transformer_path_waveform_and_position_num_tokens_match(): + domain = _make_domain() + sampler = _make_sampler_stub(domain, tokenization_settings=TOK_SETTINGS) + sampler._initialize_transforms() + + context = _make_context(domain) + x = sampler.transform_pre(context) + assert ( + x["waveform"].shape[0] + == x["position"].shape[0] + == x["drop_token_mask"].shape[0] + ) diff --git a/tests/gw/transforms/test_tokenization_transforms.py b/tests/gw/transforms/test_tokenization_transforms.py new file mode 100644 index 00000000..8421931a --- /dev/null +++ b/tests/gw/transforms/test_tokenization_transforms.py @@ -0,0 +1,321 @@ +import numpy as np +import pytest + +from dingo.gw.domains import UniformFrequencyDomain, MultibandedFrequencyDomain +from dingo.gw.transforms import StrainTokenization, DETECTOR_DICT + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def make_ufd(f_min=20.0, f_max=1024.0, T=8.0): + return UniformFrequencyDomain(f_min=f_min, f_max=f_max, delta_f=1.0 / T) + + +def make_mfd(nodes=None, f_max=1038.0, T=8.0): + if nodes is None: + nodes = [20.0, 34.0, 46.0, 62.0, 78.0, 1038.0] + base = UniformFrequencyDomain(f_min=nodes[0], f_max=f_max, delta_f=1.0 / T) + return MultibandedFrequencyDomain( + nodes=nodes, delta_f_initial=1.0 / T, base_domain=base + ) + + +def make_sample(domain, batch_size, num_channels=3): + """Build a minimal {'waveform': ..., 'asds': ...} dict. + + H1 waveform is all zeros; L1 real channel is set to [1, 2, ..., num_f] + so we can track ordering through the reshape. + """ + num_f = domain.frequency_mask_length + detectors = ["H1", "L1"] + + if batch_size is None: + # No batch dimension: shape [num_blocks, num_channels, num_f] + waveform_h1 = np.zeros([1, num_channels, num_f]) + waveform_l1 = np.ones([1, num_channels, num_f]) + waveform_l1[0, 0, :] = np.arange(1, num_f + 1) + waveform = np.concatenate([waveform_h1, waveform_l1], axis=0) + asds = {d: np.random.rand(num_f) for d in detectors} + else: + # Batched: shape [batch, num_blocks, num_channels, num_f] + waveform_h1 = np.zeros([batch_size, 1, num_channels, num_f]) + waveform_l1 = np.ones([batch_size, 1, num_channels, num_f]) + waveform_l1[0, 0, 0, :] = np.arange(1, num_f + 1) + waveform = np.concatenate([waveform_h1, waveform_l1], axis=1) + asds = {d: np.random.rand(batch_size, num_f) for d in detectors} + + return {"waveform": waveform, "asds": asds} + + +def _check_position_and_mask(out, domain, num_tokens_per_block, num_blocks): + """Shared checks for position and drop_token_mask outputs.""" + num_tokens = num_tokens_per_block * num_blocks + + # position shape + assert out["position"].shape[-2:] == (num_tokens, 3) + + # First token of each detector starts at domain.f_min + for block in range(num_blocks): + first_tok = block * num_tokens_per_block + assert np.all(out["position"][..., first_tok, 0] == domain.f_min) + + # Last token of each detector reaches at least domain.f_max - delta_f + if isinstance(domain, MultibandedFrequencyDomain): + f_max_threshold = domain.f_max - domain.delta_f[-1] + else: + f_max_threshold = domain.f_max - domain.delta_f + for block in range(num_blocks): + last_tok = block * num_tokens_per_block + num_tokens_per_block - 1 + assert np.all(out["position"][..., last_tok, 1] >= f_max_threshold) + + # f_min increases monotonically within each detector's tokens + for block in range(num_blocks): + tok_slice = slice( + block * num_tokens_per_block, (block + 1) * num_tokens_per_block + ) + f_mins = out["position"][..., tok_slice, 0] + # Take first batch element if batched + f_mins_1d = f_mins.reshape(-1, num_tokens_per_block)[0] + assert np.all( + np.diff(f_mins_1d) > 0 + ), "f_min is not monotonically increasing within a detector" + + # Each detector's tokens share the same detector index, and indices differ across detectors + unique_det_indices = set() + for block in range(num_blocks): + tok_slice = slice( + block * num_tokens_per_block, (block + 1) * num_tokens_per_block + ) + det_vals = np.unique(out["position"][..., tok_slice, 2]) + assert ( + len(det_vals) == 1 + ), "Tokens of a single detector have mixed detector indices" + unique_det_indices.add(det_vals[0]) + assert ( + len(unique_det_indices) == num_blocks + ), "Detector indices are not unique across blocks" + + # drop_token_mask: shape and default all-False + assert out["drop_token_mask"].shape[-1] == num_tokens + assert not out[ + "drop_token_mask" + ].any(), "Default mask should keep all tokens (all False)" + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def ufd_batch(): + domain = make_ufd() + # frequency_mask_length = 8032 (f_min=20, f_max=1024, delta_f=0.125). + # 8032 % 40 != 0, so these fixtures exercise the zero-padding path. + # 8032 has no convenient power-of-2 divisor, so an exact-division UFD case + # would require an artificially chosen domain; the MFD fixtures cover that. + num_tokens_per_block = 40 + return domain, num_tokens_per_block, 2, make_sample(domain, batch_size=100) + + +@pytest.fixture +def ufd_no_batch(): + domain = make_ufd() + num_tokens_per_block = 40 # non-divisible; see ufd_batch comment + return domain, num_tokens_per_block, 2, make_sample(domain, batch_size=None) + + +@pytest.fixture +def ufd_single_batch(): + domain = make_ufd() + num_tokens_per_block = 40 # non-divisible; see ufd_batch comment + return domain, num_tokens_per_block, 2, make_sample(domain, batch_size=1) + + +@pytest.fixture +def mfd_batch(): + # 43 tokens fits exactly for these nodes with T=8 + domain = make_mfd() + num_tokens_per_block = 43 + return domain, num_tokens_per_block, 2, make_sample(domain, batch_size=100) + + +@pytest.fixture +def mfd_drop_last_token(): + # Extend f_max slightly so the last token is incomplete + domain = make_mfd(f_max=1040.0) + num_tokens_per_block = 43 + return domain, num_tokens_per_block, 2, make_sample(domain, batch_size=100) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +SETUPS = ["ufd_batch", "ufd_no_batch", "ufd_single_batch", "mfd_batch"] + + +@pytest.mark.parametrize("setup", SETUPS) +def test_strain_tokenization_num_tokens(request, setup): + """Basic tokenization using num_tokens_per_block: shapes and content.""" + domain, num_tokens_per_block, num_blocks, sample = request.getfixturevalue(setup) + + transform = StrainTokenization( + domain, num_tokens_per_block=num_tokens_per_block, print_output=False + ) + out = transform(sample) + + # Output waveform shape + assert out["waveform"].shape[-2] == num_tokens_per_block * num_blocks + num_features = out["waveform"].shape[-1] + + # L1 sits in the second half of the token sequence; real is the first 1/3 of features. + # The linearly increasing values 1..num_f should be recoverable in order. + num_f = domain.frequency_mask_length + if out["waveform"].ndim == 2: + l1_real = out["waveform"][num_tokens_per_block:, : num_features // 3] + else: + l1_real = out["waveform"][0, num_tokens_per_block:, : num_features // 3] + assert np.all(l1_real.flatten()[:num_f] == np.arange(1, num_f + 1)) + + _check_position_and_mask(out, domain, num_tokens_per_block, num_blocks) + + +@pytest.mark.parametrize("setup", SETUPS) +def test_strain_tokenization_token_size(request, setup): + """Equivalent result when specifying token_size instead of num_tokens_per_block.""" + domain, num_tokens_per_block, num_blocks, sample = request.getfixturevalue(setup) + token_size = int(np.ceil(domain.frequency_mask_length / num_tokens_per_block)) + + transform = StrainTokenization(domain, token_size=token_size, print_output=False) + out = transform(sample) + + assert out["waveform"].shape[-2] == num_tokens_per_block * num_blocks + _check_position_and_mask(out, domain, num_tokens_per_block, num_blocks) + + +@pytest.mark.parametrize("setup", SETUPS + ["mfd_drop_last_token"]) +def test_strain_tokenization_drop_last_token(request, setup): + """drop_last_token removes the trailing incomplete token.""" + domain, num_tokens_per_block, num_blocks, sample = request.getfixturevalue(setup) + token_size = int(np.ceil(domain.frequency_mask_length / num_tokens_per_block)) + remainder = domain.frequency_mask_length % num_tokens_per_block + expected = (num_tokens_per_block - (1 if remainder else 0)) * num_blocks + + transform = StrainTokenization( + domain, token_size=token_size, drop_last_token=True, print_output=False + ) + out = transform(sample) + + assert out["waveform"].shape[-2] == expected + + +def test_token_bin_content(): + """Each token contains exactly the right frequency bins in order. + + Uses the MFD fixture (exact division, no padding) so token k of L1 contains + bins [k*P, ..., (k+1)*P - 1] with values [k*P+1, ..., (k+1)*P]. + """ + domain = make_mfd() + num_tokens_per_block = 43 + num_channels = 3 + P = domain.frequency_mask_length // num_tokens_per_block # bins per token + + sample = make_sample(domain, batch_size=100, num_channels=num_channels) + transform = StrainTokenization( + domain, num_tokens_per_block=num_tokens_per_block, print_output=False + ) + out = transform(sample) + + num_features = out["waveform"].shape[-1] + real_width = num_features // num_channels # = P + + # L1 tokens occupy indices [num_tokens_per_block, 2*num_tokens_per_block) + # Real channel is the first `real_width` features within each token + l1_real = out["waveform"][0, num_tokens_per_block:, :real_width] # [T, P] + + for k in range(num_tokens_per_block): + expected = np.arange(k * P + 1, (k + 1) * P + 1, dtype=float) + assert np.allclose( + l1_real[k], expected + ), f"Token {k} has wrong bin values: got {l1_real[k]}, expected {expected}" + + +def test_three_detectors(): + """Detector-index assignment and token ordering with three detectors (H1, L1, V1). + + Uses MFD with exact division so no zero-padding obscures the value checks. + """ + domain = make_mfd() # frequency_mask_length=688, 688/16=43 exactly + num_f = domain.frequency_mask_length + num_channels = 3 + batch_size = 4 + detectors = ["H1", "L1", "V1"] + + waveforms = [] + for i, det in enumerate(detectors): + w = np.full([batch_size, 1, num_channels, num_f], float(i)) + waveforms.append(w) + waveform = np.concatenate(waveforms, axis=1) # [B, 3, C, F] + asds = {d: np.ones([batch_size, num_f]) for d in detectors} + sample = {"waveform": waveform, "asds": asds} + + num_tokens_per_block = 43 + transform = StrainTokenization( + domain, num_tokens_per_block=num_tokens_per_block, print_output=False + ) + out = transform(sample) + + T = num_tokens_per_block + for block_idx, det in enumerate(detectors): + tok_slice = slice(block_idx * T, (block_idx + 1) * T) + + # All waveform values in this block's tokens equal float(block_idx) + assert np.all( + out["waveform"][0, tok_slice, :] == float(block_idx) + ), f"Wrong waveform values for detector {det}" + # Detector index in position matches DETECTOR_DICT + det_indices = out["position"][0, tok_slice, 2] + assert np.all( + det_indices == DETECTOR_DICT[det] + ), f"Wrong detector index for {det}: got {det_indices[0]}, expected {DETECTOR_DICT[det]}" + + +def test_output_dtype(): + """Output arrays preserve input dtype; drop_token_mask is always bool.""" + domain = make_ufd() + + for dtype in (np.float32, np.float64): + sample = make_sample(domain, batch_size=8) + sample["waveform"] = sample["waveform"].astype(dtype) + + transform = StrainTokenization( + domain, num_tokens_per_block=40, print_output=False + ) + out = transform(sample) + + assert out["waveform"].dtype == dtype, f"waveform dtype changed from {dtype}" + assert out["position"].dtype == dtype, f"position dtype changed from {dtype}" + assert out["drop_token_mask"].dtype == bool + + +def test_mutual_exclusivity(): + """Passing both or neither of num_tokens_per_block / token_size raises ValueError.""" + domain = make_ufd() + with pytest.raises(ValueError): + StrainTokenization(domain, print_output=False) + with pytest.raises(ValueError): + StrainTokenization( + domain, num_tokens_per_block=10, token_size=20, print_output=False + ) + + +def test_mfd_incompatible_nodes(): + """MFD node inside a token should raise ValueError.""" + # nodes=[20, 34, ...]: with token_size=200, a node will land inside a token + domain = make_mfd() + with pytest.raises(ValueError): + StrainTokenization(domain, token_size=200, print_output=False) From 87888a59cde2dec63bb18c590b811e1e6481bd4d Mon Sep 17 00:00:00 2001 From: Maximilian Dax Date: Sun, 12 Jul 2026 21:32:35 +0200 Subject: [PATCH 9/9] Port T1 augmentation and inference-time token suppression Stage 2 of the transformer merge, from origin/dingo-t1 (build-system step 6). Training-side augmentation for tokenized models, wired via data.tokenization: DropDetectors, DropFrequenciesToUpdateRange (f_cut), DropFrequencyInterval (mask_interval), DropRandomTokens, and NormalizePosition, all marking tokens in drop_token_mask that the transformer then does not attend to. Ported near-verbatim with two deviations: DropDetectors gains unbatched-sample support (our pipeline feeds single samples; the other drop transforms already handled this - covered by the ported no-batch test), and DropRandomTokens loses the increase_p_until_epoch ramp, which is dead on dingo-t1 as well (the epoch injection into samples is commented out there). Inference side: GWSampler gains a suppress property ([f_lo, f_hi] or per-detector dict) with real validation (dingo-t1 left it as an unimplemented warning stub): it requires a tokenized model trained with drop augmentation and an in-domain interval. Frequency-range updates and suppression for tokenized models go through the ported UpdateFrequencyRange transform (masks tokens instead of zeroing strain data), with NormalizePosition applied after mask updates, matching training. Unchanged frequency bounds are passed as None so the domain default does not needlessly mask the zero-padded final token. UpdateFrequencyRange gets the test dingo-t1 left as a TODO. PoolingTransformer (sinusoidal MultiPositionalEncoding + average pooling) registers as the "pooling_transformer" embedding, composing with any distribution through the generic build path. The transformer example config now demonstrates the augmentation settings. Not ported (out of step-6 scope, separate research features on dingo-t1): pretraining (enets_pretraining/pretraining_model), custom schedulers/AMP trainer changes, mixed ASD datasets, TimeShiftStrainGrid, BlockEncoding, skymaps. Co-Authored-By: Claude Fable 5 --- dingo/core/nn/transformer.py | 238 +++++ dingo/gw/gwutils.py | 15 + dingo/gw/inference/gw_samplers.py | 88 +- dingo/gw/training/train_builders.py | 68 ++ dingo/gw/transforms/__init__.py | 13 +- .../gw/transforms/tokenization_transforms.py | 988 ++++++++++++++++++ .../transformer_model/train_settings.yaml | 21 + tests/core/test_build_model.py | 47 + .../inference/test_gw_sampler_transforms.py | 64 ++ .../test_tokenization_augmentation.py | 708 +++++++++++++ 10 files changed, 2242 insertions(+), 8 deletions(-) create mode 100644 tests/gw/transforms/test_tokenization_augmentation.py diff --git a/dingo/core/nn/transformer.py b/dingo/core/nn/transformer.py index 09f01b6a..2288ad0b 100644 --- a/dingo/core/nn/transformer.py +++ b/dingo/core/nn/transformer.py @@ -370,3 +370,241 @@ def complete_settings(cls, settings: dict, sample_batch: dict) -> dict: "tokenizer_kwargs": tokenizer_kwargs, "output_dim": output_dim, } + + +class MultiPositionalEncoding(nn.Module): + """ + Sinusoidal positional encoding over several position quantities (e.g. f_min, + f_max, detector index), each allotted a share of the embedding dimension, added + to the token embeddings. Ported from the DINGO-T1 branch. + """ + + def __init__(self, d_model: int, max_vals: List[float], resolutions: List[float]): + """ + Parameters + ---------- + d_model : int + embedding size of the transformer + max_vals : List[float] + maximum value of each position quantity (sets the largest wavelength) + resolutions : List[float] + resolution of each position quantity (sets the smallest wavelength) + """ + super().__init__() + num_encodings = len(max_vals) + encoding_sizes = [((d_model // 2) // num_encodings) * 2] * num_encodings + encoding_sizes[-1] += d_model - sum(encoding_sizes) + if sum(encoding_sizes) != d_model: + raise ValueError( + f"Cannot partition d_model={d_model} into {num_encodings} encodings." + ) + + for i in range(num_encodings): + k = torch.arange(0, encoding_sizes[i], 2) + div_term = torch.exp( + torch.log(torch.tensor(max_vals[i] / resolutions[i])) + * (-k / encoding_sizes[i]) + ) + self.register_buffer("div_term_" + str(i), div_term) + self.num_encodings = num_encodings + + def forward(self, x: Tensor, position: Tensor): + """ + Parameters + ---------- + x: Tensor, shape ``[batch_size, seq_length, embedding_dim]`` + position: Tensor, shape ``[batch_size, seq_length, self.num_encodings]`` + """ + position = position.unsqueeze(-1) + start = 0 + pe = torch.zeros_like(x) + for i in range(self.num_encodings): + div_term = getattr(self, "div_term_" + str(i)) + end = start + 2 * len(div_term) + pe[:, :, start:end:2] = torch.sin(position[:, :, i, :] * div_term) + pe[:, :, start + 1 : end : 2] = torch.cos(position[:, :, i, :] * div_term) + start = end + return x + pe + + +class PoolingTransformer(nn.Module): + """ + Transformer encoder with additive sinusoidal positional encoding and average + pooling, used as an embedding network. In contrast to TransformerModel, the + tokenizer is an unconditional DenseResidualNet and position enters via + MultiPositionalEncoding rather than GLU conditioning. Ported from the DINGO-T1 + branch. + """ + + def __init__( + self, + tokenizer: nn.Module, + positional_encoder: nn.Module, + transformer_encoder: nn.Module, + final_net: Optional[nn.Module] = None, + ): + super().__init__() + self.tokenizer = tokenizer + self.positional_encoder = positional_encoder + self.transformer_encoder = transformer_encoder + self.final_net = final_net + + self.init_weights() + + def init_weights(self) -> None: + """ + Initialize parameters of the transformer encoder explicitly, due to + https://github.com/pytorch/pytorch/issues/72253. Parameters are initialized + with xavier uniform. + """ + for p in self.transformer_encoder.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + + def forward( + self, + src: Tensor, + position: Tensor = None, + src_key_padding_mask: Tensor = None, + ) -> Tensor: + x = self.tokenizer(src) + + if position is not None: + x = self.positional_encoder(x, position) + x = self.transformer_encoder(x, src_key_padding_mask=src_key_padding_mask) + + # Average over non-masked components. + if src_key_padding_mask is not None: + denominator = torch.sum(~src_key_padding_mask, -1, keepdim=True) + x = torch.sum(x * ~src_key_padding_mask.unsqueeze(-1), dim=-2) / denominator + else: + x = torch.mean(x, dim=-2) + + if self.final_net is not None: + x = self.final_net(x) + + return x + + +@EMBEDDING_NETS.register("pooling_transformer") +class PoolingTransformerEmbedding(PoolingTransformer): + """ + PoolingTransformer as a registered embedding network (see the contract in + dingo.core.nn.enets): consumes the tokenized batch entries produced by + StrainTokenization, and builds its modules from settings dicts. + """ + + input_keys = ("waveform", "position", "drop_token_mask") + + def __init__( + self, + tokenizer_kwargs: dict, + positional_encoder_kwargs: dict, + transformer_kwargs: dict, + output_dim: int, + final_net_kwargs: Optional[dict] = None, + ): + """ + Parameters + ---------- + tokenizer_kwargs : dict + Settings for the token embedding DenseResidualNet: hidden_dims, + activation (str), and optionally dropout, batch_norm, layer_norm. + input_dim is inferred from a sample batch by complete_settings; the + output_dim is transformer_kwargs["d_model"]. + positional_encoder_kwargs : dict + Settings for MultiPositionalEncoding: max_vals and resolutions, one + entry per position quantity (d_model is filled in automatically). + transformer_kwargs : dict + Settings for the transformer encoder: d_model, nhead, num_layers, and + optionally dim_feedforward, dropout. + output_dim : int + Dimension of the embedded context: the output_dim of final_net_kwargs + if given, else d_model. Inferred by complete_settings; not a user + setting. + final_net_kwargs : Optional[dict] + Settings for the DenseResidualNet applied after pooling: output_dim, + hidden_dims, activation (str), and optionally dropout, batch_norm, + layer_norm. If None, the pooled d_model-dim vector is returned + directly. + """ + d_model = transformer_kwargs["d_model"] + tokenizer_kwargs = dict(tokenizer_kwargs) + tokenizer_kwargs["activation"] = torchutils.get_activation_function_from_string( + tokenizer_kwargs["activation"] + ) + tokenizer = DenseResidualNet(output_dim=d_model, **tokenizer_kwargs) + + positional_encoder = MultiPositionalEncoding( + d_model=d_model, **positional_encoder_kwargs + ) + + transformer_layer = nn.TransformerEncoderLayer( + d_model=d_model, + dim_feedforward=transformer_kwargs.get("dim_feedforward", 2048), + nhead=transformer_kwargs["nhead"], + dropout=transformer_kwargs.get("dropout", 0.1), + batch_first=True, + ) + transformer_encoder = nn.TransformerEncoder( + transformer_layer, num_layers=transformer_kwargs["num_layers"] + ) + + final_net = None + if final_net_kwargs is not None: + final_net_kwargs = dict(final_net_kwargs) + final_net_output_dim = final_net_kwargs.pop("output_dim") + final_net_kwargs["activation"] = ( + torchutils.get_activation_function_from_string( + final_net_kwargs["activation"] + ) + ) + final_net = DenseResidualNet( + input_dim=d_model, + output_dim=final_net_output_dim, + **final_net_kwargs, + ) + else: + final_net_output_dim = d_model + if output_dim != final_net_output_dim: + raise ValueError( + f"Inconsistent settings: output_dim is {output_dim}, but the " + f"network produces {final_net_output_dim} " + f"(final_net output_dim, or d_model without a final net)." + ) + + super().__init__( + tokenizer=tokenizer, + positional_encoder=positional_encoder, + transformer_encoder=transformer_encoder, + final_net=final_net, + ) + self.output_dim = output_dim + + @classmethod + def complete_settings(cls, settings: dict, sample_batch: dict) -> dict: + """Infer the tokenizer input dim and the embedding output_dim from a sample + batch; return completed settings.""" + tokenizer_kwargs = dict(settings["tokenizer_kwargs"]) + if "input_dim" in tokenizer_kwargs: + raise ValueError( + "'input_dim' is derived from the data and must not be specified " + "in the tokenizer settings." + ) + if "output_dim" in settings: + raise ValueError( + "'output_dim' is derived from the network settings and must not " + "be specified." + ) + tokenizer_kwargs["input_dim"] = sample_batch["waveform"].shape[-1] + + final_net_kwargs = settings.get("final_net_kwargs") + if final_net_kwargs is not None: + output_dim = final_net_kwargs["output_dim"] + else: + output_dim = settings["transformer_kwargs"]["d_model"] + return { + **settings, + "tokenizer_kwargs": tokenizer_kwargs, + "output_dim": output_dim, + } diff --git a/dingo/gw/gwutils.py b/dingo/gw/gwutils.py index f5625258..68b13d01 100644 --- a/dingo/gw/gwutils.py +++ b/dingo/gw/gwutils.py @@ -144,3 +144,18 @@ def get_standardization_dict( "std": {k: std[k] for k in selected_parameters}, } return standardization_dict + + +def add_defaults_for_missing_ifos( + object_to_update: Optional[float | dict], + update_value: float, + ifos: list[str], +): + """For a per-detector settings dict, fill in update_value for any detector in + ifos that has no entry; scalars and None pass through unchanged.""" + object_to_update = deepcopy(object_to_update) + if isinstance(object_to_update, dict) and ifos is not None: + for det in ifos: + if det not in object_to_update.keys(): + object_to_update[det] = update_value + return object_to_update diff --git a/dingo/gw/inference/gw_samplers.py b/dingo/gw/inference/gw_samplers.py index d5e38aaf..8f110ca2 100644 --- a/dingo/gw/inference/gw_samplers.py +++ b/dingo/gw/inference/gw_samplers.py @@ -32,8 +32,10 @@ GetDetectorTimes, DecimateWaveformsAndASDS, MaskDataForFrequencyRangeUpdate, + NormalizePosition, SelectKeys, StrainTokenization, + UpdateFrequencyRange, ) @@ -66,6 +68,7 @@ def __init__(self: SamplerProtocol, **kwargs): # Has to be specified before init, because the information is required in _initialize_transforms() self._minimum_frequency = None self._maximum_frequency = None + self._suppress = None super().__init__(**kwargs) self.t_ref = self.base_model_metadata["train_settings"]["data"]["ref_time"] self._pesummary_package = "gw" @@ -129,6 +132,46 @@ def maximum_frequency(self: _GWMixinProtocol, value: Union[float, dict]): self._maximum_frequency = value self._initialize_transforms() + @property + def suppress(self): + """Frequency ranges whose tokens are masked out at inference: [f_lo, f_hi], + or {detector: [f_lo, f_hi]}. Only available for tokenized models trained + with drop augmentation.""" + return self._suppress + + @suppress.setter + def suppress(self: _GWMixinProtocol, value): + data_settings = self.base_model_metadata["train_settings"]["data"] + tokenization = data_settings.get("tokenization", {}) + if not tokenization: + raise ValueError( + "Token suppression requires a model trained on tokenized data." + ) + if not ( + "drop_frequency_range" in tokenization + or "drop_random_tokens" in tokenization + ): + raise ValueError( + "Token suppression requires a model trained with drop augmentation " + "(tokenization.drop_frequency_range or drop_random_tokens)." + ) + intervals = value if isinstance(value, dict) else {None: value} + if isinstance(value, dict): + unknown = set(value) - set(self.detectors) + if unknown: + raise ValueError( + f"Unknown detectors in suppress setting: {sorted(unknown)}." + ) + for interval in intervals.values(): + f_lo, f_hi = interval + if not self.domain.f_min <= f_lo < f_hi <= self.domain.f_max: + raise ValueError( + f"Suppress interval {interval} must satisfy " + f"{self.domain.f_min} <= f_lo < f_hi <= {self.domain.f_max}." + ) + self._suppress = value + self._initialize_transforms() + @property def frequency_updates(self) -> bool: def normalize(val): @@ -136,9 +179,11 @@ def normalize(val): return set(val.values()) return {val} - return normalize(self.minimum_frequency) != {self.domain.f_min} or normalize( - self.maximum_frequency - ) != {self.domain.f_max} + return ( + normalize(self.minimum_frequency) != {self.domain.f_min} + or normalize(self.maximum_frequency) != {self.domain.f_max} + or self.suppress is not None + ) @property def event_metadata(self): @@ -148,6 +193,8 @@ def event_metadata(self): metadata = {} metadata["minimum_frequency"] = self.minimum_frequency metadata["maximum_frequency"] = self.maximum_frequency + if self.suppress is not None: + metadata["suppress"] = self.suppress return metadata @event_metadata.setter @@ -158,6 +205,8 @@ def event_metadata(self, value): self.minimum_frequency = value.pop("minimum_frequency") if "maximum_frequency" in value: self.maximum_frequency = value.pop("maximum_frequency") + if value.get("suppress") is not None: + self.suppress = value.pop("suppress") self._event_metadata = value def _build_domain(self: Sampler): @@ -297,10 +346,12 @@ def _initialize_transforms(self): # * whiten and scale strain (since the inference network expects standardized # data) transform_pre.append(WhitenAndScaleStrain(self.domain.noise_std)) - if self.frequency_updates: + tokenization = self.metadata["train_settings"]["data"].get("tokenization") + if self.frequency_updates and not tokenization: # * update frequency range # Needs to happen before RepackageStrainsAndASDs since we might need to apply - # detectors specific frequency updates. + # detectors specific frequency updates. Tokenized models instead update + # the drop_token_mask after tokenization (UpdateFrequencyRange below). transform_pre.append( MaskDataForFrequencyRangeUpdate( domain=self.domain, @@ -320,7 +371,6 @@ def _initialize_transforms(self): first_index=self.domain.min_idx, ) ) - tokenization = self.metadata["train_settings"]["data"].get("tokenization") if tokenization: # StrainTokenization operates on numpy arrays, so it precedes ToTorch. transform_pre.append( @@ -331,6 +381,32 @@ def _initialize_transforms(self): drop_last_token=tokenization.get("drop_last_token", False), ) ) + if self.frequency_updates: + # Frequency-range updates / token suppression for tokenized + # models: mask out the affected tokens. Unchanged bounds are + # passed as None — the domain default as a threshold would + # needlessly mask the zero-padded final token. + transform_pre.append( + UpdateFrequencyRange( + minimum_frequency=( + self.minimum_frequency + if self.minimum_frequency != self.domain.f_min + else None + ), + maximum_frequency=( + self.maximum_frequency + if self.maximum_frequency != self.domain.f_max + else None + ), + suppress_range=self.suppress, + domain=self.domain, + ifos=self.detectors, + ) + ) + # Normalize positions after all mask updates, matching training + # (the mask transforms compare positions against frequencies in Hz). + if tokenization.get("normalize_frequency_for_positional_encoding", False): + transform_pre.append(NormalizePosition()) transform_pre.append(ToTorch(device=self.model.device)) if tokenization: # Dict of named network inputs; the model routes them by key. diff --git a/dingo/gw/training/train_builders.py b/dingo/gw/training/train_builders.py index 33190208..6a09fae4 100755 --- a/dingo/gw/training/train_builders.py +++ b/dingo/gw/training/train_builders.py @@ -19,6 +19,11 @@ GetDetectorTimes, CropMaskStrainRandom, StrainTokenization, + DropDetectors, + DropFrequenciesToUpdateRange, + DropFrequencyInterval, + DropRandomTokens, + NormalizePosition, ) from dingo.gw.noise.asd_dataset import ASDDataset from dingo.gw.prior import default_inference_parameters @@ -243,6 +248,69 @@ def set_train_transforms(wfd, data_settings, asd_dataset_path, omit_transforms=N drop_last_token=tokenization.get("drop_last_token", False), ) ) + num_tokens = transforms[-1].num_tokens_per_detector * len( + data_settings["detectors"] + ) + + # Augmentation: randomly drop detectors, frequency ranges, or random + # tokens during training. This is what enables inference-time frequency + # updates and token suppression. + if "drop_detectors" in tokenization: + drop_detectors = tokenization["drop_detectors"] + transforms.append( + DropDetectors( + num_blocks=len(data_settings["detectors"]), + p_drop_012_detectors=drop_detectors.get("p_drop_012_detectors"), + p_drop_hlv=drop_detectors.get("p_drop_hlv"), + ) + ) + if "drop_frequency_range" in tokenization: + if "f_cut" in tokenization["drop_frequency_range"]: + f_cut_settings = tokenization["drop_frequency_range"]["f_cut"] + transforms.append( + DropFrequenciesToUpdateRange( + domain=domain, + p_cut=f_cut_settings.get("p_cut", 0.2), + f_max_lower_cut=f_cut_settings.get( + "f_max_lower_cut", domain.f_min + ), + f_min_upper_cut=f_cut_settings.get( + "f_min_upper_cut", domain.f_max + ), + p_same_cut_all_detectors=f_cut_settings.get( + "p_same_cut_all_detectors", 0.2 + ), + p_lower_upper_both=f_cut_settings.get( + "p_lower_upper_both", [0.4, 0.4, 0.2] + ), + ) + ) + if "mask_interval" in tokenization["drop_frequency_range"]: + interval_settings = tokenization["drop_frequency_range"][ + "mask_interval" + ] + transforms.append( + DropFrequencyInterval( + domain=domain, + p_per_detector=interval_settings.get("p_per_detector", 0.2), + f_min=interval_settings.get("f_min", domain.f_min), + f_max=interval_settings.get("f_max", domain.f_max), + max_width=interval_settings.get("max_width", 10.0), + ) + ) + if "drop_random_tokens" in tokenization: + random_drop_settings = tokenization["drop_random_tokens"] + transforms.append( + DropRandomTokens( + p_drop=random_drop_settings.get("p_drop", 0.4), + max_num_tokens=random_drop_settings.get( + "max_num_tokens", num_tokens + ), + ) + ) + # Normalize the position entries after all drop transforms. + if tokenization.get("normalize_frequency_for_positional_encoding", False): + transforms.append(NormalizePosition()) selected_keys = ["inference_parameters", "waveform"] if "tokenization" in data_settings: diff --git a/dingo/gw/transforms/__init__.py b/dingo/gw/transforms/__init__.py index e601c190..e1be056a 100644 --- a/dingo/gw/transforms/__init__.py +++ b/dingo/gw/transforms/__init__.py @@ -6,5 +6,14 @@ from .inference_transforms import * from .utils import * from .waveform_transforms import * -from .tokenization_transforms import StrainTokenization, DETECTOR_DICT - +from .tokenization_transforms import ( + DETECTOR_DICT, + DETECTOR_DICT_INVERSE, + DropDetectors, + DropFrequenciesToUpdateRange, + DropFrequencyInterval, + DropRandomTokens, + NormalizePosition, + StrainTokenization, + UpdateFrequencyRange, +) diff --git a/dingo/gw/transforms/tokenization_transforms.py b/dingo/gw/transforms/tokenization_transforms.py index f683d850..7faef5ff 100644 --- a/dingo/gw/transforms/tokenization_transforms.py +++ b/dingo/gw/transforms/tokenization_transforms.py @@ -3,8 +3,10 @@ import numpy as np from dingo.gw.domains import UniformFrequencyDomain, MultibandedFrequencyDomain +from dingo.gw.gwutils import add_defaults_for_missing_ifos DETECTOR_DICT = {"H1": 0, "L1": 1, "V1": 2} +DETECTOR_DICT_INVERSE = {v: k for k, v in DETECTOR_DICT.items()} class StrainTokenization: @@ -230,3 +232,989 @@ def _check_mfd_node_compatibility( f"MFD nodes {mfd_nodes[~covered]} fall within a token rather than " f"between tokens. Adjust token_size or MFD nodes." ) + + +# --------------------------------------------------------------------------- +# Augmentation transforms for tokenized strain data (ported from the DINGO-T1 +# branch). These operate on the output of StrainTokenization (waveform / +# position / drop_token_mask) and mark tokens as dropped by setting +# drop_token_mask entries to True; the transformer then does not attend to +# them. Training with these augmentations is what enables inference-time +# frequency-range updates and token suppression (UpdateFrequencyRange). +# --------------------------------------------------------------------------- + + +class DropDetectors(object): + """ + Randomly drop detectors. + """ + + def __init__( + self, + num_blocks: int, + p_drop_012_detectors: list | None = None, + p_drop_hlv: dict | None = None, + print_output: bool = True, + ): + """ + Parameters + ---------- + num_blocks: int + Number of blocks (= detectors) in GW use case. + p_drop_012_detectors: list[float] + Specifies the categorical probability distribution for how many detectors to drop, in ascending order + example: [0.1, 0.6, 0.3] = [10% probability to drop 0 detectors (=3 detector setup), 60 % probability for + 2 detector setup, 30% probability for 1 detector setup] + p_drop_hlv: dict + Specifies the categorical probability distribution for which specific detectors to drop, order: H1, L1, V1 + example: {'H1': 0.1, 'L1': 0.2, 'V1': 0.7] = 10 % probability to drop H1, 20 % probability to drop L1, + 70% probability to drop V1 + print_output: bool + Whether to write print statements to the console. + """ + self.num_blocks = num_blocks + if p_drop_012_detectors is None: + p_drop_012_detectors = [1 / num_blocks for _ in range(num_blocks)] + if not np.isclose(np.sum(p_drop_012_detectors), 1.0, rtol=1e-6, atol=1e-12): + raise ValueError( + f"p_drop_012_detectors {p_drop_012_detectors} does not sum to 1." + ) + self.p_drop_012_detectors = p_drop_012_detectors + if p_drop_hlv is None: + p_drop_hlv = { + ["H1", "L1", "V1"][k]: 1 / num_blocks for k in range(num_blocks) + } + if not np.isclose( + np.sum(list(p_drop_hlv.values())), 1.0, rtol=1e-6, atol=1e-12 + ): + raise ValueError(f"p_drop_hlv {p_drop_hlv} does not sum to 1.") + # Update keys equivalently to tokenization transform + self.p_drop_hlv = {DETECTOR_DICT[k]: v for k, v in p_drop_hlv.items()} + + if len(p_drop_012_detectors) > num_blocks: + raise ValueError( + f"p_drop_num_detectors {self.p_drop_012_detectors} contains more options than" + f"detectors available: {num_blocks}. You need to specify a categorical probability" + f"value for dropping 0, ..., {num_blocks - 1} detectors." + ) + if len(self.p_drop_hlv) != num_blocks: + raise ValueError( + f"Provided values for p_drop_hlv={self.p_drop_hlv} is inconsistent with number of " + f"detectors: {num_blocks}. You need to specify a categorical probability value for each " + f"detector." + ) + if print_output: + print( + f"Transform DropDetectors activated: \n" + f" - Probabilities for dropping {[i for i in range(num_blocks)]} detectors are " + f"{self.p_drop_012_detectors}.\n" + f" - Probabilities for specific detectors are {self.p_drop_hlv}." + ) + + def __call__(self, input_sample: dict) -> dict: + """ + Parameters + ---------- + input_sample: Dict + Values for keys + - 'waveform': + Sample of shape [batch_size, num_tokens, num_features] = + [batch_size, num_blocks * num_tokens_per_block, num_channels * num_bins_per_token] + where num_blocks = number of detectors in GW use case, + num_channels>=3 (real, imag, auxiliary channels, e.g. asd), + and num_bins = number of frequency bins. + - 'position', shape [batch_size, num_tokens, 3] + contains information [f_min, f_max, block] + - 'drop_token_mask', shape [batch_size, num_tokens] + + Returns + ---------- + sample: Dict + input_sample with modified value for key + - 'drop_token_mask', shape [batch_size, num_tokens] + + """ + # The transform operates on batched arrays; the training pipeline feeds + # single samples, so temporarily add a batch axis. + batched = input_sample["drop_token_mask"].ndim > 1 + if not batched: + input_sample["drop_token_mask"] = input_sample["drop_token_mask"][ + np.newaxis + ] + blocks = input_sample["position"][..., 2] + if not batched: + blocks = blocks[np.newaxis] + num_blocks = len(np.unique(blocks)) + detectors = np.unique(blocks) + + # Convert p_drop_hlv dict to list + p_drop_hlv = [self.p_drop_hlv[k] for k in detectors] + + # Decide how many detectors to drop (either none, or one less than the number of detectors present) + # for each element in batch_size + drop_n_blocks = np.random.choice( + [i for i in range(num_blocks)], + p=self.p_drop_012_detectors, + size=[*blocks.shape[:-1]], + ) + if np.sum(drop_n_blocks) != 0: + # Treat drop 1 vs. 2 blocks separately because which detectors to drop varies + # with the number of detectors to drop + for n in [i for i in np.unique(drop_n_blocks) if i > 0]: + # Construct mask for which batch indices require updates + mask_mod = np.where(drop_n_blocks == n, True, False) + # Decide which detectors + detectors_to_drop = np.apply_along_axis( + np.random.choice, + axis=1, + arr=np.repeat( + np.expand_dims(detectors, 0), repeats=np.sum(mask_mod), axis=0 + ), + p=p_drop_hlv, + size=n, + replace=False, + ) + # Create mask such that tokens corresponding to dropped detectors are True + # (1) Drop one detector + mask_detectors = np.where( + blocks[mask_mod].T == detectors_to_drop[:, 0], True, False + ).T + if detectors_to_drop.shape[-1] > 1: + # (2) Update mask to include dropping of any further detector + for i in range(1, detectors_to_drop.shape[-1]): + mask_detectors_i = np.where( + blocks[mask_mod].T == detectors_to_drop[:, i], True, False + ).T + mask_detectors = np.logical_or(mask_detectors_i, mask_detectors) + # Keep drop=True from previous transforms with logical OR + mask_detectors = np.logical_or( + input_sample["drop_token_mask"][mask_mod], mask_detectors + ) + # Update mask + input_sample["drop_token_mask"][mask_mod] = mask_detectors + + if not batched: + input_sample["drop_token_mask"] = input_sample["drop_token_mask"][0] + return input_sample + + +class DropFrequenciesToUpdateRange(object): + """ + Randomly drop tokens such that f_min and f_max of the frequency range are updated. + + This transform does the following things: + * Decides whether to apply a cut to each element of the batch based on p_cut. + * Decides whether to treat the detectors individually or apply the same cut to all detectors. + * Decides whether to cut upper or lower end or both (potentially for each detector). + * Samples f_cut from [f_min, f_max_lower_cut] and/or [f_min_upper_cut, f_max] in UFD (potentially for each + detector). + * Converts frequency values to tokens and creates a token mask removing [f_min, f_lower_cut] and/or + [f_upper_cut, f_max] (potentially for each detector). + """ + + def __init__( + self, + domain: UniformFrequencyDomain | MultibandedFrequencyDomain, + p_cut: float, + f_max_lower_cut: float, + f_min_upper_cut: float, + p_same_cut_all_detectors: float, + p_lower_upper_both: Optional[list] = None, + print_output: bool = True, + ): + """ + Parameters + ---------- + domain: UniformFrequencyDomain | MultibandedFrequencyDomain + Domain corresponding to the data being transformed. + p_cut: float + Probability of applying a cut to each element of the batch. + f_max_lower_cut: float + Maximal frequency value to cut at the lower end of the frequency domain. f_min_lower_cut is sampled from + [f_min, f_max_lower_cut] in UFD. + f_min_upper_cut: float + Minimal frequency value to cut at the upper end of the frequency domain. f_max_upper_cut is sampled from + [f_min_upper_cut, f_max] in UFD. + p_same_cut_all_detectors: float + Probability of applying the same cut to all detectors. + p_lower_upper_both: list[float] + List of probabilities explaining with what probability we either cut at the lower, at the upper, or at both + ends. Order: [p_lower, p_upper, p_both] + print_output: bool + Whether to write print statements to the console. + """ + + self.domain = domain + self.p_cut = p_cut + self.f_max_lower_cut = f_max_lower_cut + self.f_min_upper_cut = f_min_upper_cut + self.prevent_zero_information = ( + True if self.f_max_lower_cut >= self.f_min_upper_cut else False + ) + self.p_same_cut_all_detectors = p_same_cut_all_detectors + if p_lower_upper_both is None: + p_lower_upper_both = np.array([0.4, 0.4, 0.2]) + self.p_lower_upper_both = p_lower_upper_both + if not np.isclose(np.sum(self.p_lower_upper_both), 1.0, rtol=1e-6, atol=1e-12): + raise ValueError( + f"p_lower_upper_both {self.p_lower_upper_both} does not sum to 1. " + ) + if print_output: + print( + f"Transform DropFrequencyValues activated: \n" + f" - Probability of a cut happening: {self.p_cut}\n" + f" - Lower cut sampled from [{self.domain.f_min}, {self.f_max_lower_cut}]\n" + f" - Upper cut sampled from [{self.f_min_upper_cut}, {self.domain.f_max}]\n" + f" - Probability to apply the same cut on all detectors: {self.p_same_cut_all_detectors} " + ) + if self.prevent_zero_information: + print( + f"\n - Preventing zero information is activated since [{self.domain.f_min}, {self.f_max_lower_cut}]" + f"overlaps with [{self.f_min_upper_cut}, {self.domain.f_max}] " + ) + + def __call__(self, input_sample: dict) -> dict: + """ + Parameters + ---------- + input_sample: Dict + Values for keys + - 'waveform': + Sample of shape [batch_size, num_tokens, num_features] + - 'position', shape [batch_size, num_tokens, 3] + contains information [f_min, f_max, block] + - 'drop_token_mask', shape [batch_size, num_tokens] + + Returns + ---------- + sample: Dict + input_sample with modified value for key + - 'drop_token_mask', shape [batch_size, num_tokens] + + """ + num_tokens = input_sample["waveform"].shape[-2] + blocks = input_sample["position"][..., 2] + num_blocks = len(np.unique(blocks)) + num_tokens_per_block = num_tokens // num_blocks + + # Cut in frequency domain, where we remove the upper, lower or both part(s), + # i.e. [f_min, f_cut], [f_cut, f_max], or [f_cut_min, f_cut_max] + # - Decide whether to apply a cut for each sample + # - Decide whether to treat the detectors individually or apply the same cut to all detectors + # - Decide whether to mask upper or lower range or both (potentially for each detector) + # - Sample index for f_cut from [f_min, f_max_lower_cut] and/or [f_min_upper_cut, f_max] + # in uniform frequency domain (potentially for each detector) + # - Convert frequency values to token mask + + batch_size = [*blocks.shape[:-1]] if blocks.shape[:-1] != () else [1] + # Decide whether to apply a cut for each sample + apply_cut = np.random.choice( + [True, False], p=[self.p_cut, 1 - self.p_cut], size=batch_size + ) + + # Decide whether to treat the detectors individually or apply the same cut to all detectors + same_cut_all_detectors = np.where( + apply_cut, + np.random.choice( + [True, False], + p=[self.p_same_cut_all_detectors, 1 - self.p_same_cut_all_detectors], + size=batch_size, + ), + False, + ) + batch_block_size = ( + [*blocks.shape[:-1], num_blocks] + if blocks.shape[:-1] != () + else [1, num_blocks] + ) + # (1) Different cut is applied to every detector + # Decide whether to mask upper or lower range or both (potentially for each detector) + lower_upper_both_separate = np.random.choice( + ["lower", "upper", "both"], p=self.p_lower_upper_both, size=batch_block_size + ) + mask_lower_separate = np.logical_or( + lower_upper_both_separate == "lower", lower_upper_both_separate == "both" + ) + mask_upper_separate = np.logical_or( + lower_upper_both_separate == "upper", lower_upper_both_separate == "both" + ) + # Combine with masks (a) whether we apply a cut and (b) whether we apply it to a single detector + ones_vec = np.ones((1, num_blocks), dtype=bool) + mask_lower_separate_combined = np.logical_and.reduce( + ( + mask_lower_separate, + apply_cut[..., None] * ones_vec, + ~same_cut_all_detectors[..., None] * ones_vec, + ) + ) + mask_upper_separate_combined = np.logical_and.reduce( + ( + mask_upper_separate, + apply_cut[..., None] * ones_vec, + ~same_cut_all_detectors[..., None] * ones_vec, + ) + ) + # Sample f_cut from [f_min, f_max_lower_cut] and/or [f_min_upper_cut, f_max] in UFD for each detector + if isinstance(self.domain, UniformFrequencyDomain): + f_values_base_domain = self.domain.sample_frequencies[ + self.domain.frequency_mask + ] + elif isinstance(self.domain, MultibandedFrequencyDomain): + f_values_base_domain = self.domain.base_domain.sample_frequencies[ + self.domain.base_domain.frequency_mask + ] + else: + raise ValueError(f"Unknown domain type: {self.domain}") + f_lower_separate = np.where( + mask_lower_separate_combined, + np.random.choice( + f_values_base_domain[f_values_base_domain <= self.f_max_lower_cut], + replace=True, + size=batch_block_size, + ), + -1, + ) + f_upper_separate = np.where( + mask_upper_separate_combined, + np.random.choice( + f_values_base_domain[f_values_base_domain >= self.f_min_upper_cut], + replace=True, + size=batch_block_size, + ), + np.inf, + ) + + # Construct mask: f_cut_lower >= f_min_per_token and f_cut_upper <= f_max_per_token + token_mask_separate_lower = ( + np.repeat(f_lower_separate, repeats=num_tokens_per_block, axis=-1) + >= input_sample["position"][..., 0] + ) + token_mask_separate_upper = ( + np.repeat(f_upper_separate, repeats=num_tokens_per_block, axis=-1) + <= input_sample["position"][..., 1] + ) + + # Combine into one mask + token_mask_separate = np.logical_or( + token_mask_separate_lower, token_mask_separate_upper + ) + if self.prevent_zero_information: + # If all tokens are masked in one sample, only apply upper or lower mask + replace_mask = np.where( + np.sum(token_mask_separate, axis=-1) == num_tokens, True, False + ) + repl_mask = np.repeat( + replace_mask[..., np.newaxis], repeats=num_tokens, axis=-1 + ) + # Decide whether to choose lower or upper instead of both + lower_upper_probs = self.p_lower_upper_both[:2] / np.sum( + self.p_lower_upper_both[:2] + ) + lower_upper_global = np.random.choice( + ["lower", "upper"], p=lower_upper_probs, size=batch_size + ) + mask_lower_separate_replace = np.where( + lower_upper_global == "lower", True, False + ) + mask_lower_sep_repl = np.repeat( + mask_lower_separate_replace[..., np.newaxis], + repeats=num_tokens, + axis=-1, + ) + # Create replace mask + mask_combined_separate_replace = np.where( + mask_lower_sep_repl, + token_mask_separate_lower, + token_mask_separate_upper, + ) + # Combine with token_mask_separate + token_mask_separate = np.where( + repl_mask, mask_combined_separate_replace, token_mask_separate + ) + + # (2) Same cut is applied to all detectors + # Decide whether to mask upper or lower or both + lower_upper_both_same = np.random.choice( + ["lower", "upper", "both"], p=self.p_lower_upper_both, size=batch_size + ) + mask_lower_same = np.logical_or( + lower_upper_both_same == "lower", lower_upper_both_same == "both" + ) + mask_upper_same = np.logical_or( + lower_upper_both_same == "upper", lower_upper_both_same == "both" + ) + # Combine with masks (a) whether we apply a cut and (b) whether we apply it to all detectors + mask_lower_combined = np.logical_and.reduce( + (mask_lower_same, apply_cut, same_cut_all_detectors) + ) + mask_upper_combined = np.logical_and.reduce( + (mask_upper_same, apply_cut, same_cut_all_detectors) + ) + # Sample f_cut from [f_min, f_max_lower_cut] and/or [f_min_upper_cut, f_max] in UFD + f_lower_same = np.where( + mask_lower_combined, + np.random.choice( + f_values_base_domain[f_values_base_domain <= self.f_max_lower_cut], + replace=True, + size=batch_size, + ), + -1, + ) + f_upper_same = np.where( + mask_upper_combined, + np.random.choice( + f_values_base_domain[f_values_base_domain >= self.f_min_upper_cut], + replace=True, + size=batch_size, + ), + np.inf, + ) + # Construct mask: f_cut_lower >= f_min_per_token and f_cut_upper <= f_max_per_token + # (Assume that all detectors have same f_min and f_max values) + f_mins = input_sample["position"][..., 0:num_tokens_per_block, 0] + f_maxs = input_sample["position"][..., 0:num_tokens_per_block, 1] + token_mask_same_lower = f_lower_same[:, np.newaxis] >= f_mins + token_mask_same_upper = f_upper_same[:, np.newaxis] <= f_maxs + + # Combine into one mask + token_mask_same_one_detector = np.logical_or( + token_mask_same_lower, token_mask_same_upper + ) + if self.prevent_zero_information: + # If all tokens are masked in one block, only apply upper or lower mask + replace_mask = np.where( + np.sum(token_mask_same_one_detector, axis=-1) == num_tokens_per_block, + True, + False, + ) + repl_mask = np.repeat( + replace_mask[..., np.newaxis], repeats=num_tokens_per_block, axis=-1 + ) + # Decide whether to choose lower or upper instead of both + lower_upper_probs = self.p_lower_upper_both[:2] / np.sum( + self.p_lower_upper_both[:2] + ) + lower_upper_global = np.random.choice( + ["lower", "upper"], p=lower_upper_probs, size=batch_size + ) + mask_lower_same_replace = np.where( + lower_upper_global == "lower", True, False + ) + mask_lower_same_repl = np.repeat( + mask_lower_same_replace[..., np.newaxis], + repeats=num_tokens_per_block, + axis=-1, + ) + # Create replace mask + mask_combined_same_replace = np.where( + mask_lower_same_repl, token_mask_same_lower, token_mask_same_upper + ) + # Combine with token_mask_same_one_detector + token_mask_same_one_detector = np.where( + repl_mask, mask_combined_same_replace, token_mask_same_one_detector + ) + + # Duplicate for number of detectors + token_mask_same = np.tile(token_mask_same_one_detector, reps=num_blocks) + + # Modify mask + if len(input_sample["drop_token_mask"].shape) == 1: + token_mask_separate = token_mask_separate.squeeze() + token_mask_same = token_mask_same.squeeze() + input_sample["drop_token_mask"] = np.logical_or.reduce( + (input_sample["drop_token_mask"], token_mask_separate, token_mask_same) + ) + + return input_sample + + +class DropFrequencyInterval(object): + """ + Randomly drop tokens corresponding to specific frequency interval. + + This transform does the following things: + * Decides whether to mask a frequency interval per detector based on p_per_detector. + * Samples f_lower from [f_min, f_max - max_width]. + * Samples f_upper from [f_lower, f_lower + max_width]. + * Converts f_lower and f_upper to tokens and creates a token mask removing all tokens in [f_lower, f_upper]. + """ + + def __init__( + self, + domain: UniformFrequencyDomain | MultibandedFrequencyDomain, + p_per_detector: float, + f_min: float, + f_max: float, + max_width: float, + print_output: bool = True, + ): + """ + Parameters + ---------- + domain: UniformFrequencyDomain | MultibandedFrequencyDomain + Domain corresponding to the data being transformed. + p_per_detector: float + Probability of dropping tokens (corresponding to a frequency interval) independently per detector. + f_min: float + Minimal frequency value for which we allow tokens to be dropped. + f_max: float + Maximum frequency value for which we allow tokens to be dropped. + max_width: float + Maximal width of frequency interval that can be dropped. + print_output: bool + Whether to write print statements to the console. + """ + self.domain = domain + self.p_per_detector = p_per_detector + self.interval_f_min = f_min if domain.f_min < f_min else domain.f_min + self.interval_f_max = f_max if domain.f_max > f_max else domain.f_max + interval_width = self.interval_f_max - self.interval_f_min + self.interval_max_width = ( + max_width if max_width < interval_width else interval_width + ) + if print_output: + print( + f"Transform DropFrequencyInterval activated:" + f" Settings: \n" + f" - Probability of dropping interval per detector: {self.p_per_detector}\n" + f" - Interval range sampled from [{self.interval_f_min}, {self.interval_f_max}]\n" + f" - Maximal width of interval: {self.interval_max_width}, but the effective interval can be larger " + f"if {self.interval_f_min} and {self.interval_f_max} fall in the middle of a token." + ) + + def __call__(self, input_sample: dict) -> dict: + """ + Parameters + ---------- + input_sample: Dict + Values for keys + - 'waveform': + Sample of shape [batch_size, num_tokens, num_features] + - 'position', shape [batch_size, num_tokens, 3] + contains information [f_min, f_max, block] + - 'drop_token_mask', shape [batch_size, num_tokens] + + Returns + ---------- + sample: Dict + input_sample with modified value for key + - 'drop_token_mask', shape [batch_size, num_tokens] + + """ + num_tokens = input_sample["waveform"].shape[-2] + blocks = input_sample["position"][..., 2] + num_blocks = len(np.unique(blocks)) + num_tokens_per_block = num_tokens // num_blocks + + # Mask frequency range: + # - Decide whether to apply a mask for each detector + # - Sample f_mask_lower and f_mask_upper in uniform frequency domain + # - Get tokens corresponding to frequency values + # - Mask everything in between, i.e. [f_mask_lower, f_mask_upper] + + batch_block_size = ( + [*blocks.shape[:-1], num_blocks] + if blocks.shape[:-1] != () + else [1, num_blocks] + ) + # Decide whether to cut or mask frequency range for each block + mask_interval = np.random.choice( + [True, False], + p=[self.p_per_detector, 1 - self.p_per_detector], + size=batch_block_size, + ) + + # Sample f_lower and f_upper in UFD + if isinstance(self.domain, UniformFrequencyDomain): + f_values_base_domain = self.domain.sample_frequencies[ + self.domain.frequency_mask + ] + elif isinstance(self.domain, MultibandedFrequencyDomain): + f_values_base_domain = self.domain.base_domain.sample_frequencies[ + self.domain.base_domain.frequency_mask + ] + else: + raise ValueError(f"Unknown domain type: {self.domain}") + # f_lower from [interval_f_min, interval_f_max - interval_max_width] + mask_f_vals_lower = np.logical_and( + self.interval_f_min <= f_values_base_domain, + f_values_base_domain <= self.interval_f_max - self.interval_max_width, + ) + possible_f_vals_lower = f_values_base_domain[mask_f_vals_lower] + f_lower_full = np.random.choice( + possible_f_vals_lower, replace=True, size=batch_block_size + ) + f_lower = np.where(mask_interval, f_lower_full, np.inf) + + # f_upper from [f_lower, f_lower + interval_max_width] + # Sampling f_upper is more complicated because it depends on the f_lower sampled for each batch index and + # detector + mask_f_vals_upper = np.logical_and( + f_lower_full[:, :, np.newaxis] + <= f_values_base_domain[np.newaxis, np.newaxis, :], + f_values_base_domain[np.newaxis, np.newaxis, :] + <= f_lower_full[:, :, np.newaxis] + self.interval_max_width, + ) + possible_indices_upper = np.stack( + [ + np.apply_along_axis( + np.argwhere, arr=mask_f_vals_upper[:, b, :], axis=-1 + ).squeeze() + for b in range(num_blocks) + ], + axis=-2, + ) + possible_f_vals_upper = f_values_base_domain[possible_indices_upper] + f_upper_no_mask = np.stack( + [ + np.apply_along_axis( + np.random.choice, arr=possible_f_vals_upper[..., b, :], axis=-1 + ) + for b in range(num_blocks) + ], + axis=-1, + ) + f_upper = np.where(mask_interval, f_upper_no_mask, -1.0) + + # Construct mask: f_lower <= f_maxs AND f_upper >= f_mins + f_mins = input_sample["position"][..., 0] + f_maxs = input_sample["position"][..., 1] + token_mask_lower = ( + np.repeat(f_lower, repeats=num_tokens_per_block, axis=-1) <= f_maxs + ) + token_mask_upper = ( + np.repeat(f_upper, repeats=num_tokens_per_block, axis=-1) >= f_mins + ) + + # Combine into one mask + token_mask = np.logical_and(token_mask_lower, token_mask_upper) + + # Modify mask + if len(input_sample["drop_token_mask"].shape) == 1: + token_mask = token_mask.squeeze() + input_sample["drop_token_mask"] = np.logical_or( + input_sample["drop_token_mask"], token_mask + ) + + return input_sample + + +class DropRandomTokens(object): + """ + Randomly drop tokens for data points. Whether tokens will be dropped depends on the drop probability p_drop. + The number of tokens that will be dropped is sampled uniformly from [1, max_num_tokens], disregarding any domain + information. + """ + + def __init__( + self, + p_drop: float, + max_num_tokens: int, + print_output: bool = True, + ): + """ + Parameters + ---------- + p_drop: float + Probability of dropping tokens from a data point. + max_num_tokens: int + Maximum number of tokens that can be dropped. + print_output: bool + Whether to write print statements to the console. + """ + self.p_drop = p_drop + self.max_num_tokens = max_num_tokens + if print_output: + print( + f"Transform DropRandomTokens activated:\n" + f" - Probability of dropping tokens for each data point: {self.p_drop}\n" + f" - Maximal number of tokens that can be dropped: {self.max_num_tokens}" + ) + + def __call__(self, input_sample: dict) -> dict: + """ + Parameters + ---------- + input_sample: Dict + Values for keys + - 'waveform': + Sample of shape [batch_size, num_tokens, num_features] + - 'position', shape [batch_size, num_tokens, 3] + contains information [f_min, f_max, block] + - 'drop_token_mask', shape [batch_size, num_tokens] + + Returns + ---------- + sample: Dict + input_sample with modified value for key + - 'position', shape [batch_size, num_tokens, 3] + + """ + sample_without_channel = input_sample["waveform"][..., 0] + num_tokens = sample_without_channel.shape[-1] + + batch_size = ( + [*sample_without_channel.shape[:-1]] + if sample_without_channel.shape[:-1] != () + else [1] + ) + probs = [self.p_drop, 1 - self.p_drop] + drop_mask = np.random.choice( + [True, False], + p=probs, + replace=True, + size=batch_size, + ) + num_tokens_to_drop = np.random.choice( + np.arange(1, self.max_num_tokens + 1), size=batch_size + ) + + batch_token_size = ( + [*sample_without_channel.shape] + if sample_without_channel.shape[:-1] != () + else [1, num_tokens] + ) + # Generate random values for all tokens + random_scores = np.random.uniform(size=batch_token_size) + # Sort the scores in ascending order, and get indices + sorted_indices = np.argsort(random_scores, axis=-1) + # Create an index mask for selecting top-k per row + row_indices = np.arange(batch_size[0])[:, np.newaxis] + token_ranks = np.arange(num_tokens) + # For each row, get threshold index + thresholds = num_tokens_to_drop[:, np.newaxis] > token_ranks + # Build boolean mask + token_mask = np.zeros(batch_token_size, dtype=bool) + token_mask[row_indices, sorted_indices] = thresholds + + # Combine masks + token_mask = np.logical_and( + np.repeat(drop_mask[..., np.newaxis], repeats=num_tokens, axis=-1), + token_mask, + ) + + # Modify mask + if len(input_sample["drop_token_mask"].shape) == 1: + token_mask = token_mask.squeeze() + input_sample["drop_token_mask"] = np.logical_or( + input_sample["drop_token_mask"], token_mask + ) + + return input_sample + + +class NormalizePosition(object): + """ + Normalize f_min and f_max in position + """ + + def __call__(self, input_sample: dict) -> dict: + """ + Parameters + ---------- + input_sample: Dict + Values for keys + - 'waveform': + Sample of shape [batch_size, num_tokens, num_features] + - 'position', shape [batch_size, num_tokens, 3] + contains information [f_min, f_max, block] + - 'drop_token_mask', shape [batch_size, num_tokens] + + Returns + ---------- + sample: Dict + input_sample with modified value for key + - 'position', shape [batch_size, num_tokens, 3] + + """ + position = input_sample["position"] + f_min = position[..., 0].min() + f_max = position[..., 1].max() + position[..., 0] = (position[..., 0] - f_min) / (f_max - f_min) + position[..., 1] = (position[..., 1] - f_min) / (f_max - f_min) + input_sample["position"] = position + + return input_sample + + +class UpdateFrequencyRange(object): + """ + # TODO: Rename MaskDataForFrequencyRangeUpdateTransformer to distinguish from similar transform applied for cropping + Update token mask according to frequency range update + """ + + def __init__( + self, + minimum_frequency: Optional[float | dict[str, float]] = None, + maximum_frequency: Optional[float | dict[str, float]] = None, + suppress_range: Optional[ + list[float, float] | dict[str, list[float, float]] + ] = None, + domain: Optional[UniformFrequencyDomain | MultibandedFrequencyDomain] = None, + ifos: Optional[list[str]] = None, + print_output: bool = True, + ): + """ + Parameters + ---------- + minimum_frequency: Optional[float | dict[str, float]] + Update of f_min, if float, the same value will be used for all detectors. + maximum_frequency: Optional[float | dict[str, float]] + Update of f_max, if float, the same value will be used for all detectors. + suppress_range: list[float, float] | dict[str, list[float, float]] | None + Suppress ranges [f_min, f_max], either for all detectors or for individual detectors. + domain: UniformFrequencyDomain | MultibandedFrequencyDomain + ifos: list[str] + List of detectors. + print_output: bool + Whether to write print statements to the console. + """ + # Include defaults in case of missing minimum-/maximum frequency values per detector + self.minimum_frequency = add_defaults_for_missing_ifos( + object_to_update=minimum_frequency, update_value=domain.f_min, ifos=ifos + ) + self.maximum_frequency = add_defaults_for_missing_ifos( + object_to_update=maximum_frequency, update_value=domain.f_max, ifos=ifos + ) + self.suppress_range = suppress_range + self.print_output = print_output + if print_output: + print( + f"Transform UpdateFrequencyRange activated:" + f" Settings: \n" + f" - Minimum_frequency update: {self.minimum_frequency}\n" + f" - Maximum_frequency update: {self.maximum_frequency}\n" + f" - Suppress range: {self.suppress_range}\n" + ) + + def __call__(self, input_sample: dict) -> dict: + """ + Parameters + ---------- + input_sample: Dict + Values for keys + - 'waveform': + Sample of shape [batch_size, num_tokens, num_features] + - 'position', shape [batch_size, num_tokens, 3] + contains information [f_min, f_max, block] + - 'drop_token_mask', shape [batch_size, num_tokens] + + Returns + ---------- + sample: Dict + input_sample with modified value for key + - 'drop_token_mask', shape [batch_size, num_tokens] + + """ + # TODO: Write test for transform. Vectorize (not required for inference) + sample = input_sample.copy() + blocks = np.unique(sample["position"][..., 2]) + num_blocks = len(blocks) + num_tokens_per_block = sample["position"].shape[-2] // num_blocks + + # Assume that f_min is the same for all detectors + f_min_per_token = sample["position"][..., 0] + f_max_per_token = sample["position"][..., 1] + f_min_per_token_single = f_min_per_token[:num_tokens_per_block] + f_max_per_token_single = f_max_per_token[:num_tokens_per_block] + + # Start with empty mask that we will successively update + mask = np.zeros_like(sample["drop_token_mask"], dtype=bool) + # Update minimum_frequency + if self.minimum_frequency is not None: + # Same for all detectors + if isinstance(self.minimum_frequency, float) or isinstance( + self.minimum_frequency, int + ): + # Do not mask token if f_min_per_token = minimum_frequency + mask_min = np.where( + f_min_per_token < self.minimum_frequency, True, False + ) + mask = np.logical_or(mask, mask_min) + # Different for each detector + elif isinstance(self.minimum_frequency, dict): + for b in blocks: + if DETECTOR_DICT_INVERSE[b] in self.minimum_frequency: + # Do not mask token if f_min_per_token = minimum_frequency + mask_min = np.where( + f_min_per_token_single + < self.minimum_frequency[DETECTOR_DICT_INVERSE[b]], + True, + False, + ) + mask_b = np.where(sample["position"][..., 2] == b, True, False) + mask[mask_b] = np.logical_or(mask_min, mask[mask_b]) + else: + raise TypeError( + f"self.minimum_frequency is type {type(self.minimum_frequency)} but must be either a float, an int or a dict." + ) + if self.print_output: + print(f"Updated f_min with {self.minimum_frequency}.") + + # Update maximum_frequency + if self.maximum_frequency is not None: + # Same for all detectors + if isinstance(self.maximum_frequency, float) or isinstance( + self.maximum_frequency, int + ): + # Do not mask token if f_max_per_token = maximum_frequency + mask_max = np.where( + f_max_per_token > self.maximum_frequency, True, False + ) + mask = np.logical_or(mask, mask_max) + # Different for each detector + elif isinstance(self.maximum_frequency, dict): + for b in blocks: + if DETECTOR_DICT_INVERSE[b] in self.maximum_frequency: + # Do not mask token if f_max_per_token = maximum_frequency + mask_max = np.where( + f_max_per_token_single + > self.maximum_frequency[DETECTOR_DICT_INVERSE[b]], + True, + False, + ) + mask_b = np.where(sample["position"][..., 2] == b, True, False) + mask[mask_b] = np.logical_or(mask_max, mask[mask_b]) + else: + raise TypeError( + f"self.maximum_frequency is type {type(self.maximum_frequency)} but must be either a float, an int or a dict." + ) + if self.print_output: + print(f"Updated f_max with {self.maximum_frequency}.") + + # Update suppress_range + if self.suppress_range is not None: + # Same for all detectors + if isinstance(self.suppress_range, list): + f_min_lower, f_max_upper = self.suppress_range + mask_lower = np.where(f_max_per_token >= f_min_lower, True, False) + mask_upper = np.where(f_min_per_token <= f_max_upper, True, False) + mask_interval = np.logical_and(mask_lower, mask_upper) + mask = np.logical_or(mask, mask_interval) + # Different for each detector + elif isinstance(self.suppress_range, dict): + for b in blocks: + if DETECTOR_DICT_INVERSE[b] in self.suppress_range: + f_min_lower, f_max_upper = self.suppress_range[ + DETECTOR_DICT_INVERSE[b] + ] + mask_lower = np.where( + f_max_per_token_single >= f_min_lower, True, False + ) + mask_upper = np.where( + f_min_per_token_single <= f_max_upper, True, False + ) + mask_interval = np.logical_and(mask_lower, mask_upper) + mask_b = np.where(sample["position"][..., 2] == b, True, False) + mask[mask_b] = np.logical_or(mask_interval, mask[mask_b]) + else: + raise TypeError( + f"self.suppress is type {type(self.suppress_range)} but must be either a list or a dict." + ) + if self.print_output: + print(f"Updated suppress_range with {self.suppress_range}.") + + # Update drop_token_mask + sample["drop_token_mask"] = np.logical_or(mask, sample["drop_token_mask"]) + + return sample diff --git a/examples/transformer_model/train_settings.yaml b/examples/transformer_model/train_settings.yaml index 98ea5bd6..1d0560a1 100644 --- a/examples/transformer_model/train_settings.yaml +++ b/examples/transformer_model/train_settings.yaml @@ -34,8 +34,29 @@ data: - psi # Tokenized strain representation: the frequency axis is split into fixed-size # tokens with per-token position information (f_min, f_max, detector). + # The drop_* augmentations mask random tokens during training, which is what + # enables inference-time frequency-range updates, token suppression, and + # dropping detectors. tokenization: token_size: 16 + drop_detectors: + p_drop_012_detectors: [0.6, 0.3, 0.1] + p_drop_hlv: + H1: 0.3 + L1: 0.3 + V1: 0.4 + drop_frequency_range: + f_cut: + p_cut: 0.25 + f_max_lower_cut: 180. + f_min_upper_cut: 80. + p_same_cut_all_detectors: 0.7 + p_lower_upper_both: [0.1, 0.7, 0.2] + mask_interval: + p_per_detector: 0.1 + f_min: 20. + f_max: 1800. + max_width: 10. model: distribution: diff --git a/tests/core/test_build_model.py b/tests/core/test_build_model.py index f521634b..c6fb94cc 100644 --- a/tests/core/test_build_model.py +++ b/tests/core/test_build_model.py @@ -715,3 +715,50 @@ def test_transformer_embedding_with_context_parameters(): ) theta, context = tokenized_batch(with_context_parameters=True) assert torch.isfinite(pm.loss(theta, context)) + + +def test_pooling_transformer_embedding_through_build_path(): + """The second tokenized architecture from the T1 branch (sinusoidal positional + encoding + average pooling) builds through the same generic path.""" + settings = model_settings("normalizing_flow", completed=False) + model = settings["train_settings"]["model"] + model["embedding_net"] = { + "type": "pooling_transformer", + "kwargs": { + "tokenizer_kwargs": { + "hidden_dims": [16], + "activation": "elu", + "batch_norm": False, + "layer_norm": True, + }, + "positional_encoder_kwargs": { + "max_vals": [1024.0, 1024.0, 3.0], + "resolutions": [0.25, 0.25, 1.0], + }, + "transformer_kwargs": { + "d_model": 16, + "dim_feedforward": 32, + "nhead": 4, + "dropout": 0.0, + "num_layers": 1, + }, + "final_net_kwargs": { + "activation": "elu", + "hidden_dims": [16], + "output_dim": EMBEDDING_OUTPUT_DIM, + "batch_norm": False, + }, + }, + } + + model = complete_model_settings(model, tokenized_data_sample()) + kwargs = model["embedding_net"]["kwargs"] + assert kwargs["tokenizer_kwargs"]["input_dim"] == NUM_FEATURES + assert kwargs["output_dim"] == EMBEDDING_OUTPUT_DIM + assert model["distribution"]["kwargs"]["context_dim"] == EMBEDDING_OUTPUT_DIM + + pm = build_model_from_kwargs( + settings={"train_settings": {"model": model}}, device="cpu" + ) + theta, context = tokenized_batch() + assert torch.isfinite(pm.loss(theta, context)) diff --git a/tests/gw/inference/test_gw_sampler_transforms.py b/tests/gw/inference/test_gw_sampler_transforms.py index 9018bf32..48684e88 100644 --- a/tests/gw/inference/test_gw_sampler_transforms.py +++ b/tests/gw/inference/test_gw_sampler_transforms.py @@ -1,6 +1,7 @@ """Tests for GWSampler._initialize_transforms with and without tokenization.""" import numpy as np +import pytest import torch from unittest.mock import MagicMock @@ -60,6 +61,7 @@ def _make_sampler_stub(domain, tokenization_settings=None): sampler.inference_parameters = INFERENCE_PARAMS sampler._minimum_frequency = None sampler._maximum_frequency = None + sampler._suppress = None return sampler @@ -199,3 +201,65 @@ def test_transformer_path_waveform_and_position_num_tokens_match(): == x["position"].shape[0] == x["drop_token_mask"].shape[0] ) + + +# --------------------------------------------------------------------------- +# Token suppression (inference-side frequency updates for tokenized models) +# --------------------------------------------------------------------------- + +TOK_WITH_DROP = { + **TOK_SETTINGS, + "drop_frequency_range": {"f_cut": {"p_cut": 0.25}}, +} + + +def test_suppress_requires_tokenized_model_with_drop_augmentation(): + domain = _make_domain() + sampler = _make_sampler_stub(domain, tokenization_settings=None) + with pytest.raises(ValueError, match="tokenized"): + sampler.suppress = [50.0, 60.0] + + sampler = _make_sampler_stub(domain, tokenization_settings=TOK_SETTINGS) + with pytest.raises(ValueError, match="drop augmentation"): + sampler.suppress = [50.0, 60.0] + + +def test_suppress_validates_interval(): + domain = _make_domain() + sampler = _make_sampler_stub(domain, tokenization_settings=TOK_WITH_DROP) + with pytest.raises(ValueError, match="f_lo < f_hi"): + sampler.suppress = [60.0, 50.0] + with pytest.raises(ValueError, match="f_lo < f_hi"): + sampler.suppress = [5.0, 60.0] # below domain f_min + with pytest.raises(ValueError, match="Unknown detectors"): + sampler.suppress = {"V1": [50.0, 60.0]} + + +def test_suppress_masks_tokens(): + from dingo.gw.transforms import UpdateFrequencyRange + + domain = _make_domain() + sampler = _make_sampler_stub(domain, tokenization_settings=TOK_WITH_DROP) + sampler.suppress = [50.0, 60.0] + assert sampler.frequency_updates + + types = [type(t) for t in sampler.transform_pre.transforms] + assert UpdateFrequencyRange in types + + context = _make_context(domain) + x = sampler.transform_pre(context) + mask = x["drop_token_mask"].numpy() + position = x["position"].numpy() + overlaps = (position[..., 1] >= 50.0) & (position[..., 0] <= 60.0) + assert np.array_equal(mask, overlaps) + assert mask.any() and not mask.all() + + +def test_no_frequency_update_has_no_update_transform(): + from dingo.gw.transforms import UpdateFrequencyRange + + domain = _make_domain() + sampler = _make_sampler_stub(domain, tokenization_settings=TOK_WITH_DROP) + sampler._initialize_transforms() + types = [type(t) for t in sampler.transform_pre.transforms] + assert UpdateFrequencyRange not in types diff --git a/tests/gw/transforms/test_tokenization_augmentation.py b/tests/gw/transforms/test_tokenization_augmentation.py new file mode 100644 index 00000000..a4fc1893 --- /dev/null +++ b/tests/gw/transforms/test_tokenization_augmentation.py @@ -0,0 +1,708 @@ +"""Tests for the tokenized-strain augmentation transforms (ported from the +DINGO-T1 branch). The StrainTokenization tests live in +test_tokenization_transforms.py.""" + +import numpy as np +import pytest + +from dingo.gw.domains import UniformFrequencyDomain, MultibandedFrequencyDomain +from dingo.gw.transforms import ( + StrainTokenization, + DropDetectors, + DropFrequenciesToUpdateRange, + DropFrequencyInterval, + DropRandomTokens, +) +from dingo.gw.transforms import NormalizePosition + + +@pytest.fixture +def strain_tokenization_setup(): + num_tokens_per_block = 40 # needs to be larger than for MFD + f_min = 20.0 + f_max = 1024.0 + T = 8.0 + domain = UniformFrequencyDomain(f_min, f_max, delta_f=1 / T) + num_f = domain.frequency_mask_length + + batch_size = 100 + waveform_h1 = np.zeros([batch_size, 1, 3, num_f]) + waveform_l1 = np.ones([batch_size, 1, 3, num_f]) + # Set real part of second detector to linearly increasing values + waveform_l1[0, 0, 0, :] *= np.arange(1, num_f + 1) + waveform = np.concatenate([waveform_h1, waveform_l1], axis=-3) + num_blocks = waveform.shape[-3] + asds = { + "H1": np.random.random([batch_size, num_f]), + "L1": np.random.random([batch_size, num_f]), + } + + sample = {"waveform": waveform, "asds": asds} + + return domain, num_tokens_per_block, num_blocks, sample + + +@pytest.fixture +def strain_tokenization_setup_no_batch(): + num_tokens_per_block = 40 # needs to be larger than for MFD + f_min = 20.0 + f_max = 1024.0 + T = 8.0 + domain = UniformFrequencyDomain(f_min, f_max, delta_f=1 / T) + num_f = domain.frequency_mask_length + + waveform_h1 = np.zeros([1, 3, num_f]) + waveform_l1 = np.ones([1, 3, num_f]) + # Set real part of second detector to linearly increasing values + waveform_l1[0, 0, :] *= np.arange(1, num_f + 1) + waveform = np.concatenate([waveform_h1, waveform_l1], axis=-3) + num_blocks = waveform.shape[-3] + asds = {"H1": np.random.random([num_f]), "L1": np.random.random([num_f])} + + sample = {"waveform": waveform, "asds": asds} + + return domain, num_tokens_per_block, num_blocks, sample + + +@pytest.fixture +def strain_tokenization_setup_single_batch(): + num_tokens_per_block = 40 # needs to be larger than for MFD + f_min = 20.0 + f_max = 1024.0 + T = 8.0 + domain = UniformFrequencyDomain(f_min, f_max, delta_f=1 / T) + num_f = domain.frequency_mask_length + + waveform_h1 = np.zeros([1, 3, num_f]) + waveform_l1 = np.ones([1, 3, num_f]) + # Set real part of second detector to linearly increasing values + waveform_l1[0, 0, :] *= np.arange(1, num_f + 1) + waveform = np.concatenate([waveform_h1, waveform_l1], axis=-3) + num_blocks = waveform.shape[-3] + asds = {"H1": np.random.random([1, num_f]), "L1": np.random.random([1, num_f])} + + sample = {"waveform": np.expand_dims(waveform, axis=0), "asds": asds} + + return domain, num_tokens_per_block, num_blocks, sample + + +@pytest.fixture +def strain_tokenization_setup_mfd(): + num_tokens_per_block = 43 # Fits exactly, no truncation/extrapolation necessary + nodes = [20.0, 34.0, 46.0, 62.0, 78.0, 1038.0] + f_min = 20.0 + f_max = 1038.0 + T = 8.0 + base_domain = UniformFrequencyDomain(f_min=f_min, f_max=f_max, delta_f=1 / T) + domain = MultibandedFrequencyDomain( + nodes=nodes, delta_f_initial=1 / T, base_domain=base_domain + ) + num_f = domain.frequency_mask_length + + batch_size = 100 + waveform_h1 = np.zeros([batch_size, 1, 3, num_f]) + waveform_l1 = np.ones([batch_size, 1, 3, num_f]) + # Set real part of second detector to linearly increasing values + waveform_l1[0, 0, 0, :] *= np.arange(1, num_f + 1) + waveform = np.concatenate([waveform_h1, waveform_l1], axis=-3) + num_blocks = waveform.shape[-3] + asds = { + "H1": np.random.random([batch_size, num_f]), + "L1": np.random.random([batch_size, num_f]), + } + + sample = {"waveform": waveform, "asds": asds} + + return domain, num_tokens_per_block, num_blocks, sample + + +@pytest.fixture +def strain_tokenization_setup_mfd_drop_last_token(): + num_tokens_per_block = 43 # Fits exactly, no truncation/extrapolation necessary + nodes = [20.0, 34.0, 46.0, 62.0, 78.0, 1038.0] + f_min = 20.0 + f_max = 1040.0 + T = 8.0 + base_domain = UniformFrequencyDomain(f_min=f_min, f_max=f_max, delta_f=1 / T) + domain = MultibandedFrequencyDomain( + nodes=nodes, delta_f_initial=1 / T, base_domain=base_domain + ) + num_f = domain.frequency_mask_length + + batch_size = 100 + waveform_h1 = np.zeros([batch_size, 1, 3, num_f]) + waveform_l1 = np.ones([batch_size, 1, 3, num_f]) + # Set real part of second detector to linearly increasing values + waveform_l1[0, 0, 0, :] *= np.arange(1, num_f + 1) + waveform = np.concatenate([waveform_h1, waveform_l1], axis=-3) + num_blocks = waveform.shape[-3] + asds = { + "H1": np.random.random([batch_size, num_f]), + "L1": np.random.random([batch_size, num_f]), + } + + sample = {"waveform": waveform, "asds": asds} + + return domain, num_tokens_per_block, num_blocks, sample + + +@pytest.mark.parametrize( + "setup", + [ + "strain_tokenization_setup", + "strain_tokenization_setup_no_batch", + "strain_tokenization_setup_single_batch", + "strain_tokenization_setup_mfd", + ], +) +def test_DropDetectors(request, setup): + domain, num_tokens_per_block, num_blocks, sample = request.getfixturevalue(setup) + + # Initialize StrainTokenization transform + token_transformation = StrainTokenization( + domain, + num_tokens_per_block=num_tokens_per_block, + ) + # Initialize DropDetectors transform + drop_transformation = DropDetectors( + num_blocks=num_blocks, + p_drop_012_detectors=[0.0, 1.0], + p_drop_hlv={"H1": 1.0, "L1": 0.0}, + ) + + # Evaluate StrainTokenization transform + out = token_transformation(sample) + # Evaluate DropDetectors transform + out = drop_transformation(out) + + # Check that mask has expected shape + assert out["drop_token_mask"].shape[-1] == num_tokens_per_block * num_blocks + # Check that mask only contains True for tokens of one detector + assert np.all(np.sum(out["drop_token_mask"], axis=-1) == num_tokens_per_block) + + trafo_dict = { + "p_drop_012_detectors": [0.3, 0.7], + "p_drop_hlv": {"H1": 0.4, "L1": 0.6}, + } + + # Initialize DropDetectors transform + drop_transformation = DropDetectors( + num_blocks=num_blocks, + p_drop_012_detectors=trafo_dict["p_drop_012_detectors"], + p_drop_hlv=trafo_dict["p_drop_hlv"], + ) + out = token_transformation(sample) + out = drop_transformation(out) + + # Check whether probability for dropping one detector aligns with p_drop_012_detectors[1] + count_dropped_tokens = np.sum(out["drop_token_mask"], axis=-1) + assert np.all(np.isin(count_dropped_tokens, [0, num_tokens_per_block])) + + # Only run tests involving probabilities if we can average over the batch dimension + if len(out["position"].shape) > 2 and out["position"].shape[0] > 1: + prob_drop_1_detector = np.mean(np.where(count_dropped_tokens > 0, 1, 0)) + assert np.isclose( + prob_drop_1_detector, + trafo_dict["p_drop_012_detectors"][1], + atol=0.1, + rtol=0.1, + ) + + # Check whether probabilities for individual detectors are consistent with p_drop_hlv + detectors = [det for det in out["asds"].keys()] + for b in range(num_blocks): + b_min, b_max = b * num_tokens_per_block, (b + 1) * num_tokens_per_block + vals = out["drop_token_mask"][..., b_min:b_max] + # Check that either 0 or num_tokens_per_block values are dropped + count_dropped_tokens = np.sum(vals, axis=-1) + assert np.all(np.isin(count_dropped_tokens, [0, num_tokens_per_block])) + prob_drop_detector = np.mean(np.where(count_dropped_tokens > 0, 1, 0)) + prob_expected = ( + trafo_dict["p_drop_012_detectors"][1] + * trafo_dict["p_drop_hlv"][detectors[b]] + ) + assert np.isclose(prob_drop_detector, prob_expected, atol=0.1, rtol=0.1) + + +@pytest.mark.parametrize( + "setup", + [ + "strain_tokenization_setup", + "strain_tokenization_setup_no_batch", + "strain_tokenization_setup_single_batch", + "strain_tokenization_setup_mfd", + ], +) +def test_DropFrequenciesToUpdateRange(request, setup): + domain, num_tokens_per_block, num_blocks, sample = request.getfixturevalue(setup) + # (1) Test cuts in frequency domain + + # Initialize StrainTokenization transform + token_transformation = StrainTokenization( + domain, + num_tokens_per_block=num_tokens_per_block, + ) + drop_dict = { + "p_cut": 0.2, + "f_max_lower_cut": 100.0, + "f_min_upper_cut": 800.0, + "p_lower_upper_both": [0.4, 0.4, 0.2], + "p_same_cut_all_detectors": 0.7, + } + drop_transformation = DropFrequenciesToUpdateRange( + domain=domain, + p_cut=drop_dict["p_cut"], + f_max_lower_cut=drop_dict["f_max_lower_cut"], + f_min_upper_cut=drop_dict["f_min_upper_cut"], + p_lower_upper_both=drop_dict["p_lower_upper_both"], + p_same_cut_all_detectors=drop_dict["p_same_cut_all_detectors"], + ) + # Evaluate transforms + out = token_transformation(sample) + out = drop_transformation(out) + + # Check that dropped tokens are either at frequencies lower than f_max_lower_cut and larger than f_min_upper_cut + dropped_f_mins = out["position"][..., 0][out["drop_token_mask"]] + dropped_f_maxs = out["position"][..., 1][out["drop_token_mask"]] + assert np.all( + np.logical_or( + dropped_f_mins < drop_dict["f_max_lower_cut"], + dropped_f_maxs > drop_dict["f_min_upper_cut"], + ) + ) + + # Only run tests involving probabilities if we can average over the batch dimension + if len(out["position"].shape) > 2 and out["position"].shape[0] > 1: + # Check that p_cut is correct + num_removed_tokens = np.sum(out["drop_token_mask"], axis=-1) + prob_cut = np.mean(np.where(num_removed_tokens > 0.0, 1.0, 0.0)) + # Very noisy because we are just averaging over 100 samples + assert np.isclose(prob_cut, drop_dict["p_cut"], atol=0.1, rtol=0.1) + + # Check that p_upper_lower_both is correct + # Find token indices which correspond to f_max_lower_cut and f_min_upper_cut + mask_lower = ( + out["position"][..., :num_tokens_per_block, 0] + <= drop_dict["f_max_lower_cut"] + ) + mask_upper = ( + out["position"][..., :num_tokens_per_block, 1] + >= drop_dict["f_min_upper_cut"] + ) + lower_blocks = [] + upper_blocks = [] + both_blocks = [] + for b in range(num_blocks): + b_min, b_max = b * num_tokens_per_block, (b + 1) * num_tokens_per_block + vals = out["drop_token_mask"][:, b_min:b_max] + num_masked_lower = np.sum(np.where(mask_lower, vals, False), axis=-1) + num_masked_upper = np.sum(np.where(mask_upper, vals, False), axis=-1) + lower = np.where(num_masked_lower > 0.0, 1.0, 0.0) + upper = np.where(num_masked_upper > 0.0, 1.0, 0.0) + both = np.logical_and(lower, upper) + lower = np.where(np.logical_and(lower, ~both), 1.0, 0.0) + upper = np.where(np.logical_and(upper, ~both), 1.0, 0.0) + assert np.isclose( + np.mean(lower), + drop_dict["p_cut"] * drop_dict["p_lower_upper_both"][0], + atol=0.1, + rtol=0.1, + ) + assert np.isclose( + np.mean(upper), + drop_dict["p_cut"] * drop_dict["p_lower_upper_both"][1], + atol=0.1, + rtol=0.1, + ) + assert np.isclose( + np.mean(both), + drop_dict["p_cut"] * drop_dict["p_lower_upper_both"][2], + atol=0.1, + rtol=0.1, + ) + + lower_blocks.append(lower) + upper_blocks.append(upper) + both_blocks.append(both) + + # Check that p_cut_all_detectors is correct + all_lower = np.logical_and(*lower_blocks) + all_upper = np.logical_and(*upper_blocks) + all_both = np.logical_and(*both_blocks) + p_cut_all_detectors = np.mean( + np.logical_or.reduce((all_lower, all_upper, all_both)) + ) + assert np.isclose( + p_cut_all_detectors, + drop_dict["p_cut"] * drop_dict["p_same_cut_all_detectors"], + atol=0.1, + rtol=0.1, + ) + + # Check that we sample the cut frequencies uniformly in UFD / uniformly in MFD bands + # Make sure to only consider bins that are completely in [f_min, f_max_lower] and [f_min_upper, f_max] + # => remove tokens at boundary + edge_mask_lower = mask_lower[..., :-1] & ~mask_lower[..., 1:] + mask_lower_strict = mask_lower.copy() + mask_lower_strict[..., :-1][edge_mask_lower] = False + edge_mask_upper = ~mask_upper[..., :-1] & mask_upper[..., 1:] + mask_upper_strict = mask_upper.copy() + mask_upper_strict[..., 1:][edge_mask_upper] = False + # Combine detectors as well as lower & both and upper & both to get better stats + masked_lower_blocks, masked_upper_blocks = [], [] + edge_mask_lower_blocks, edge_mask_upper_blocks = [], [] + for b in range(num_blocks): + b_min, b_max = b * num_tokens_per_block, (b + 1) * num_tokens_per_block + vals = out["drop_token_mask"][:, b_min:b_max] + masked_lower_blocks.append(np.where(mask_lower_strict, vals, False)) + masked_upper_blocks.append(np.where(mask_upper_strict, vals, False)) + edge_mask_lower_blocks.append( + masked_lower_blocks[-1][..., :-1] & ~masked_lower_blocks[-1][..., 1:] + ) + edge_mask_upper_blocks.append( + ~masked_upper_blocks[-1][..., :-1] & masked_upper_blocks[-1][..., 1:] + ) + num_tokens_masked_lower = np.apply_over_axes( + np.sum, np.array(masked_lower_blocks), [0, 1] + ).squeeze() # (num_tokens) + num_tokens_masked_upper = np.apply_over_axes( + np.sum, np.array(masked_upper_blocks), [0, 1] + ).squeeze() # (num_tokens) + + # Since we mask from f_min to a random f_lower and from a random f_upper to f_max, we expect the count of masked + # tokens to decrease at the lower end and increase at the upper end. + assert np.all(num_tokens_masked_lower[1:] <= num_tokens_masked_lower[:-1]) + assert np.all(num_tokens_masked_upper[1:] >= num_tokens_masked_upper[:-1]) + + num_cuts_lower = np.apply_over_axes( + np.sum, np.array(edge_mask_lower_blocks), [0, 1] + ).squeeze() # (num_tokens-1) + num_cuts_upper = np.apply_over_axes( + np.sum, np.array(edge_mask_upper_blocks), [0, 1] + ).squeeze() # (num_tokens-1) + + if isinstance(domain, UniformFrequencyDomain): + # We sample f_max_lower and f_min_upper in UFD, so we expect the masked edge tokens to be uniformly + # distributed. + non_zero_lower = num_cuts_lower[num_cuts_lower > 0.0] + non_zero_upper = num_cuts_upper[num_cuts_upper > 0.0] + if not non_zero_lower.size == 0: + assert np.isclose( + np.mean(non_zero_lower), non_zero_lower, atol=5, rtol=5 + ).all() + if not non_zero_upper.size == 0: + assert np.isclose( + np.mean(non_zero_upper), non_zero_upper, atol=5, rtol=5 + ).all() + + elif isinstance(domain, MultibandedFrequencyDomain): + # We expect tokens completely within [f_min, f_max_lower] and [f_min_upper, f_max] + # AND with the same compression factor (i.e., tokens between the same nodes) to be masked with equal + # probability. + # Lower cut + first_band_and_below_lower_cut = np.logical_and( + out["position"][0, :num_tokens_per_block, 1] < domain.nodes[1], + mask_lower[0, :], + ) + non_zero_lower = num_cuts_lower[first_band_and_below_lower_cut[:-1]] + if not non_zero_lower.size == 0: + assert np.isclose( + np.mean(non_zero_lower), non_zero_lower, atol=5, rtol=5 + ).all() + second_band_and_below_lower_cut = np.logical_and( + out["position"][0, :num_tokens_per_block, 0] > domain.nodes[1], + out["position"][0, :num_tokens_per_block, 1] < domain.nodes[2], + mask_lower[0, :], + ) + non_zero_lower_2 = num_cuts_lower[second_band_and_below_lower_cut[:-1]] + if not non_zero_lower_2.size == 0: + assert np.isclose( + np.mean(non_zero_lower_2), non_zero_lower_2, atol=5, rtol=5 + ).all() + # Upper cut + last_band_and_above_upper_cut = np.logical_and( + out["position"][0, :num_tokens_per_block, 1] > domain.nodes[-2], + mask_upper[0, :], + ) + non_zero_upper = num_cuts_upper[last_band_and_above_upper_cut[:-1]] + if not non_zero_upper.size == 0: + assert np.isclose( + np.mean(non_zero_upper), non_zero_upper, atol=5, rtol=5 + ).all() + + # Check that we never mask all tokens if f_max_lower > f_min_upper + drop_dict = { + "p_cut": 1.0, + "f_max_lower_cut": 900.0, + "f_min_upper_cut": 100.0, + "p_lower_upper_both": [0.1, 0.1, 0.8], + "p_same_cut_all_detectors": 0.7, + } + drop_transformation = DropFrequenciesToUpdateRange( + domain=domain, + p_cut=drop_dict["p_cut"], + f_max_lower_cut=drop_dict["f_max_lower_cut"], + f_min_upper_cut=drop_dict["f_min_upper_cut"], + p_lower_upper_both=drop_dict["p_lower_upper_both"], + p_same_cut_all_detectors=drop_dict["p_same_cut_all_detectors"], + ) + # Evaluate transforms + out = token_transformation(sample) + out = drop_transformation(out) + + # Check that we do not drop all tokens in one sample + num_dropped_tokens = np.sum(out["drop_token_mask"], axis=-1) + assert np.all(num_dropped_tokens < num_tokens_per_block * num_blocks) + # Check that dropped tokens are either at frequencies lower than f_max_lower_cut and larger than f_min_upper_cut + dropped_f_mins = out["position"][..., 0][out["drop_token_mask"]] + dropped_f_maxs = out["position"][..., 1][out["drop_token_mask"]] + assert np.all( + np.logical_or( + dropped_f_mins < drop_dict["f_max_lower_cut"], + dropped_f_maxs > drop_dict["f_min_upper_cut"], + ) + ) + + +@pytest.mark.parametrize( + "setup", + [ + "strain_tokenization_setup", + "strain_tokenization_setup_no_batch", + "strain_tokenization_setup_single_batch", + "strain_tokenization_setup_mfd", + ], +) +def test_DropFrequencyInterval(request, setup): + domain, num_tokens_per_block, num_blocks, sample = request.getfixturevalue(setup) + # (2) Test masking frequency interval + + # Initialize StrainTokenization transform + token_transformation = StrainTokenization( + domain, + num_tokens_per_block=num_tokens_per_block, + ) + # Test mask_glitch + drop_dict = { + "p_glitch_per_detector": 0.4, + "f_min": 100.0, + "f_max": 500.0, + "max_width": 100.0, + } + drop_transformation = DropFrequencyInterval( + domain=domain, + p_per_detector=drop_dict["p_glitch_per_detector"], + f_min=drop_dict["f_min"], + f_max=drop_dict["f_max"], + max_width=drop_dict["max_width"], + ) + # Evaluate transforms + out = token_transformation(sample) + out = drop_transformation(out) + + # Check that no tokens are masked outside [f_min, f_max] + # It can happen that drop_dict['f_min'] falls in the middle of a token. Such a token might be dropped, resulting in + # the f_min of the token being lower than drop_dict['f_min']. The same can happen for drop_dict['f_max']. + # To exclude this case, we select f_max (f_min) of the dropped tokens when comparing to drop_dict['f_min'] + # (drop_dict['f_max']) + dropped_f_mins = out["position"][..., 1][out["drop_token_mask"]] + dropped_f_maxs = out["position"][..., 0][out["drop_token_mask"]] + assert np.all( + np.logical_and( + dropped_f_mins >= drop_dict["f_min"], dropped_f_maxs <= drop_dict["f_max"] + ) + ) + + # Check that masked ranges are not wider than max_width + mask_has_true = [] + for b in range(num_blocks): + b_min, b_max = b * num_tokens_per_block, (b + 1) * num_tokens_per_block + vals = out["drop_token_mask"][..., b_min:b_max] + first_true_idx = np.argmax(vals, axis=-1) + last_true_idx = vals.shape[-1] - 1 - np.argmax(vals[..., ::-1], axis=-1) + + # Initialize edge mask + edge_mask_lower = np.zeros_like(vals, dtype=bool) + edge_mask_upper = np.zeros_like(vals, dtype=bool) + if len(out["position"].shape) > 2: + batch_indices = np.arange(vals.shape[0]) + edge_mask_lower[batch_indices, first_true_idx] = True + edge_mask_upper[batch_indices, last_true_idx] = True + else: + edge_mask_lower[first_true_idx] = True + edge_mask_upper[last_true_idx] = True + + # If a row has no True values, clear accidentally masked values at the beginning or end + has_true = np.any(vals, axis=-1) + edge_mask_lower[~has_true] = False + edge_mask_upper[~has_true] = False + + # Depending on the position of the tokens and the values of drop_dict['f_min'] and drop_dict['f_max'], it can + # happen that we mask a larger frequency range than drop_dict['max_width']. This is the case when + # drop_dict['f_min'] is located in the middle/at the upper end of a token and drop_dict['f_max'] is located in + # the middle/at the lower end of a token. + # Similar to the case of testing that tokens fall in the correct range [drop_dict['f_min'], drop_dict['f_max']], + # we select f_max (f_min) of the dropped tokens for the lower edge (upper edge) + dropped_f_mins_lower = out["position"][..., b_min:b_max, 1][edge_mask_lower] + dropped_f_maxs_upper = out["position"][..., b_min:b_max, 0][edge_mask_upper] + # If we only dropped one token, dropped_f_mins_lower (based on f_max) is larger than dropped_f_maxs_upper + # (based on f_min). We mask these tokens during the check + diff_f = dropped_f_maxs_upper - dropped_f_mins_lower + assert np.all(np.where(diff_f > 0.0, diff_f <= drop_dict["max_width"], True)) + + # Save has_true for next check + mask_has_true.append(has_true) + + # Only run tests involving probabilities if we can average over the batch dimension + if len(out["position"].shape) > 2 and out["position"].shape[0] > 1: + # Check that p_glitch_per_detector is correct + masked_blocks = np.concatenate(mask_has_true) + prob_glitch = np.mean(masked_blocks) + assert np.isclose( + prob_glitch, drop_dict["p_glitch_per_detector"], atol=0.1, rtol=0.1 + ) + + mask_tokens = np.logical_and( + out["position"][0, :num_tokens_per_block, 0] > drop_dict["f_min"], + out["position"][0, :num_tokens_per_block, 1] < drop_dict["f_max"], + ) + # Check that f_lower and f_upper are sampled uniformly in UFD / uniformly in MFD bands between f_min and f_max + if isinstance(domain, UniformFrequencyDomain): + # We sample the lower edge uniformly between f_min and f_max, so we expect the lower edge to be uniformly + # distributed. + lower_edge_counts = np.sum(edge_mask_lower, axis=0) + non_zero_lower = lower_edge_counts[mask_tokens] + if not non_zero_lower.size == 0: + assert np.isclose( + np.mean(non_zero_lower), non_zero_lower, atol=5, rtol=5 + ).all() + + # Since the sampling range of f_max depends on the sampled f_min, it isn't straight forward to assume + # something about the statistics. + + +@pytest.mark.parametrize( + "setup", + [ + "strain_tokenization_setup", + "strain_tokenization_setup_no_batch", + "strain_tokenization_setup_single_batch", + "strain_tokenization_setup_mfd", + ], +) +def test_DropRandomTokens(request, setup): + domain, num_tokens_per_block, num_blocks, sample = request.getfixturevalue(setup) + # Drop random tokens + # Initialize StrainTokenization transform + token_transformation = StrainTokenization( + domain, + num_tokens_per_block=num_tokens_per_block, + ) + drop_dict = { + "p_drop": 0.2, + "max_num_tokens": num_tokens_per_block, + } + drop_trafo = DropRandomTokens( + p_drop=drop_dict["p_drop"], + max_num_tokens=drop_dict["max_num_tokens"], + ) + + # Evaluate transforms + out = token_transformation(sample) + out = drop_trafo(out) + + # Check that not more than max_num_tokens are masked per sample + num_removed_tokens = np.sum(out["drop_token_mask"], axis=-1) + assert np.all(num_removed_tokens <= drop_dict["max_num_tokens"]) + + # Only check probabilities if we can average over the batch dimension + if len(out["position"].shape) > 2 and out["position"].shape[0] > 1: + # Check that p_drop is correct + prob_drop = np.mean(np.where(num_removed_tokens > 0.0, 1.0, 0.0)) + assert np.isclose(prob_drop, drop_dict["p_drop"], atol=0.1, rtol=0.1) + + # Check that dropped tokens are uniformly distributed + hist_removed_tokens = np.sum(out["drop_token_mask"], axis=0) + mean_removed_tokens = np.mean(hist_removed_tokens) + assert np.all( + np.isclose(mean_removed_tokens, hist_removed_tokens, atol=5, rtol=5) + ) + + +@pytest.mark.parametrize( + "setup", + [ + "strain_tokenization_setup", + "strain_tokenization_setup_no_batch", + "strain_tokenization_setup_single_batch", + "strain_tokenization_setup_mfd", + ], +) +def test_NormalizePosition(request, setup): + domain, num_tokens_per_block, num_blocks, sample = request.getfixturevalue(setup) + + # Initialize StrainTokenization transform + token_transformation = StrainTokenization( + domain, + num_tokens_per_block=num_tokens_per_block, + ) + trafo = NormalizePosition() + # Evaluate transforms + out = token_transformation(sample) + original_positions = out["position"].copy() + out = trafo(out) + + # Check that blocks remain the same + assert np.all(out["position"][..., 2] == original_positions[..., 2]) + + # Check that f_min and f_max are correctly normalized + f_min = np.min(original_positions[..., 0]) + f_max = np.max(original_positions[..., 1]) + f_min_norm = (original_positions[..., 0] - f_min) / (f_max - f_min) + f_max_norm = (original_positions[..., 1] - f_min) / (f_max - f_min) + assert np.all(f_min_norm == out["position"][..., 0]) + assert np.all(f_max_norm == out["position"][..., 1]) + + +def test_UpdateFrequencyRange(strain_tokenization_setup_no_batch): + """UpdateFrequencyRange (inference side, unbatched) masks tokens outside the + updated frequency range and inside suppressed intervals; T1 left it untested.""" + from dingo.gw.transforms import UpdateFrequencyRange + + domain, num_tokens_per_block, num_blocks, sample = ( + strain_tokenization_setup_no_batch + ) + tokenize = StrainTokenization(domain, num_tokens_per_block=num_tokens_per_block) + + # f_min update: tokens starting below the new minimum are masked. + out = tokenize(dict(sample)) + trafo = UpdateFrequencyRange( + minimum_frequency=100.0, domain=domain, ifos=["H1", "L1"] + ) + out = trafo(out) + expected = out["position"][..., 0] < 100.0 + assert np.array_equal(out["drop_token_mask"], expected) + + # f_max update per detector: only that detector's tokens are affected. + out = tokenize(dict(sample)) + trafo = UpdateFrequencyRange( + maximum_frequency={"L1": 500.0}, domain=domain, ifos=["H1", "L1"] + ) + out = trafo(out) + l1 = out["position"][..., 2] == 1 + # H1 gets the default threshold domain.f_max, which still masks the final + # (zero-padded) token whose extrapolated f_max exceeds it. + expected_h1 = out["position"][~l1][..., 1] > domain.f_max + assert np.array_equal(out["drop_token_mask"][~l1], expected_h1) + expected_l1 = out["position"][l1][..., 1] > 500.0 + assert np.array_equal(out["drop_token_mask"][l1], expected_l1) + + # Suppressed interval: tokens overlapping [f_lo, f_hi] are masked. + out = tokenize(dict(sample)) + trafo = UpdateFrequencyRange( + suppress_range=[100.0, 200.0], domain=domain, ifos=["H1", "L1"] + ) + out = trafo(out) + overlaps = (out["position"][..., 1] >= 100.0) & (out["position"][..., 0] <= 200.0) + assert np.array_equal(out["drop_token_mask"], overlaps) + assert out["drop_token_mask"].any() and not out["drop_token_mask"].all()