diff --git a/dingo/core/samplers.py b/dingo/core/samplers.py index efc4fc1fa..f76f0a0d6 100644 --- a/dingo/core/samplers.py +++ b/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 diff --git a/dingo/gw/likelihood.py b/dingo/gw/likelihood.py index 49b29941f..c111498d3 100644 --- a/dingo/gw/likelihood.py +++ b/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): diff --git a/tests/core/test_samplers.py b/tests/core/test_samplers.py new file mode 100644 index 000000000..926cee286 --- /dev/null +++ b/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,) diff --git a/tests/gw/test_likelihood.py b/tests/gw/test_likelihood.py new file mode 100644 index 000000000..b027929f2 --- /dev/null +++ b/tests/gw/test_likelihood.py @@ -0,0 +1,152 @@ +import numpy as np +import pandas as pd +import pytest + +from dingo.gw.domains import UniformFrequencyDomain +from dingo.gw.likelihood import ( + StationaryGaussianGWLikelihood, + inner_product, + inner_product_complex, +) + + +@pytest.fixture +def likelihood_and_theta(): + """Set up a simple likelihood for testing.""" + f_min = 20.0 + f_max = 1024.0 + T = 8.0 + domain = UniformFrequencyDomain(f_min, f_max, delta_f=1 / T) + + waveform = { + "H1": np.where(domain.frequency_mask, (1.0 + 1j) * 1e-20, 0.0), + "L1": np.where(domain.frequency_mask, (1.0 + 1j) * 1e-20, 0.0), + } + asds = { + "H1": np.where(domain.frequency_mask, 1e-20, 1.0), + "L1": np.where(domain.frequency_mask, 1e-20, 1.0), + } + + sample = {"waveform": waveform, "asds": asds} + + likelihood = StationaryGaussianGWLikelihood( + wfg_kwargs={ + "approximant": "IMRPhenomXPHM", + "f_ref": 20.0, + "spin_conversion_phase": 0.0, + }, + wfg_domain=domain, + data_domain=domain, + event_data=sample, + t_ref=1248242632.0, + ) + + theta = { + "chirp_mass": 50.0, + "mass_ratio": 0.8, + "a_1": 0.3, + "a_2": 0.4, + "tilt_1": 1.3, + "tilt_2": 1.4, + "phi_12": 0.3, + "phi_jl": 3.3, + "theta_jn": 1.2, + "luminosity_distance": 4000.0, + "geocent_time": -0.03, + "dec": -0.3, + "ra": 1.5, + "psi": 1.2, + "phase": 1.4, + } + + return likelihood, theta + + +class TestLikelihoodDecomposition: + """Test the likelihood decomposition log L = log_Zn + kappa2 - 0.5*rho2opt.""" + + def test_log_likelihood_decomposition(self, likelihood_and_theta): + """log L = log_Zn + Re(d|h) - 0.5*(h|h)""" + likelihood, theta = likelihood_and_theta + log_l = likelihood.log_likelihood(theta) + d_inner_h, rho2opt = likelihood._d_inner_h_complex(theta) + kappa2 = d_inner_h.real + reconstructed = likelihood.log_Zn + kappa2 - 0.5 * rho2opt + assert np.isclose(log_l, reconstructed) + + def test_rho2opt_is_positive(self, likelihood_and_theta): + """rho2opt = (h|h) should be non-negative.""" + likelihood, theta = likelihood_and_theta + _, rho2opt = likelihood._d_inner_h_complex(theta) + assert rho2opt >= 0 + + def test_rho2opt_matches_direct_computation(self, likelihood_and_theta): + """rho2opt from _d_inner_h_complex should match direct (h|h).""" + likelihood, theta = likelihood_and_theta + _, rho2opt = likelihood._d_inner_h_complex(theta) + + mu = likelihood.signal(theta)["waveform"] + rho2opt_direct = sum( + [inner_product(mu_ifo, mu_ifo) for mu_ifo in mu.values()] + ) + assert np.isclose(rho2opt, rho2opt_direct) + + +class TestDInnerHComplex: + """Test the complex inner product methods.""" + + def test_d_inner_h_complex_returns_complex(self, likelihood_and_theta): + """d_inner_h_complex returns a complex scalar (backward compat).""" + likelihood, theta = likelihood_and_theta + result = likelihood.d_inner_h_complex(theta) + assert np.isscalar(result) or isinstance( + result, (complex, np.complexfloating) + ) + + def test_private_returns_tuple(self, likelihood_and_theta): + """_d_inner_h_complex returns (d_inner_h, rho2opt) tuple.""" + likelihood, theta = likelihood_and_theta + result = likelihood._d_inner_h_complex(theta) + assert isinstance(result, tuple) + assert len(result) == 2 + + def test_public_matches_private(self, likelihood_and_theta): + """d_inner_h_complex matches _d_inner_h_complex[0].""" + likelihood, theta = likelihood_and_theta + public = likelihood.d_inner_h_complex(theta) + private_tuple = likelihood._d_inner_h_complex(theta) + assert np.isclose(public, private_tuple[0]) + + def test_multi_without_rho2opt(self, likelihood_and_theta): + """d_inner_h_complex_multi without return_rho2opt returns array only.""" + likelihood, theta = likelihood_and_theta + theta_df = pd.DataFrame([theta, theta]) + result = likelihood.d_inner_h_complex_multi(theta_df) + assert isinstance(result, np.ndarray) + assert result.shape == (2,) + assert np.iscomplexobj(result) + + def test_multi_with_rho2opt(self, likelihood_and_theta): + """d_inner_h_complex_multi with return_rho2opt returns tuple.""" + likelihood, theta = likelihood_and_theta + theta_df = pd.DataFrame([theta, theta, theta]) + d_inner_h, rho2opt = likelihood.d_inner_h_complex_multi( + theta_df, return_rho2opt=True + ) + assert d_inner_h.shape == (3,) + assert rho2opt.shape == (3,) + assert np.iscomplexobj(d_inner_h) + assert not np.iscomplexobj(rho2opt) + assert np.all(rho2opt >= 0) + + def test_multi_rho2opt_consistency(self, likelihood_and_theta): + """rho2opt from multi matches single-sample computation.""" + likelihood, theta = likelihood_and_theta + d_inner_h_single, rho2opt_single = likelihood._d_inner_h_complex(theta) + + theta_df = pd.DataFrame([theta]) + d_inner_h_multi, rho2opt_multi = likelihood.d_inner_h_complex_multi( + theta_df, return_rho2opt=True + ) + assert np.isclose(d_inner_h_single, d_inner_h_multi[0]) + assert np.isclose(rho2opt_single, rho2opt_multi[0])