NN build system: registry, dict batches, generic build path, transformer merge (steps 1-6)

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()}

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,

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)},

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