segmtree 0.3.0__py3-none-any.whl

segmtree/__init__.py

@@ -0,0 +1,30 @@
+"""segmtree: interpretable unsupervised decision trees for segmentation.
+
+Partitions rows of a feature matrix into segments described by short,
+human-readable rule chains, without any target variable. See
+:class:`SegTree` for the estimator and the README for a full walkthrough.
+"""
+
+from __future__ import annotations
+
+from .rules import Condition, Rule, extract_rules, profile_leaves, replay_check, replay_labels
+from .stability import StabilityResult, adjusted_rand_score, bootstrap_stability
+from .tree import Cut, Leaf, SegTree
+
+__version__ = "0.3.0"
+
+__all__ = [
+    "Condition",
+    "Cut",
+    "Leaf",
+    "Rule",
+    "SegTree",
+    "StabilityResult",
+    "__version__",
+    "adjusted_rand_score",
+    "bootstrap_stability",
+    "extract_rules",
+    "profile_leaves",
+    "replay_check",
+    "replay_labels",
+]

segmtree/_base.py

@@ -0,0 +1,94 @@
+"""Shared infrastructure: sklearn-compatible base class and input validation.
+
+The package's only runtime dependency is numpy. When scikit-learn is
+installed we reuse its ``BaseEstimator`` for full ecosystem compatibility
+(``clone``, pipelines, hyperparameter search); otherwise a minimal,
+API-compatible fallback is used.
+"""
+
+from __future__ import annotations
+
+import inspect
+
+import numpy as np
+
+__all__ = ["BaseEstimator", "as_2d_array"]
+
+try:  # pragma: no cover - exercised implicitly depending on environment
+    from sklearn.base import BaseEstimator
+
+    _HAS_SKLEARN = True
+except ImportError:  # pragma: no cover
+
+    class BaseEstimator:
+        """Minimal stand-in for :class:`sklearn.base.BaseEstimator`."""
+
+        def get_params(self, deep=True):
+            params = {}
+            for name, p in inspect.signature(type(self).__init__).parameters.items():
+                if name == "self" or p.kind in (
+                    inspect.Parameter.VAR_POSITIONAL,
+                    inspect.Parameter.VAR_KEYWORD,
+                ):
+                    continue
+                params[name] = getattr(self, name)
+            return params
+
+        def set_params(self, **params):
+            valid = self.get_params(deep=False)
+            for key, value in params.items():
+                if key not in valid:
+                    raise ValueError(
+                        f"Invalid parameter {key!r} for {type(self).__name__}. "
+                        f"Valid parameters are: {sorted(valid)}."
+                    )
+                setattr(self, key, value)
+            return self
+
+        def __repr__(self):
+            parts = []
+            for name in self.get_params(deep=False):
+                try:
+                    default = inspect.signature(type(self).__init__).parameters[name].default
+                    current = getattr(self, name)
+                    equal = (
+                        bool(np.all(current == default))
+                        if isinstance(current, np.ndarray)
+                        else (current == default)
+                    )
+                except Exception:
+                    equal = True
+                if not equal:
+                    parts.append(f"{name}={getattr(self, name)!r}")
+            return f"{type(self).__name__}({', '.join(parts)})"
+
+
+def as_2d_array(X, *, context="X"):
+    """Coerce tabular input to a 2-D finite numeric ndarray.
+
+    Accepts anything ``np.asarray`` understands plus objects exposing
+    ``to_numpy`` (e.g. pandas DataFrames). Returns ``(array, column_names_or_None)``.
+    """
+    columns = None
+    if hasattr(X, "to_numpy") and hasattr(X, "columns"):
+        columns = [str(c) for c in X.columns]
+        X = X.to_numpy()
+    arr = np.asarray(X)
+    if arr.dtype == object:
+        try:
+            arr = arr.astype(np.float64)
+        except (TypeError, ValueError) as e:
+            raise TypeError(f"{context} must be numeric; got dtype=object ({e})") from e
+    if not np.issubdtype(arr.dtype, np.number):
+        raise TypeError(f"{context} must have a numeric dtype, got {arr.dtype}.")
+    if arr.ndim != 2:
+        raise ValueError(f"{context} must be 2-D (n_samples, n_features), got shape {arr.shape}.")
+    if arr.shape[0] == 0 or arr.shape[1] == 0:
+        raise ValueError(f"{context} must not be empty, got shape {arr.shape}.")
+    if not np.isfinite(arr).all():
+        n_bad = int((~np.isfinite(arr)).sum())
+        raise ValueError(
+            f"{context} contains {n_bad} NaN/infinite value(s). "
+            "Impute or drop them before fitting (segmtree does no imputation)."
+        )
+    return arr, columns

segmtree/rules.py

@@ -0,0 +1,215 @@
+"""Rule extraction, replay validation and segment profiling."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from ._base import as_2d_array
+
+__all__ = [
+    "Condition",
+    "Rule",
+    "extract_rules",
+    "profile_leaves",
+    "replay_check",
+    "replay_labels",
+    "rule_mask",
+]
+
+
+def _fmt(value):
+    if float(value).is_integer() and abs(value) < 2**53:
+        return str(int(value))
+    return f"{value:g}"
+
+
+def _fmt_codes(categories):
+    return "{" + ", ".join(_fmt(float(c)) for c in sorted(categories)) + "}"
+
+
+@dataclass(frozen=True)
+class Condition:
+    """One axis-aligned test.
+
+    ``op`` is ``"<="``/``">"`` with a float ``threshold``, or ``"in"`` /
+    ``"not in"`` with integer category codes in ``categories``.
+    """
+
+    feature: str
+    op: str
+    threshold: float = float("nan")
+    categories: tuple = field(default=())
+
+    def __post_init__(self):
+        if self.op not in ("<=", ">", "in", "not in"):
+            raise ValueError(f"Unknown condition op: {self.op!r}.")
+
+    def __str__(self):
+        if self.op in ("in", "not in"):
+            return f"{self.feature} {self.op} {_fmt_codes(self.categories)}"
+        return f"{self.feature} {self.op} {_fmt(self.threshold)}"
+
+
+@dataclass(frozen=True)
+class Rule:
+    """The full rule chain defining one segment (leaf)."""
+
+    leaf_id: int
+    conditions: tuple
+    n_samples: int
+    fraction: float
+    depth: int
+
+    @property
+    def text(self):
+        return " AND ".join(str(c) for c in self.conditions) if self.conditions else "TRUE"
+
+    def __str__(self):
+        return f"#{self.leaf_id}: {self.text}"
+
+
+def extract_rules(tree):
+    """Return the :class:`Rule` for every fitted segment, in leaf-id order."""
+    if not hasattr(tree, "labels_"):
+        raise RuntimeError("This SegTree instance is not fitted yet; call fit() first.")
+    rules = []
+    for lf in tree.leaves_:
+        conds = []
+        for name, op, payload in lf.conditions:
+            if op in ("in", "not in"):
+                conds.append(Condition(name, op, categories=tuple(sorted(payload))))
+            else:
+                conds.append(Condition(name, op, threshold=float(payload)))
+        rules.append(
+            Rule(
+                leaf_id=lf.id,
+                conditions=tuple(conds),
+                n_samples=lf.n_samples,
+                fraction=lf.fraction,
+                depth=lf.depth,
+            )
+        )
+    return rules
+
+
+def rule_mask(rule, X, feature_names):
+    """Boolean mask of rows satisfying ``rule``; evaluated independently of the tree."""
+    mask = np.ones(X.shape[0], dtype=bool)
+    index = {nm: i for i, nm in enumerate(feature_names)}
+    for cond in rule.conditions:
+        j = index[cond.feature]
+        col = X[:, j]
+        if cond.op == "<=":
+            mask &= col <= cond.threshold
+        elif cond.op == ">":
+            mask &= col > cond.threshold
+        elif cond.op == "in":
+            mask &= np.isin(col, cond.categories)
+        else:  # "not in"
+            mask &= ~np.isin(col, cond.categories)
+    return mask
+
+
+def replay_labels(tree, X):
+    """Assign labels by replaying every rule chain on ``X``.
+
+    Independent of :meth:`SegTree.predict` (no tree traversal), so agreement
+    between the two validates that stored rules reproduce the partition.
+    Rows matching no rule receive -1. Honors ``tree.handle_unknown``: with
+    ``"error"`` unseen category codes raise ``ValueError``.
+    """
+    Xarr, _ = as_2d_array(X, context="X")
+    if (
+        getattr(tree, "handle_unknown", "complement") == "error"
+        and getattr(tree, "categorical_indices_", np.array([])).size
+    ):
+        cat_codes = tree._encode_input_categories(Xarr)
+        tree._check_unknown(Xarr, cat_codes)
+    out = np.full(Xarr.shape[0], -1, dtype=np.intp)
+    covered = np.zeros(Xarr.shape[0], dtype=bool)
+    for rule in extract_rules(tree):
+        m = rule_mask(rule, Xarr, tree.feature_names_)
+        out[m & ~covered] = rule.leaf_id
+        covered |= m
+    return out
+
+
+def replay_check(tree, X):
+    """True if rule replay reproduces ``tree.predict(X)`` exactly."""
+    return bool(np.array_equal(replay_labels(tree, X), tree.predict(X)))
+
+
+def profile_leaves(tree, X, top_k=3):
+    """Profile each segment by its most distinctive features.
+
+    For every leaf, computes the centroid of ``X`` rows assigned to it.
+    Continuous/binary features are ranked by absolute z-score versus the
+    whole dataset (positive = above average). Categorical columns are *not*
+    z-scored; instead each leaf receives ``category_enrichment``: the most
+    enriched (feature, code) pairs by lift of in-leaf share over global share
+    (only pairs covering >= 5% of the leaf, lift > 1, capped at ``top_k``,
+    sorted by descending lift).
+
+    Returns a list of dicts with keys ``leaf_id``, ``n_samples``, ``fraction``,
+    ``top_features`` (tuple of ``(feature_name, z_score)`` for non-categorical
+    features only) and ``category_enrichment`` (tuple of
+    ``(feature_name, code, leaf_share, lift)``; empty when the tree has no
+    categorical features).
+    """
+    Xarr, _ = as_2d_array(X, context="X")
+    labels = tree.predict(Xarr)
+    cat_idx = set(int(j) for j in getattr(tree, "categorical_indices_", []))
+    num_cols = [j for j in range(Xarr.shape[1]) if j not in cat_idx]
+    mu = Xarr.mean(axis=0)
+    sd = Xarr.std(axis=0)
+    sd[sd < 1e-12] = 1e-12
+
+    # global category shares for lift computation
+    global_dist = {}
+    for j in cat_idx:
+        known = np.asarray(tree.categories_[tree.feature_names_[j]])
+        local = np.searchsorted(known, np.rint(Xarr[:, j]).astype(np.int64))
+        counts = np.bincount(local, minlength=known.size + 1)[: known.size]
+        global_dist[j] = counts / max(counts.sum(), 1)
+
+    profiles = []
+    for lf in tree.leaves_:
+        rows_mask = labels == lf.id
+        sub = Xarr[rows_mask]
+        enrichment = []
+        for j in cat_idx:
+            known = np.asarray(tree.categories_[tree.feature_names_[j]])
+            local = np.searchsorted(known, np.rint(sub[:, j]).astype(np.int64))
+            counts = np.bincount(local, minlength=known.size + 1)[: known.size]
+            if counts.sum() == 0:
+                continue
+            leaf_share = counts / counts.sum()
+            lift = leaf_share / np.maximum(global_dist[j], 1e-12)
+            eligible = np.flatnonzero((leaf_share >= 0.05) & (lift > 1.0))
+            for pos in eligible:
+                enrichment.append(
+                    (
+                        tree.feature_names_[j],
+                        int(known[pos]),
+                        float(leaf_share[pos]),
+                        float(lift[pos]),
+                    )
+                )
+        enrichment.sort(key=lambda item: -item[3])
+        zcols = sub[:, num_cols] if num_cols else np.empty((len(sub), 0))
+        znum = (zcols.mean(axis=0) - mu[num_cols]) / sd[num_cols] if num_cols else ()
+        top_order = np.argsort(-np.abs(znum))[:top_k] if num_cols else []
+        profiles.append(
+            {
+                "leaf_id": lf.id,
+                "n_samples": int(sub.shape[0]),
+                "fraction": sub.shape[0] / Xarr.shape[0],
+                "top_features": tuple(
+                    (tree.feature_names_[num_cols[p]], float(znum[p])) for p in top_order
+                ),
+                "category_enrichment": tuple(enrichment[:top_k]),
+            }
+        )
+    return profiles

segmtree/stability.py

@@ -0,0 +1,100 @@
+"""Bootstrap stability assessment for fitted trees."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from ._base import as_2d_array
+from .tree import SegTree
+
+__all__ = ["StabilityResult", "adjusted_rand_score", "bootstrap_stability"]
+
+
+@dataclass(frozen=True)
+class StabilityResult:
+    """Bootstrap stability of a segmentation.
+
+    ``scores[i]`` is the Adjusted Rand Index between the original partition
+    and the partition induced on the full data by a tree refit on bootstrap
+    replicate ``i``. Values above ~0.7 indicate a stable structure.
+    """
+
+    scores: np.ndarray
+    mean: float
+    std: float
+    n_replicates: int
+
+
+def adjusted_rand_score(labels_a, labels_b):
+    """Adjusted Rand Index between two labelings (pure numpy).
+
+    Matches :func:`sklearn.metrics.adjusted_rand_score` numerically; provided
+    to keep the package dependency-free.
+    """
+    a = np.asarray(labels_a).ravel()
+    b = np.asarray(labels_b).ravel()
+    if a.shape[0] != b.shape[0]:
+        raise ValueError("Label arrays must have the same length.")
+    _, ai = np.unique(a, return_inverse=True)
+    _, bi = np.unique(b, return_inverse=True)
+    na, nb, n = int(ai.max()) + 1, int(bi.max()) + 1, a.shape[0]
+    contingency = np.bincount(ai * nb + bi, minlength=na * nb).astype(np.int64).reshape(na, nb)
+
+    def comb2(x):
+        return x * (x - 1) // 2
+
+    sum_comb = comb2(contingency.astype(np.float64)).sum()
+    sum_a = comb2(contingency.sum(axis=1).astype(np.float64)).sum()
+    sum_b = comb2(contingency.sum(axis=0).astype(np.float64)).sum()
+    expected = sum_a * sum_b / comb2(float(n))
+    max_index = 0.5 * (sum_a + sum_b)
+    denom = max_index - expected
+    if denom == 0.0:
+        return 1.0
+    return float((sum_comb - expected) / denom)
+
+
+def bootstrap_stability(tree, X, *, n_replicates=12, random_state=None):
+    """Refit the tree on bootstrap resamples and measure agreement with the fit.
+
+    Parameters
+    ----------
+    tree : SegTree
+        A fitted tree; its hyperparameters are reused for every replicate.
+    X : array-like, shape (n_samples, n_features)
+        The training matrix used to fit ``tree``.
+    n_replicates : int, default=12
+        Number of bootstrap replicates.
+    random_state : int or None, default=None
+        Seed for reproducibility.
+
+    Returns
+    -------
+    StabilityResult
+    """
+    if not hasattr(tree, "labels_"):
+        raise RuntimeError("tree must be fitted before calling bootstrap_stability().")
+    if isinstance(n_replicates, bool) or not isinstance(n_replicates, int) or n_replicates < 1:
+        raise ValueError("n_replicates must be an int >= 1.")
+    Xarr, _ = as_2d_array(X, context="X")
+    if Xarr.shape != (tree.n_samples_in_, tree.n_features_in_):
+        raise ValueError(
+            f"X must be the training matrix with shape "
+            f"{(tree.n_samples_in_, tree.n_features_in_)}, got {Xarr.shape}."
+        )
+    params = tree.get_params()
+    base_labels = tree.labels_
+    rng = np.random.default_rng(random_state)
+    scores = np.empty(n_replicates, dtype=np.float64)
+    for b in range(n_replicates):
+        idx = rng.integers(0, Xarr.shape[0], Xarr.shape[0])
+        replica = SegTree(**params).fit(Xarr[idx])
+        scores[b] = adjusted_rand_score(base_labels, replica.predict(Xarr))
+    return StabilityResult(
+        scores=scores,
+        mean=float(scores.mean()),
+        std=float(scores.std()),
+        n_replicates=n_replicates,
+    )