Core GNPE extensions for BNS support

dingo/core/samplers.py

@@ -366,6 +366,7 @@ def __init__(
         model: BasePosteriorModel,
         init_sampler: Sampler,
         num_iterations: int = 1,
+        fixed_context_parameters: Optional[dict] = None,
     ):
         """
         Parameters
@@ -375,8 +376,16 @@ def __init__(
             Used for generating initial samples
         num_iterations : int
             Number of GNPE iterations to be performed by sampler.
+        fixed_context_parameters : dict, optional
+            Dictionary of context parameters to hold fixed (not iterated) during
+            GNPE sampling. Keys are parameter names (e.g., "chirp_mass_proxy"),
+            values are the fixed float values. These are inserted into the
+            extrinsic_parameters at each iteration without being updated by the
+            Gibbs loop. Useful for single-step GNPE where a proxy parameter
+            (e.g., chirp mass) is known ahead of time.
         """
         self.gnpe_parameters = []  # Should be set in subclass _initialize_transform()
+        self.fixed_context_parameters = fixed_context_parameters or {}
 
         super().__init__(model)
         self.init_sampler = init_sampler
@@ -415,6 +424,28 @@ def gnpe_proxy_parameters(self):
     def _kernel_log_prob(self, samples):
         raise NotImplementedError("To be implemented in subclass.")
 
+    def _preprocess_data(self, context):
+        """
+        One-time data preprocessing before the GNPE Gibbs loop. Converts raw
+        context data (e.g., strain, ASDs) into the tensor format expected by the
+        network.
+
+        By default, delegates to init_sampler.transform_pre. Subclasses can
+        override to add additional preprocessing (e.g., heterodyning, decimation)
+        that should happen once rather than every iteration.
+
+        Parameters
+        ----------
+        context : dict
+            Raw context data.
+
+        Returns
+        -------
+        torch.Tensor
+            Preprocessed data tensor.
+        """
+        return self.init_sampler.transform_pre(context)
+
     def _run_sampler(
         self,
         num_samples: int,
@@ -423,7 +454,7 @@ def _run_sampler(
         if context is None:
             raise ValueError("self.context must be set to run sampler.")
 
-        data_ = self.init_sampler.transform_pre(context)
+        data_ = self._preprocess_data(context)
 
         # TODO: Reimplement outlier removal in IterationTracker? Save setting somewhere.
         if self.remove_init_outliers == 0.0:
@@ -442,6 +473,14 @@ def _run_sampler(
             inds = torch.where(init_samples["log_prob"] >= thr)[0][:num_samples]
             init_samples = {k: v[inds] for k, v in init_samples.items()}
 
+        # Insert fixed context parameters into initialization samples.
+        init_samples.update(
+            {
+                k: torch.ones(num_samples) * v
+                for k, v in self.fixed_context_parameters.items()
+            }
+        )
+
         # We could be starting with either the GNPE parameters *or* their proxies,
         # depending on the nature of the initialization network.
 
@@ -473,8 +512,17 @@ def _run_sampler(
                 x["extrinsic_parameters"] = proxies.copy()
             else:
                 x["extrinsic_parameters"] = {
-                    k: x["extrinsic_parameters"][k] for k in self.gnpe_parameters
+                    k: x["extrinsic_parameters"][k]
+                    for k in self.gnpe_parameters
+                    if k + "_proxy" not in self.fixed_context_parameters
+                }
+            # Add fixed context parameters (constant across iterations).
+            x["extrinsic_parameters"].update(
+                {
+                    k: torch.ones(num_samples, dtype=torch.float32) * v
+                    for k, v in self.fixed_context_parameters.items()
                 }
+            )
 
             # TODO: Depending on whether start_with_proxies is True, this might end up
             #  comparing proxies vs gnpe_parameters for the first iteration.
@@ -510,10 +558,14 @@ def _run_sampler(
             # Extract the proxy parameters from x["extrinsic_parameters"]. These have
             # not been standardized. They are persistent from before sampling took place,
             # since this is when they were placed here and their values should not have
-            # changed.
-            proxies = {
-                p: x["extrinsic_parameters"][p] for p in self.gnpe_proxy_parameters
-            }
+            # changed. Include both iterated proxies and fixed context parameters.
+            proxies = {}
+            for p in self.gnpe_proxy_parameters:
+                if p in x["extrinsic_parameters"]:
+                    proxies[p] = x["extrinsic_parameters"][p]
+            for k in self.fixed_context_parameters:
+                if k in x["extrinsic_parameters"]:
+                    proxies[k] = x["extrinsic_parameters"][k]
 
             print(
                 f"it {i}.\tmin pvalue: {self.iteration_tracker.pvalue_min:.3f}"
@@ -535,7 +587,9 @@ def _run_sampler(
             # In this case it makes sense to save the log_prob and the proxy parameters.
 
             samples = x["parameters"]
-            samples["log_prob"] = x["log_prob"] + proxy_log_prob
+            samples["log_prob"] = x["log_prob"] + proxy_log_prob.to(
+                x["log_prob"].device
+            )
 
             # The log_prob returned by gnpe is not just the log_prob over parameters
             # theta, but instead the log_prob in the *joint* space q(theta,theta^|x),
@@ -567,7 +621,39 @@ def _run_sampler(
 
         # Safety check for unconditional flows. Make sure the proxies haven't changed.
         if start_with_proxies and self.num_iterations == 1:
-            for k in proxies:
+            for k in init_proxies:
                 assert torch.equal(proxies[k], init_proxies[k])
 
         return samples
+
+
+class FixedInitSampler:
+    """
+    Lightweight initialization sampler that provides fixed parameter values.
+
+    Used in place of a trained init network when proxy parameters are known
+    ahead of time (e.g., chirp_mass_proxy for single-step GNPE). Returns the
+    same fixed values for every sample.
+
+    Parameters
+    ----------
+    init_parameters : dict
+        Parameter names and their fixed values.
+    log_prob : float
+        Log probability to assign to each sample. Default is 0 (uniform).
+    """
+
+    def __init__(self, init_parameters: dict, log_prob: float = 0.0):
+        self.init_parameters = init_parameters
+        self._log_prob = log_prob
+        self.unconditional_model = False
+        self.metadata = {}
+        self.transform_pre = Compose([])
+
+    def _run_sampler(self, num_samples, *args, **kwargs):
+        samples = {
+            k: torch.ones(num_samples) * v
+            for k, v in self.init_parameters.items()
+        }
+        samples["log_prob"] = torch.ones(num_samples) * self._log_prob
+        return samples

dingo/gw/likelihood.py

@@ -640,7 +640,10 @@ def _log_likelihood_calibration_marginalized(self, theta):
         return logsumexp(likelihoods) - np.log(len(likelihoods))
 
     def d_inner_h_complex_multi(
-        self, theta: pd.DataFrame, num_processes: int = 1
+        self,
+        theta: pd.DataFrame,
+        num_processes: int = 1,
+        return_rho2opt: bool = False,
     ) -> np.ndarray:
         """
         Calculate the complex inner product (d | h(theta)) between the stored data d
@@ -652,26 +655,31 @@ def d_inner_h_complex_multi(
             Parameters at which to evaluate h.
         num_processes : int
             Number of parallel processes to use.
+        return_rho2opt : bool
+            If True, also return rho2opt = (h|h) for each sample. Useful for
+            synthetic phase sampling where the full likelihood decomposition
+            log L = log_Zn + Re[(d|h)*exp(2i*phase)] - 0.5*(h|h) is needed.
 
         Returns
         -------
-        complex : Inner product
+        np.ndarray or tuple of np.ndarray
+            Complex inner products. If return_rho2opt is True, returns
+            (d_inner_h_complex, rho2opt).
         """
         with threadpool_limits(limits=1, user_api="blas"):
-            # Generator object for theta rows. For idx this yields row idx of
-            # theta dataframe, converted to dict, ready to be passed to
-            # self.log_likelihood.
             theta_generator = (d[1].to_dict() for d in theta.iterrows())
 
             if num_processes > 1:
                 with Pool(processes=num_processes) as pool:
-                    d_inner_h_complex = pool.map(
-                        self.d_inner_h_complex, theta_generator
-                    )
+                    results = pool.map(self._d_inner_h_complex, theta_generator)
             else:
-                d_inner_h_complex = list(map(self.d_inner_h_complex, theta_generator))
+                results = list(map(self._d_inner_h_complex, theta_generator))
 
-        return np.array(d_inner_h_complex)
+        d_inner_h = np.array([r[0] for r in results])
+        if return_rho2opt:
+            rho2opt = np.array([r[1] for r in results]).real
+            return d_inner_h, rho2opt
+        return d_inner_h
 
     def d_inner_h_complex(self, theta):
         """
@@ -685,20 +693,26 @@ def d_inner_h_complex(self, theta):
 
         Returns
         -------
-        complex : Inner product
+        complex : Inner product (d|h)
         """
         # TODO: Implement for time marginalization.
-        return self._d_inner_h_complex(theta)
+        return self._d_inner_h_complex(theta)[0]
 
     def _d_inner_h_complex(self, theta):
+        """Return (d|h) and (h|h) as a tuple. Computing both is essentially
+        free since the waveform mu is already generated."""
         mu = self.signal(theta)["waveform"]
         d = self.whitened_strains
-        return sum(
+        d_inner_h = sum(
             [
                 inner_product_complex(d_ifo, mu_ifo)
                 for d_ifo, mu_ifo in zip(d.values(), mu.values())
             ]
         )
+        rho2opt = sum(
+            [inner_product(mu_ifo, mu_ifo) for mu_ifo in mu.values()]
+        )
+        return d_inner_h, rho2opt
 
 
 def inner_product(a, b, min_idx=0, delta_f=None, psd=None):

tests/core/test_samplers.py

@@ -0,0 +1,67 @@
+import numpy as np
+import pytest
+import torch
+from unittest.mock import MagicMock, patch
+from torchvision.transforms import Compose
+
+from dingo.core.samplers import FixedInitSampler
+
+
+class TestFixedInitSampler:
+    """Tests for the FixedInitSampler class."""
+
+    def test_basic_construction(self):
+        params = {"chirp_mass_proxy": 1.2, "foo": 3.0}
+        sampler = FixedInitSampler(init_parameters=params)
+        assert sampler.init_parameters == params
+        assert sampler._log_prob == 0.0
+        assert sampler.unconditional_model is False
+        assert sampler.metadata == {}
+        assert isinstance(sampler.transform_pre, Compose)
+
+    def test_custom_log_prob(self):
+        sampler = FixedInitSampler(
+            init_parameters={"x": 1.0}, log_prob=-5.0
+        )
+        assert sampler._log_prob == -5.0
+
+    def test_run_sampler_returns_correct_shape(self):
+        params = {"chirp_mass_proxy": 1.2, "mass_ratio_proxy": 0.8}
+        sampler = FixedInitSampler(init_parameters=params)
+        num_samples = 50
+        samples = sampler._run_sampler(num_samples)
+
+        assert set(samples.keys()) == {"chirp_mass_proxy", "mass_ratio_proxy", "log_prob"}
+        for key in samples:
+            assert samples[key].shape == (num_samples,)
+
+    def test_run_sampler_returns_correct_values(self):
+        params = {"chirp_mass_proxy": 1.2, "mass_ratio_proxy": 0.8}
+        sampler = FixedInitSampler(init_parameters=params, log_prob=-3.0)
+        samples = sampler._run_sampler(10)
+
+        assert torch.allclose(samples["chirp_mass_proxy"], torch.ones(10) * 1.2)
+        assert torch.allclose(samples["mass_ratio_proxy"], torch.ones(10) * 0.8)
+        assert torch.allclose(samples["log_prob"], torch.ones(10) * (-3.0))
+
+    def test_run_sampler_single_sample(self):
+        params = {"x": 42.0}
+        sampler = FixedInitSampler(init_parameters=params)
+        samples = sampler._run_sampler(1)
+        assert samples["x"].shape == (1,)
+        assert samples["x"].item() == 42.0
+
+    def test_run_sampler_ignores_extra_args(self):
+        """_run_sampler should accept and ignore extra positional/keyword args
+        (matching the Sampler._run_sampler signature which receives context)."""
+        params = {"x": 1.0}
+        sampler = FixedInitSampler(init_parameters=params)
+        samples = sampler._run_sampler(5, "ignored_context", extra_kwarg=True)
+        assert samples["x"].shape == (5,)
+
+    def test_empty_parameters(self):
+        """FixedInitSampler with no parameters should still return log_prob."""
+        sampler = FixedInitSampler(init_parameters={})
+        samples = sampler._run_sampler(3)
+        assert set(samples.keys()) == {"log_prob"}
+        assert samples["log_prob"].shape == (3,)