nullcal 0.2.1__tar.gz → 0.3.0__tar.gz

{nullcal-0.2.1 → nullcal-0.3.0}/.github/workflows/scheduled_release.yml

@@ -31,7 +31,7 @@ jobs:
 
             - name: Check for new commits
               id: check_commits
-              uses: isaac-cf-wong/detect-git-changes-action@a8b0d224a4fdff017607079f83f01f013b9d66a6 # v0.0.20
+              uses: isaac-cf-wong/detect-git-changes-action@051225b2ef2faf2a02d61f16d1a027cdf941498d # v0.0.21
               with:
                   use-latest-tag: true
 

{nullcal-0.2.1 → nullcal-0.3.0}/.github/workflows/support_floor_update.yml

@@ -22,7 +22,7 @@ jobs:
 
             - name: Apply support policy (spec0, from pyproject tool table)
               id: policy
-              uses: isaac-cf-wong/dependency-support-policy-action@dfc24273d70d9730ede96b5fdef09a8d7887cf04 # v0.2.10
+              uses: isaac-cf-wong/dependency-support-policy-action@af9e20cfdcb7d14a4c08614e981e8c2339d3da0f # v0.2.11
               with:
                   mode: update
 

{nullcal-0.2.1 → nullcal-0.3.0}/.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
           - id: check-json
 
     - repo: https://github.com/astral-sh/ruff-pre-commit
-      rev: v0.16.9
+      rev: v0.16.10
       hooks:
           - id: ruff
             args: [--fix, --config=pyproject.toml]
@@ -54,7 +54,7 @@ repos:
           - id: actionlint
 
     - repo: https://github.com/astral-sh/uv-pre-commit
-      rev: 0.12.20
+      rev: 0.12.23
       hooks:
           # Update the uv lockfile
           - id: uv-lock

{nullcal-0.2.1 → nullcal-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: nullcal
-Version: 0.2.1
+Version: 0.3.0
 Summary: A package to constrain calibration errors of a closed-geometry network of gravitational-wave detectors.
 Project-URL: Source, https://github.com/Leuven-Gravity-Institute/nullcal
 Project-URL: Tracker, https://github.com/Leuven-Gravity-Institute/nullcal/issues
@@ -37,7 +37,7 @@ Classifier: Programming Language :: Python :: 3.14
 Requires-Python: >=3.12
 Requires-Dist: blackjax>=1.6.2
 Requires-Dist: jax>=0.11.1
-Requires-Dist: numpy>=1.26.4
+Requires-Dist: numpy>=2.2.0
 Requires-Dist: pandas>=2.3.3
 Requires-Dist: pyyaml>=6.0.3
 Requires-Dist: scipy>=1.16.3

{nullcal-0.2.1 → nullcal-0.3.0}/pyproject.toml

@@ -24,7 +24,7 @@ classifiers = [
 requires-python = ">=3.12"
 dynamic = ["version"]
 dependencies = [
-  "numpy>=1.26.4",
+  "numpy>=2.2.0",
   "blackjax>=1.6.2",
   "jax>=0.11.1",
   "scipy>=1.16.3",
@@ -53,15 +53,15 @@ Tracker = "https://github.com/Leuven-Gravity-Institute/nullcal/issues"
 
 [dependency-groups]
 dev = [
-  "ruff>=0.16.9",
-  "prek>=0.5.3",
+  "ruff>=0.16.10",
+  "prek>=0.5.4",
   "pytest-github-actions-annotate-failures>=0.4.2",
   "pytest-cov>=7.1.0",
-  "pytest-mock>=3.15.1",
+  "pytest-mock>=3.16.0",
   "pytest>=9.1.1",
 ]
 
-docs = ["zensical>=0.0.65", "mkdocstrings-python>=2.0.9"]
+docs = ["zensical>=0.0.67", "mkdocstrings-python>=2.0.9"]
 build = ["hatch>=1.18.1"]
 release = ["git-cliff>=2.14.2"]
 

{nullcal-0.2.1 → nullcal-0.3.0}/src/nullcal/likelihood/recalibration_likelihood.py

@@ -169,22 +169,50 @@ class RecalibrationLikelihood:
             phase,
         )
 
-    def _frequency_domain_null_stream(self, amplitude: jax.Array, phase: jax.Array) -> jax.Array:
+    def _validated_strain(self, whitened_strain) -> jax.Array:
+        strain = jnp.asarray(whitened_strain)
+        expected_shape = (self.parameter_shape[0], self._frequency_count)
+        if tuple(strain.shape[-2:]) != expected_shape:
+            raise ValueError(
+                f"whitened_frequency_domain_strain must have trailing shape {expected_shape}; "
+                f"received {tuple(strain.shape)}"
+            )
+        return strain
+
+    def _calibrated_projector(self, amplitude: jax.Array, phase: jax.Array) -> jax.Array:
+        """Build the per-frequency null projector for the calibrated response."""
         factors = self._calibration_factor(amplitude, phase)
         response = self._whitened_antenna_response * jnp.swapaxes(factors, 0, 1)[:, :, None]
         response_dagger = jnp.swapaxes(jnp.conj(response), 1, 2)
         gram = response_dagger @ response
         projected_response = response @ jnp.linalg.solve(gram, response_dagger)
-        projector = jnp.eye(response.shape[1], dtype=response.dtype)[None, :, :] - projected_response
-        masked_null_stream = jnp.einsum("fij,jf->if", projector, self._whitened_frequency_domain_strain)
+        return jnp.eye(response.shape[1], dtype=response.dtype)[None, :, :] - projected_response
+
+    def _project(self, projector: jax.Array, selected_strain: jax.Array) -> jax.Array:
+        """Apply the null projector to a masked frequency-domain strain."""
+        return jnp.einsum("fij,...jf->...if", projector, selected_strain)
+
+    def _frequency_domain_null_stream(
+        self, amplitude: jax.Array, phase: jax.Array, whitened_strain: jax.Array | None = None
+    ) -> jax.Array:
+        if whitened_strain is None:
+            selected_strain = self._whitened_frequency_domain_strain
+            realisation_shape = selected_strain.shape[:-1]
+        else:
+            strain = self._validated_strain(whitened_strain)
+            selected_strain = strain[..., self._frequency_indices]
+            realisation_shape = strain.shape[:-1]
+        projector = self._calibrated_projector(amplitude, phase)
+        masked_null_stream = self._project(projector, selected_strain)
+        full_shape = (*realisation_shape, self._frequency_count)
         return (
-            jnp.zeros((self.parameter_shape[0], self._frequency_count), dtype=masked_null_stream.dtype)
-            .at[:, self._frequency_indices]
+            jnp.zeros(full_shape, dtype=masked_null_stream.dtype)
+            .at[..., self._frequency_indices]
             .set(masked_null_stream)
         )
 
-    def _time_frequency_null_stream(self, amplitude: jax.Array, phase: jax.Array) -> jax.Array:
-        frequency_domain = self._frequency_domain_null_stream(amplitude, phase)
+    def _transform_to_time_frequency(self, frequency_domain: jax.Array) -> jax.Array:
+        """Map a ``(..., detector, frequency)`` array through the WDM transform."""
 
         def transform(detector_data):
             return (
@@ -197,14 +225,52 @@ class RecalibrationLikelihood:
                 * self._wavelet_scale
             )
 
-        time_frequency = jax.vmap(transform)(frequency_domain)
-        return jnp.where(self.time_frequency_filter[None, :, :], time_frequency, 0.0)
+        leading_shape = frequency_domain.shape[:-1]
+        flattened = frequency_domain.reshape((-1, frequency_domain.shape[-1]))
+        transformed = jax.vmap(transform)(flattened)
+        return transformed.reshape(leading_shape + transformed.shape[-2:])
+
+    def _time_frequency_null_stream(
+        self, amplitude: jax.Array, phase: jax.Array, whitened_strain: jax.Array | None = None
+    ) -> jax.Array:
+        frequency_domain = self._frequency_domain_null_stream(amplitude, phase, whitened_strain)
+        return jnp.where(self.time_frequency_filter, self._transform_to_time_frequency(frequency_domain), 0.0)
+
+    def _log_likelihood_from_stream(self, null_stream: jax.Array) -> jax.Array:
+        # The detector axis always sits immediately before the time and
+        # frequency axes, so summing the trailing three leaves exactly one
+        # value per leading realisation (and a scalar for a single realisation).
+        return -0.5 * jnp.sum(jnp.abs(null_stream) ** 2, axis=(-1, -2, -3))
 
     def log_likelihood_fn(self, params: Mapping[str, jax.Array]) -> jax.Array:
         """Return the pure JAX null-stream log likelihood for ``params``."""
         amplitude, phase = self._validated_parameters(params)
-        null_stream = self._time_frequency_null_stream(amplitude, phase)
-        return -0.5 * jnp.sum(jnp.abs(null_stream) ** 2)
+        return self._log_likelihood_from_stream(self._time_frequency_null_stream(amplitude, phase))
+
+    def log_likelihood_for_strain(self, params: Mapping[str, jax.Array], whitened_frequency_domain_strain) -> jax.Array:
+        """Evaluate the likelihood on a supplied whitened strain realisation.
+
+        ``whitened_frequency_domain_strain`` has trailing shape
+        ``(detector, frequency)`` and may carry any number of leading
+        realisation axes. The return value drops those leading axes and keeps
+        one log likelihood per realisation, so a ``vmap`` over the leading axis
+        yields one posterior target per realisation.
+        """
+        amplitude, phase = self._validated_parameters(params)
+        null_stream = self._time_frequency_null_stream(amplitude, phase, whitened_frequency_domain_strain)
+        return self._log_likelihood_from_stream(null_stream)
+
+    def logdensity_for_strain(self, params: Mapping[str, jax.Array], whitened_frequency_domain_strain) -> jax.Array:
+        """Return log prior plus the per-realisation log likelihood on supplied strain."""
+        amplitude, phase = self._validated_parameters(params)
+        return self.log_likelihood_for_strain(params, whitened_frequency_domain_strain) + calibration_log_prior(
+            amplitude,
+            phase,
+            self.amplitude_prior_mean,
+            self.amplitude_prior_sigma,
+            self.phase_prior_mean,
+            self.phase_prior_sigma,
+        )
 
     def logdensity_fn(self, params: Mapping[str, jax.Array]) -> jax.Array:
         """Return normalized Gaussian log prior plus null-stream log likelihood."""

nullcal-0.3.0/src/nullcal/sampler.py

@@ -0,0 +1,228 @@
+"""BlackJAX NUTS sampling with convergence diagnostics."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from typing import Any
+
+import blackjax
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from .result import Result
+
+MINIMUM_DIAGNOSTIC_CHAINS = 2
+
+
+def _jitter_position(position, key, scale: float):
+    leaves, structure = jax.tree.flatten(position)
+    keys = jax.random.split(key, len(leaves))
+    jittered = [
+        jnp.asarray(leaf) + scale * jax.random.normal(leaf_key, shape=jnp.asarray(leaf).shape)
+        for leaf, leaf_key in zip(leaves, keys, strict=True)
+    ]
+    return jax.tree.unflatten(structure, jittered)
+
+
+def _run_chain(
+    logdensity_fn: Callable[[Any], jax.Array],
+    initial_position: Mapping[str, Any],
+    chain_key: jax.Array,
+    *,
+    num_warmup: int,
+    num_samples: int,
+    target_acceptance_rate: float,
+    initial_position_jitter: float,
+):
+    """Adapt and sample one NUTS chain from a single jittered start."""
+    jitter_key, warmup_key, sample_key = jax.random.split(chain_key, 3)
+    chain_initial_position = _jitter_position(initial_position, jitter_key, initial_position_jitter)
+    adaptation = blackjax.window_adaptation(
+        blackjax.nuts,
+        logdensity_fn,
+        target_acceptance_rate=target_acceptance_rate,
+    )
+    adapted, _ = adaptation.run(warmup_key, chain_initial_position, num_steps=num_warmup)
+    kernel = blackjax.nuts(logdensity_fn, **adapted.parameters)
+    sample_keys = jax.random.split(sample_key, num_samples)
+
+    def step(state, key):
+        next_state, info = kernel.step(key, state)
+        return next_state, (next_state.position, next_state.logdensity, info.is_divergent, info.num_integration_steps)
+
+    _, (positions, logdensity, divergences, integration_steps) = jax.lax.scan(step, adapted.state, sample_keys)
+    return positions, logdensity, divergences, jnp.sum(integration_steps)
+
+
+def sample_nuts(
+    logdensity_fn: Callable[[Any], jax.Array],
+    initial_position: Mapping[str, Any],
+    *,
+    seed: int = 0,
+    num_chains: int = 4,
+    num_warmup: int = 1_000,
+    num_samples: int = 1_000,
+    target_acceptance_rate: float = 0.8,
+    initial_position_jitter: float = 0.01,
+) -> Result:
+    """Adapt and run independent NUTS chains, returning mandatory diagnostics.
+
+    Chains are executed independently rather than treating detector data as a
+    realization batch. Returned parameter arrays flatten chain and draw into a
+    single leading sample dimension; chain-aware R-hat and ESS are retained in
+    ``Result.metadata``.
+    """
+    if num_chains < MINIMUM_DIAGNOSTIC_CHAINS:
+        raise ValueError("at least two chains are required to compute R-hat")
+    if num_warmup <= 0 or num_samples <= 0:
+        raise ValueError("num_warmup and num_samples must be positive")
+    if initial_position_jitter < 0.0:
+        raise ValueError("initial_position_jitter must be non-negative")
+
+    root_key = jax.random.key(seed)
+    chain_keys = jax.random.split(root_key, num_chains)
+    chain_positions = []
+    chain_logdensity = []
+    chain_divergences = []
+    chain_integration_steps = []
+
+    for chain_key in chain_keys:
+        positions, logdensity, divergences, integration_steps = _run_chain(
+            logdensity_fn,
+            initial_position,
+            chain_key,
+            num_warmup=num_warmup,
+            num_samples=num_samples,
+            target_acceptance_rate=target_acceptance_rate,
+            initial_position_jitter=initial_position_jitter,
+        )
+        chain_positions.append(positions)
+        chain_logdensity.append(logdensity)
+        chain_divergences.append(divergences)
+        chain_integration_steps.append(integration_steps)
+
+    positions_by_chain = jax.tree.map(lambda *values: jnp.stack(values), *chain_positions)
+    logdensity_by_chain = jnp.stack(chain_logdensity)
+    divergence_by_chain = jnp.stack(chain_divergences)
+    integration_steps_by_chain = jnp.stack(chain_integration_steps)
+
+    diagnostic_leaves = [value.reshape((num_chains, num_samples, -1)) for value in jax.tree.leaves(positions_by_chain)]
+    diagnostic_array = jnp.concatenate(diagnostic_leaves, axis=-1)
+    rhat = blackjax.diagnostics.rhat(diagnostic_array)
+    ess_bulk = blackjax.diagnostics.ess_bulk(diagnostic_array)
+    ess_tail = blackjax.diagnostics.ess_tail(diagnostic_array)
+
+    flattened_positions = jax.tree.map(
+        lambda value: value.reshape((num_chains * num_samples, *value.shape[2:])), positions_by_chain
+    )
+    flattened_logdensity = logdensity_by_chain.reshape(num_chains * num_samples)
+    flattened_divergence = divergence_by_chain.reshape(num_chains * num_samples)
+    metadata = {
+        "seed": seed,
+        "num_chains": num_chains,
+        "num_warmup": num_warmup,
+        "num_samples_per_chain": num_samples,
+        "target_acceptance_rate": target_acceptance_rate,
+        "max_rhat": float(np.asarray(jnp.max(rhat))),
+        "min_ess_bulk": float(np.asarray(jnp.min(ess_bulk))),
+        "min_ess_tail": float(np.asarray(jnp.min(ess_tail))),
+        "divergences": int(np.asarray(jnp.sum(divergence_by_chain))),
+        "integration_steps": int(np.asarray(jnp.sum(integration_steps_by_chain))),
+    }
+    return Result(
+        samples=flattened_positions,
+        logdensity=flattened_logdensity,
+        info={"is_divergent": flattened_divergence},
+        metadata=metadata,
+    )
+
+
+def sample_nuts_batched(
+    logdensity_fn: Callable[[Any, Any], jax.Array],
+    initial_position: Mapping[str, Any],
+    realisation_data: Any,
+    *,
+    seed: int = 0,
+    num_chains_per_realisation: int = 1,
+    num_warmup: int = 1_000,
+    num_samples: int = 1_000,
+    target_acceptance_rate: float = 0.8,
+    initial_position_jitter: float = 0.01,
+) -> Result:
+    """Run one NUTS chain per (realisation, chain) pair under a single ``vmap``.
+
+    ``logdensity_fn(position, realisation)`` evaluates one posterior target:
+    the leading axis of every leaf of ``realisation_data`` is the realisation
+    index. All chains advance inside one ``vmap``, so a campaign pays one
+    dispatch and one accelerator launch per batch instead of one per
+    realisation. The returned sample arrays flatten ``(realisation, chain,
+    draw)`` into the leading axis; the split sizes are in ``Result.metadata``.
+
+    The single-realisation :func:`sample_nuts` is the diagnostic path and
+    computes R-hat and ESS; this batched path trades those diagnostics for
+    throughput and is the cost-measurement target.
+    """
+    if num_chains_per_realisation < 1:
+        raise ValueError("num_chains_per_realisation must be at least one")
+    if num_warmup <= 0 or num_samples <= 0:
+        raise ValueError("num_warmup and num_samples must be positive")
+    if initial_position_jitter < 0.0:
+        raise ValueError("initial_position_jitter must be non-negative")
+
+    leaves = jax.tree.leaves(realisation_data)
+    if not leaves:
+        raise ValueError("realisation_data must contain at least one array")
+    realisation_count = int(jnp.shape(leaves[0])[0]) if jnp.ndim(leaves[0]) > 0 else 0
+    if realisation_count < 1:
+        raise ValueError("realisation_data must contain at least one realisation on a leading axis")
+
+    total_chains = realisation_count * num_chains_per_realisation
+    realisation_index = jnp.repeat(jnp.arange(realisation_count), num_chains_per_realisation)
+    batched_realisations = jax.tree.map(lambda value: jnp.asarray(value)[realisation_index], realisation_data)
+    broadcast_positions = jax.tree.map(
+        lambda value: jnp.broadcast_to(jnp.asarray(value, dtype=jnp.float64), (total_chains, *jnp.shape(value))),
+        initial_position,
+    )
+    chain_keys = jax.random.split(jax.random.key(seed), total_chains)
+
+    def chain(position, chain_key, realisation):
+        return _run_chain(
+            lambda sample_position: logdensity_fn(sample_position, realisation),
+            position,
+            chain_key,
+            num_warmup=num_warmup,
+            num_samples=num_samples,
+            target_acceptance_rate=target_acceptance_rate,
+            initial_position_jitter=initial_position_jitter,
+        )
+
+    positions_by_chain, logdensity_by_chain, divergence_by_chain, integration_steps_by_chain = jax.vmap(chain)(
+        broadcast_positions, chain_keys, batched_realisations
+    )
+
+    flattened_positions = jax.tree.map(
+        lambda value: value.reshape((total_chains * num_samples, *value.shape[2:])), positions_by_chain
+    )
+    flattened_logdensity = logdensity_by_chain.reshape(total_chains * num_samples)
+    flattened_divergence = divergence_by_chain.reshape(total_chains * num_samples)
+    metadata = {
+        "seed": seed,
+        "num_realisations": realisation_count,
+        "num_chains_per_realisation": num_chains_per_realisation,
+        "num_warmup": num_warmup,
+        "num_samples_per_chain": num_samples,
+        "target_acceptance_rate": target_acceptance_rate,
+        "initial_position_jitter": initial_position_jitter,
+        "divergences": int(np.asarray(jnp.sum(divergence_by_chain))),
+        "integration_steps": int(np.asarray(jnp.sum(integration_steps_by_chain))),
+    }
+    return Result(
+        samples=flattened_positions,
+        logdensity=flattened_logdensity,
+        info={"is_divergent": flattened_divergence},
+        metadata=metadata,
+    )
+
+
+__all__ = ["sample_nuts", "sample_nuts_batched"]