diff --git a/dingo/core/SVD.py b/dingo/core/SVD.py new file mode 100644 index 000000000..a15dabb3b --- /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/density/__init__.py b/dingo/core/density/__init__.py index e465636a5..af76a40dd 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 53cd08858..725359205 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 @@ -13,10 +14,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 +26,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( @@ -71,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/__init__.py b/dingo/core/nn/__init__.py index e69de29bb..ed46a48c9 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 77529f79c..c58cf4ef9 100644 --- a/dingo/core/nn/cfnets.py +++ b/dingo/core/nn/cfnets.py @@ -1,13 +1,11 @@ -import copy +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 create_enet_with_projection_layer_and_dense_resnet - -from dingo.core.nn.enets import DenseResidualNet +from dingo.core.nn.resnet import DenseResidualNet class ContinuousFlow(nn.Module): @@ -29,37 +27,67 @@ 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",), ): """ Parameters ---------- 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 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 - 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) 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 @@ -157,55 +185,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 @@ -213,16 +233,15 @@ 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"], + layer_norm=kwargs.get("layer_norm", False), context_features=glu_dim, ) @@ -230,8 +249,9 @@ def create_cf( 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 @@ -257,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), ) @@ -271,14 +292,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 @@ -298,7 +320,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. @@ -306,12 +328,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 206730f84..56d432d91 100755 --- a/dingo/core/nn/enets.py +++ b/dingo/core/nn/enets.py @@ -1,11 +1,38 @@ -"""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. +* ``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, Union, List -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.nn.resnet import DenseResidualNet +from dingo.core.registry import CONTEXT_MERGERS, EMBEDDING_NETS from dingo.core.utils import torchutils @@ -157,84 +184,175 @@ def forward(self, x, **_): return x -class DenseResidualNet(nn.Module): +@EMBEDDING_NETS.register("dense_svd") +class DenseSVDEmbedding(nn.Sequential): """ - 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. + 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. - Module specs - -------- - input dimension: (batch_size, input_dim) - output dimension: (batch_size, output_dim) + 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_dim: int, + input_dims: List[int], output_dim: int, hidden_dims: Tuple, - activation: Callable = F.elu, + svd: dict, + activation: str = "elu", dropout: float = 0.0, batch_norm: bool = True, - context_features: int = None, + V_rb_list: Union[Tuple, None] = None, ): """ Parameters ---------- - input_dim : int - dimension of the input to this module + 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 of this module + output dimension (dimension of the embedded context) 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. + 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 directly. Usually None: the projection is seeded via the + initialize_weights hook instead. """ - - 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) - ] + 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, ) - 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)] + super().__init__(projection, resnet) + self.output_dim = output_dim + self.svd_settings = svd + + @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)} + + 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. - 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 + 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): @@ -277,6 +395,113 @@ 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 + + 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.""" + 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, @@ -347,17 +572,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 b0727df13..70e82c7f8 100644 --- a/dingo/core/nn/nsf.py +++ b/dingo/core/nn/nsf.py @@ -3,16 +3,13 @@ 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.nn.resnet import DenseResidualNet 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): @@ -42,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, @@ -82,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. @@ -95,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: @@ -103,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, @@ -116,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, @@ -124,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, @@ -189,56 +230,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( @@ -246,8 +294,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). @@ -264,82 +310,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) - model = FlowWrapper(flow, embedding_net) - return model - - if __name__ == "__main__": pass diff --git a/dingo/core/nn/resnet.py b/dingo/core/nn/resnet.py new file mode 100644 index 000000000..2856c9e5a --- /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 000000000..2288ad0b0 --- /dev/null +++ b/dingo/core/nn/transformer.py @@ -0,0 +1,610 @@ +"""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, + } + + +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/core/posterior_models/__init__.py b/dingo/core/posterior_models/__init__.py index 148346f98..8281c1fa4 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 68508d199..a2b66b2b0 100755 --- a/dingo/core/posterior_models/base_model.py +++ b/dingo/core/posterior_models/base_model.py @@ -18,27 +18,32 @@ 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 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__( self, model_filename: str = None, metadata: dict = None, - initial_weights: dict = None, device: str = "cuda", load_training_info: bool = True, ): @@ -51,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 """ @@ -63,10 +66,10 @@ 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: + 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(). @@ -94,8 +97,61 @@ 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. 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). + """ + embedding_settings = self.model_kwargs.get("embedding_net") + if embedding_settings is None: + return None + 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: + 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 + # 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): + def sample(self, context: Optional[dict] = None, num_samples: int = 1): """ Sample parameters theta from the posterior model, @@ -103,9 +159,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. @@ -117,7 +174,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, @@ -129,9 +186,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. @@ -143,7 +201,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, @@ -153,9 +211,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 ------- @@ -165,7 +225,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. @@ -174,9 +234,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 ------- @@ -232,6 +293,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(), @@ -333,6 +414,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"] @@ -469,6 +553,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( @@ -482,15 +572,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() @@ -509,12 +600,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 2f7571032..60c1a83a4 100644 --- a/dingo/core/posterior_models/build_model.py +++ b/dingo/core/posterior_models/build_model.py @@ -1,7 +1,15 @@ -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 +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 ( + CONTEXT_MERGERS, + EMBEDDING_NETS, + NEURAL_DISTRIBUTIONS, +) from dingo.core.utils.backward_compatibility import ( torch_load_with_fallback, update_model_config, @@ -11,12 +19,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' distribution 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 +38,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: @@ -50,56 +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"] - 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(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: list): +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. - * 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 + 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). + + 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: 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. + 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[1].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["embedding_kwargs"]["added_context"] = True - model_kwargs["posterior_kwargs"]["context_dim"] = ( - model_kwargs["embedding_kwargs"]["output_dim"] + gnpe_proxy_dim + 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)})." ) - except IndexError: - model_kwargs["embedding_kwargs"]["added_context"] = False - model_kwargs["posterior_kwargs"]["context_dim"] = model_kwargs[ - "embedding_kwargs" - ]["output_dim"] + embedding_settings["kwargs"] = embedding_cls.complete_settings( + embedding_settings.get("kwargs", {}), sample_batch + ) + 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: + 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 2f9b09bbc..0cbe4de8b 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). @@ -49,12 +49,9 @@ class ContinuousFlowPosteriorModel(BasePosteriorModel): 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,14 +119,12 @@ 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) - - def sample(self, *context: torch.Tensor, num_samples: int = None): + 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): """ Sample parameters theta from the posterior model, @@ -139,9 +134,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 +145,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 +170,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 +188,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 +198,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 +221,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 +236,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 +247,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 3e53a61c8..2302ec7ff 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) @@ -34,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): """ @@ -59,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. @@ -69,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 ------- @@ -87,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 8481aeda9..d04cd56d1 100644 --- a/dingo/core/posterior_models/normalizing_flow.py +++ b/dingo/core/posterior_models/normalizing_flow.py @@ -1,12 +1,11 @@ -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, - create_nsf_wrapped, -) +from dingo.core.nn.nsf import FlowWrapper, create_nsf_model -class NormalizingFlowPosteriorModel(BasePosteriorModel): +@NEURAL_DISTRIBUTIONS.register("normalizing_flow") +class NormalizingFlowPosteriorModel(NeuralDistribution): """ Posterior model based on a (discrete) normalizing flow. @@ -38,25 +37,27 @@ 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): - 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 5fbaca74e..eece3a2dd 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) @@ -36,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: @@ -48,7 +51,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. @@ -56,8 +59,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 ------- @@ -65,7 +69,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/registry.py b/dingo/core/registry.py new file mode 100644 index 000000000..3cddca3b0 --- /dev/null +++ b/dingo/core/registry.py @@ -0,0 +1,152 @@ +""" +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/result.py b/dingo/core/result.py index 67017271a..5b905a9f8 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/samplers.py b/dingo/core/samplers.py index efc4fc1fa..0479fb98c 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 @@ -151,21 +151,25 @@ 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 = x.unsqueeze(0) - x = [x] + x = {k: v.unsqueeze(0) for k, v in x.items()} 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 +279,16 @@ 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] + 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 = [] + 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)) @@ -363,14 +369,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 @@ -490,12 +496,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/core/transforms.py b/dingo/core/transforms.py index 7bb208048..09c10da82 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 af48f34ef..a362ac4d5 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"] @@ -147,3 +157,48 @@ 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"] + + # 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 1b0e2bbeb..0711b36c4 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 81077046a..bf785798c 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 da856d9b1..dded52a14 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 06e6aa96c..a41963ce6 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 2d109d7bb..5d6c33124 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/SVD.py b/dingo/gw/SVD.py index e77b7761d..df637d778 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 a52ca9865..6e369337c 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 e8911f87c..bad3a3869 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 af750f5be..86997c409 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/gwutils.py b/dingo/gw/gwutils.py index f56252581..68b13d017 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 2a1b63d89..8f110ca2a 100644 --- a/dingo/gw/inference/gw_samplers.py +++ b/dingo/gw/inference/gw_samplers.py @@ -32,14 +32,17 @@ GetDetectorTimes, DecimateWaveformsAndASDS, MaskDataForFrequencyRangeUpdate, + NormalizePosition, + SelectKeys, + StrainTokenization, + UpdateFrequencyRange, ) class SamplerProtocol(Protocol): base_model_metadata: dict - def _initialize_transforms(self) -> None: - ... + def _initialize_transforms(self) -> None: ... class _GWMixinProtocol(SamplerProtocol): @@ -65,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" @@ -128,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): @@ -135,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): @@ -147,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 @@ -157,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): @@ -173,7 +223,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 ): @@ -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, @@ -309,19 +360,61 @@ 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"), - ] + ) + ) + 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), + ) + ) + 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. + 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/ls_cli.py b/dingo/gw/ls_cli.py index 5c401bee7..7ce09d5ed 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 ca62856bb..6a09fae48 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 ( @@ -17,11 +13,17 @@ AddWhiteNoiseComplex, SelectStandardizeRepackageParameters, RepackageStrainsAndASDS, - UnpackDict, + SelectKeys, GNPECoalescenceTimes, SampleExtrinsicParameters, GetDetectorTimes, CropMaskStrainRandom, + StrainTokenization, + DropDetectors, + DropFrequenciesToUpdateRange, + DropFrequencyInterval, + DropRandomTokens, + NormalizePosition, ) from dingo.gw.noise.asd_dataset import ASDDataset from dingo.gw.prior import default_inference_parameters @@ -29,6 +31,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 +168,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 +211,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 @@ -176,12 +238,87 @@ 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), + ) + ) + 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: + 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(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. @@ -191,151 +328,72 @@ 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, - UnpackDict, + SelectKeys, CropMaskStrainRandom, - ], + StrainTokenization, + ] + 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, + 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 ac1420b0f..3710eff6f 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 @@ -12,13 +13,14 @@ 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, - build_svd_for_embedding_network, + initialization_dataloader, ) from dingo.core.utils.trainutils import RuntimeLimits from dingo.core.utils import ( @@ -28,7 +30,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,13 +84,12 @@ 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. - 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 ---------- @@ -101,7 +102,7 @@ def prepare_training_new( Returns ------- - (BasePosteriorModel, WaveformDataset) + (NeuralDistribution, WaveformDataset) """ data_settings = deepcopy(train_settings["data"]) # Optionally copy files to local and update path @@ -115,23 +116,8 @@ 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 = {} - # 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. - - if train_settings["model"].get("embedding_kwargs", None): - # First, build the SVD for seeding the embedding network. - 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, - **train_settings["model"]["embedding_kwargs"]["svd"], - ) + update_model_config(train_settings["model"]) # Map old schemas forward. # 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 @@ -140,14 +126,10 @@ 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) - # 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, @@ -159,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 @@ -180,7 +192,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 +209,7 @@ def prepare_training_resume( Returns ------- - (BasePosteriorModel, WaveformDataset) + (NeuralDistribution, WaveformDataset) """ pm = build_model_from_kwargs( @@ -232,7 +244,7 @@ def prepare_training_resume( def initialize_stage( - pm: BasePosteriorModel, + pm: NeuralDistribution, wfd: WaveformDataset, stage: dict, num_workers: int, @@ -248,7 +260,7 @@ def initialize_stage( Parameters ---------- - pm : BasePosteriorModel + pm : NeuralDistribution wfd : WaveformDataset stage : dict Settings specific to current stage of training @@ -301,7 +313,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 +321,7 @@ def train_stages( Parameters ---------- - pm : BasePosteriorModel + pm : NeuralDistribution wfd : WaveformDataset train_dir : str Directory for saving checkpoints and train history. diff --git a/dingo/gw/transforms/__init__.py b/dingo/gw/transforms/__init__.py index 75e3785e8..e1be056ab 100644 --- a/dingo/gw/transforms/__init__.py +++ b/dingo/gw/transforms/__init__.py @@ -5,4 +5,15 @@ 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 ( + DETECTOR_DICT, + DETECTOR_DICT_INVERSE, + DropDetectors, + DropFrequenciesToUpdateRange, + DropFrequencyInterval, + DropRandomTokens, + NormalizePosition, + StrainTokenization, + UpdateFrequencyRange, +) diff --git a/dingo/gw/transforms/detector_transforms.py b/dingo/gw/transforms/detector_transforms.py index 4daf4de0f..da8c45a53 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/dingo/gw/transforms/general_transforms.py b/dingo/gw/transforms/general_transforms.py index 94eeba7dd..cef3d8abf 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/dingo/gw/transforms/tokenization_transforms.py b/dingo/gw/transforms/tokenization_transforms.py new file mode 100644 index 000000000..7faef5ff7 --- /dev/null +++ b/dingo/gw/transforms/tokenization_transforms.py @@ -0,0 +1,1220 @@ +from typing import Optional + +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: + """ + 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." + ) + + +# --------------------------------------------------------------------------- +# 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/dingo/pipe/default_settings.py b/dingo/pipe/default_settings.py index 53edc4415..870f76967 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 9138066bc..184832d49 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 97b2025ea..674a19d5b 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 4aef4a518..801522fe3 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 747352d69..697cc1c0c 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 7bdf87ee5..c4a044019 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 40ead7e5e..6b3593090 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 1095f6d54..9bc5af2eb 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 d3e44f089..5eee2aa27 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/examples/transformer_model/train_settings.yaml b/examples/transformer_model/train_settings.yaml new file mode 100644 index 000000000..1d0560a13 --- /dev/null +++ b/examples/transformer_model/train_settings.yaml @@ -0,0 +1,123 @@ +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). + # 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: + 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 25a0c060e..c6fb94cc2 100644 --- a/tests/core/test_build_model.py +++ b/tests/core/test_build_model.py @@ -1,87 +1,764 @@ +""" +Tests for the model build path. + +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 +settings/checkpoints through the update_model_config boundary (see +hackathon/NN_Build_System_Design.md). +""" + +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, + complete_model_settings, ) +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(completed=True): + """Embedding-net kwargs; completed=False gives user-style settings (no dims).""" + kwargs = { + "svd": {"size": 10}, + "output_dim": EMBEDDING_OUTPUT_DIM, + "hidden_dims": [32, 16, 8], + "activation": "elu", + "dropout": 0.0, + "batch_norm": True, + } + if completed: + kwargs["input_dims"] = list(DATA_SHAPE) + return kwargs + + +def nsf_distribution_kwargs(completed=True): + kwargs = { + "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", + }, + } + if completed: + kwargs["theta_dim"] = NUM_PARAMETERS + kwargs["context_dim"] = EMBEDDING_OUTPUT_DIM + return kwargs -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_distribution_kwargs(completed=True): + kwargs = { + "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}, + }, + } + if completed: + kwargs["theta_dim"] = NUM_PARAMETERS + kwargs["context_dim"] = EMBEDDING_OUTPUT_DIM + return kwargs -def _settings(posterior_model_type="normalizing_flow"): +def model_settings(model_type, completed=True): + if model_type == "normalizing_flow": + distribution_kwargs = nsf_distribution_kwargs(completed) + else: + 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": { - "input_dim": 3, - "context_dim": None, - "num_flow_steps": 2, - "base_transform_kwargs": BASE_TRANSFORM_KWARGS, + "distribution": {"type": model_type, "kwargs": distribution_kwargs}, + "embedding_net": { + "type": "dense_svd", + "kwargs": embedding_kwargs(completed), }, } } } -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 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["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 = {"waveform": torch.rand(BATCH_SIZE, *DATA_SHAPE)} + if with_gnpe_proxies: + context["context_parameters"] = 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"]["distribution"]["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) + + +# ----------------------------------------------------------------------------------- +# complete_model_settings: per-architecture dimension inference +# ----------------------------------------------------------------------------------- + + +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)) + + # 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_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 completed["context_merger"] == { + "type": "concat", + "kwargs": {"num_context_parameters": GNPE_PROXY_DIM}, + } + assert ( + 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 +# ----------------------------------------------------------------------------------- -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) -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(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() -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") + 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_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_normalizing_flow_with_gnpe_context(): + """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["context_merger"] = { + "type": "concat", + "kwargs": {"num_context_parameters": GNPE_PROXY_DIM}, + } + model["distribution"]["kwargs"]["context_dim"] = ( + EMBEDDING_OUTPUT_DIM + GNPE_PROXY_DIM ) - 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 + 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,) + + # 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 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_net"] + model["distribution"]["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) + + 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,) + + +# ----------------------------------------------------------------------------------- +# Parameter-contract accessors (interface consumed by the factorized sampler) +# ----------------------------------------------------------------------------------- -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_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 == [] + + +# ----------------------------------------------------------------------------------- +# Data-driven weight initialization (init_data_spec / initialize_weights hooks) +# ----------------------------------------------------------------------------------- + + +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 + 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, 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, + } + + # 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])) - 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 + +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 +# ----------------------------------------------------------------------------------- + + +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": old["posterior_kwargs"], + "embedding_net_kwargs": old["embedding_kwargs"], + } + 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.""" + 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 + + +@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_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) + + +# ----------------------------------------------------------------------------------- +# 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) + + 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) + + +# ----------------------------------------------------------------------------------- +# 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)) + + +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/core/test_nsf.py b/tests/core/test_nsf.py index 0e3f73b0a..df87fcb67 100644 --- a/tests/core/test_nsf.py +++ b/tests/core/test_nsf.py @@ -2,12 +2,13 @@ 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, + MLPContextMerger, + 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 @@ -116,15 +117,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 +158,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 +170,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 +187,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 +214,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 +224,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 ( @@ -230,27 +237,150 @@ 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"]) - model = create_nsf_with_rb_projection_embedding_net( - d.nde_kwargs, d.embedding_net_kwargs + embedding_net = DenseSVDEmbedding(**completed) + assert embedding_net.input_keys == ("waveform",) + assert embedding_net.output_dim == kwargs["output_dim"] + + 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 ) - loss = -model(d.y, d.x, d.z) + # State-dict layout matches the historic builder (old checkpoints). + legacy = create_enet_with_projection_layer_and_dense_resnet( + **d.embedding_net_kwargs + ) + 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." assert torch.all(loss > 0) and torch.all(loss < 40), ( "Unexpected log prob encountered. Network initialization or " "normalization seems broken." ) - with pytest.raises(ValueError): - model(d.y, d.z, d.x) - with pytest.raises(ValueError): - model(d.y, d.x, d.z, d.z) - with pytest.raises(RuntimeError): - model(d.y, d.x, d.x) + +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 + ) + + +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_registry.py b/tests/core/test_registry.py new file mode 100644 index 000000000..0b346b033 --- /dev/null +++ b/tests/core/test_registry.py @@ -0,0 +1,132 @@ +"""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 the distribution 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"]["distribution"][ + "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 diff --git a/tests/core/test_resnet.py b/tests/core/test_resnet.py new file mode 100644 index 000000000..5051b55ba --- /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 000000000..e0d098010 --- /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/core/test_unconditional_density_estimation.py b/tests/core/test_unconditional_density_estimation.py index 38b6a9306..4660106b3 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): diff --git a/tests/gw/inference/__init__.py b/tests/gw/inference/__init__.py new file mode 100644 index 000000000..e69de29bb 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 000000000..48684e889 --- /dev/null +++ b/tests/gw/inference/test_gw_sampler_transforms.py @@ -0,0 +1,265 @@ +"""Tests for GWSampler._initialize_transforms with and without tokenization.""" + +import numpy as np +import pytest +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 + sampler._suppress = 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] + ) + + +# --------------------------------------------------------------------------- +# 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/test_time_alignment_validation.py b/tests/gw/test_time_alignment_validation.py new file mode 100644 index 000000000..c1ecedc63 --- /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 267dc014f..e71edd72c 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: diff --git a/tests/gw/transforms/test_general_transforms.py b/tests/gw/transforms/test_general_transforms.py index 19be7ec33..be8c654f5 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 diff --git a/tests/gw/transforms/test_tokenization_augmentation.py b/tests/gw/transforms/test_tokenization_augmentation.py new file mode 100644 index 000000000..a4fc1893f --- /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() diff --git a/tests/gw/transforms/test_tokenization_transforms.py b/tests/gw/transforms/test_tokenization_transforms.py new file mode 100644 index 000000000..8421931a6 --- /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)