antecedent.data

  1"""Data constructors and conversion probes (Arrow / float64 columns)."""
  2
  3from __future__ import annotations
  4
  5from dataclasses import dataclass
  6from typing import Any, Mapping, Sequence
  7
  8import numpy as np
  9from numpy.typing import NDArray
 10
 11from ._data import as_columns, as_multi_env_columns, to_f64
 12from ._native import ArrowLoadInfo, load_float64_arrow_c_columns, load_float64_columns
 13
 14
 15@dataclass(frozen=True)
 16class EventFrame:
 17    """Irregular event marks + timestamps; aligned via ``align_interval_ns`` before temporal algos."""
 18
 19    names: list[str]
 20    columns: list[NDArray[np.float64]]
 21    event_times_ns: NDArray[np.int64]
 22    align_interval_ns: int
 23
 24
 25@dataclass(frozen=True)
 26class PanelFrame:
 27    """Multi-unit time series sharing one schema.
 28
 29    Discovery: ``JPCMCIPlus`` (multi-env context) or pooled-units ``PCMCI`` /
 30    ``PCMCIPlus`` / ``LPCMCI``. Estimation uses stacked cluster-HAC SE
 31    (frequentist) or ``BayesianTemporalGcomp`` when ``inference=Bayesian``.
 32    """
 33
 34    names: list[str]
 35    unit_columns: list[list[NDArray[np.float64]]]
 36    unit_ids: list[int]
 37
 38
 39@dataclass(frozen=True)
 40class MultiEnvFrame:
 41    """Multi-environment series (J-PCMCI+ discovery)."""
 42
 43    names: list[str]
 44    env_columns: list[list[NDArray[np.float64]]]
 45
 46
 47def event(
 48    data: Mapping[str, Any] | Any,
 49    event_times_ns: Sequence[int] | NDArray[np.int64],
 50    *,
 51    align_interval_ns: int,
 52) -> EventFrame:
 53    """Build an [`EventFrame`] for ``analyze`` (duration-bin align → temporal path)."""
 54    if align_interval_ns <= 0:
 55        raise ValueError("align_interval_ns must be > 0")
 56    names, columns = as_columns(data)
 57    times = np.asarray(event_times_ns, dtype=np.int64)
 58    if times.ndim != 1:
 59        raise ValueError(f"event_times_ns must be 1-d, got shape {times.shape}")
 60    if len(times) != len(columns[0]):
 61        raise ValueError(
 62            f"event_times_ns length {len(times)} != column length {len(columns[0])}"
 63        )
 64    return EventFrame(
 65        names=names,
 66        columns=columns,
 67        event_times_ns=times,
 68        align_interval_ns=int(align_interval_ns),
 69    )
 70
 71
 72def panel(
 73    units: Sequence[Mapping[str, Any] | Any] | Mapping[Any, Mapping[str, Any] | Any],
 74) -> PanelFrame:
 75    """Build a [`PanelFrame`] from a sequence of unit frames or ``{unit_id: frame}``."""
 76    if isinstance(units, Mapping):
 77        ids = [int(k) for k in units.keys()]
 78        frames = list(units.values())
 79    else:
 80        frames = list(units)
 81        ids = list(range(len(frames)))
 82    if not frames:
 83        raise ValueError("panel needs ≥1 unit")
 84    names, first = as_columns(frames[0])
 85    unit_columns = [first]
 86    for i, frame in enumerate(frames[1:], start=1):
 87        n, cols = as_columns(frame)
 88        if n != names:
 89            raise ValueError(
 90                f"unit {i} column names {n!r} do not match unit 0 {names!r}"
 91            )
 92        unit_columns.append(cols)
 93    return PanelFrame(names=names, unit_columns=unit_columns, unit_ids=ids)
 94
 95
 96def multi_env(
 97    envs: Sequence[Mapping[str, Any] | Any],
 98) -> MultiEnvFrame:
 99    """Build a [`MultiEnvFrame`] (sequence of environment frames for J-PCMCI+)."""
100    names, env_columns = as_multi_env_columns(envs)
101    return MultiEnvFrame(names=names, env_columns=env_columns)
102
103
104__all__ = [
105    "ArrowLoadInfo",
106    "EventFrame",
107    "MultiEnvFrame",
108    "PanelFrame",
109    "event",
110    "load_float64_arrow_c_columns",
111    "load_float64_columns",
112    "multi_env",
113    "panel",
114    "to_f64",
115]