diff --git a/dingo/core/utils/trainutils.py b/dingo/core/utils/trainutils.py index bc7d12154..62d62ca2d 100644 --- a/dingo/core/utils/trainutils.py +++ b/dingo/core/utils/trainutils.py @@ -56,7 +56,7 @@ def __init__( self.counter = 0 self.best_score = None self.early_stop = False - self.val_loss_min = np.Inf + self.val_loss_min = np.inf self.delta = delta if metric not in ["training", "validation"]: raise ValueError( diff --git a/dingo/gw/gwutils.py b/dingo/gw/gwutils.py index f2b0730c0..f56252581 100644 --- a/dingo/gw/gwutils.py +++ b/dingo/gw/gwutils.py @@ -33,7 +33,7 @@ def get_extrinsic_prior_dict(extrinsic_prior): TODO: Move to dingo.gw.prior.py?""" extrinsic_prior_dict = default_extrinsic_dict.copy() for k, v in extrinsic_prior.items(): - if v.lower() != "default": + if not isinstance(v, str) or v.lower() != "default": extrinsic_prior_dict[k] = v return extrinsic_prior_dict diff --git a/dingo/gw/result.py b/dingo/gw/result.py index 51616032d..97503e942 100644 --- a/dingo/gw/result.py +++ b/dingo/gw/result.py @@ -415,13 +415,18 @@ def sample_calibration_parameters(self, calibration_sampling_kwargs: dict): self.calibration_sampling_kwargs = calibration_sampling_kwargs # Handle correction_type defaults - correction_type = self.calibration_sampling_kwargs.get("correction_type", "data") + correction_type = self.calibration_sampling_kwargs.get( + "correction_type", "data" + ) if correction_type is None: correction_type_dict = { - ifo: CALIBRATION_CORRECTION_TYPE_LOOKUP[ifo] for ifo in self.interferometers + ifo: CALIBRATION_CORRECTION_TYPE_LOOKUP[ifo] + for ifo in self.interferometers } elif correction_type == "data" or correction_type == "template": - correction_type_dict = {ifo: correction_type for ifo in self.interferometers} + correction_type_dict = { + ifo: correction_type for ifo in self.interferometers + } elif isinstance(correction_type, dict): correction_type_dict = correction_type else: @@ -440,8 +445,8 @@ def sample_calibration_parameters(self, calibration_sampling_kwargs: dict): ) # Removing the delta function priors on the frequency nodes, amplitude and phase. - # Usually the frequency nodes are set to delta functions, but we also remove the - # the amplitude and phase delta functions if present. + # Usually the frequency nodes are set to delta functions, but we also remove the + # amplitude and phase delta functions if present. # This avoids large log probs and log priors, since the density of a delta function # at the sampled point is infinite. The delta functions do not affect the sampling, # since they just fix certain parameters to constant values. @@ -457,9 +462,9 @@ def sample_calibration_parameters(self, calibration_sampling_kwargs: dict): delta_log_prob = np.zeros(num_samples) # Here we will sample the calibration parameters from the prior. - # We treat the *prior as the proposal* distribution and - # therefore add the log_prob of the sampled calibration parameters - # to the existing log_prob. We also will update the prior + # We treat the *prior as the proposal* distribution and + # therefore add the log_prob of the sampled calibration parameters + # to the existing log_prob. We also will update the prior # to include the calibration priors using the importance_sampling_metadata prior_update = self.importance_sampling_metadata.get("prior_update", {}) for ifo, prior in calibration_priors.items(): @@ -475,6 +480,12 @@ def sample_calibration_parameters(self, calibration_sampling_kwargs: dict): # Update prior_update dict with calibration parameters. This is for # persistence when saving to hdf5 for param_name, prior_obj in prior.items(): + # bilby's Prior.__repr__ isn't parseable for numpy scalars on numpy>2.0 + # Upstream fix: https://github.com/bilby-dev/bilby/pull/1108 + # Can be removed once dingo requires a bilby release that includes it. + for attr, value in prior_obj.get_instantiation_dict().items(): + if isinstance(value, np.generic): + setattr(prior_obj, attr, value.item()) prior_update[param_name] = repr(prior_obj) # Store prior_update for persistence on save/reload @@ -560,8 +571,25 @@ def sample_synthetic_phase( # Restrict to samples that are within the prior. param_keys = [k for k, v in self.prior.items() if not isinstance(v, Constraint)] theta = self.samples[param_keys] - log_prior = self.prior.ln_prob(theta, axis=0) - constraints = self.prior.evaluate_constraints(theta) + + # Compute log_prior only for non-DeltaFunction parameters. DeltaFunction + # priors return ln_prob = +inf at the peak, which causes check_ln_prob to + # skip constraint evaluation and return +inf for every sample, so + # np.isfinite(log_prior) would be False for all samples. Additionally, RA + # corrections (trigger_time vs model ref_time) can shift fixed parameters by + # tiny amounts, making DeltaFunction ln_prob = -inf for all samples. + prior_keys_for_lp = [ + k for k, v in self.prior.items() + if not isinstance(v, Constraint) and not isinstance(v, DeltaFunction) + ] + log_prior = self.prior.ln_prob(self.samples[prior_keys_for_lp], axis=0) + # Pass a plain dict so bilby's evaluate_constraints handles the argument + # correctly. bilby's evaluate_constraints mishandles a DataFrame argument: + # its internal .values() call raises TypeError (DataFrame.values is a + # property, not a method), causing the try/except inside bilby to fall + # through to ``np.ones_like(out_sample)``, which returns a 2-D array and + # causes a shape-broadcast error in the subsequent element-wise multiplication. + constraints = self.prior.evaluate_constraints(dict(theta)) np.putmask(log_prior, constraints == 0, -np.inf) within_prior = np.isfinite(log_prior) @@ -690,7 +718,9 @@ def get_samples_bilby_phase(self, num_processes=1): num_processes=num_processes, ) - def get_pesummary_samples(self, num_processes=1, resampling_method="clip+rejection"): + def get_pesummary_samples( + self, num_processes=1, resampling_method="clip+rejection" + ): """Samples in a form suitable for PESummary. These samples are adjusted to undo certain conventions used internally by diff --git a/pyproject.toml b/pyproject.toml index 7eef484df..26be4fada 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,8 +61,10 @@ write_to = "dingo/_version.py" markers = [ "slow: marks tests as slow", "asimov: marks tests for Asimov/LIGO pipeline integration", + "heavy: marks heavy end-to-end tests (train + importance sample; GPU; minutes-long). Deselected by default; run with -m heavy.", ] -addopts = "-m 'not asimov'" +addopts = "-m 'not asimov and not heavy'" +pythonpath = ["."] [project.urls] homepage = "https://github.com/dingo-gw/dingo" diff --git a/tests/core/test_backward_compatibility.py b/tests/core/test_backward_compatibility.py new file mode 100644 index 000000000..56a6d6b69 --- /dev/null +++ b/tests/core/test_backward_compatibility.py @@ -0,0 +1,24 @@ +import pytest + +from dingo.core.utils.backward_compatibility import check_minimum_version + + +def test_check_minimum_version_passes_for_recent_version(): + # A version well above the compatibility threshold must not raise. + assert check_minimum_version("dingo=99.0.0", raise_exception=True) is None + + +def test_check_minimum_version_raises_for_old_version(): + with pytest.raises(ValueError, match="backwards compatibility"): + check_minimum_version("dingo=0.0.1", raise_exception=True) + + +def test_check_minimum_version_treats_none_as_oldest(): + # A "None" version string is treated as 0.0.0, i.e. below the threshold. + with pytest.raises(ValueError): + check_minimum_version("None", raise_exception=True) + + +def test_check_minimum_version_warns_but_does_not_raise_by_default(): + # With raise_exception=False (default), an old version only warns. + assert check_minimum_version("dingo=0.0.1") is None diff --git a/tests/core/test_build_model.py b/tests/core/test_build_model.py new file mode 100644 index 000000000..25a0c060e --- /dev/null +++ b/tests/core/test_build_model.py @@ -0,0 +1,87 @@ +import numpy as np +import pytest + +from dingo.core.posterior_models.build_model import ( + autocomplete_model_kwargs, + build_model_from_kwargs, +) +from dingo.core.posterior_models.normalizing_flow import NormalizingFlowPosteriorModel + + +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 _settings(posterior_model_type="normalizing_flow"): + 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, + }, + } + } + } + + +def test_build_model_dispatches_to_normalizing_flow(): + model = build_model_from_kwargs(settings=_settings(), device="cpu") + assert isinstance(model, NormalizingFlowPosteriorModel) + + +def test_build_model_dispatch_is_case_insensitive(): + model = build_model_from_kwargs( + settings=_settings("Normalizing_Flow"), device="cpu" + ) + assert isinstance(model, NormalizingFlowPosteriorModel) + + +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()) + + +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") + + +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))] + ) + + 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 + + +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)] + ) + + 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 diff --git a/tests/core/test_density_interpolation.py b/tests/core/test_density_interpolation.py new file mode 100644 index 000000000..d8f01e2b9 --- /dev/null +++ b/tests/core/test_density_interpolation.py @@ -0,0 +1,61 @@ +import numpy as np + +from dingo.core.density.interpolation import ( + interpolated_log_prob, + interpolated_log_prob_multi, + interpolated_sample_and_log_prob, + interpolated_sample_and_log_prob_multi, +) + + +# A uniform distribution on [0, W] discretized on a grid: density = 1/W, so the +# (normalized) log prob at any interior point is -log(W). The Interped wrapper +# normalizes internally, so the input `values` need not be normalized. +WIDTH = 2.0 +SAMPLE_POINTS = np.linspace(0.0, WIDTH, 500) +UNIFORM_VALUES = np.ones_like(SAMPLE_POINTS) + + +def test_interpolated_log_prob_normalizes_uniform_distribution(): + log_prob = interpolated_log_prob(SAMPLE_POINTS, UNIFORM_VALUES, WIDTH / 2) + assert log_prob == np.float64(log_prob) # scalar + # A constant density integrates exactly under the trapezoidal normalization, so + # ln_prob = -log(W) holds to floating-point precision (observed residual 0.0). + np.testing.assert_allclose(log_prob, -np.log(WIDTH), atol=1e-10) + + +def test_interpolated_sample_and_log_prob_in_range_and_consistent(): + np.random.seed(0) + sample, log_prob = interpolated_sample_and_log_prob(SAMPLE_POINTS, UNIFORM_VALUES) + assert 0.0 <= sample <= WIDTH + # Both calls build the same Interped from the same data, so the returned log_prob + # equals interpolated_log_prob at the drawn sample to floating-point precision. + np.testing.assert_allclose( + log_prob, + interpolated_log_prob(SAMPLE_POINTS, UNIFORM_VALUES, sample), + atol=1e-10, + ) + + +def test_interpolated_log_prob_multi_matches_single(): + batch_values = np.stack([UNIFORM_VALUES, UNIFORM_VALUES]) + eval_points = np.array([0.5, 1.5]) + result = interpolated_log_prob_multi( + SAMPLE_POINTS, batch_values, eval_points, num_processes=1 + ) + assert result.shape == (2,) + expected = [ + interpolated_log_prob(SAMPLE_POINTS, UNIFORM_VALUES, p) for p in eval_points + ] + np.testing.assert_allclose(result, expected) + + +def test_interpolated_sample_and_log_prob_multi_shapes(): + np.random.seed(0) + batch_values = np.stack([UNIFORM_VALUES] * 4) + samples, log_probs = interpolated_sample_and_log_prob_multi( + SAMPLE_POINTS, batch_values, num_processes=1 + ) + assert samples.shape == (4,) + assert log_probs.shape == (4,) + assert np.all((samples >= 0.0) & (samples <= WIDTH)) diff --git a/tests/core/test_gnpeutils.py b/tests/core/test_gnpeutils.py new file mode 100644 index 000000000..902d64ea6 --- /dev/null +++ b/tests/core/test_gnpeutils.py @@ -0,0 +1,48 @@ +import numpy as np + +from dingo.core.utils.gnpeutils import IterationTracker + + +def test_pvalue_min_is_minus_inf_before_two_updates(): + tracker = IterationTracker() + assert tracker.pvalue_min == -np.inf # no ks_result yet + + tracker.update({"a": np.random.default_rng(0).normal(size=100)}) + # After a single update there is still no KS comparison. + assert tracker.pvalue_min == -np.inf + + +def test_first_update_stores_data_with_leading_axis(): + tracker = IterationTracker() + tracker.update({"a": np.zeros(50), "b": np.ones(50)}) + assert tracker.data["a"].shape == (1, 50) + assert tracker.data["b"].shape == (1, 50) + + +def test_second_update_computes_ks_pvalue(): + rng = np.random.default_rng(0) + tracker = IterationTracker() + tracker.update({"a": rng.normal(size=200)}) + tracker.update({"a": rng.normal(size=200)}) + # Two samples from the same distribution -> a finite, valid p-value. + assert tracker.ks_result is not None + assert np.isfinite(tracker.pvalue_min) + assert 0.0 <= tracker.pvalue_min <= 1.0 + + +def test_store_data_true_accumulates_rows(): + rng = np.random.default_rng(0) + tracker = IterationTracker(store_data=True) + for _ in range(3): + tracker.update({"a": rng.normal(size=20)}) + # store_data=True keeps every iteration along axis 0. + assert tracker.data["a"].shape == (3, 20) + + +def test_store_data_false_keeps_only_latest(): + rng = np.random.default_rng(0) + tracker = IterationTracker(store_data=False) + for _ in range(3): + tracker.update({"a": rng.normal(size=20)}) + # store_data=False replaces the stored data each iteration. + assert tracker.data["a"].shape == (1, 20) diff --git a/tests/core/test_misc.py b/tests/core/test_misc.py new file mode 100644 index 000000000..c3e595ea2 --- /dev/null +++ b/tests/core/test_misc.py @@ -0,0 +1,55 @@ +import numpy as np + +from dingo.core.utils.misc import get_version, recursive_check_dicts_are_equal + + +def test_get_version_returns_string(): + version = get_version() + assert isinstance(version, str) and version + + +def test_recursive_check_equal_nested_dicts(): + a = {"x": 1, "nested": {"y": 2}} + assert recursive_check_dicts_are_equal(a, {"x": 1, "nested": {"y": 2}}) is True + + +def test_recursive_check_different_keys(): + assert recursive_check_dicts_are_equal({"a": 1}, {"b": 1}) is False + + +def test_recursive_check_different_types(): + # 1 (int) vs 1.0 (float) differ by type. + assert recursive_check_dicts_are_equal({"a": 1}, {"a": 1.0}) is False + + +def test_recursive_check_nested_value_differs(): + assert recursive_check_dicts_are_equal({"n": {"y": 2}}, {"n": {"y": 3}}) is False + + +def test_recursive_check_arrays(): + assert recursive_check_dicts_are_equal( + {"a": np.array([1, 2])}, {"a": np.array([1, 2])} + ) + assert not recursive_check_dicts_are_equal( + {"a": np.array([1, 2])}, {"a": np.array([1, 3])} + ) + + +def test_recursive_check_string_numbers_within_tolerance(): + # Numeric literals in strings are compared with a tolerance (they can drift by + # float-precision across machines); non-numeric characters must match exactly. + a = {"p": "Uniform(minimum=1.0000000, maximum=2.0)"} + b = {"p": "Uniform(minimum=1.0000001, maximum=2.0)"} + assert recursive_check_dicts_are_equal(a, b) is True + + +def test_recursive_check_string_numbers_beyond_tolerance(): + a = {"p": "Uniform(minimum=1.0, maximum=2.0)"} + b = {"p": "Uniform(minimum=9.0, maximum=2.0)"} + assert recursive_check_dicts_are_equal(a, b) is False + + +def test_recursive_check_string_non_numeric_differs(): + a = {"p": "Uniform(minimum=1.0)"} + b = {"p": "Normal(minimum=1.0)"} + assert recursive_check_dicts_are_equal(a, b) is False diff --git a/tests/core/test_result.py b/tests/core/test_result.py index 752fd0227..74b519e4c 100644 --- a/tests/core/test_result.py +++ b/tests/core/test_result.py @@ -3,6 +3,8 @@ import numpy as np import pandas as pd import pytest +from bilby.core.prior import Constraint, DeltaFunction, PriorDict, Uniform +from scipy.special import logsumexp from dingo.core.result import Result, _clip_weights @@ -150,9 +152,9 @@ def test_higher_max_increases_output(): ] # Expected output grows monotonically with k (stochastically, but very reliably). for i in range(len(counts) - 1): - assert counts[i] <= counts[i + 1], ( - f"k={i+1} gave {counts[i]} samples but k={i+2} gave {counts[i+1]}" - ) + assert ( + counts[i] <= counts[i + 1] + ), f"k={i+1} gave {counts[i]} samples but k={i+2} gave {counts[i+1]}" # --------------------------------------------------------------------------- @@ -233,3 +235,303 @@ def test_clip_weights_reproducible(): out1 = result.rejection_sample(clip_weights=True, random_state=42) out2 = result.rejection_sample(clip_weights=True, random_state=42) pd.testing.assert_frame_equal(out1, out2) + + +# --------------------------------------------------------------------------- +# Statistical / metadata properties +# --------------------------------------------------------------------------- + + +def make_result(samples=None, settings=None, event_metadata=None): + """Construct a minimal core Result, optionally with settings/event_metadata.""" + dictionary = {} + if samples is not None: + dictionary["samples"] = samples + if settings is not None: + dictionary["settings"] = settings + result = Result(dictionary=dictionary) + if event_metadata is not None: + result.event_metadata = event_metadata + return result + + +def test_num_samples(): + result = make_result(pd.DataFrame({"x": np.arange(7.0)})) + assert result.num_samples == 7 + + +def test_num_samples_zero_when_no_samples(): + result = make_result() + assert result.num_samples == 0 + + +def test_effective_sample_size_uniform_weights_equals_n(): + result = make_result_with_weights(np.ones(10)) + assert result.effective_sample_size == pytest.approx(10.0) + assert result.n_eff == result.effective_sample_size + + +def test_effective_sample_size_matches_formula(): + weights = np.array([1.0, 2.0, 3.0, 4.0]) + result = make_result_with_weights(weights) + expected = np.sum(weights) ** 2 / np.sum(weights**2) + assert result.effective_sample_size == pytest.approx(expected) + + +def test_effective_sample_size_none_without_weights(): + result = make_result(pd.DataFrame({"x": [1.0, 2.0]})) + assert result.effective_sample_size is None + + +def test_sample_efficiency(): + result = make_result_with_weights(np.ones(8)) + assert result.sample_efficiency == pytest.approx(1.0) # uniform -> n_eff / N = 1 + + +def test_sample_efficiency_none_without_weights(): + result = make_result(pd.DataFrame({"x": [1.0, 2.0]})) + assert result.sample_efficiency is None + + +def test_log_evidence_std_requires_weights_and_log_evidence(): + result = make_result_with_weights([1.0, 2.0, 3.0]) + # No log_evidence set yet. + assert result.log_evidence_std is None + result.log_evidence = -5.0 + assert result.log_evidence_std is not None + assert result.log_evidence_std > 0 + + +def test_log_bayes_factor(): + result = make_result(pd.DataFrame({"x": [1.0]})) + assert result.log_bayes_factor is None + result.log_evidence = -3.0 + result.log_noise_evidence = -10.0 + assert result.log_bayes_factor == pytest.approx(7.0) + + +def test_injection_parameters(): + result = make_result(pd.DataFrame({"x": [1.0]})) + assert result.injection_parameters is None + result.event_metadata = {"injection_parameters": {"mass": 30.0}} + assert result.injection_parameters == {"mass": 30.0} + + +def test_metadata_and_base_metadata(): + settings = {"train_settings": {"data": {}}, "value": 1} + result = make_result(pd.DataFrame({"x": [1.0]}), settings=settings) + assert result.metadata is result.settings + # Non-unconditional model: base_metadata is the full metadata. + assert result.base_metadata is result.metadata + + +def test_base_metadata_unconditional_returns_base(): + base = {"some": "precursor"} + settings = {"train_settings": {"data": {"unconditional": True}}, "base": base} + result = make_result(pd.DataFrame({"x": [1.0]}), settings=settings) + assert result.base_metadata is base + + +def test_parameter_subset_keeps_only_requested_columns(): + samples = pd.DataFrame({"a": [1.0, 2.0], "b": [3.0, 4.0], "log_prob": [0.0, 0.0]}) + result = make_result(samples) + subset = result.parameter_subset(["a"]) + assert list(subset.samples.columns) == ["a"] + assert isinstance(subset, Result) + + +# --------------------------------------------------------------------------- +# _calculate_evidence +# --------------------------------------------------------------------------- + + +def test_calculate_evidence_matches_logsumexp(): + n = 6 + samples = pd.DataFrame( + { + "x": np.arange(n, dtype=float), + "log_prob": np.linspace(0.0, 1.0, n), + "log_prior": np.full(n, -2.0), + "log_likelihood": np.linspace(-1.0, 1.0, n), + } + ) + result = make_result(samples) + result._calculate_evidence() + + log_weights = samples["log_prior"] + samples["log_likelihood"] - samples["log_prob"] + expected = logsumexp(log_weights) - np.log(n) + assert result.log_evidence == pytest.approx(expected) + # Stored weights are normalized to mean 1. + assert result.samples["weights"].mean() == pytest.approx(1.0) + + +def test_calculate_evidence_includes_delta_log_prob_target(): + n = 4 + delta = np.array([0.0, 0.5, -0.5, 1.0]) + samples = pd.DataFrame( + { + "log_prob": np.zeros(n), + "log_prior": np.zeros(n), + "log_likelihood": np.zeros(n), + "delta_log_prob_target": delta, + } + ) + result = make_result(samples) + result._calculate_evidence() + expected = logsumexp(delta) - np.log(n) + assert result.log_evidence == pytest.approx(expected) + + +# --------------------------------------------------------------------------- +# importance_sample orchestration (stubbed prior + likelihood) +# --------------------------------------------------------------------------- + + +class _StubResult(Result): + """Result with an analytic prior and a trivial likelihood, to exercise the + importance_sample orchestration without any GW machinery.""" + + def _build_prior(self): + self.prior = PriorDict( + { + "a": Uniform(0.0, 1.0, name="a"), + "b": Uniform(0.0, 1.0, name="b"), + "c": DeltaFunction(0.5, name="c"), # fixed parameter + } + ) + + def _build_likelihood(self, **likelihood_kwargs): + class _StubLikelihood: + log_Zn = -10.0 + + def log_likelihood_multi(self, theta, num_processes=1): + return np.zeros(len(theta)) + + self.likelihood = _StubLikelihood() + + +def _stub_samples(n=5, with_log_prob=True): + data = { + "a": np.linspace(0.1, 0.9, n), + "b": np.linspace(0.1, 0.9, n), + "c": np.full(n, 0.5), + } + if with_log_prob: + data["log_prob"] = np.zeros(n) + return pd.DataFrame(data) + + +def test_importance_sample_populates_columns_and_evidence(): + result = _StubResult(dictionary={"samples": _stub_samples()}) + result.importance_sample(num_processes=1) + for col in ("log_prior", "log_likelihood", "weights"): + assert col in result.samples.columns + assert result.log_evidence is not None + assert result.log_noise_evidence == -10.0 + + +def test_importance_sample_requires_samples(): + result = _StubResult(dictionary={"samples": _stub_samples()}) + result.samples = None + with pytest.raises(KeyError): + result.importance_sample() + + +def test_importance_sample_requires_log_prob(): + result = _StubResult(dictionary={"samples": _stub_samples(with_log_prob=False)}) + with pytest.raises(KeyError, match="log probability"): + result.importance_sample() + + +# --------------------------------------------------------------------------- +# split / merge +# --------------------------------------------------------------------------- + + +def _result_with_evidence_columns(n=10): + samples = pd.DataFrame( + { + "x": np.arange(n, dtype=float), + "log_prob": np.zeros(n), + "log_prior": np.zeros(n), + "log_likelihood": np.zeros(n), + } + ) + return Result( + dictionary={"samples": samples, "settings": {"train_settings": {"data": {}}}} + ) + + +def test_split_partitions_samples(): + result = _result_with_evidence_columns(10) + parts = result.split(3) + assert len(parts) == 3 + assert sum(p.num_samples for p in parts) == 10 + assert all(isinstance(p, Result) for p in parts) + + +def test_split_then_merge_round_trips(): + result = _result_with_evidence_columns(10) + merged = Result.merge(result.split(3)) + np.testing.assert_array_equal( + merged.samples["x"].to_numpy(), result.samples["x"].to_numpy() + ) + + +def test_merge_incompatible_metadata_raises(): + result = _result_with_evidence_columns(10) + parts = result.split(2) + parts[1].settings = {"train_settings": {"data": {"changed": True}}} + with pytest.raises(ValueError, match="same metadata"): + Result.merge(parts) + + +# --------------------------------------------------------------------------- +# sampling_importance_resampling +# --------------------------------------------------------------------------- + + +def test_sampling_importance_resampling_count_and_drops_weights(): + result = make_result_with_weights([1.0, 2.0, 3.0, 0.5, 4.0], extra_col=False) + out = result.sampling_importance_resampling(num_samples=3, random_state=0) + assert len(out) == 3 + assert "weights" not in out.columns + + +def test_sampling_importance_resampling_too_many_raises(): + result = make_result_with_weights([1.0, 2.0, 3.0]) + with pytest.raises(ValueError, match="Cannot sample more"): + result.sampling_importance_resampling(num_samples=100) + + +# --------------------------------------------------------------------------- +# prior-derived parameter-key properties +# --------------------------------------------------------------------------- + + +def test_parameter_key_properties_partition_the_prior(): + result = Result(dictionary={"samples": pd.DataFrame({"a": [0.5]})}) + # core Result._build_prior leaves prior=None; set a prior with one of each kind. + result.prior = PriorDict( + { + "a": Uniform(0.0, 1.0, name="a"), # search parameter + "b": Constraint(0.0, 1.0, name="b"), # constraint + "c": DeltaFunction(0.5, name="c"), # fixed parameter + } + ) + assert result.search_parameter_keys == ["a"] + assert result.constraint_parameter_keys == ["b"] + assert result.fixed_parameter_keys == ["c"] + + +# --------------------------------------------------------------------------- +# print_summary (smoke) +# --------------------------------------------------------------------------- + + +def test_print_summary_runs_with_and_without_evidence(capsys): + result = make_result_with_weights([1.0, 2.0, 3.0]) + result.print_summary() # no log_evidence yet + result.log_evidence = -3.0 + result.print_summary() # with evidence / n_eff / efficiency + assert "Number of samples" in capsys.readouterr().out diff --git a/tests/core/test_samplers.py b/tests/core/test_samplers.py new file mode 100644 index 000000000..b86918826 --- /dev/null +++ b/tests/core/test_samplers.py @@ -0,0 +1,261 @@ +import numpy as np +import pandas as pd +import pytest +import torch + +from dingo.core.posterior_models.normalizing_flow import NormalizingFlowPosteriorModel +from dingo.core.result import Result +from dingo.core.samplers import Sampler + + +INFERENCE_PARAMETERS = ["a", "b", "c"] + +# Non-trivial standardization so that the log_prob change-of-variables correction +# (-sum(log(std))) is actually exercised by the round-trip test below. +STANDARDIZATION = { + "mean": {"a": 1.0, "b": -2.0, "c": 0.5}, + "std": {"a": 2.0, "b": 0.5, "c": 3.0}, +} + +BASE_TRANSFORM_KWARGS = { + "hidden_dim": 16, + "num_transform_blocks": 1, + "activation": "elu", + "dropout_probability": 0.0, + "batch_norm": False, + "num_bins": 4, + "base_transform_type": "rq-coupling", +} + + +def _build_model(unconditional, context_dim=None, embedding_kwargs=None): + """Build a small normalizing-flow posterior model for use with a Sampler.""" + posterior_kwargs = { + "input_dim": len(INFERENCE_PARAMETERS), + "context_dim": context_dim, + "num_flow_steps": 2, + "base_transform_kwargs": BASE_TRANSFORM_KWARGS, + } + model_kwargs = { + "posterior_model_type": "normalizing_flow", + "posterior_kwargs": posterior_kwargs, + } + if embedding_kwargs is not None: + model_kwargs["embedding_kwargs"] = embedding_kwargs + + metadata = { + "train_settings": { + "model": model_kwargs, + "data": { + "inference_parameters": INFERENCE_PARAMETERS, + "standardization": STANDARDIZATION, + "unconditional": unconditional, + }, + } + } + if unconditional: + # For unconditional models the Sampler keeps the precursor ("base") metadata + # separately; it just needs to exist. + metadata["base"] = {} + + return NormalizingFlowPosteriorModel(metadata=metadata, device="cpu") + + +@pytest.fixture() +def unconditional_sampler(): + """Sampler wrapping an unconditional flow. + + This takes the context-free (x = []) path through run_sampler / log_prob, which + isolates the count, batching, and log_prob logic with minimal setup. + """ + return Sampler(model=_build_model(unconditional=True)) + + +@pytest.fixture() +def conditional_sampler(): + """Sampler wrapping a conditional flow, without an embedding network. + + Sufficient for exercising the context-required error paths, which are checked + before the network is ever called. + """ + return Sampler(model=_build_model(unconditional=False, context_dim=5)) + + +@pytest.fixture() +def conditional_sampler_with_context(): + """Sampler wrapping a conditional flow *with* an embedding network. + + The full conditional sampling path runs here. The base Sampler's default + transform_pre is the identity (a subclass would replace it); we substitute a + minimal callable that extracts the context tensor so that data reaches the + embedding network. + """ + embedding_kwargs = { + "input_dims": (2, 3, 20), + "svd": {"size": 10}, + "V_rb_list": None, + "output_dim": 8, + "hidden_dims": [32, 16, 8], + "activation": "elu", + "dropout": 0.0, + "batch_norm": False, + "added_context": False, + } + sampler = Sampler( + model=_build_model( + unconditional=False, context_dim=8, embedding_kwargs=embedding_kwargs + ) + ) + sampler.transform_pre = lambda context: context["data"] + return sampler + + +def test_run_sampler_returns_correct_number_of_samples(unconditional_sampler): + unconditional_sampler.run_sampler(num_samples=17) + assert len(unconditional_sampler.samples) == 17 + + +def test_samples_columns_are_inference_parameters_plus_log_prob(unconditional_sampler): + unconditional_sampler.run_sampler(num_samples=10) + assert list(unconditional_sampler.samples.columns) == INFERENCE_PARAMETERS + [ + "log_prob" + ] + + +@pytest.mark.parametrize("batch_size", [None, 3, 5, 7, 10, 25]) +def test_run_sampler_count_invariant_to_batch_size(unconditional_sampler, batch_size): + """The number of samples must not depend on batch_size. + + Covers the divmod branch in run_sampler: exact divisors (5, 10), divisors with a + remainder (3, 7), a batch larger than the request (25 -> one short batch), and the + default single-batch case (None). + """ + num_samples = 10 + unconditional_sampler.run_sampler(num_samples=num_samples, batch_size=batch_size) + assert len(unconditional_sampler.samples) == num_samples + + +def test_conditional_run_sampler_without_context_raises(conditional_sampler): + with pytest.raises(ValueError, match="Context must be set"): + conditional_sampler.run_sampler(num_samples=5) + + +def test_conditional_log_prob_without_context_raises(conditional_sampler): + samples = pd.DataFrame({p: [0.0] for p in INFERENCE_PARAMETERS}) + with pytest.raises(ValueError, match="Context must be set"): + conditional_sampler.log_prob(samples) + + +def test_unconditional_run_sampler_needs_no_context(unconditional_sampler): + # Unconditional models carry no context and require none in order to sample. + assert unconditional_sampler.context is None + unconditional_sampler.run_sampler(num_samples=8) + assert len(unconditional_sampler.samples) == 8 + + +def test_log_prob_round_trip_matches_sampling_log_prob(unconditional_sampler): + """log_prob recomputed at the sampled points reproduces the sampling log_prob. + + This exercises the full standardize -> de-standardize round trip, including the + change-of-variables correction (-sum(log(std))). The stored log_prob column must + be dropped first, since log_prob() adds any log_prob already present (see + test_log_prob_adds_existing_log_prob_column). + """ + unconditional_sampler.run_sampler(num_samples=20) + samples = unconditional_sampler.samples + recomputed = unconditional_sampler.log_prob(samples.drop(columns="log_prob")) + # The network runs in float32, and sampling vs. log_prob use the flow's forward + # vs. inverse transforms, which round-trip only to the float32 floor (observed + # ~1e-6); atol=1e-5 leaves ~10x margin. + np.testing.assert_allclose(recomputed, samples["log_prob"].to_numpy(), atol=1e-5) + + +def test_log_prob_adds_existing_log_prob_column(unconditional_sampler): + """A log_prob column in the input is added to the recomputed network log_prob. + + This is relied upon by GNPE / density-recovery, where the carried log_prob is a + proposal contribution. The base Sampler does not drop it. + """ + unconditional_sampler.run_sampler(num_samples=6) + samples = unconditional_sampler.samples + stored = samples["log_prob"].to_numpy() + recomputed = unconditional_sampler.log_prob(samples.drop(columns="log_prob")) + combined = unconditional_sampler.log_prob(samples) + # float32 round-trip floor, as in test_log_prob_round_trip_matches_sampling_log_prob. + np.testing.assert_allclose(combined, recomputed + stored, atol=1e-5) + + +def test_log_prob_accepts_dataframe_and_dict_inputs(unconditional_sampler): + unconditional_sampler.run_sampler(num_samples=6) + samples = unconditional_sampler.samples.drop(columns="log_prob") + + # DataFrame -> one log_prob per row. + lp_dataframe = unconditional_sampler.log_prob(samples) + assert lp_dataframe.shape == (6,) + + # dict of arrays -> one log_prob per row. + lp_dict = unconditional_sampler.log_prob( + {p: samples[p].to_numpy() for p in INFERENCE_PARAMETERS} + ) + assert lp_dict.shape == (6,) + + # dict of scalars -> a single log_prob. + lp_scalar = unconditional_sampler.log_prob( + {p: samples[p].iloc[0] for p in INFERENCE_PARAMETERS} + ) + assert lp_scalar.shape == (1,) + + +def test_to_result_round_trips_samples_and_settings(unconditional_sampler): + unconditional_sampler.run_sampler(num_samples=10) + result = unconditional_sampler.to_result() + assert isinstance(result, Result) + pd.testing.assert_frame_equal(result.samples, unconditional_sampler.samples) + assert result.settings == unconditional_sampler.metadata + + +def test_to_hdf5_round_trips_samples(unconditional_sampler, tmp_path): + unconditional_sampler.run_sampler(num_samples=10) + unconditional_sampler.to_hdf5(label="result", outdir=str(tmp_path)) + + reloaded = Result(file_name=str(tmp_path / "result.hdf5")) + np.testing.assert_allclose( + reloaded.samples[INFERENCE_PARAMETERS].to_numpy(), + unconditional_sampler.samples[INFERENCE_PARAMETERS].to_numpy(), + ) + + +def test_context_setter_moves_parameters_to_event_metadata(conditional_sampler): + conditional_sampler.context = {"waveform": [1, 2, 3], "parameters": {"mass": 30.0}} + assert conditional_sampler.event_metadata["injection_parameters"] == {"mass": 30.0} + # The injected truth parameters are removed from the context itself. + assert "parameters" not in conditional_sampler.context + + +def test_conditional_samples_depend_on_context(conditional_sampler_with_context): + """For a conditional model the samples depend on the context. + + With the torch RNG fixed, the same context reproduces samples exactly, while a + different context changes them -- isolating the context as the source of variation. + """ + sampler = conditional_sampler_with_context + context_a = {"data": torch.randn(2, 3, 20)} + context_b = {"data": torch.randn(2, 3, 20)} + + torch.manual_seed(0) + sampler.context = context_a + sampler.run_sampler(num_samples=8) + samples_a = sampler.samples.to_numpy().copy() + + torch.manual_seed(0) + sampler.context = context_a + sampler.run_sampler(num_samples=8) + samples_a_repeat = sampler.samples.to_numpy().copy() + + torch.manual_seed(0) + sampler.context = context_b + sampler.run_sampler(num_samples=8) + samples_b = sampler.samples.to_numpy().copy() + + assert np.array_equal(samples_a, samples_a_repeat) + assert not np.array_equal(samples_a, samples_b) diff --git a/tests/core/test_torchutils.py b/tests/core/test_torchutils.py new file mode 100644 index 000000000..f58735d7f --- /dev/null +++ b/tests/core/test_torchutils.py @@ -0,0 +1,73 @@ +import pytest +import torch +import torch.nn as nn + +from dingo.core.utils import torchutils + + +def test_get_activation_function_from_string(): + activation = torchutils.get_activation_function_from_string("elu") + # Returns a callable activation (torch.nn.functional.elu). + assert callable(activation) + out = activation(torch.tensor([-1.0, 1.0])) + assert out.shape == (2,) + + +def test_get_activation_function_unknown_raises(): + with pytest.raises(ValueError): + torchutils.get_activation_function_from_string("not_an_activation") + + +def test_get_optimizer_from_kwargs(): + model = nn.Linear(2, 2) + optimizer = torchutils.get_optimizer_from_kwargs( + model.parameters(), type="adam", lr=0.01 + ) + assert isinstance(optimizer, torch.optim.Adam) + + +def test_get_scheduler_from_kwargs(): + model = nn.Linear(2, 2) + optimizer = torchutils.get_optimizer_from_kwargs( + model.parameters(), type="adam", lr=0.01 + ) + scheduler = torchutils.get_scheduler_from_kwargs(optimizer, type="cosine", T_max=10) + assert isinstance(scheduler, torch.optim.lr_scheduler.CosineAnnealingLR) + + +def test_split_dataset_into_train_and_test(): + dataset = torch.utils.data.TensorDataset(torch.arange(100.0)) + train, test = torchutils.split_dataset_into_train_and_test(dataset, 0.8) + assert len(train) == 80 + assert len(test) == 20 + # The split is a partition (no shared indices, full coverage). + assert len(train) + len(test) == len(dataset) + + +def test_get_number_of_model_parameters(): + model = nn.Linear(2, 3) # weight 2*3 + bias 3 = 9 + assert torchutils.get_number_of_model_parameters(model) == 9 + + +def test_get_lr_returns_optimizer_learning_rates(): + optimizer = torchutils.get_optimizer_from_kwargs( + nn.Linear(2, 2).parameters(), type="adam", lr=0.007 + ) + assert torchutils.get_lr(optimizer) == [0.007] + + +def test_torch_detach_to_cpu(): + tensor = torch.ones(3, requires_grad=True) + detached = torchutils.torch_detach_to_cpu(tensor) + assert not detached.requires_grad + # Non-tensor inputs pass through unchanged. + assert torchutils.torch_detach_to_cpu(5) == 5 + + +def test_set_requires_grad_flag_by_name(): + model = nn.Linear(2, 3) + torchutils.set_requires_grad_flag( + model, name_contains="weight", requires_grad=False + ) + assert model.weight.requires_grad is False + assert model.bias.requires_grad is True diff --git a/tests/core/test_trainutils.py b/tests/core/test_trainutils.py new file mode 100644 index 000000000..12702f042 --- /dev/null +++ b/tests/core/test_trainutils.py @@ -0,0 +1,172 @@ +import math +import os + +import pytest + +from dingo.core.utils.trainutils import ( + AvgTracker, + EarlyStopping, + LossInfo, + RuntimeLimits, + write_history, +) + + +# --------------------------------------------------------------------------- +# AvgTracker +# --------------------------------------------------------------------------- + + +def test_avg_tracker_empty_returns_nan(): + tracker = AvgTracker() + assert math.isnan(tracker.get_avg()) + + +def test_avg_tracker_weighted_average_and_last_value(): + tracker = AvgTracker() + tracker.update(2.0) + tracker.update(4.0) + assert tracker.get_avg() == 3.0 + assert tracker.x == 4.0 + + # With explicit counts the average is sum / total-count. + tracker = AvgTracker() + tracker.update(10.0, n=2) + tracker.update(2.0, n=3) + assert tracker.get_avg() == 12.0 / 5 + + +# --------------------------------------------------------------------------- +# EarlyStopping +# --------------------------------------------------------------------------- + + +def test_early_stopping_first_call_sets_best_and_returns_true(): + es = EarlyStopping(patience=3) + assert es(1.0) is True + assert es.best_score == -1.0 + assert es.counter == 0 + assert es.early_stop is False + + +def test_early_stopping_improving_loss_keeps_counter_zero(): + es = EarlyStopping(patience=3) + es(5.0) + for loss in (4.0, 3.0, 2.0): + assert es(loss) is True + assert es.counter == 0 + assert es.early_stop is False + + +def test_early_stopping_triggers_after_patience_non_improving(): + es = EarlyStopping(patience=3) + es(1.0) # best + # Three consecutive non-improving losses reach the patience limit. + assert es(2.0) is False and es.counter == 1 + assert es(2.0) is False and es.counter == 2 + assert es(2.0) is False and es.counter == 3 + assert es.early_stop is True + + +def test_early_stopping_delta_makes_small_improvement_count_as_non_improving(): + # An improvement must exceed `delta` to reset the counter. + es = EarlyStopping(patience=5, delta=1.0) + es(10.0) # best_score = -10 + # New loss 9.5 -> score -9.5, which is < best_score + delta = -9.0, so non-improving. + assert es(9.5) is False + assert es.counter == 1 + + +def test_early_stopping_invalid_metric_raises(): + with pytest.raises(ValueError, match="training.*validation|validation.*training"): + EarlyStopping(metric="not_a_metric") + + +# --------------------------------------------------------------------------- +# RuntimeLimits +# --------------------------------------------------------------------------- + + +def test_runtime_limits_none_never_exceeded(): + limits = RuntimeLimits() + assert limits.limits_exceeded(epoch=10_000) is False + + +def test_runtime_limits_total_epochs(): + limits = RuntimeLimits(max_epochs_total=10) + assert limits.limits_exceeded(epoch=9) is False + assert limits.limits_exceeded(epoch=10) is True + + +def test_runtime_limits_epochs_per_run(): + limits = RuntimeLimits(max_epochs_per_run=5, epoch_start=3) + assert limits.limits_exceeded(epoch=7) is False # 7 - 3 = 4 < 5 + assert limits.limits_exceeded(epoch=8) is True # 8 - 3 = 5 >= 5 + + +def test_runtime_limits_epochs_per_run_requires_epoch(): + limits = RuntimeLimits(max_epochs_per_run=5, epoch_start=0) + with pytest.raises(ValueError, match="epoch required"): + limits.limits_exceeded(epoch=None) + + +def test_runtime_limits_per_run_requires_epoch_start(): + with pytest.raises(ValueError, match="epoch_start required"): + RuntimeLimits(max_epochs_per_run=5) + + +def test_runtime_limits_time_limit_zero_is_exceeded_immediately(): + # Any elapsed time >= 0 trips a zero time limit. + limits = RuntimeLimits(max_time_per_run=0.0) + assert limits.limits_exceeded(epoch=0) is True + + +def test_local_limits_ignore_total_epoch_limit(): + # local_limits_exceeded honours per-run/time limits but not the total epoch limit. + limits = RuntimeLimits(max_epochs_total=5) + assert limits.local_limits_exceeded(epoch=100) is False + + limits = RuntimeLimits(max_epochs_per_run=2, epoch_start=0) + assert limits.local_limits_exceeded(epoch=1) is False + assert limits.local_limits_exceeded(epoch=2) is True + + +# --------------------------------------------------------------------------- +# LossInfo +# --------------------------------------------------------------------------- + + +def test_loss_info_weighted_average_across_batches(): + info = LossInfo(epoch=1, len_dataset=100, batch_size=10) + info.update(2.0, n=4) + info.update(3.0, n=2) + # Weighted: (2*4 + 3*2) / (4 + 2) = 14 / 6. + assert info.get_avg() == pytest.approx(14.0 / 6) + assert info.loss == 3.0 + + +# --------------------------------------------------------------------------- +# write_history +# --------------------------------------------------------------------------- + + +def test_write_history_appends_rows(tmp_path): + log_dir = str(tmp_path) + write_history(log_dir, 1, -1.0, -0.5, [0.001]) + write_history(log_dir, 2, -2.0, -1.5, [0.0005]) + + history_file = os.path.join(log_dir, "history.txt") + with open(history_file) as f: + rows = [line.strip().split("\t") for line in f if line.strip()] + + assert len(rows) == 2 + assert rows[0] == ["1", "-1.0", "-0.5", "0.001"] + assert rows[1] == ["2", "-2.0", "-1.5", "0.0005"] + + +def test_write_history_refuses_to_overwrite_on_first_epoch(tmp_path): + log_dir = str(tmp_path) + write_history(log_dir, 1, -1.0, -0.5, [0.001]) + # Writing epoch 1 again would clobber the existing file. + with pytest.raises(AssertionError): + write_history(log_dir, 1, -1.0, -0.5, [0.001]) diff --git a/tests/core/test_unconditional_density_estimation.py b/tests/core/test_unconditional_density_estimation.py new file mode 100644 index 000000000..38b6a9306 --- /dev/null +++ b/tests/core/test_unconditional_density_estimation.py @@ -0,0 +1,113 @@ +import numpy as np +import pandas as pd +import pytest + +from dingo.core.density.unconditional_density_estimation import ( + train_unconditional_density_estimator, +) +from dingo.core.posterior_models.normalizing_flow import NormalizingFlowPosteriorModel +from dingo.core.result import Result + + +PARAMETERS = ["x", "y"] + + +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. + """ + return { + "data": {}, + "model": { + "posterior_model_type": "normalizing_flow", + "posterior_kwargs": { + "num_flow_steps": 2, + "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", + }, + }, + }, + "training": { + "device": "cpu", + "num_workers": 0, + "train_fraction": 0.9, + "batch_size": 64, + "epochs": 1, + "optimizer": {"type": "adam", "lr": 0.01}, + "scheduler": {"type": "cosine", "T_max": 1}, + }, + } + + +@pytest.fixture() +def result(): + """A Result holding Gaussian samples over three parameters.""" + samples = pd.DataFrame( + { + "x": np.random.normal(1.0, 2.0, 256), + "y": np.random.normal(-3.0, 0.5, 256), + "z": np.random.normal(0.0, 1.0, 256), + } + ) + return Result( + dictionary={"samples": samples, "settings": {"train_settings": {"data": {}}}} + ) + + +def test_train_unconditional_density_estimator_basic(result, tmp_path): + settings = _nde_settings() + settings["data"]["parameters"] = PARAMETERS + model = train_unconditional_density_estimator(result, settings, str(tmp_path)) + + assert isinstance(model, NormalizingFlowPosteriorModel) + + # 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 + + # Standardization is computed from the training samples. + expected_mean = result.samples[PARAMETERS].to_numpy().mean(axis=0) + expected_std = result.samples[PARAMETERS].to_numpy().std(axis=0) + stored = settings["data"]["standardization"] + for i, p in enumerate(PARAMETERS): + assert stored["mean"][p] == pytest.approx(expected_mean[i]) + assert stored["std"][p] == pytest.approx(expected_std[i]) + + # The trained model can sample with the right shapes. + theta, log_prob = model.sample_and_log_prob(num_samples=5) + assert tuple(theta.shape) == (5, len(PARAMETERS)) + assert tuple(log_prob.shape) == (5,) + + +def test_train_unconditional_density_estimator_uses_all_parameters_by_default( + result, tmp_path +): + # 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] + + +def test_train_unconditional_flow_end_to_end(result, tmp_path): + model = result.train_unconditional_flow( + PARAMETERS, _nde_settings(), train_dir=str(tmp_path), threshold_std=np.inf + ) + assert isinstance(model, NormalizingFlowPosteriorModel) + # Trained over the requested subset only. + assert model.model_kwargs["posterior_kwargs"]["input_dim"] == len(PARAMETERS) + + +def test_train_unconditional_flow_rejects_too_many_outliers(result): + # A tight threshold removes far more than 5% of Gaussian samples, so the outlier + # guard raises before any training happens (no train_dir needed). + with pytest.raises(ValueError, match="Too many proxy samples"): + result.train_unconditional_flow(PARAMETERS, _nde_settings(), threshold_std=1.0) diff --git a/tests/gw/inference/test_gw_samplers.py b/tests/gw/inference/test_gw_samplers.py new file mode 100644 index 000000000..dbeb0bc89 --- /dev/null +++ b/tests/gw/inference/test_gw_samplers.py @@ -0,0 +1,341 @@ +import numpy as np +import pytest +from astropy.time import Time +from astropy.utils import iers + +from dingo.core.posterior_models.normalizing_flow import NormalizingFlowPosteriorModel +from dingo.gw.domains import UniformFrequencyDomain +from dingo.gw.inference.gw_samplers import ( + GWSampler, + check_frequency_updates, + _validate_maximum_frequency, + _validate_minimum_frequency, +) + +# Avoid network access (and the associated timeout) when astropy computes sidereal +# time in _correct_reference_time; the bundled IERS data is sufficient for tests. +iers.conf.auto_download = False + + +DETECTORS = ["H1", "L1"] +INFERENCE_PARAMETERS = ["chirp_mass", "mass_ratio", "ra", "dec"] + +DOMAIN_SETTINGS = { + "type": "UniformFrequencyDomain", + "f_min": 20.0, + "f_max": 1024.0, + "delta_f": 0.25, +} + + +# --------------------------------------------------------------------------- +# Frequency-range validators (pure functions, only need a domain). +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def domain(): + return UniformFrequencyDomain(f_min=20.0, f_max=1024.0, delta_f=0.25) + + +@pytest.mark.parametrize( + "validate, valid_change, beyond_bound", + [ + (_validate_minimum_frequency, 40.0, 10.0), # raise f_min; below hard f_min + (_validate_maximum_frequency, 512.0, 2048.0), # lower f_max; above hard f_max + ], +) +def test_frequency_validator_no_op_when_unchanged( + domain, validate, valid_change, beyond_bound +): + # Value equal to the domain bound is a no-op and is allowed even without cropping. + bound = domain.f_min if validate is _validate_minimum_frequency else domain.f_max + assert validate(bound, DETECTORS, domain, None) is None + + +@pytest.mark.parametrize( + "validate, valid_change, beyond_bound", + [ + (_validate_minimum_frequency, 40.0, 10.0), + (_validate_maximum_frequency, 512.0, 2048.0), + ], +) +def test_frequency_validator_expands_float_to_all_detectors( + domain, validate, valid_change, beyond_bound +): + # A float applies to every detector; a valid change passes with cropping on. + # The cap/floor must be given explicitly, else it defaults to the domain bound. + crop = {"cropping_probability": 0.5, "f_min_upper": 100.0, "f_max_lower": 400.0} + assert validate(valid_change, DETECTORS, domain, crop) is None + + +@pytest.mark.parametrize( + "validate, valid_change, beyond_bound", + [ + (_validate_minimum_frequency, 40.0, 10.0), + (_validate_maximum_frequency, 512.0, 2048.0), + ], +) +def test_frequency_validator_rejects_value_beyond_hard_bound( + domain, validate, valid_change, beyond_bound +): + crop = {"cropping_probability": 0.5} + with pytest.raises(ValueError, match="domain.f_"): + validate(beyond_bound, DETECTORS, domain, crop) + + +@pytest.mark.parametrize( + "validate, valid_change", + [(_validate_minimum_frequency, 40.0), (_validate_maximum_frequency, 512.0)], +) +def test_frequency_validator_rejects_detector_key_mismatch( + domain, validate, valid_change +): + crop = {"cropping_probability": 0.5} + with pytest.raises(ValueError, match="exactly detectors"): + validate({"H1": valid_change}, DETECTORS, domain, crop) + + +@pytest.mark.parametrize( + "validate, valid_change", + [(_validate_minimum_frequency, 40.0), (_validate_maximum_frequency, 512.0)], +) +def test_frequency_validator_rejects_change_when_cropping_disabled( + domain, validate, valid_change +): + # No crop settings at all. + with pytest.raises(ValueError, match="[Cc]ropping"): + validate(valid_change, DETECTORS, domain, None) + # Crop settings present but probability zero. + with pytest.raises(ValueError, match="[Cc]ropping"): + validate(valid_change, DETECTORS, domain, {"cropping_probability": 0.0}) + + +def test_validate_minimum_frequency_rejects_value_above_cap(domain): + crop = {"cropping_probability": 0.5, "f_min_upper": 60.0} + assert _validate_minimum_frequency(50.0, DETECTORS, domain, crop) is None + with pytest.raises(ValueError, match="upper bound"): + _validate_minimum_frequency(80.0, DETECTORS, domain, crop) + + +def test_validate_maximum_frequency_rejects_value_below_floor(domain): + crop = {"cropping_probability": 0.5, "f_max_lower": 400.0} + assert _validate_maximum_frequency(500.0, DETECTORS, domain, crop) is None + with pytest.raises(ValueError, match="lower bound"): + _validate_maximum_frequency(300.0, DETECTORS, domain, crop) + + +def test_validate_minimum_frequency_rejects_differing_values_when_not_independent( + domain, +): + crop = { + "cropping_probability": 0.5, + "independent_detectors": False, + "f_min_upper": 100.0, + } + with pytest.raises(ValueError, match="[Ii]ndependent"): + _validate_minimum_frequency({"H1": 40.0, "L1": 50.0}, DETECTORS, domain, crop) + + +def test_check_frequency_updates_accepts_valid_and_rejects_invalid(): + model_metadata = { + "train_settings": { + "data": { + "detectors": DETECTORS, + "random_strain_cropping": { + "cropping_probability": 0.5, + "f_min_upper": 100.0, + "f_max_lower": 400.0, + }, + } + }, + "dataset_settings": {"domain": DOMAIN_SETTINGS}, + } + # Valid frequency updates pass without raising. + assert check_frequency_updates(model_metadata, f_min=40.0, f_max=512.0) is None + # Beyond the hard bound raises. + with pytest.raises(ValueError, match="domain.f_min"): + check_frequency_updates(model_metadata, f_min=10.0) + + +# --------------------------------------------------------------------------- +# GWSamplerMixin methods (lightweight GWSampler; network not exercised). +# --------------------------------------------------------------------------- + + +def _build_gw_sampler(unconditional=False, domain_update=None): + """Build a GWSampler around a tiny flow plus minimal but valid GW metadata. + + The network is never run by the methods under test here; it only needs to exist. + """ + standardization = { + "mean": {p: 0.0 for p in INFERENCE_PARAMETERS}, + "std": {p: 1.0 for p in INFERENCE_PARAMETERS}, + } + posterior_kwargs = { + "input_dim": len(INFERENCE_PARAMETERS), + "context_dim": None, + "num_flow_steps": 2, + "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", + }, + } + data_settings = { + "unconditional": unconditional, + "inference_parameters": INFERENCE_PARAMETERS, + "standardization": standardization, + "detectors": DETECTORS, + "ref_time": 1126259462.4, + "extrinsic_prior": { + "dec": "default", + "ra": "default", + "geocent_time": "default", + "luminosity_distance": "default", + "psi": "default", + }, + } + if domain_update is not None: + data_settings["domain_update"] = domain_update + + metadata = { + "train_settings": { + "model": { + "posterior_model_type": "normalizing_flow", + "posterior_kwargs": posterior_kwargs, + }, + "data": data_settings, + }, + "dataset_settings": { + "domain": DOMAIN_SETTINGS, + "intrinsic_prior": { + "mass_1": "bilby.core.prior.Constraint(minimum=10, maximum=80)", + "mass_2": "bilby.core.prior.Constraint(minimum=10, maximum=80)", + "chirp_mass": "bilby.gw.prior.UniformInComponentsChirpMass(" + "minimum=25, maximum=31)", + "mass_ratio": "bilby.gw.prior.UniformInComponentsMassRatio(" + "minimum=0.125, maximum=1)", + "phase": "default", + "a_1": 0.0, + "a_2": 0.0, + }, + }, + } + if unconditional: + metadata["base"] = metadata + model = NormalizingFlowPosteriorModel(metadata=metadata, device="cpu") + return GWSampler(model=model) + + +@pytest.fixture() +def gw_sampler(): + return _build_gw_sampler() + + +def test_build_domain_from_metadata(gw_sampler): + assert isinstance(gw_sampler.domain, UniformFrequencyDomain) + assert gw_sampler.domain.f_min == DOMAIN_SETTINGS["f_min"] + assert gw_sampler.domain.f_max == DOMAIN_SETTINGS["f_max"] + + +def test_build_domain_applies_domain_update(): + sampler = _build_gw_sampler(domain_update={"f_min": 30.0}) + assert sampler.domain.f_min == 30.0 + + +def test_correct_reference_time_round_trip(gw_sampler): + gw_sampler._event_metadata = {"time_event": gw_sampler.t_ref + 3600.0} + samples = {"ra": np.array([0.5, 1.5, 2.5]), "dec": np.array([0.1, 0.2, 0.3])} + original_ra = samples["ra"].copy() + + gw_sampler._correct_reference_time(samples, inverse=False) + assert not np.allclose(samples["ra"], original_ra) + assert np.all((samples["ra"] >= 0) & (samples["ra"] < 2 * np.pi)) + + gw_sampler._correct_reference_time(samples, inverse=True) + np.testing.assert_allclose(samples["ra"], original_ra) + + +def test_correct_reference_time_matches_sidereal_shift(gw_sampler): + """The RA shift must equal the difference in apparent sidereal time.""" + t_event = gw_sampler.t_ref + 3600.0 + gw_sampler._event_metadata = {"time_event": t_event} + samples = {"ra": np.array([0.5, 1.5, 2.5])} + original_ra = samples["ra"].copy() + + ra_correction = ( + Time(t_event, format="gps", scale="utc").sidereal_time("apparent", "greenwich") + - Time(gw_sampler.t_ref, format="gps", scale="utc").sidereal_time( + "apparent", "greenwich" + ) + ).rad + + gw_sampler._correct_reference_time(samples, inverse=False) + np.testing.assert_allclose( + samples["ra"], (original_ra + ra_correction) % (2 * np.pi) + ) + + +def test_correct_reference_time_noop_when_time_matches_reference(gw_sampler): + gw_sampler._event_metadata = {"time_event": gw_sampler.t_ref} + samples = {"ra": np.array([0.5, 1.5])} + original_ra = samples["ra"].copy() + gw_sampler._correct_reference_time(samples, inverse=False) + np.testing.assert_allclose(samples["ra"], original_ra) + + +def test_correct_reference_time_noop_without_ra(gw_sampler): + gw_sampler._event_metadata = {"time_event": gw_sampler.t_ref + 3600.0} + samples = {"dec": np.array([0.1, 0.2])} + # No "ra" key: nothing to correct, and no error. + gw_sampler._correct_reference_time(samples, inverse=False) + assert "ra" not in samples + + +def test_post_process_forward_adds_fixed_prior_parameters(gw_sampler): + samples = {p: np.zeros(5) for p in INFERENCE_PARAMETERS} + gw_sampler._event_metadata = None + gw_sampler._post_process(samples, inverse=False) + # a_1 and a_2 are DeltaFunctions (0.0) in the intrinsic prior; they get added. + for fixed in ("a_1", "a_2"): + assert fixed in samples + np.testing.assert_array_equal(samples[fixed], np.zeros(5)) + + +def test_post_process_inverse_drops_non_inference_parameters(gw_sampler): + samples = { + "chirp_mass": np.array([28.0]), + "ra": np.array([1.0]), + "log_prob": np.array([0.5]), + "extra": np.array([9.0]), + } + gw_sampler._event_metadata = None + gw_sampler._post_process(samples, inverse=True) + assert set(samples) <= set(INFERENCE_PARAMETERS) + assert "log_prob" not in samples + assert "extra" not in samples + + +def test_frequency_updates_flag(gw_sampler): + # By default the requested range equals the domain, so no updates are flagged and + # the min/max frequencies report the domain bounds. + assert gw_sampler.minimum_frequency == gw_sampler.domain.f_min + assert gw_sampler.maximum_frequency == gw_sampler.domain.f_max + assert gw_sampler.frequency_updates is False + + # A requested minimum frequency that differs from the domain flags an update. + # (Set the private attribute directly to bypass the validating setter, which + # would also rebuild the transforms.) + gw_sampler._minimum_frequency = 40.0 + assert gw_sampler.frequency_updates is True + + +def test_event_metadata_injects_frequency_bounds(gw_sampler): + metadata = gw_sampler.event_metadata + assert metadata["minimum_frequency"] == gw_sampler.domain.f_min + assert metadata["maximum_frequency"] == gw_sampler.domain.f_max diff --git a/tests/gw/noise/test_asd_parameterization.py b/tests/gw/noise/test_asd_parameterization.py new file mode 100644 index 000000000..7374ab332 --- /dev/null +++ b/tests/gw/noise/test_asd_parameterization.py @@ -0,0 +1,150 @@ +import numpy as np +import pytest + +from dingo.gw.domains import UniformFrequencyDomain +from dingo.gw.noise.synthetic.asd_parameterization import ( + curve_fit, + fit_broadband_noise, + parameterize_single_psd, +) +from dingo.gw.noise.synthetic.utils import ( + get_index_for_elem, + lorentzian_eval, + reconstruct_psds_from_parameters, +) + + +@pytest.fixture() +def domain(): + return UniformFrequencyDomain(f_min=0.0, f_max=1024.0, delta_f=1.0) + + +# --------------------------------------------------------------------------- +# utils: get_index_for_elem, lorentzian_eval +# --------------------------------------------------------------------------- + + +def test_get_index_for_elem_returns_nearest(): + arr = np.array([0.0, 1.0, 2.0, 3.0]) + assert get_index_for_elem(arr, 2.0) == 2 # exact + assert get_index_for_elem(arr, 2.4) == 2 # nearest below + assert get_index_for_elem(arr, 2.6) == 3 # nearest above + + +def test_lorentzian_eval_returns_zeros_for_degenerate_params(domain): + x = domain.sample_frequencies + np.testing.assert_array_equal(lorentzian_eval(x, 0.0, 5.0, 100.0), np.zeros_like(x)) + np.testing.assert_array_equal( + lorentzian_eval(x, 200.0, -1.0, 100.0), np.zeros_like(x) + ) + + +def test_lorentzian_eval_peaks_at_f0(domain): + x = domain.sample_frequencies + f0 = 200.0 + line = lorentzian_eval(x, f0, 5.0, 100.0) + assert line.shape == x.shape + assert line.max() > 0 + assert x[np.argmax(line)] == pytest.approx(f0, abs=domain.delta_f) + + +def test_lorentzian_eval_truncation_suppresses_tails(domain): + x = domain.sample_frequencies + f0 = 200.0 + full = lorentzian_eval(x, f0, 5.0, 100.0) + truncated = lorentzian_eval(x, f0, 5.0, 100.0, delta_f=5.0) + far_from_peak = np.abs(x - f0) > 50 + # The exponential truncation makes the tails strictly smaller. + assert truncated[far_from_peak].sum() < full[far_from_peak].sum() + + +# --------------------------------------------------------------------------- +# fit_broadband_noise +# --------------------------------------------------------------------------- + + +def test_fit_broadband_noise_recovers_flat_level(domain): + # A PSD that is constant in log space should be reproduced by the spline nodes. + const = np.log(1e-40) + log_psd = np.full_like(domain.sample_frequencies, const) + xs, ys = fit_broadband_noise(domain, log_psd, num_spline_positions=8, sigma=1.0) + + assert len(xs) == len(ys) == 8 + # ys are accumulated in float32, so rounding limits agreement to ~1e-5 + np.testing.assert_allclose(ys, const, atol=1e-4) + # Node positions are increasing and lie within the frequency range. + assert np.all(np.diff(xs) > 0) + assert xs[0] >= domain.sample_frequencies[0] + assert xs[-1] <= domain.sample_frequencies[-1] + + +# --------------------------------------------------------------------------- +# curve_fit +# --------------------------------------------------------------------------- + + +def test_curve_fit_recovers_injected_lorentzian(domain): + x = domain.sample_frequencies + segment = (x >= 150) & (x <= 250) + frequencies = x[segment] + f0_true, A_true, Q_true = 200.0, 5.0, 100.0 + line = lorentzian_eval(frequencies, f0_true, A_true, Q_true) + + data = { + "psd": line, + "broadband_noise": np.zeros_like(frequencies), + "frequencies": frequencies, + "lower_freq": frequencies[0], + "upper_freq": frequencies[-1], + } + f0, A, Q = curve_fit(data, std=1e-3) + + assert f0 == pytest.approx(f0_true, abs=2 * domain.delta_f) + # Fitted parameters respect the optimizer bounds. + assert frequencies[0] <= f0 <= frequencies[-1] + assert 0 <= A <= 12 + assert 10 <= Q <= 1000 + + +# --------------------------------------------------------------------------- +# parameterize_single_psd (integration) +# --------------------------------------------------------------------------- + + +def test_parameterize_single_psd_shapes(domain): + psd = np.full_like(domain.sample_frequencies, 1e-40) + settings = { + "sigma": 1.0, + "num_spline_positions": 8, + "num_spectral_segments": 5, + "delta_f": -1, # non-positive -> no Lorentzian truncation + } + out = parameterize_single_psd(psd, domain, settings) + + assert set(out.keys()) == {"x_positions", "y_values", "spectral_features"} + assert len(out["y_values"]) == settings["num_spline_positions"] + assert len(out["x_positions"]) == settings["num_spline_positions"] + assert out["spectral_features"].shape == (settings["num_spectral_segments"], 3) + + +# --------------------------------------------------------------------------- +# reconstruct_psds_from_parameters +# --------------------------------------------------------------------------- + + +def test_reconstruct_psds_from_parameters_shape_and_positivity(domain): + num_spline, num_segments = 8, 5 + parameters = { + "x_positions": np.linspace(20.0, domain.f_max, num_spline), + "y_values": np.full(num_spline, np.log(1e-40)), + "spectral_features": np.zeros((num_segments, 3)), # no spectral lines + } + # smoothen=True -> deterministic exp(base_noise) reconstruction (no added noise). + psds = reconstruct_psds_from_parameters( + parameters, domain, {"sigma": 1.0, "smoothen": True} + ) + assert psds.shape == (1, len(domain)) + assert np.all(psds > 0) + # A cubic spline through a constant is exactly that constant, so exp(base_noise) + # reconstructs 1e-40 to (float64) machine precision (observed rel. error ~1e-15). + np.testing.assert_allclose(psds[0], 1e-40, rtol=1e-12) diff --git a/tests/gw/noise/test_asd_sampling.py b/tests/gw/noise/test_asd_sampling.py new file mode 100644 index 000000000..7f2e2e9cc --- /dev/null +++ b/tests/gw/noise/test_asd_sampling.py @@ -0,0 +1,57 @@ +import numpy as np +import pytest + +from dingo.gw.noise.synthetic.asd_sampling import KDE + + +NUM_ASDS = 50 +NUM_SPLINE = 8 +NUM_SEGMENTS = 5 +DETECTORS = ["H1", "L1"] + + +def _parameter_dict(): + """Synthetic parameterization of a set of ASDs (the output format of + parameterize_asd_dataset): spline y-values + Lorentzian spectral features.""" + + def per_detector(): + return { + "x_positions": np.linspace(20.0, 1000.0, NUM_SPLINE), + "y_values": np.random.normal(-90.0, 1.0, size=(NUM_ASDS, NUM_SPLINE)), + "spectral_features": np.random.normal( + [100.0, -1.0, 500.0], [1.0, 0.1, 5.0], size=(NUM_ASDS, NUM_SEGMENTS, 3) + ), + } + + return {det: per_detector() for det in DETECTORS} + + +@pytest.fixture() +def kde(): + settings = { + "bandwidth_spectral": 0.1, + "bandwidth_spline": 0.1, + "split_frequencies": [200.0, 500.0], + } + kde = KDE(_parameter_dict(), settings) + kde.fit() + return kde + + +def test_kde_sample_shapes(kde): + out = kde.sample(num_samples=10) + assert set(out.keys()) == set(DETECTORS) + for det in DETECTORS: + assert out[det]["spectral_features"].shape == (10, NUM_SEGMENTS, 3) + assert out[det]["y_values"].shape == (10, NUM_SPLINE) + + +def test_kde_sample_rescaling_shifts_base_noise_mean(kde): + rescaling_ys = {det: np.full(NUM_SPLINE, -85.0) for det in DETECTORS} + out = kde.sample(num_samples=30, rescaling_ys=rescaling_ys) + # Rescaling subtracts the per-node sample mean and adds the target, so the + # post-rescaling per-node sample mean equals the target *exactly* (a mean-shift), + # independent of sample size. Only floating-point error remains (observed ~1e-14). + np.testing.assert_allclose( + out["H1"]["y_values"].mean(axis=0), rescaling_ys["H1"], atol=1e-10 + ) diff --git a/tests/gw/test_build_domain.py b/tests/gw/test_build_domain.py new file mode 100644 index 000000000..4d1cfdcb7 --- /dev/null +++ b/tests/gw/test_build_domain.py @@ -0,0 +1,72 @@ +import pytest + +from dingo.gw.domains import UniformFrequencyDomain +from dingo.gw.domains.build_domain import ( + build_domain, + build_domain_from_model_metadata, +) + + +def test_build_domain_uniform_frequency(): + settings = { + "type": "UniformFrequencyDomain", + "f_min": 20.0, + "f_max": 1024.0, + "delta_f": 0.25, + } + domain = build_domain(settings) + assert isinstance(domain, UniformFrequencyDomain) + assert domain.f_min == 20.0 and domain.f_max == 1024.0 + + +def test_build_domain_frequency_domain_alias(): + # "FD" and "FrequencyDomain" are aliases for UniformFrequencyDomain. + domain = build_domain({"type": "FD", "f_min": 20.0, "f_max": 512.0, "delta_f": 0.5}) + assert isinstance(domain, UniformFrequencyDomain) + + +def test_build_domain_missing_type_raises(): + with pytest.raises(ValueError, match='"type"'): + build_domain({"f_min": 20.0, "f_max": 1024.0, "delta_f": 0.25}) + + +def test_build_domain_unknown_type_raises(): + with pytest.raises(NotImplementedError, match="not implemented"): + build_domain({"type": "NotADomain"}) + + +# --------------------------------------------------------------------------- +# build_domain_from_model_metadata +# --------------------------------------------------------------------------- + +_METADATA = { + "dataset_settings": { + "domain": { + "type": "UniformFrequencyDomain", + "f_min": 20.0, + "f_max": 1024.0, + "delta_f": 0.25, + } + }, + "train_settings": {"data": {}}, +} + + +def test_build_domain_from_model_metadata(): + domain = build_domain_from_model_metadata(_METADATA) + assert isinstance(domain, UniformFrequencyDomain) + assert domain.f_min == 20.0 + + +def test_build_domain_from_model_metadata_applies_domain_update(): + metadata = { + **_METADATA, + "train_settings": {"data": {"domain_update": {"f_min": 30.0}}}, + } + assert build_domain_from_model_metadata(metadata).f_min == 30.0 + + +def test_build_domain_from_model_metadata_base_is_noop_for_uniform_domain(): + # A UniformFrequencyDomain has no base_domain, so base=True returns it unchanged. + domain = build_domain_from_model_metadata(_METADATA, base=True) + assert isinstance(domain, UniformFrequencyDomain) diff --git a/tests/gw/test_gw_result.py b/tests/gw/test_gw_result.py new file mode 100644 index 000000000..b0aec44af --- /dev/null +++ b/tests/gw/test_gw_result.py @@ -0,0 +1,300 @@ +import numpy as np +import pandas as pd +import pytest + +from dingo.gw.domains import UniformFrequencyDomain +from dingo.gw.gwutils import get_extrinsic_prior_dict +from dingo.gw.likelihood import StationaryGaussianGWLikelihood +from dingo.gw.prior import build_prior_with_defaults +from dingo.gw.result import Result + + +DETECTORS = ["H1", "L1"] +REF_TIME = 1126259462.391 + +DOMAIN_SETTINGS = { + "type": "UniformFrequencyDomain", + "f_min": 20.0, + "f_max": 256.0, # small (T=2s) so the likelihood is cheap + "delta_f": 0.5, +} +WAVEFORM_GENERATOR = {"approximant": "IMRPhenomD", "f_ref": 20.0} + +INTRINSIC_PRIOR = { + "mass_1": "bilby.core.prior.Constraint(minimum=10.0, maximum=80.0, name='mass_1')", + "mass_2": "bilby.core.prior.Constraint(minimum=10.0, maximum=80.0, name='mass_2')", + "chirp_mass": "bilby.gw.prior.UniformInComponentsChirpMass(" + "minimum=15.0, maximum=100.0, name='chirp_mass')", + "mass_ratio": "bilby.gw.prior.UniformInComponentsMassRatio(" + "minimum=0.125, maximum=1.0, name='mass_ratio')", + "phase": "default", + "chi_1": "bilby.gw.prior.AlignedSpin(name='chi_1', a_prior=Uniform(minimum=0, maximum=0.9))", + "chi_2": "bilby.gw.prior.AlignedSpin(name='chi_2', a_prior=Uniform(minimum=0, maximum=0.9))", + "theta_jn": "default", + "luminosity_distance": 100.0, + "geocent_time": 0.0, +} +EXTRINSIC_PRIOR = { + "dec": "default", + "ra": "default", + "geocent_time": "bilby.core.prior.Uniform(" + "minimum=-0.10, maximum=0.10, name='geocent_time')", + "psi": "default", + "luminosity_distance": "bilby.core.prior.Uniform(" + "minimum=100.0, maximum=1000.0, name='luminosity_distance')", +} + + +def _metadata(): + return { + "dataset_settings": { + "domain": DOMAIN_SETTINGS, + "waveform_generator": WAVEFORM_GENERATOR, + "intrinsic_prior": INTRINSIC_PRIOR, + }, + "train_settings": { + "data": { + "detectors": DETECTORS, + "ref_time": REF_TIME, + "extrinsic_prior": EXTRINSIC_PRIOR, + } + }, + } + + +def _context(): + """Synthetic strain + ASD context (constant above f_min, masked below).""" + domain = UniformFrequencyDomain( + DOMAIN_SETTINGS["f_min"], DOMAIN_SETTINGS["f_max"], DOMAIN_SETTINGS["delta_f"] + ) + mask = domain.frequency_mask + waveform = {d: np.where(mask, (1.0 + 1j) * 1e-21, 0.0) for d in DETECTORS} + asds = {d: np.where(mask, 1e-21, 1.0) for d in DETECTORS} + return {"waveform": waveform, "asds": asds} + + +def make_gw_result(n=5, drop_phase=False, event_metadata=None): + """Build a gw Result with `n` rows drawn from its own prior (so columns/ranges are + consistent with the waveform generator), plus a synthetic context.""" + full_prior = build_prior_with_defaults( + {**INTRINSIC_PRIOR, **get_extrinsic_prior_dict(EXTRINSIC_PRIOR)} + ) + samples = pd.DataFrame(full_prior.sample(n)) + samples["log_prob"] = 0.0 + if drop_phase: + samples = samples.drop(columns="phase") + return Result( + dictionary={ + "samples": samples, + "context": _context(), + "event_metadata": {} if event_metadata is None else event_metadata, + "settings": _metadata(), + } + ) + + +# --------------------------------------------------------------------------- +# Simple property accessors +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def gw_result(): + return make_gw_result() + + +@pytest.mark.parametrize( + "attr, key", + [ + ("synthetic_phase_kwargs", "synthetic_phase"), + ("time_marginalization_kwargs", "time_marginalization"), + ("phase_marginalization_kwargs", "phase_marginalization"), + ("calibration_marginalization_kwargs", "calibration_marginalization"), + ("calibration_sampling_kwargs", "calibration_sampling"), + ], +) +def test_kwargs_round_trip_via_importance_sampling_metadata(gw_result, attr, key): + assert getattr(gw_result, attr) is None # not set yet + value = {"some": "setting"} + setattr(gw_result, attr, value) + assert getattr(gw_result, attr) == value + assert gw_result.importance_sampling_metadata[key] == value + + +def test_use_base_domain_default_and_noop_on_uniform_domain(gw_result): + # UniformFrequencyDomain has no base_domain, so the setter is a no-op. + assert gw_result.use_base_domain is False + gw_result.use_base_domain = True + assert gw_result.use_base_domain is False + + +def test_f_ref_and_approximant(gw_result): + assert gw_result.f_ref == WAVEFORM_GENERATOR["f_ref"] + assert gw_result.approximant == WAVEFORM_GENERATOR["approximant"] + + +def test_t_ref_defaults_to_ref_time(gw_result): + assert gw_result.t_ref == REF_TIME + + +def test_t_ref_uses_event_time_when_present(): + result = make_gw_result(event_metadata={"time_event": REF_TIME + 100.0}) + assert result.t_ref == REF_TIME + 100.0 + + +def test_minimum_maximum_frequency_default_to_domain(gw_result): + assert gw_result.minimum_frequency == gw_result.domain.f_min + assert gw_result.maximum_frequency == gw_result.domain.f_max + + +def test_minimum_maximum_frequency_override_and_setter(gw_result): + result = make_gw_result( + event_metadata={"minimum_frequency": 30.0, "maximum_frequency": 200.0} + ) + assert result.minimum_frequency == 30.0 + assert result.maximum_frequency == 200.0 + + gw_result.minimum_frequency = 25.0 + assert gw_result.event_metadata["minimum_frequency"] == 25.0 + assert gw_result.minimum_frequency == 25.0 + + +def test_interferometers(gw_result): + assert gw_result.interferometers == DETECTORS + + +def test_build_domain_creates_uniform_frequency_domain(gw_result): + assert isinstance(gw_result.domain, UniformFrequencyDomain) + assert gw_result.domain.f_min == DOMAIN_SETTINGS["f_min"] + assert gw_result.domain.f_max == DOMAIN_SETTINGS["f_max"] + + +# --------------------------------------------------------------------------- +# Likelihood tests +# --------------------------------------------------------------------------- + + +def test_build_likelihood(gw_result): + gw_result._build_likelihood() + assert isinstance(gw_result.likelihood, StationaryGaussianGWLikelihood) + assert np.isfinite(gw_result.likelihood.log_Zn) + + +def test_importance_sample_populates_columns_and_evidence(gw_result): + gw_result.importance_sample(num_processes=1) + for col in ("log_prior", "log_likelihood", "weights"): + assert col in gw_result.samples.columns + assert np.isfinite(gw_result.log_evidence) + assert np.isfinite(gw_result.log_noise_evidence) + # Normalized weights have mean 1. + assert gw_result.samples["weights"].mean() == pytest.approx(1.0) + + +def test_importance_sample_requires_log_prob(): + result = make_gw_result() + result.samples = result.samples.drop(columns="log_prob") + with pytest.raises(KeyError, match="log probability"): + result.importance_sample() + + +def test_sample_synthetic_phase_adds_phase_column(): + result = make_gw_result(drop_phase=True) + assert "phase" not in result.samples.columns + log_prob_before = result.samples["log_prob"].to_numpy().copy() + + result.sample_synthetic_phase({"n_grid": 16, "approximation_22_mode": True}) + + assert "phase" in result.samples.columns + phase = result.samples["phase"].to_numpy() + assert np.all((phase >= 0) & (phase < 2 * np.pi)) + # log_prob is updated with the synthetic-phase conditional. + assert not np.array_equal(result.samples["log_prob"].to_numpy(), log_prob_before) + + +def test_sample_synthetic_phase_requires_uniform_phase_prior(): + # When `phase` is in the samples, the phase prior is not split off (it is None), + # so synthetic phase sampling is not applicable and must raise. + result = make_gw_result(drop_phase=False) + with pytest.raises(ValueError, match="[Pp]hase prior"): + result.sample_synthetic_phase({"n_grid": 16}) + + +# --------------------------------------------------------------------------- +# sample_calibration_parameters +# --------------------------------------------------------------------------- + + +def _write_envelope(path): + """Write a minimal LVC-format calibration envelope file. + + Columns: frequency, median-amp, median-phase, -1sigma-amp, -1sigma-phase, + +1sigma-amp, +1sigma-phase. At least 4 rows are needed (cubic spline), spanning + the domain's [f_min, f_max]. + """ + freqs = np.geomspace(15.0, 300.0, 8) + data = np.column_stack( + [ + freqs, + np.ones_like(freqs), # median amplitude ~ 1 + np.zeros_like(freqs), # median phase ~ 0 + np.full_like(freqs, 0.99), # -1 sigma amplitude + np.full_like(freqs, -0.01), # -1 sigma phase + np.full_like(freqs, 1.01), # +1 sigma amplitude + np.full_like(freqs, 0.01), # +1 sigma phase + ] + ) + np.savetxt(path, data) + + +def _calibration_kwargs(tmp_path, correction_type="data", num_nodes=5): + envelopes = {} + for ifo in DETECTORS: + path = tmp_path / f"{ifo}.txt" + _write_envelope(path) + envelopes[ifo] = str(path) + return { + "calibration_envelope": envelopes, + "num_calibration_nodes": num_nodes, + "correction_type": correction_type, + } + + +def test_sample_calibration_parameters_adds_recalib_columns(tmp_path): + result = make_gw_result() + log_prob_before = result.samples["log_prob"].to_numpy().copy() + n_nodes = 5 + + result.sample_calibration_parameters( + _calibration_kwargs(tmp_path, num_nodes=n_nodes) + ) + + # Amplitude + phase nodes per detector (the frequency nodes are delta functions + # and are dropped before sampling). + recalib_cols = [c for c in result.samples.columns if c.startswith("recalib_")] + assert len(recalib_cols) == 2 * n_nodes * len(DETECTORS) + # The calibration prior log_prob is folded into the proposal log_prob. + assert not np.array_equal(result.samples["log_prob"].to_numpy(), log_prob_before) + # The calibration priors are recorded for persistence and added to the prior. + assert len(result.importance_sampling_metadata["prior_update"]) == len(recalib_cols) + assert any("recalib" in key for key in result.prior.keys()) + + +def test_sample_calibration_parameters_invalid_correction_type(): + result = make_gw_result() + # Parsed before any envelope file is read, so no files are needed. + with pytest.raises(ValueError, match="not understood"): + result.sample_calibration_parameters({"correction_type": "bogus"}) + + +@pytest.mark.parametrize( + "correction_type", + ["data", "template", {"H1": "data", "L1": "template"}, None], +) +def test_sample_calibration_parameters_correction_type_variants( + tmp_path, correction_type +): + result = make_gw_result() + result.sample_calibration_parameters( + _calibration_kwargs(tmp_path, correction_type=correction_type) + ) + assert any(c.startswith("recalib_") for c in result.samples.columns) diff --git a/tests/gw/test_gwutils.py b/tests/gw/test_gwutils.py new file mode 100644 index 000000000..a7b2055d7 --- /dev/null +++ b/tests/gw/test_gwutils.py @@ -0,0 +1,150 @@ +import numpy as np +import pandas as pd +import pytest + +from dingo.gw.domains import UniformFrequencyDomain +from dingo.gw.gwutils import ( + get_extrinsic_prior_dict, + get_mismatch, + get_standardization_dict, + get_window, +) + + +# --------------------------------------------------------------------------- +# get_mismatch +# +# mismatch = 1 - overlap, with overlap = / sqrt( ). +# Properties checked here mirror bilby's overlap test +# (bilby/test/gw/utils_test.py::TestGWUtils::test_overlap), adapted to dingo's +# get_mismatch. +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def domain(): + return UniformFrequencyDomain(20.0, 256.0, delta_f=0.5) + + +@pytest.fixture() +def waveforms(domain): + rng = np.random.default_rng(0) + n = len(domain) + a = rng.normal(size=n) + 1j * rng.normal(size=n) + b = rng.normal(size=n) + 1j * rng.normal(size=n) + return a, b + + +def test_mismatch_of_identical_waveforms_is_zero(domain, waveforms): + a, _ = waveforms + assert get_mismatch(a, a, domain) == pytest.approx(0.0, abs=1e-12) + + +def test_mismatch_is_scale_invariant(domain, waveforms): + # Overlap is normalized, so a rescaling of one waveform leaves the mismatch at 0. + a, _ = waveforms + assert get_mismatch(a, 3.0 * a, domain) == pytest.approx(0.0, abs=1e-12) + + +def test_mismatch_is_symmetric(domain, waveforms): + a, b = waveforms + assert get_mismatch(a, b, domain) == pytest.approx(get_mismatch(b, a, domain)) + + +def test_mismatch_is_in_valid_range(domain, waveforms): + # overlap in [-1, 1] => mismatch = 1 - overlap in [0, 2]. + a, b = waveforms + assert 0.0 <= get_mismatch(a, b, domain) <= 2.0 + + +# --------------------------------------------------------------------------- +# get_window +# --------------------------------------------------------------------------- + + +def test_get_window_tukey_length_and_range(): + T, f_s = 4.0, 1024 + window = get_window({"type": "tukey", "roll_off": 0.4, "T": T, "f_s": f_s}) + assert len(window) == int(T * f_s) + assert np.all((window >= 0.0) & (window <= 1.0)) + # A Tukey window tapers to (near) zero at the edges. + assert window[0] < 1e-6 and window[-1] < 1e-6 + + +def test_get_window_unknown_type_raises(): + with pytest.raises(NotImplementedError, match="window type"): + get_window({"type": "not_a_window"}) + + +# --------------------------------------------------------------------------- +# get_extrinsic_prior_dict +# --------------------------------------------------------------------------- + + +def test_get_extrinsic_prior_dict_expands_default_and_keeps_override(): + override = "bilby.core.prior.Uniform(minimum=100, maximum=1000)" + out = get_extrinsic_prior_dict({"ra": "default", "luminosity_distance": override}) + # "default" is replaced by the package default prior (no longer the literal string). + assert out["ra"] != "default" + # A non-default value is passed through unchanged. + assert out["luminosity_distance"] == override + + +# --------------------------------------------------------------------------- +# get_standardization_dict +# --------------------------------------------------------------------------- + + +class _StubWaveformDataset: + """Minimal stand-in exposing only what get_standardization_dict needs: + parameter_mean_std() for intrinsic params (extrinsic ones come from the prior).""" + + def __init__(self, luminosity_distance_std=0.0): + self._ld_std = luminosity_distance_std + self.parameters = pd.DataFrame({"chirp_mass": [30.0]}) + + def parameter_mean_std(self): + mean = {"chirp_mass": 30.0, "luminosity_distance": 100.0} + std = {"chirp_mass": 5.0, "luminosity_distance": self._ld_std} + return mean, std + + +@pytest.fixture() +def extrinsic_prior(): + return get_extrinsic_prior_dict( + { + "ra": "default", + "dec": "default", + "psi": "default", + "luminosity_distance": ( + "bilby.core.prior.Uniform(" + "minimum=100, maximum=1000, name='luminosity_distance')" + ), + "geocent_time": ( + "bilby.core.prior.Uniform(" + "minimum=-0.1, maximum=0.1, name='geocent_time')" + ), + } + ) + + +def test_get_standardization_dict_combines_intrinsic_and_extrinsic(extrinsic_prior): + selected = ["chirp_mass", "ra", "luminosity_distance"] + out = get_standardization_dict(extrinsic_prior, _StubWaveformDataset(), selected) + + assert set(out["mean"]) == set(selected) == set(out["std"]) + # Intrinsic parameter values come straight from the dataset. + assert out["mean"]["chirp_mass"] == 30.0 + assert out["std"]["chirp_mass"] == 5.0 + # Extrinsic parameter standardization is analytic / from the prior. + assert out["std"]["ra"] > 0 + + +def test_get_standardization_dict_rejects_nonzero_intrinsic_std_for_extrinsic( + extrinsic_prior, +): + # luminosity_distance is sampled as an extrinsic parameter, so the dataset must + # hold it at a fixed (std 0) value; a non-zero intrinsic std is an error. + wfd = _StubWaveformDataset(luminosity_distance_std=5.0) + with pytest.raises(ValueError, match="fixed value"): + get_standardization_dict(extrinsic_prior, wfd, ["chirp_mass"]) diff --git a/tests/gw/test_likelihood.py b/tests/gw/test_likelihood.py new file mode 100644 index 000000000..4679a8aa7 --- /dev/null +++ b/tests/gw/test_likelihood.py @@ -0,0 +1,213 @@ +import numpy as np +import pytest +from bilby.core.prior import Uniform +from scipy.integrate import trapezoid + +from dingo.gw.domains import UniformFrequencyDomain +from dingo.gw.likelihood import ( + StationaryGaussianGWLikelihood, + inner_product, + inner_product_complex, +) + + +# --------------------------------------------------------------------------- +# Pure inner-product helpers +# --------------------------------------------------------------------------- + +A = np.array([1 + 2j, 3 - 1j, 0 + 1j]) +B = np.array([2 + 0j, 1 + 1j, 1 - 1j]) + + +def test_inner_product_whitened(): + result = inner_product(A, B) + assert result == pytest.approx(np.sum(A.conj() * B).real) + assert result == pytest.approx(3.0) # hand-computed + + +def test_inner_product_min_idx_truncates_leading_bins(): + result = inner_product(A, B, min_idx=1) + assert result == pytest.approx(np.sum((A.conj() * B)[1:]).real) + assert result == pytest.approx(1.0) # hand-computed + + +def test_inner_product_unwhitened(): + psd = np.array([1.0, 2.0, 4.0]) + delta_f = 0.5 + result = inner_product(A, B, delta_f=delta_f, psd=psd) + assert result == pytest.approx(4 * delta_f * np.sum(A.conj() * B / psd).real) + assert result == pytest.approx(5.5) # hand-computed + + +def test_inner_product_psd_without_delta_f_raises(): + with pytest.raises(ValueError, match="delta_f and psd"): + inner_product(A, B, psd=np.ones(3)) + + +def test_inner_product_sums_only_axis_0(): + # A trailing axis (e.g., a phase grid) is preserved; the sum is over axis 0 only. + a = np.ones((3, 2)) + b = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) + result = inner_product(a, b) + assert result.shape == (2,) + np.testing.assert_allclose(result, [9.0, 12.0]) + + +def test_inner_product_self_is_sum_of_squared_magnitudes(): + result = inner_product(A, A) + assert result == pytest.approx(np.sum(np.abs(A) ** 2)) + assert result >= 0 + + +def test_inner_product_complex_retains_imaginary_part(): + # The complex variant does not take the real part. + result = inner_product_complex(A, B) + assert result == pytest.approx(np.sum(A.conj() * B)) + assert np.iscomplexobj(result) + # Its real part equals the (real) inner_product. + assert result.real == pytest.approx(inner_product(A, B)) + + +def test_inner_product_complex_psd_without_delta_f_raises(): + with pytest.raises(ValueError, match="delta_f and psd"): + inner_product_complex(A, B, psd=np.ones(3)) + + +# --------------------------------------------------------------------------- +# StationaryGaussianGWLikelihood +# --------------------------------------------------------------------------- + +THETA = { + "chirp_mass": 30.0, + "mass_ratio": 0.8, + "chi_1": 0.1, + "chi_2": -0.1, + "theta_jn": 1.0, + "phase": 1.3, + "ra": 1.5, + "dec": -0.3, + "psi": 1.2, + "luminosity_distance": 500.0, + "geocent_time": 0.0, +} + + +@pytest.fixture() +def domain(): + return UniformFrequencyDomain(20.0, 256.0, delta_f=0.5) + + +@pytest.fixture() +def event_data(domain): + mask = domain.frequency_mask + waveform = {d: np.where(mask, (1.0 + 1j) * 1e-21, 0.0) for d in ("H1", "L1")} + asds = {d: np.where(mask, 1e-21, 1.0) for d in ("H1", "L1")} + return {"waveform": waveform, "asds": asds} + + +def make_likelihood(domain, event_data, **extra): + return StationaryGaussianGWLikelihood( + wfg_kwargs={"approximant": "IMRPhenomD", "f_ref": 20.0}, + wfg_domain=domain, + data_domain=domain, + event_data=event_data, + t_ref=1126259462.4, + **extra, + ) + + +def test_log_Zn_is_minus_half_inner_product_of_data(domain, event_data): + likelihood = make_likelihood(domain, event_data) + expected = sum( + -0.5 * inner_product(d, d) for d in likelihood.whitened_strains.values() + ) + assert likelihood.log_Zn == pytest.approx(expected) + + +def test_log_likelihood_decomposition_identity(domain, event_data): + """log L = log_Zn + - 1/2 .""" + likelihood = make_likelihood(domain, event_data) + mu = likelihood.signal(dict(THETA))["waveform"] + d = likelihood.whitened_strains + rho2opt = sum(inner_product(m, m) for m in mu.values()) + kappa2 = sum(inner_product(di, mi) for di, mi in zip(d.values(), mu.values())) + expected = likelihood.log_Zn + kappa2 - 0.5 * rho2opt + assert likelihood.log_likelihood(dict(THETA)) == pytest.approx(expected) + + +def test_multiple_marginalizations_raise(domain, event_data): + likelihood = make_likelihood( + domain, + event_data, + time_marginalization_kwargs={"t_lower": -0.05, "t_upper": 0.05, "n_fft": 1}, + phase_marginalization_kwargs={}, + ) + with pytest.raises(NotImplementedError): + likelihood.log_likelihood(dict(THETA)) + + +def _brute_force_marginal_log_likelihood(plain_likelihood, key, prior, n=1000): + """Numerically marginalize the non-marginalized likelihood over ``key``. + + Direct adaptation of bilby's marginalization-correctness check + (bilby/test/gw/likelihood/marginalization_test.py::TestMarginalizations._template): + evaluate the non-marginalized likelihood on a grid of the marginalized parameter + and integrate it against the parameter's prior with the trapezoidal rule. + """ + values = np.linspace(prior.minimum, prior.maximum, n) + prior_values = prior.prob(values) + ln_likes = np.array( + [plain_likelihood._log_likelihood({**THETA, key: float(v)}) for v in values] + ) + like = np.exp(ln_likes - ln_likes.max()) + return np.log(trapezoid(like * prior_values, values)) + ln_likes.max() + + +def test_phase_marginalization_matches_brute_force_integral(domain, event_data): + """Analytic phase-marginalized likelihood matches the brute-force integral. + + Same comparison as bilby's marginalization test (``_template``), specialized to + the phase parameter with a uniform [0, 2*pi) prior. + """ + marginalized = make_likelihood( + domain, event_data, phase_marginalization_kwargs={"approximation_22_mode": True} + )._log_likelihood_phase_marginalized(dict(THETA)) + + brute_force = _brute_force_marginal_log_likelihood( + make_likelihood(domain, event_data), + key="phase", + prior=Uniform(minimum=0.0, maximum=2 * np.pi, name="phase"), + ) + # The analytic (Bessel) form is exact for a (2,2)-dominated waveform, and the + # brute-force integrand is smooth and periodic on [0, 2*pi), so the trapezoidal + # rule is spectrally accurate -> the two agree to the integration floor + # (observed residual ~0). (bilby's own test uses a much looser delta=0.5.) + assert marginalized == pytest.approx(brute_force, abs=1e-3) + + +def test_time_marginalization_matches_brute_force_integral(domain, event_data): + """FFT-based time-marginalized likelihood matches the brute-force integral. + + Same comparison as bilby's marginalization test (``_template``), specialized to + the geocent_time parameter with a uniform prior over [t_lower, t_upper]. + """ + t_lower, t_upper = -0.02, 0.02 + marginalized = make_likelihood( + domain, + event_data, + time_marginalization_kwargs={ + "t_lower": t_lower, + "t_upper": t_upper, + "n_fft": 5, + }, + )._log_likelihood_time_marginalized(dict(THETA)) + + brute_force = _brute_force_marginal_log_likelihood( + make_likelihood(domain, event_data), + key="geocent_time", + prior=Uniform(minimum=t_lower, maximum=t_upper, name="geocent_time"), + ) + # The residual is dominated by the FFT time-grid discretization (resolution + # delta_t / n_fft = 1 / (f_max * n_fft)); observed ~0.019 for n_fft=5. The + # tolerance is set just above that. (bilby's own test uses a looser delta=0.5.) + assert marginalized == pytest.approx(brute_force, abs=0.05) diff --git a/tests/integration/config/asd_dataset_settings.yaml b/tests/integration/config/asd_dataset_settings.yaml new file mode 100644 index 000000000..0ce647568 --- /dev/null +++ b/tests/integration/config/asd_dataset_settings.yaml @@ -0,0 +1,13 @@ +dataset_settings: + f_s: 4096 + time_psd: 1024 + T: 4.0 + window: + roll_off: 0.4 + type: tukey + time_gap: 0 + num_psds_max: 1 + detectors: + - H1 + - L1 + observing_run: O1 diff --git a/tests/integration/config/gw150914.ini b/tests/integration/config/gw150914.ini new file mode 100644 index 000000000..282df87f3 --- /dev/null +++ b/tests/integration/config/gw150914.ini @@ -0,0 +1,22 @@ +local = True +device = cuda + +# GW150914 trigger: 'GW150914' is resolved to GPS 1126259462.4 via the gwosc package. +trigger-time = GW150914 +label = gw150914 +outdir = gw150914_out +detectors = [H1, L1] +channel-dict = {H1:GWOSC, L1:GWOSC} + +importance-sample = True +recover-log-prob = False +# Match the synthetic-phase settings used for the injection stage. +importance-sampling-settings = {synthetic_phase: {approximation_22_mode: True, n_grid: 5001, uniform_weight: 0.01}} + +num_samples = 20000 +batch_size = 20000 + +model = model_latest.pt + +# Duration must match the training domain: 1 / delta_f = 1 / 0.25 = 4 s. +duration = 4 diff --git a/tests/integration/config/injection.ini b/tests/integration/config/injection.ini new file mode 100644 index 000000000..88beb3f06 --- /dev/null +++ b/tests/integration/config/injection.ini @@ -0,0 +1,23 @@ +local = True +device = cuda + +trigger-time = 1126259462.4 +label = heavy_ci +outdir = inference_out +detectors = [H1, L1] + +importance-sample = True +recover-log-prob = False +importance-sampling-settings = {synthetic_phase: {approximation_22_mode: True, n_grid: 5001, uniform_weight: 0.01}} + +num_samples = 20000 # SEARCH-TUNED (IS proposal count) +batch_size = 20000 + +model = model_latest.pt + +gaussian-noise = True +n-simulation = 1 +injection = True +dingo-injection = True +asd-dataset = asds_O1.hdf5 +injection-dict = {chirp_mass: 31.2, mass_ratio: 0.864, phase: 0.0, a_1: 0.0, a_2: 0.0, tilt_1: 0.0, tilt_2: 0.0, phi_12: 0.0, phi_jl: 0.0, theta_jn: 2.68, luminosity_distance: 439.0, geocent_time: 1126259462.4, dec: -1.21, ra: 1.68, psi: 0.0} diff --git a/tests/integration/config/train_settings.yaml b/tests/integration/config/train_settings.yaml new file mode 100644 index 000000000..9c49ad468 --- /dev/null +++ b/tests/integration/config/train_settings.yaml @@ -0,0 +1,61 @@ +data: + waveform_dataset_path: waveform_dataset.hdf5 + train_fraction: 0.95 + detectors: + - H1 + - L1 + # All extrinsic parameters fixed (delta) so only chirp_mass/mass_ratio vary. + extrinsic_prior: + dec: -1.21 + ra: 1.68 + geocent_time: 0.0 + psi: 0.0 + luminosity_distance: 439.0 + ref_time: 1126259462.391 + inference_parameters: + - chirp_mass + - mass_ratio + +model: + posterior_model_type: normalizing_flow + posterior_kwargs: + num_flow_steps: 3 # SEARCH-TUNED + base_transform_kwargs: + hidden_dim: 32 # SEARCH-TUNED + num_transform_blocks: 3 + activation: elu + dropout_probability: 0.0 + batch_norm: True + num_bins: 8 + base_transform_type: rq-coupling + embedding_kwargs: + output_dim: 64 + hidden_dims: [256, 128, 64] # SEARCH-TUNED + activation: elu + dropout: 0.0 + batch_norm: True + svd: + num_training_samples: 1000 + num_validation_samples: 100 + size: 50 # SEARCH-TUNED + +training: + stage_0: + epochs: 10 # SEARCH-TUNED + asd_dataset_path: asds_O1.hdf5 + freeze_rb_layer: True + optimizer: + type: adam + lr: 0.0001 + scheduler: + type: cosine + T_max: 10 # keep equal to epochs + batch_size: 64 + +local: + device: cuda + num_workers: 8 + runtime_limits: + max_time_per_run: 3600000 + max_epochs_per_run: 1000 + checkpoint_epochs: 1000 diff --git a/tests/integration/config/waveform_dataset_settings.yaml b/tests/integration/config/waveform_dataset_settings.yaml new file mode 100644 index 000000000..2c2d09bd9 --- /dev/null +++ b/tests/integration/config/waveform_dataset_settings.yaml @@ -0,0 +1,29 @@ +domain: + type: UniformFrequencyDomain + f_min: 20.0 + f_max: 1024.0 + delta_f: 0.25 + +waveform_generator: + approximant: IMRPhenomXPHM + f_ref: 20.0 + +# Only chirp_mass, mass_ratio and phase vary. Spins zero; everything else fixed (delta). +intrinsic_prior: + mass_1: bilby.core.prior.Constraint(minimum=5.0, maximum=100.0, name='mass_1') + mass_2: bilby.core.prior.Constraint(minimum=5.0, maximum=100.0, name='mass_2') + chirp_mass: bilby.gw.prior.UniformInComponentsChirpMass(minimum=20.0, maximum=40.0, name='chirp_mass') + mass_ratio: bilby.gw.prior.UniformInComponentsMassRatio(minimum=0.5, maximum=1.0, name='mass_ratio') + phase: default + a_1: 0.0 + a_2: 0.0 + tilt_1: 0.0 + tilt_2: 0.0 + phi_12: 0.0 + phi_jl: 0.0 + theta_jn: 2.68 + luminosity_distance: 439.0 + geocent_time: 0.0 + +num_samples: 50000 # SEARCH-TUNED +compression: None diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py new file mode 100644 index 000000000..ae6465c4f --- /dev/null +++ b/tests/integration/conftest.py @@ -0,0 +1,27 @@ +"""Shared skip conditions for heavy integration tests.""" +import shutil +import subprocess + + +def _detect_apptainer(): + for cmd in ("apptainer", "singularity"): + if shutil.which(cmd): + return cmd + return None + + +def _detect_gpu(): + if not shutil.which("nvidia-smi"): + return False + try: + out = subprocess.run( + ["nvidia-smi", "-L"], capture_output=True, text=True, timeout=30 + ) + return out.returncode == 0 and "GPU 0" in out.stdout + except (subprocess.SubprocessError, OSError): + return False + + +APPTAINER_CMD = _detect_apptainer() +HAS_APPTAINER = APPTAINER_CMD is not None +HAS_GPU = _detect_gpu() diff --git a/tests/integration/run_pipeline.py b/tests/integration/run_pipeline.py new file mode 100644 index 000000000..5388ba4d7 --- /dev/null +++ b/tests/integration/run_pipeline.py @@ -0,0 +1,164 @@ +"""Run the heavy end-to-end DINGO smoke pipeline and print the sample efficiency. + +Stages: waveform dataset -> ASD dataset (fixed GPS) -> train -> dingo_pipe local +injection with importance sampling. Reads the IS sample efficiency from the +result file and prints it for the pytest wrapper to parse. +""" +import argparse +import glob +import os +import pickle +import shutil +import subprocess +import time + +HERE = os.path.dirname(os.path.abspath(__file__)) + + +def run(cmd, cwd): + print(f"+ {' '.join(cmd)}", flush=True) + subprocess.run(cmd, cwd=cwd, check=True) + + +def timed_stage(name, fn): + """Run fn(), return elapsed seconds, and return the result.""" + t0 = time.time() + result = fn() + elapsed = time.time() - t0 + return elapsed, result + + +def find_is_efficiency(outdir): + """Locate the importance-sampled result and return sample_efficiency (0-1).""" + from dingo.gw.result import Result + + candidates = sorted( + glob.glob(os.path.join(outdir, "**", "*.hdf5"), recursive=True) + ) + last_exc_info = None # (path, exception) from the most recent load failure + for path in candidates: + try: + result = Result(file_name=path) + except Exception as exc: + last_exc_info = (path, exc) + continue + eff = getattr(result, "sample_efficiency", None) + if eff is not None: + return eff, path + last_exc_msg = ( + f" Last load error: {last_exc_info[0]}: {last_exc_info[1]}" + if last_exc_info + else " No files raised a load error (files present but lacked sample_efficiency)." + ) + raise RuntimeError( + f"No importance-sampled result with sample_efficiency found under {outdir}. " + f"Scanned: {candidates}\n{last_exc_msg}" + ) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--config-dir", default=os.path.join(HERE, "config")) + parser.add_argument("--workdir", default=None) + args = parser.parse_args() + + # Default to /tmp so the script works inside a read-only SIF container. + # /tmp is always writable (apptainer mounts a tmpfs there). + # Always start from a clean workdir: apptainer bind-mounts the host /tmp, + # so stale files from a previous (possibly failed) run may be present. + workdir = args.workdir or os.path.join("/tmp", "dingo_run") + shutil.rmtree(workdir, ignore_errors=True) + os.makedirs(workdir, exist_ok=True) + for name in ( + "waveform_dataset_settings.yaml", + "asd_dataset_settings.yaml", + "train_settings.yaml", + "injection.ini", + "gw150914.ini", + ): + shutil.copy(os.path.join(args.config_dir, name), os.path.join(workdir, name)) + + stage_times = {} + pipeline_start = time.time() + + # 1. Waveform dataset + def _stage_waveform(): + run( + ["dingo_generate_dataset", "--settings_file", "waveform_dataset_settings.yaml", + "--num_processes", "8", "--out_file", "waveform_dataset.hdf5"], + cwd=workdir, + ) + + stage_times["waveform_dataset"], _ = timed_stage("waveform_dataset", _stage_waveform) + + # 2. ASD dataset at a fixed GPS time (GWOSC, no auth) + def _stage_asd(): + # Deterministic time segments: one fixed O1 segment (GWOSC open data), + # clear of GW150914 (GPS 1126259462.4); length >= asd time_psd. + ts_path = os.path.join(workdir, "time_segments.pkl") + gps_start, seg_len = 1126257000, 1024 + segments = {det: [(gps_start, gps_start + seg_len)] for det in ("H1", "L1")} + with open(ts_path, "wb") as f: + pickle.dump(segments, f) + run( + ["dingo_generate_asd_dataset", "--settings_file", "asd_dataset_settings.yaml", + "--data_dir", workdir, "--time_segments_file", ts_path, + "--out_name", "asds_O1.hdf5"], + cwd=workdir, + ) + + stage_times["asd_dataset"], _ = timed_stage("asd_dataset", _stage_asd) + + # 3. Train + def _stage_train(): + run(["dingo_train", "--settings_file", "train_settings.yaml", "--train_dir", workdir], + cwd=workdir) + + stage_times["train"], _ = timed_stage("train", _stage_train) + + # 4. Inference + importance sampling via dingo_pipe (local mode). + # dingo_pipe generates submit/bash_