nullcal 0.2.0__py3-none-any.whl

nullcal/null_stream/whiten.py

@@ -0,0 +1,106 @@
+"""
+Functions for whitening.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def compute_whitened_frequency_domain_strain(
+    frequency_domain_strain_array: np.ndarray,
+    power_spectral_density_array: np.ndarray,
+    delta_f: float,
+    frequency_mask: np.ndarray,
+) -> np.ndarray:
+    """Compute the whitened frequency domain strain.
+
+    Args:
+        frequency_domain_strain_array (np.ndarray): Frequency domain strain array.
+            Dimensions: (detector, frequency).
+        power_spectral_density_array (np.ndarray): Power spectral density array.
+            Dimensions: (detector, frequency).
+        delta_f (float): Frequency resolution in Hz.
+        frequency_mask (np.ndarray): A frequency mask.
+
+    Raises:
+        ValueError: delta_f must be positive.
+        ValueError: Shape mismatch. frequency_domain_strain_array and power_spectral_density_array
+            must have the same dimensions.
+        ValueError: Shape mismatch. frequency_domain_strain_array and frequency_mask
+            must have the same frequency dimension.
+
+    Returns:
+        np.ndarray: Whitened frequency domain strain.
+    """
+    if delta_f <= 0.0:
+        raise ValueError("delta_f must be positive.")
+    if frequency_domain_strain_array.shape != power_spectral_density_array.shape:
+        raise ValueError(
+            "Shape mismatch. frequency_domain_strain_array:"
+            f"{frequency_domain_strain_array.shape}"
+            f"and power_spectral_density_array: {power_spectral_density_array.shape}"
+            "must have the same dimensions."
+        )
+    if frequency_domain_strain_array.shape[1] != frequency_mask.shape[0]:
+        raise ValueError(
+            "Shape mismatch. frequency_domain_strain_array:"
+            f"{frequency_domain_strain_array.shape} and frequency_mask:"
+            f"{frequency_mask.shape} must have the same frequency dimension."
+        )
+    return np.divide(
+        frequency_domain_strain_array,
+        np.sqrt(power_spectral_density_array / (2 * delta_f)),
+        where=(frequency_mask & np.all(power_spectral_density_array, axis=0)),
+        out=np.zeros_like(frequency_domain_strain_array),
+    )
+
+
+def compute_whitened_antenna_response(
+    antenna_response_matrix: np.ndarray,
+    power_spectral_density_array: np.ndarray,
+    delta_f: float,
+    frequency_mask: np.ndarray,
+) -> np.ndarray:
+    """Compute the whitened antenna response function.
+
+    Args:
+        antenna_response_matrix (np.ndarray): Antenna response matrix.
+        power_spectral_density_array (np.ndarray): Power spectral density array.
+        delta_f (float): Frequency resolution in Hz.
+        frequency_mask (np.ndarray): Frequency mask.
+
+    Raises:
+        ValueError: delta_f must be positive.
+        ValueError: Shape mismatch. antenna_response_matrix and power_spectral_density_array
+            must have the same detector dimension.
+        ValueError: Shape mismatch. power_spectral_density_array and frequency_mask
+            must have the same frequency dimensions.
+
+    Returns:
+        np.ndarray: Whitened antenna response function.
+    """
+    if delta_f <= 0:
+        raise ValueError("delta_f must be positive.")
+    n_det_1, n_mode_1 = antenna_response_matrix.shape
+    n_det_2, n_freq_2 = power_spectral_density_array.shape
+    n_freq_3 = frequency_mask.shape[0]
+    if n_det_1 != n_det_2:
+        raise ValueError(
+            f"Shape mismatch. antenna_response_matrix: (detector={n_det_1}, mode={n_mode_1})"
+            f"and power_spectral_density_array: (detector={n_det_2}, frequency={n_freq_2})"
+            "must have the same detector dimension."
+        )
+    if n_freq_2 != n_freq_3:
+        raise ValueError(
+            f"Shape mismatch. power_spectral_density_array: (detector={n_det_2}, freq={n_freq_2})"
+            f"and frequency_mask: (frequency: {n_freq_3}) must have the same frequency dimension."
+        )
+    output = np.zeros((n_freq_2, n_det_1, n_mode_1), dtype=antenna_response_matrix.dtype)
+    frequency_mask = frequency_mask & np.all(power_spectral_density_array, axis=0)
+    output[frequency_mask, :, :] = np.einsum(
+        "dm,df->fdm",
+        antenna_response_matrix,
+        1 / np.sqrt(power_spectral_density_array[:, frequency_mask] / (2 * delta_f)),
+    )
+    return output

nullcal/result/__init__.py

@@ -0,0 +1,5 @@
+"""Result containers."""
+
+from .result import Result
+
+__all__ = ["Result"]

nullcal/result/result.py

@@ -0,0 +1,50 @@
+"""Sampler result types."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass(frozen=True)
+class Result:
+    """Posterior samples extracted from a sequence of sampler states."""
+
+    samples: Mapping[str, Any]
+    logdensity: Any
+    info: Any | None = None
+    metadata: Mapping[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if not self.samples:
+            raise ValueError("samples must contain at least one parameter array")
+        sample_counts = {value.shape[0] for value in self.samples.values() if value.ndim > 0}
+        if len(sample_counts) != 1 or any(value.ndim == 0 for value in self.samples.values()):
+            raise ValueError("all parameter arrays must have the same leading sample dimension")
+        if self.logdensity.ndim != 1 or self.logdensity.shape[0] != self.sample_count:
+            raise ValueError("logdensity must contain one value per posterior sample")
+
+    @property
+    def sample_count(self) -> int:
+        """Number of retained posterior samples."""
+        return next(iter(self.samples.values())).shape[0]
+
+    @classmethod
+    def from_states(
+        cls,
+        states: Any,
+        *,
+        info: Any | None = None,
+        metadata: Mapping[str, Any] | None = None,
+    ) -> Result:
+        """Extract positions and log densities from stacked sampler states."""
+        return cls(
+            samples=states.position,
+            logdensity=states.logdensity,
+            info=info,
+            metadata={} if metadata is None else metadata,
+        )
+
+
+__all__ = ["Result"]

nullcal/result/utils.py

@@ -0,0 +1,13 @@
+"""Small result-processing helpers."""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def spline_percentage_xform(delta_a: float | np.ndarray) -> float | np.ndarray:
+    """Convert fractional amplitude error to percent."""
+    return delta_a * 100
+
+
+__all__ = ["spline_percentage_xform"]

nullcal/sampler.py

@@ -0,0 +1,113 @@
+"""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 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_states = []
+    chain_infos = []
+
+    adaptation = blackjax.window_adaptation(
+        blackjax.nuts,
+        logdensity_fn,
+        target_acceptance_rate=target_acceptance_rate,
+    )
+
+    for chain_key in chain_keys:
+        jitter_key, warmup_key, sample_key = jax.random.split(chain_key, 3)
+        chain_initial_position = _jitter_position(initial_position, jitter_key, initial_position_jitter)
+        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, kernel=kernel):
+            next_state, info = kernel.step(key, state)
+            return next_state, (next_state, info)
+
+        _, (states, infos) = jax.lax.scan(step, adapted.state, sample_keys)
+        chain_states.append(states)
+        chain_infos.append(infos)
+
+    positions_by_chain = jax.tree.map(lambda *values: jnp.stack(values), *(state.position for state in chain_states))
+    logdensity_by_chain = jnp.stack([state.logdensity for state in chain_states])
+    divergence_by_chain = jnp.stack([info.is_divergent for info in chain_infos])
+
+    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))),
+    }
+    return Result(
+        samples=flattened_positions,
+        logdensity=flattened_logdensity,
+        info={"is_divergent": flattened_divergence},
+        metadata=metadata,
+    )
+
+
+__all__ = ["sample_nuts"]

nullcal/studies/__init__.py

@@ -0,0 +1 @@
+"""Numerical studies that inform nullcal model choices."""