ergodic 0.1.0__py3-none-any.whl

ergodic/__init__.py

@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Ergodic: Causal AI.
+
+Tools for causal discovery, causal inference, and process intelligence. The
+shared objects live at the top level: the mixed-graph family (:class:`MixedGraph`
+with the :func:`dag`, :func:`admg`, :func:`cpdag`, :func:`pag` builders) and
+:class:`DomainKnowledge` for prior constraints. Identification turns a graph into
+an estimand with :func:`identify_effect`.
+
+Homepage: https://ergodic.ai
+"""
+
+from __future__ import annotations
+
+from . import datasets
+from .data import (
+    Dataset,
+    EventLog,
+    HierarchicalTSData,
+    Hierarchy,
+    Ordinal,
+    PanelData,
+    Schema,
+    SchemaInferenceWarning,
+    TabularData,
+    TimeSeriesData,
+    VariableType,
+    event_log,
+    hts,
+    panel,
+    tabular,
+    timeseries,
+)
+from .discovery import discover, granger
+
+# `forecast` names both the pillar and its facade, and the function is what the
+# top level exports: `from ergodic import forecast` gives the callable, while
+# `from ergodic.forecast import backtest` still reaches the module. Import the
+# function last so the binding is the deliberate one.
+from .forecast import backtest, forecast
+from .graph import (
+    Edge,
+    GraphKind,
+    InvalidGraphError,
+    Mark,
+    MixedGraph,
+    admg,
+    cpdag,
+    d_separated,
+    dag,
+    m_separated,
+    mixed,
+    pag,
+)
+from .identification import (
+    IDA,
+    Functional,
+    Hedge,
+    Identification,
+    IDResult,
+    Marginal,
+    MediationID,
+    Prob,
+    Product,
+    Ratio,
+    adjustment_set,
+    frontdoor_set,
+    ida,
+    identify_distribution,
+    identify_effect,
+    identify_mediation,
+    instruments,
+    is_frontdoor_set,
+    is_instrument,
+    is_valid_adjustment_set,
+    minimal_adjustment_set,
+)
+from .inference import (
+    CATEResult,
+    CausalForest,
+    EffectEstimate,
+    EffectTarget,
+    ITEResult,
+    MediationEstimate,
+    PerUnitEffect,
+    PolicyTree,
+    Refutation,
+    SelectionReport,
+    VariableImportance,
+    cate,
+    did,
+    estimate,
+    estimate_effect,
+    estimate_mediation,
+    event_study,
+    interrupted_time_series,
+    policy_tree,
+    refute_effect,
+    select_cate,
+    synthetic_control,
+)
+from .knowledge import DomainKnowledge, InconsistentKnowledgeError
+from .process import (
+    case_table,
+    conformance,
+    decision_points,
+    decision_table,
+    discover_model,
+    kpi_panel,
+    process_map,
+)
+from .scm import (
+    SCM,
+    BinaryGumbelMax,
+    CategoricalGumbelMax,
+    GaussianAdditive,
+    SCMResult,
+    Sensitivity,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "IDA",
+    "SCM",
+    "BinaryGumbelMax",
+    "CATEResult",
+    "CategoricalGumbelMax",
+    "CausalForest",
+    "Dataset",
+    "DomainKnowledge",
+    "Edge",
+    "EffectEstimate",
+    "EffectTarget",
+    "EventLog",
+    "Functional",
+    "GaussianAdditive",
+    "GraphKind",
+    "Hedge",
+    "HierarchicalTSData",
+    "Hierarchy",
+    "IDResult",
+    "ITEResult",
+    "Identification",
+    "InconsistentKnowledgeError",
+    "InvalidGraphError",
+    "Marginal",
+    "Mark",
+    "MediationEstimate",
+    "MediationID",
+    "MixedGraph",
+    "Ordinal",
+    "PanelData",
+    "PerUnitEffect",
+    "PolicyTree",
+    "Prob",
+    "Product",
+    "Ratio",
+    "Refutation",
+    "SCMResult",
+    "Schema",
+    "SchemaInferenceWarning",
+    "SelectionReport",
+    "Sensitivity",
+    "TabularData",
+    "TimeSeriesData",
+    "VariableImportance",
+    "VariableType",
+    "__version__",
+    "adjustment_set",
+    "admg",
+    "backtest",
+    "case_table",
+    "cate",
+    "conformance",
+    "cpdag",
+    "d_separated",
+    "dag",
+    "datasets",
+    "decision_points",
+    "decision_table",
+    "did",
+    "discover",
+    "discover_model",
+    "estimate",
+    "estimate_effect",
+    "estimate_mediation",
+    "event_log",
+    "event_study",
+    "forecast",
+    "frontdoor_set",
+    "granger",
+    "hts",
+    "ida",
+    "identify_distribution",
+    "identify_effect",
+    "identify_mediation",
+    "instruments",
+    "interrupted_time_series",
+    "is_frontdoor_set",
+    "is_instrument",
+    "is_valid_adjustment_set",
+    "kpi_panel",
+    "m_separated",
+    "minimal_adjustment_set",
+    "mixed",
+    "pag",
+    "panel",
+    "policy_tree",
+    "process_map",
+    "refute_effect",
+    "select_cate",
+    "synthetic_control",
+    "tabular",
+    "timeseries",
+]

ergodic/data/__init__.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Typed datasets: the data objects discovery and inference read.
+
+Each shape embeds a column :class:`Schema` (continuous, categorical, ordinal,
+datetime, or excluded) and shares the accessors on :class:`Dataset`. The shapes
+are tabular, time-series, panel, and hierarchical time-series; the event log
+follows.
+
+Build a dataset with :func:`tabular`, :func:`timeseries`, :func:`panel`, or
+:func:`hts`, each of which takes a pandas or polars dataframe, a dict of columns,
+or a numpy array.
+"""
+
+from __future__ import annotations
+
+from ._dataset import Dataset
+from ._eventlog import EventLog, event_log
+from ._hierarchy import Hierarchy
+from ._hts import HierarchicalTSData, hts
+from ._panel import PanelData, panel
+from ._schema import (
+    ColumnType,
+    Ordinal,
+    Schema,
+    SchemaInferenceWarning,
+    VariableType,
+    infer_schema,
+)
+from ._tabular import TabularData, tabular
+from ._timeseries import TimeSeriesData, timeseries
+
+__all__ = [
+    "ColumnType",
+    "Dataset",
+    "EventLog",
+    "HierarchicalTSData",
+    "Hierarchy",
+    "Ordinal",
+    "PanelData",
+    "Schema",
+    "SchemaInferenceWarning",
+    "TabularData",
+    "TimeSeriesData",
+    "VariableType",
+    "event_log",
+    "hts",
+    "infer_schema",
+    "panel",
+    "tabular",
+    "timeseries",
+]

ergodic/data/_dataset.py

@@ -0,0 +1,223 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The dataset base: a typed, immutable column store the algorithms read.
+
+``Dataset`` holds the normalized columns, the :class:`Schema`, and the accessors
+that discovery and inference need: a float matrix, integer codes for categorical
+columns, missing-value reporting, and a name check against a graph. It is
+immutable, so retyping a column, selecting columns, or dropping rows returns a
+new dataset. The concrete shapes (tabular, time-series, panel, and the rest)
+extend it with their index columns.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field, replace
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from ._schema import (
+    Ordinal,
+    Schema,
+    VariableType,
+    _encode,
+    _missing_mask,
+    _resolve_type,
+    _unique_nonnull,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Mapping, Sequence
+    from typing import Self
+
+    import numpy.typing as npt
+    import pandas as pd
+
+    from ..graph import MixedGraph
+
+
+@dataclass(frozen=True, eq=False)
+class Dataset:
+    """A typed column store. Built through the shape factories, not directly."""
+
+    schema: Schema
+    n_rows: int
+    _store: Mapping[str, npt.NDArray[Any]] = field(repr=False)
+
+    # -- shape ---------------------------------------------------------------
+
+    @property
+    def names(self) -> tuple[str, ...]:
+        return self.schema.names
+
+    @property
+    def n_columns(self) -> int:
+        return len(self.schema)
+
+    @property
+    def shape(self) -> tuple[int, int]:
+        return (self.n_rows, self.n_columns)
+
+    def __len__(self) -> int:
+        return self.n_rows
+
+    @property
+    def index_columns(self) -> tuple[str, ...]:
+        """Columns that index the rows (time, entity, case), not modeled variables."""
+        return ()
+
+    @property
+    def variables(self) -> tuple[str, ...]:
+        """Column names the algorithms see: not the excluded ones, not the index."""
+        index = set(self.index_columns)
+        return tuple(name for name in self.schema.variables if name not in index)
+
+    # -- column access -------------------------------------------------------
+
+    def column(self, name: str) -> npt.NDArray[Any]:
+        """A read-only copy of one column's raw values."""
+        if name not in self.schema:
+            raise KeyError(name)
+        arr = np.array(self._store[name], copy=True)
+        arr.setflags(write=False)
+        return arr
+
+    def codes(self, name: str) -> npt.NDArray[np.int64]:
+        """Integer category codes for a categorical or ordinal column, -1 for missing.
+
+        A non-missing value outside the column's categories raises a ``ValueError``.
+        """
+        col = self.schema[name]
+        if not col.is_discrete:
+            raise ValueError(f"column {name!r} is {col.kind.value}, not categorical or ordinal")
+        return _encode(self._store[name], col.categories, col.ordered, name)
+
+    def categories(self, name: str) -> tuple[Any, ...] | None:
+        return self.schema[name].categories
+
+    def matrix(self, names: Sequence[str] | None = None) -> npt.NDArray[np.float64]:
+        """A float matrix of the named columns. Discrete columns become their codes,
+        and a missing value (an NA category or NaT) becomes NaN."""
+        chosen = list(names) if names is not None else list(self.variables)
+        if not chosen:
+            return np.empty((self.n_rows, 0), dtype=np.float64)
+        cols: list[npt.NDArray[Any]] = []
+        for name in chosen:
+            col = self.schema[name]
+            if col.is_continuous:
+                cols.append(np.asarray(self._store[name], dtype=np.float64))
+            elif col.is_discrete:
+                encoded = self.codes(name).astype(np.float64)
+                encoded[encoded == -1.0] = np.nan
+                cols.append(encoded)
+            elif col.is_datetime:
+                stamps = np.asarray(self._store[name], dtype="datetime64[ns]")
+                floats = stamps.view(np.int64).astype(np.float64)
+                floats[np.isnat(stamps)] = np.nan
+                cols.append(floats)
+            else:
+                raise ValueError(f"column {name!r} is excluded; it has no numeric form")
+        return np.asarray(np.column_stack(cols), dtype=np.float64)
+
+    def to_numpy(self) -> npt.NDArray[np.float64]:
+        return self.matrix()
+
+    def to_pandas(self) -> pd.DataFrame:
+        import pandas as pd
+
+        return pd.DataFrame({name: self._store[name] for name in self.schema.names})
+
+    # -- missing data --------------------------------------------------------
+
+    def missing(self) -> dict[str, int]:
+        """Missing-value count per column."""
+        return {name: int(_missing_mask(self._store[name]).sum()) for name in self.schema.names}
+
+    @property
+    def has_missing(self) -> bool:
+        return any(count > 0 for count in self.missing().values())
+
+    def dropna(self) -> Self:
+        """Drop every row with a missing value in any column."""
+        keep = np.ones(self.n_rows, dtype=bool)
+        for arr in self._store.values():
+            keep &= ~_missing_mask(arr)
+        store = {name: arr[keep] for name, arr in self._store.items()}
+        return replace(self, _store=store, n_rows=int(keep.sum()))
+
+    # -- immutable edits -----------------------------------------------------
+
+    def take(self, rows: Sequence[int] | npt.NDArray[np.int64]) -> Self:
+        """A new dataset with the rows at the given positions, in the given order.
+
+        Positions may repeat (a bootstrap resample) and may be negative
+        (Python indexing). Row identity is all this knows about; shapes with
+        row semantics (time series, panel) keep their index columns but it is
+        the caller's job to resample in a way their structure survives.
+        """
+        idx = np.asarray(rows, dtype=np.int64)
+        if idx.ndim != 1:
+            raise ValueError(f"rows must be one-dimensional, got shape {idx.shape}")
+        if idx.size and (int(idx.min()) < -self.n_rows or int(idx.max()) >= self.n_rows):
+            raise IndexError(f"row positions out of range for {self.n_rows} rows")
+        store = {name: arr[idx] for name, arr in self._store.items()}
+        return replace(self, _store=store, n_rows=int(idx.size))
+
+    def select(self, names: Iterable[str]) -> Self:
+        """A new dataset with the named columns, plus any index columns, in order."""
+        chosen = [str(n) for n in names]
+        absent = [n for n in chosen if n not in self.schema]
+        if absent:
+            raise KeyError(f"no such columns: {absent}")
+        for index_name in self.index_columns:
+            if index_name not in chosen:
+                chosen.append(index_name)
+        schema = Schema(tuple(self.schema[n] for n in chosen))
+        store = {n: self._store[n] for n in chosen}
+        return replace(self, schema=schema, _store=store)
+
+    def with_type(
+        self,
+        name: str,
+        kind: VariableType | str | Ordinal,
+        categories: Sequence[Any] | None = None,
+    ) -> Self:
+        """A new dataset with ``name`` retyped, categories read from data when needed."""
+        if name not in self.schema:
+            raise KeyError(name)
+        new = _resolve_type(name, _unique_nonnull(self._store[name]), kind, categories)
+        if new.is_discrete and (categories is not None or isinstance(kind, Ordinal)):
+            # encode once now, so declared categories that miss a value fail here
+            _encode(self._store[name], new.categories, new.ordered, name)
+        cols = tuple(new if c.name == name else c for c in self.schema.columns)
+        return replace(self, schema=Schema(cols))
+
+    def exclude(self, *names: str) -> Self:
+        """A new dataset with the named columns marked excluded."""
+        out = self
+        for name in names:
+            out = out.with_type(name, VariableType.EXCLUDED)
+        return out
+
+    # -- graph alignment -----------------------------------------------------
+
+    def alignment(self, graph: MixedGraph) -> list[str]:
+        """Where the variables and the graph's nodes disagree, empty when they match."""
+        nodes = set(graph.nodes)
+        cols = set(self.variables)
+        out = [f"graph node {n!r} has no column" for n in sorted(nodes - cols)]
+        out += [f"column {c!r} is not a graph node" for c in sorted(cols - nodes)]
+        return out
+
+    def aligns_with(self, graph: MixedGraph) -> bool:
+        """True when the variables are exactly the graph's nodes."""
+        return not self.alignment(graph)
+
+    def __repr__(self) -> str:
+        parts = [
+            f"{len(names)} {kind.value}"
+            for kind in VariableType
+            if (names := self.schema.of_kind(kind))
+        ]
+        body = ", ".join(parts)
+        return f"{type(self).__name__}(rows={self.n_rows}, cols={self.n_columns}: {body})"

ergodic/data/_eventlog.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Event log: variable-length traces of discrete events, keyed by a case.
+
+Build it with :func:`event_log`, naming the case, activity, and timestamp
+columns. Rows are events, grouped by case and ordered in time. The log reads as
+the process-mining primitives the discovery algorithms need: the trace per case,
+the distinct variants, and the directly-follows counts.
+
+The projections to the grid shapes (case features to tabular, event rates to a
+time-series) land with the process pillar. This is the container and its core
+relations.
+"""
+
+from __future__ import annotations
+
+from collections import Counter
+from dataclasses import dataclass
+from itertools import pairwise
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from ._dataset import Dataset
+from ._frame import to_columns
+from ._schema import infer_schema
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Mapping, Sequence
+
+    from ._schema import Ordinal, VariableType
+
+
+@dataclass(frozen=True, eq=False)
+class EventLog(Dataset):
+    """Events keyed by case, activity, and timestamp; one row per event."""
+
+    case: str
+    activity: str
+    timestamp: str
+
+    @property
+    def index_columns(self) -> tuple[str, ...]:
+        return (self.case, self.activity, self.timestamp)
+
+    @property
+    def cases(self) -> tuple[Any, ...]:
+        return tuple(sorted(set(self._store[self.case].tolist())))
+
+    @property
+    def n_cases(self) -> int:
+        return len(set(self._store[self.case].tolist()))
+
+    @property
+    def activities(self) -> tuple[str, ...]:
+        return tuple(sorted(set(self._store[self.activity].tolist())))
+
+    @property
+    def n_activities(self) -> int:
+        return len(set(self._store[self.activity].tolist()))
+
+    def traces(self) -> dict[Any, tuple[str, ...]]:
+        """The ordered activity sequence per case."""
+        cases = self._store[self.case].tolist()
+        acts = self._store[self.activity].tolist()
+        out: dict[Any, list[str]] = {}
+        for case_id, act in zip(cases, acts, strict=True):
+            out.setdefault(case_id, []).append(act)
+        return {case_id: tuple(seq) for case_id, seq in out.items()}
+
+    def variants(self) -> Counter[tuple[str, ...]]:
+        """How many cases follow each distinct activity sequence."""
+        return Counter(self.traces().values())
+
+    def directly_follows(self) -> Counter[tuple[str, str]]:
+        """Counts of each ``(a, b)`` where ``b`` immediately follows ``a`` in a case."""
+        out: Counter[tuple[str, str]] = Counter()
+        for trace in self.traces().values():
+            for first, second in pairwise(trace):
+                out[(first, second)] += 1
+        return out
+
+    def __repr__(self) -> str:
+        return (
+            f"EventLog(cases={self.n_cases}, events={self.n_rows}, activities={self.n_activities})"
+        )
+
+
+def event_log(
+    data: Any,
+    *,
+    case: str,
+    activity: str,
+    timestamp: str,
+    columns: Sequence[str] | None = None,
+    types: Mapping[str, VariableType | str | Ordinal] | None = None,
+    exclude: Iterable[str] | None = None,
+    categorical_threshold: int = 10,
+) -> EventLog:
+    """Build an :class:`EventLog`, grouped by case and ordered by timestamp."""
+    store, n_rows, declared = to_columns(data, columns)
+    for role, name in (("case", case), ("activity", activity), ("timestamp", timestamp)):
+        if name not in store:
+            raise KeyError(f"{role} column {name!r} not found in {sorted(store)}")
+    overrides: dict[str, VariableType | str | Ordinal] = dict(types or {})
+    overrides.setdefault(case, "categorical")
+    overrides.setdefault(activity, "categorical")
+    if timestamp not in overrides:
+        overrides[timestamp] = "datetime" if store[timestamp].dtype.kind == "M" else "continuous"
+    schema = infer_schema(
+        store,
+        n_rows=n_rows,
+        declared=declared,
+        overrides=overrides,
+        exclude=exclude,
+        categorical_threshold=categorical_threshold,
+    )
+    order = np.lexsort((store[timestamp], store[case]))
+    ordered = {name: arr[order] for name, arr in store.items()}
+    return EventLog(
+        schema=schema,
+        n_rows=n_rows,
+        _store=ordered,
+        case=case,
+        activity=activity,
+        timestamp=timestamp,
+    )