nullcal 0.2.0__tar.gz → 0.3.0__tar.gz

{nullcal-0.2.0 → nullcal-0.3.0}/.github/workflows/codeql.yml

@@ -22,13 +22,13 @@ jobs:
               uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6
 
             - name: Initialize CodeQL
-              uses: github/codeql-action/init@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4
+              uses: github/codeql-action/init@2892aa5e19bbd11bc0cff5427e3b750a04d9e3c2 # v4
               with:
                   languages: python
                   queries: security-and-quality
 
             - name: Autobuild
-              uses: github/codeql-action/autobuild@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4
+              uses: github/codeql-action/autobuild@2892aa5e19bbd11bc0cff5427e3b750a04d9e3c2 # v4
 
             - name: Perform CodeQL Analysis
-              uses: github/codeql-action/analyze@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4
+              uses: github/codeql-action/analyze@2892aa5e19bbd11bc0cff5427e3b750a04d9e3c2 # v4

{nullcal-0.2.0 → nullcal-0.3.0}/.github/workflows/draft_release.yml

@@ -22,7 +22,7 @@ jobs:
 
             - name: Generate changelog
               id: changelog
-              uses: orhun/git-cliff-action@3d96a18cc4ec17e9dc69ddcc424ccafaf1f78ce2 # v4.9.0
+              uses: orhun/git-cliff-action@a9a95522b26fe6403f7bb24031f21fb573d0f5ff # v4.9.1
               with:
                   config: cliff.toml
                   args: --unreleased --verbose

{nullcal-0.2.0 → nullcal-0.3.0}/.github/workflows/release.yml

@@ -30,7 +30,7 @@ jobs:
 
             - name: Generate changelog
               id: changelog
-              uses: orhun/git-cliff-action@3d96a18cc4ec17e9dc69ddcc424ccafaf1f78ce2 # v4.9.0
+              uses: orhun/git-cliff-action@a9a95522b26fe6403f7bb24031f21fb573d0f5ff # v4.9.1
               with:
                   config: cliff.toml
                   args: --current --verbose

{nullcal-0.2.0 → 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@10d1e3ac7c5ee0997d7b370e9929cacac9c87013 # v0.0.19
+              uses: isaac-cf-wong/detect-git-changes-action@051225b2ef2faf2a02d61f16d1a027cdf941498d # v0.0.21
               with:
                   use-latest-tag: true
 

{nullcal-0.2.0 → 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@74dfcfcc079cd22b26713242940f65ba2b3c70c7 # v0.2.9
+              uses: isaac-cf-wong/dependency-support-policy-action@af9e20cfdcb7d14a4c08614e981e8c2339d3da0f # v0.2.11
               with:
                   mode: update
 

{nullcal-0.2.0 → 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.8
+      rev: v0.16.10
       hooks:
           - id: ruff
             args: [--fix, --config=pyproject.toml]
@@ -29,13 +29,13 @@ repos:
           - id: taplo-format
 
     - repo: https://github.com/crate-ci/typos
-      rev: v1.50.2
+      rev: v1.50.3
       hooks:
           - id: typos
             args: []
 
     - repo: https://github.com/rbubley/mirrors-prettier
-      rev: v3.9.8
+      rev: v3.9.9
       hooks:
           - id: prettier
             exclude: ^tests/e2e/diagnostics/likelihood_agreement\.json$
@@ -54,7 +54,7 @@ repos:
           - id: actionlint
 
     - repo: https://github.com/astral-sh/uv-pre-commit
-      rev: 0.12.17
+      rev: 0.12.23
       hooks:
           # Update the uv lockfile
           - id: uv-lock

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: nullcal
-Version: 0.2.0
+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.0 → nullcal-0.3.0}/docs/user-guide/quickstart.md

@@ -77,6 +77,29 @@ log_posterior = likelihood.logdensity_fn(params)
 Use `nullcal.sampler.sample_nuts` when the standard window-adapted NUTS run and
 R-hat, ESS, and divergence diagnostics are desired together.
 
+## 5. Summarize and Apply a Calibration Posterior
+
+```python
+from nullcal.calibration import posterior_median_calibration_factor
+
+# result is the output of nullcal.sampler.sample_nuts(...).
+frequencies = data.frequency_array[np.all(data.mask, axis=0)]
+factor = posterior_median_calibration_factor(
+    frequencies,
+    likelihood.knot_frequencies,
+    result.samples["amplitude"],
+    result.samples["phase"],
+)
+corrected_strain = np.asarray(data.strain).copy()
+corrected_strain[:, np.all(data.mask, axis=0)] /= np.asarray(factor)
+```
+
+The likelihood multiplies the whitened antenna response by the calibration
+factor. Correcting observed strain uses its inverse, as above. The point
+estimate evaluates the complex factor for every posterior sample at every
+frequency, then takes separate medians of its real and imaginary parts.
+Evaluating a spline from median knot values is a different estimator.
+
 ## Next Steps
 
 - See [Installation](installation.md) for environment setup

{nullcal-0.2.0 → 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.8",
-  "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.62", "mkdocstrings-python>=2.0.8"]
+docs = ["zensical>=0.0.67", "mkdocstrings-python>=2.0.9"]
 build = ["hatch>=1.18.1"]
 release = ["git-cliff>=2.14.2"]
 
@@ -84,7 +84,8 @@ parallel = true
 source = ["src"]
 
 [tool.coverage.report]
-fail_under = 55
+# Policy target: keep at least three-fifths of the package covered while retaining regression headroom.
+fail_under = 60
 omit = ["*/tests/*", "*/site-packages/*"]
 
 [tool.coverage.paths]

{nullcal-0.2.0 → nullcal-0.3.0}/src/nullcal/calibration.py

@@ -7,6 +7,12 @@ uniform log-frequency grid and also permits nonuniform knot placement.
 Importing this module enables JAX's process-wide ``jax_enable_x64`` setting.
 Calibration inference requires float64 precision, so this side effect is an
 intentional part of the module's public behavior.
+
+For a posterior point estimate, evaluate the complex calibration factor for
+each sample at each frequency, then take the componentwise median of its real
+and imaginary parts. Interpolating median knot values gives a different curve.
+The factor multiplies the antenna response; to correct observed strain, divide
+the strain by the estimated factor.
 """
 
 from __future__ import annotations
@@ -21,6 +27,7 @@ jax.config.update("jax_enable_x64", True)
 import jax.numpy as jnp  # noqa: E402
 
 MINIMUM_CUBIC_SPLINE_KNOTS = 4
+POSTERIOR_ARRAY_NDIM = 3
 
 
 def _float64(values):
@@ -109,6 +116,32 @@ def calibration_factor(frequencies, knot_frequencies, amplitude, phase):
     return (1.0 + delta_amplitude) * (2.0 + imaginary_phase) / (2.0 - imaginary_phase)
 
 
+def posterior_median_calibration_factor(frequencies, knot_frequencies, amplitudes, phases):
+    """Return the per-frequency posterior median of the complex factor.
+
+    ``amplitudes`` and ``phases`` have shape ``(sample, detector, knot)``;
+    knots may have shape ``(knot,)`` or ``(detector, knot)``. The complex median
+    is defined componentwise in real and imaginary coordinates. This evaluates
+    every posterior sample before taking the median, rather than evaluating a
+    curve from median knot values.
+    """
+    amplitudes = _float64(amplitudes)
+    phases = _float64(phases)
+    knots = _float64(knot_frequencies)
+    if amplitudes.ndim != POSTERIOR_ARRAY_NDIM or phases.shape != amplitudes.shape or amplitudes.shape[0] == 0:
+        raise ValueError("amplitudes and phases must have matching nonempty (sample, detector, knot) shapes")
+    if knots.ndim == 1:
+        knots = jnp.broadcast_to(knots, amplitudes.shape[1:])
+    if knots.shape != amplitudes.shape[1:]:
+        raise ValueError("knot_frequencies must match the detector-by-knot shape")
+
+    factors = jax.vmap(
+        jax.vmap(calibration_factor, in_axes=(None, 0, 0, 0)),
+        in_axes=(None, None, 0, 0),
+    )(frequencies, knots, amplitudes, phases)
+    return jnp.median(jnp.real(factors), axis=0) + 1j * jnp.median(jnp.imag(factors), axis=0)
+
+
 def calibration_log_prior(
     amplitude,
     phase,
@@ -185,5 +218,6 @@ __all__ = [
     "calibration_factor",
     "calibration_log_prior",
     "calibration_parameters_to_unconstrained",
+    "posterior_median_calibration_factor",
     "unconstrained_to_calibration_parameters",
 ]

{nullcal-0.2.0 → 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"]