statespacecheck 0.1.0__py3-none-any.whl

statespacecheck/__init__.py

@@ -0,0 +1,53 @@
+"""State space model goodness of fit diagnostics for neuroscience.
+
+This package provides tools to assess the consistency between state
+distributions and their component likelihood distributions in Bayesian
+state space models.
+"""
+
+from statespacecheck._validation import DistributionArray
+from statespacecheck.highest_density import DEFAULT_COVERAGE, highest_density_region
+from statespacecheck.periods import (
+    aggregate_over_period,
+    combine_flags,
+    find_low_overlap_intervals,
+    flag_extreme_kl,
+    flag_extreme_pvalues,
+    flag_low_overlap,
+)
+from statespacecheck.predictive_checks import (
+    log_predictive_density,
+    predictive_density,
+    predictive_pvalue,
+)
+from statespacecheck.state_consistency import (
+    hpd_overlap,
+    kl_divergence,
+)
+from statespacecheck.viz import plot_diagnostics
+
+__all__ = [
+    "highest_density_region",
+    "kl_divergence",
+    "hpd_overlap",
+    "predictive_density",
+    "log_predictive_density",
+    "aggregate_over_period",
+    "predictive_pvalue",
+    "DEFAULT_COVERAGE",
+    "DistributionArray",
+    "find_low_overlap_intervals",
+    "flag_extreme_kl",
+    "flag_extreme_pvalues",
+    "flag_low_overlap",
+    "combine_flags",
+    "plot_diagnostics",
+]
+
+try:
+    from ._version import __version__
+except ImportError:
+    # Fallback for development installs
+    from importlib.metadata import version
+
+    __version__ = version("statespacecheck")

statespacecheck/_validation.py

@@ -0,0 +1,189 @@
+"""Validation utilities for distributions and parameters."""
+
+import numpy as np
+from numpy.typing import NDArray
+
+# Type aliases for distribution arrays
+DistributionArray = NDArray[np.floating]
+
+
+def validate_coverage(coverage: float) -> None:
+    """Validate that coverage is in the valid range (0, 1).
+
+    Parameters
+    ----------
+    coverage : float
+        Coverage value to validate
+
+    Raises
+    ------
+    ValueError
+        If coverage is not in (0, 1)
+    """
+    if not (0.0 < coverage < 1.0):
+        raise ValueError(
+            f"coverage must be in (0, 1), got {coverage}. "
+            f"Coverage represents the probability mass of the highest density region "
+            f"and must be a value between 0 and 1 (exclusive). "
+            f"For example, use 0.95 for a 95% credible region."
+        )
+
+
+def validate_distribution(
+    distribution: DistributionArray,
+    name: str = "distribution",
+    min_ndim: int = 1,
+    allow_nan: bool = True,
+) -> DistributionArray:
+    """Validate and clean distribution array.
+
+    Parameters
+    ----------
+    distribution : np.ndarray
+        Distribution to validate
+    name : str
+        Name for error messages
+    min_ndim : int
+        Minimum number of dimensions required
+    allow_nan : bool
+        Whether to allow NaN values (converted to 0 if True)
+
+    Returns
+    -------
+    clean : np.ndarray
+        Original shape array with NaN/inf converted to 0 if allow_nan=True
+
+    Raises
+    ------
+    ValueError
+        If validation fails
+    """
+    arr = np.asarray(distribution, dtype=float)
+
+    if arr.ndim < min_ndim:
+        if min_ndim == 1:
+            expected_shape = "(n_time,)"
+        elif min_ndim == 2:
+            expected_shape = "(n_time, n_position)"
+        else:
+            expected_shape = f"{min_ndim}D"
+        raise ValueError(
+            f"{name} must be at least {min_ndim}D with shape {expected_shape}, "
+            f"got shape {arr.shape}. "
+            f"State space diagnostics require time-series data where the first "
+            f"dimension is time. "
+            f"For 1D spatial data use shape (n_time, n_position_bins), "
+            f"for 2D spatial data use shape (n_time, n_x_bins, n_y_bins). "
+            f"Did you forget to add the time dimension?"
+        )
+
+    # Handle non-finite values
+    clean: DistributionArray
+    if allow_nan:
+        # Use standard NumPy idiom: convert NaN/inf to 0
+        clean = np.nan_to_num(arr, nan=0.0, posinf=0.0, neginf=0.0)
+    else:
+        clean = arr.copy()
+        if not np.all(np.isfinite(clean)):
+            raise ValueError(
+                f"{name} contains non-finite values (NaN or inf). "
+                f"Probability distributions must have finite values. "
+                f"If you have invalid spatial bins (e.g., inaccessible locations), "
+                f"consider setting them to 0 instead of NaN, or ensure "
+                f"allow_nan=True in the validation."
+            )
+
+    # Check for negative values
+    finite_mask = np.isfinite(arr)
+    if np.any(clean[finite_mask] < 0):
+        raise ValueError(
+            f"{name} must be non-negative (probability or weight). "
+            f"Found negative values in the distribution. "
+            f"Probability distributions and weights must be >= 0. "
+            f"Check your data for errors or ensure proper normalization."
+        )
+
+    return clean
+
+
+def flatten_time_spatial(arr: DistributionArray) -> DistributionArray:
+    """Flatten array to (n_time, n_spatial) shape.
+
+    Parameters
+    ----------
+    arr : np.ndarray, shape (n_time, ...)
+        Array where ... represents arbitrary spatial dimensions.
+
+    Returns
+    -------
+    flat : np.ndarray, shape (n_time, n_spatial)
+        Flattened array.
+    """
+    n_time = arr.shape[0]
+    # Use numpy's automatic dimension calculation with -1
+    return arr.reshape(n_time, -1)
+
+
+def validate_paired_distributions(
+    dist1: DistributionArray,
+    dist2: DistributionArray,
+    name1: str = "state_dist",
+    name2: str = "likelihood",
+    min_ndim: int = 2,
+) -> tuple[DistributionArray, DistributionArray]:
+    """Validate two distributions have matching shapes.
+
+    Parameters
+    ----------
+    dist1 : np.ndarray
+        First distribution
+    dist2 : np.ndarray
+        Second distribution
+    name1 : str
+        Name for first distribution (error messages)
+    name2 : str
+        Name for second distribution (error messages)
+    min_ndim : int
+        Minimum number of dimensions required
+
+    Returns
+    -------
+    clean1 : np.ndarray
+        First distribution, cleaned
+    clean2 : np.ndarray
+        Second distribution, cleaned
+
+    Raises
+    ------
+    ValueError
+        If shapes don't match or validation fails
+    """
+    clean1 = validate_distribution(dist1, name1, min_ndim=min_ndim)
+    clean2 = validate_distribution(dist2, name2, min_ndim=min_ndim)
+
+    if clean1.shape != clean2.shape:
+        raise ValueError(
+            f"{name1} and {name2} must have same shape, got {clean1.shape} vs {clean2.shape}. "
+            f"Both distributions must cover the same time points and spatial bins. "
+            f"Common causes: different spatial discretization, mismatched time periods, "
+            f"or one distribution missing time/spatial dimensions. "
+            f"Ensure both arrays use consistent binning and time indexing."
+        )
+
+    return clean1, clean2
+
+
+def get_spatial_axes(arr: DistributionArray) -> tuple[int, ...]:
+    """Get tuple of spatial dimension axes (all except time axis 0).
+
+    Parameters
+    ----------
+    arr : np.ndarray, shape (n_time, ...)
+        Array where ... are spatial dimensions.
+
+    Returns
+    -------
+    spatial_axes : tuple[int, ...]
+        Tuple of axis indices for spatial dimensions.
+    """
+    return tuple(range(1, arr.ndim))

statespacecheck/_version.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = None

statespacecheck/highest_density.py

@@ -0,0 +1,138 @@
+"""Functions for computing highest density regions."""
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ._validation import (
+    DistributionArray,
+    flatten_time_spatial,
+    validate_coverage,
+    validate_distribution,
+)
+
+# Default coverage probability for highest density regions
+DEFAULT_COVERAGE = 0.95
+
+
+def highest_density_region(
+    distribution: DistributionArray, *, coverage: float = DEFAULT_COVERAGE
+) -> NDArray[np.bool_]:
+    """Compute boolean mask indicating highest density region membership.
+
+    Vectorized HPD mask for arrays shaped (n_time, *spatial). For each time t,
+    includes all bins with value >= threshold_t, where threshold_t is chosen so
+    cumulative mass >= coverage * total_t.
+
+    Parameters
+    ----------
+    distribution : np.ndarray, shape (n_time, ...)
+        Probability distributions over position at each time point where
+        ... represents arbitrary spatial dimensions.
+    coverage : float, optional
+        Desired coverage probability for the highest density region. Must be between 0 and 1.
+        Default is 0.95 for 95% coverage.
+
+    Returns
+    -------
+    isin_hd : np.ndarray, shape (n_time, ...)
+        Boolean mask indicating which positions are in the highest density region at each
+        time point, matching input shape.
+
+    Raises
+    ------
+    ValueError
+        If coverage is not in the range (0, 1).
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from statespacecheck import highest_density_region
+    >>> # Simple 1D example with peaked distribution
+    >>> distribution = np.array([[0.1, 0.6, 0.3], [0.2, 0.5, 0.3]])
+    >>> region = highest_density_region(distribution, coverage=0.9)
+    >>> region.shape
+    (2, 3)
+    >>> region.dtype
+    dtype('bool')
+
+    See Also
+    --------
+    hpd_overlap : Compute overlap between HPD regions of two distributions
+    kl_divergence : Measure information divergence between distributions
+
+    Notes
+    -----
+    - NaNs are ignored (treated as 0 mass).
+    - If total mass at time t <= 0 or not finite, returns all-False for that t.
+    - Works in unnormalized space to avoid numerical issues.
+    - Fully vectorized with no Python loops for efficiency.
+    - Uses `>=` threshold: all bins with value equal to cutoff are included.
+    - Due to ties, actual coverage may slightly exceed requested coverage.
+    - This ensures consistent behavior across equivalent distributions.
+
+    References
+    ----------
+    .. [1] https://stats.stackexchange.com/questions/240749/how-to-find-95-credible-interval
+
+    """
+    validate_coverage(coverage)
+
+    # Use centralized validation: handles NaN/inf → 0, checks non-negativity, validates dimensions
+    clean = validate_distribution(
+        distribution,
+        name="distribution",
+        min_ndim=2,  # Require at least (n_time, n_spatial)
+        allow_nan=True,
+    )
+
+    # Flatten to (n_time, n_spatial) for vectorized operations
+    flat = flatten_time_spatial(clean)
+
+    n_time = clean.shape[0]
+    n_spatial = flat.shape[1]
+
+    # Compute total mass and target mass for each time point
+    # Shape: (n_time,)
+    totals = flat.sum(axis=1)
+    target = coverage * totals
+
+    # Identify rows with no mass -> empty HPD (all False)
+    empty = ~np.isfinite(totals) | (totals <= 0)
+
+    # Sort each row descending (vectorized)
+    # Shape: (n_time, n_spatial)
+    flat_sorted = np.sort(flat, axis=1)[:, ::-1]
+
+    # Row-wise cumulative sums
+    # Shape: (n_time, n_spatial)
+    csum = np.cumsum(flat_sorted, axis=1)
+
+    # Find the first index where cumulative >= target (per row)
+    # Shape: (n_time, n_spatial) boolean
+    ge = csum >= target[:, None]
+
+    # Check if each row has at least one True value
+    # Shape: (n_time,)
+    has_true = ge.any(axis=1)
+
+    # argmax gives first True index; if none True, returns 0 (we fix below)
+    # Shape: (n_time,)
+    idx = ge.argmax(axis=1)
+
+    # If a row never reaches target but has positive mass (rare numeric case),
+    # choose the last index. If it's truly empty, handle later.
+    idx = np.where(has_true, idx, n_spatial - 1)
+
+    # Per-row cutoff (unnormalized)
+    # Shape: (n_time,)
+    cutoff = np.take_along_axis(flat_sorted, idx[:, None], axis=1).squeeze(1)
+
+    # Empty rows -> set cutoff to +inf so mask is all False
+    cutoff = np.where(empty, np.inf, cutoff)
+
+    # Broadcast cutoff back to spatial shape and build mask
+    # Use the **clean** array for the comparison to keep behavior consistent
+    # Broadcasting: reshape cutoff from (n_time,) to (n_time, 1, 1, ...) to match spatial dims
+    # Using tuple unpacking for clarity
+    broadcast_shape = (n_time,) + (1,) * (clean.ndim - 1)
+    return clean >= cutoff.reshape(broadcast_shape)