Mondrian Prediction

The Mondrian block groups everything built on the Mondrian process — a stochastic process that recursively partitions feature space into axis-aligned boxes. Because the partition depends only on the bag of feature values (never their order), a tree grown from the augmented bag {train ∪ (x_test, y)} yields an exact conformal predictor, and its leaves form a valid Venn taxonomy (ALRW2 §2.2.9, §4.6, Thm 6.4).

Two distinct meanings of Mondrian

• Mondrian tree / forest predictors (this page) use the Mondrian process to discover an adaptive partition and score nonconformity in the leaves.
• The group-conditional wrappers (MondrianConformalClassifier / MondrianConformalRegressor) are the general Mondrian taxonomy framework — a user-supplied category function applied on top of any base predictor. They are documented at the bottom of this page.

MondrianTree (core)

The standalone partition object. Grow it from a data bag with an external RNG, then traverse it, read leaf statistics, or visualise it. Conformal and Venn predictors are thin adapters over this class.

online_cp.mondrian.MondrianTree

A single sampled Mondrian partition of a data bag X.

MondrianTree owns one realisation of the Mondrian process at a given lifetime and exposes growing, traversal, leaf statistics and visualisation. Conformal and Venn predictors are thin adapters that grow a tree from an augmented bag and read off leaf statistics.

The partition is a bag function: it depends only on the feature values in X, never on their order. Growing with a string lifetime yields a deterministic truncation of the master tree grown with the same RNG.

Parameters:

Name	Type	Description	Default
root	_MondrianNode	Root node of the sampled tree. Prefer :meth:grow over constructing directly.	required
X	ndarray of shape (n, d)	The data bag the tree was grown from.	required
lifetime	float	The concrete (resolved) lifetime the tree was grown at.	required
feature_weights	ndarray of shape (d,) or None	The resolved per-dimension split weights (None for isotropic).	None

Examples:

>>> import numpy as np
>>> from online_cp.mondrian import MondrianTree
>>> rng = np.random.default_rng(0)
>>> X = rng.uniform(0, 1, size=(20, 2))
>>> tree = MondrianTree.grow(X, rng, lifetime=2.0)
>>> leaf = tree.find_leaf(X[0])
>>> leaf.is_leaf()
True
>>> all(node.is_leaf() for node in tree.collect_leaves())
True

Source code in src/online_cp/mondrian/tree.py

class MondrianTree:
    """A single sampled Mondrian partition of a data bag ``X``.

    ``MondrianTree`` owns one realisation of the Mondrian process at a given
    ``lifetime`` and exposes growing, traversal, leaf statistics and
    visualisation. Conformal and Venn predictors are thin adapters that grow a
    tree from an augmented bag and read off leaf statistics.

    The partition is a *bag function*: it depends only on the feature values in
    ``X``, never on their order. Growing with a string ``lifetime`` yields a
    deterministic truncation of the master tree grown with the same RNG.

    Parameters
    ----------
    root : _MondrianNode
        Root node of the sampled tree. Prefer :meth:`grow` over constructing
        directly.
    X : ndarray of shape (n, d)
        The data bag the tree was grown from.
    lifetime : float
        The concrete (resolved) lifetime the tree was grown at.
    feature_weights : ndarray of shape (d,) or None
        The resolved per-dimension split weights (``None`` for isotropic).

    Examples
    --------
    >>> import numpy as np
    >>> from online_cp.mondrian import MondrianTree
    >>> rng = np.random.default_rng(0)
    >>> X = rng.uniform(0, 1, size=(20, 2))
    >>> tree = MondrianTree.grow(X, rng, lifetime=2.0)
    >>> leaf = tree.find_leaf(X[0])
    >>> leaf.is_leaf()
    True
    >>> all(node.is_leaf() for node in tree.collect_leaves())
    True
    """

    def __init__(
        self,
        root: _MondrianNode,
        X: NDArray,
        *,
        lifetime: float,
        feature_weights: NDArray | None = None,
    ) -> None:
        self.root = root
        self.X = np.asarray(X)
        self.lifetime = lifetime
        self.feature_weights = feature_weights

    @classmethod
    def grow(
        cls,
        X: NDArray,
        rng: np.random.Generator,
        *,
        lifetime: float | str = 1.0,
        x_test: NDArray | None = None,
        max_depth: int | None = None,
        min_samples_leaf: int = 1,
        feature_weights: str | NDArray | None = None,
        verbose: int = 0,
    ) -> MondrianTree:
        """Grow a Mondrian tree from the data bag ``X`` using ``rng``.

        The generator ``rng`` is advanced in place. String ``lifetime``
        specifications (``'sqrt_n'`` / ``'density'``) grow a master tree to
        resolve the budget and then restore the RNG state, so the returned tree
        is a deterministic truncation of that master tree.

        Parameters
        ----------
        X : ndarray of shape (n, d)
            Feature bag; the partition is a function of these values only.
        rng : numpy.random.Generator
            Random generator (advanced in place).
        lifetime : float or {'sqrt_n', 'density'}
            Mondrian time budget, or an unsupervised auto-tuning rule.
        x_test : ndarray of shape (d,), optional
            Reference point for the ``'sqrt_n'`` rule (defaults to the last row
            of ``X``, i.e. the augmented test point).
        max_depth : int or None
            Hard depth cap (``None`` = no cap).
        min_samples_leaf : int
            Minimum number of points allowed in a leaf.
        feature_weights : {'variance'} or array-like of shape (d,) or None
            Unsupervised per-dimension split weights.
        verbose : int
            Verbosity level.

        Returns
        -------
        MondrianTree
        """
        X = np.asarray(X)
        fw_arr = _resolve_feature_weights(X, feature_weights)
        ref = X[-1] if x_test is None else np.asarray(x_test).ravel()
        lt_val = _resolve_lifetime(X, ref, lifetime, rng, fw_arr)
        root = _sample_mondrian_tree(
            rng,
            X,
            np.arange(X.shape[0]),
            parent_time=0.0,
            lifetime=lt_val,
            verbose=verbose,
            max_depth=max_depth,
            min_samples_leaf=min_samples_leaf,
            feature_weights=fw_arr,
        )
        return cls(root, X, lifetime=lt_val, feature_weights=fw_arr)

    def extend(self, x: NDArray, rng: np.random.Generator) -> MondrianTree:
        """Return a new tree with ``x`` projected in (online ``ExtendMondrianBlock``).

        The result has the same law as :meth:`grow` on the augmented point set
        but costs only ``O(depth)`` (Roy & Teh 2009). This tree is left unchanged;
        ``x`` is appended as the last row of the returned tree's ``X``.

        Requires a fixed float ``lifetime`` (the projective regime); string
        lifetimes are not supported online.
        """
        if not isinstance(self.lifetime, (int, float)):
            raise ValueError(
                "extend() requires a fixed float lifetime; got "
                f"{self.lifetime!r}. String/auto-tuned lifetimes are batch-only."
            )
        x = np.asarray(x, dtype=float).ravel()
        x_idx = self.X.shape[0]
        new_root = _extend_mondrian(
            self.root, x, x_idx, rng, float(self.lifetime), self.feature_weights
        )
        return MondrianTree(
            new_root,
            np.vstack([self.X, x]),
            lifetime=self.lifetime,
            feature_weights=self.feature_weights,
        )

    def find_leaf(self, x: NDArray) -> _MondrianNode:
        """Return the leaf node whose cell contains ``x``."""
        return _find_leaf(self.root, np.asarray(x).ravel())

    def collect_leaves(self) -> list[_MondrianNode]:
        """Return all leaf nodes, left-to-right."""
        return _collect_leaves(self.root)

    def iter_nodes(self) -> list[tuple]:
        """Depth-first pre-order over ``(node, depth, parent_id, node_id)`` tuples."""
        return _iter_nodes(self.root)

    def struct_stats(self) -> dict:
        """Return structural stats: ``n_nodes``, ``n_leaves``, ``n_branches``, ``height``."""
        return _tree_struct_stats(self.root)

    def assign_counts(
        self,
        y: NDArray,
        label_to_idx: dict,
        K: int,
        n_train: int | None = None,
    ) -> None:
        """Attach per-class count vectors to every leaf (modifies the tree in place)."""
        _assign_counts(self.root, y, label_to_idx, K, n_train=n_train)

    def to_dataframe(self):
        """Export the tree structure as a :class:`pandas.DataFrame` (one row per node)."""
        import pandas as pd

        rows = []
        node_id_map: dict[str, int] = {}
        next_id = 0
        for node, depth, parent_id_str, node_id_str in _iter_nodes(self.root):
            if node_id_str not in node_id_map:
                node_id_map[node_id_str] = next_id
                next_id += 1
            node_id_int = node_id_map[node_id_str]
            if parent_id_str is None:
                parent_id_int = None
            else:
                if parent_id_str not in node_id_map:
                    node_id_map[parent_id_str] = next_id
                    next_id += 1
                parent_id_int = node_id_map[parent_id_str]
            rows.append({
                "node_id": node_id_int,
                "parent_id": parent_id_int,
                "is_leaf": node.is_leaf(),
                "depth": depth,
                "split_dim": node.split_dim,
                "split_loc": float(node.split_loc) if not node.is_leaf() else None,
                "split_time": float(node.split_time) if not node.is_leaf() else None,
                "parent_time": float(node.parent_time),
                "n_points": node.n_points(),
            })
        return pd.DataFrame(rows)

    def draw(self, ax=None, max_depth=None, backend="auto"):
        """Draw a node-link diagram of the tree.

        See :meth:`ConformalMondrianTreeClassifier.draw
        <online_cp.classifiers.ConformalMondrianTreeClassifier.draw>` for the
        backend semantics. Uses a structural label (split test / leaf size).

        Notes
        -----
        For publication-quality vector output:

        * **Graphviz** (when ``ax=None``): Returns a :class:`graphviz.Digraph`
          which can be saved as SVG/PDF without quality loss::

              >>> dg = tree.draw(backend="graphviz")  # doctest: +SKIP
              >>> dg.render("mytree", format="svg", cleanup=True)  # doctest: +SKIP

        * **Matplotlib**: Use ``backend="matplotlib"`` and save with
          ``bbox_inches='tight'``::

              >>> ax = tree.draw(backend="matplotlib")  # doctest: +SKIP
              >>> ax.figure.savefig("tree.svg", dpi=300, bbox_inches="tight")  # doctest: +SKIP
        """
        def label_fn(node, depth, *, collapsed=False):
            if node.is_leaf() or collapsed:
                n = node.n_points()
                return f"leaf\nn={n}" if n else "empty"
            return f"x[{node.split_dim}] ≤\n{node.split_loc:.3g}"

        return _render_tree(
            self.root, ax, label_fn, max_depth, backend, title="Mondrian Tree",
        )

    def draw_partition(self, ax=None, scatter=True):
        """Draw the 2-D box partition (requires ``d == 2``), one colour per leaf."""
        if self.X.shape[1] != 2:
            raise ValueError(
                f"draw_partition() requires exactly 2 features, got {self.X.shape[1]}."
            )
        from matplotlib import colormaps

        ax = _get_ax(ax)
        leaves = self.collect_leaves()
        cmap = colormaps["tab20"].resampled(max(len(leaves), 1))
        leaf_order = {id(leaf): i for i, leaf in enumerate(leaves)}

        def leaf_color_fn(leaf):
            return cmap(leaf_order[id(leaf)] % cmap.N)

        _draw_partition(
            self.root, self.X, ax, leaf_color_fn,
            scatter_X=self.X if scatter else None,
        )
        return ax

    def __len__(self) -> int:
        return _tree_struct_stats(self.root)["n_leaves"]

    def __repr__(self) -> str:
        stats = _tree_struct_stats(self.root)
        return (
            f"MondrianTree(lifetime={self.lifetime!r}, "
            f"n_leaves={stats['n_leaves']}, height={stats['height']})"
        )

grow(X: NDArray, rng: np.random.Generator, *, lifetime: float | str = 1.0, x_test: NDArray | None = None, max_depth: int | None = None, min_samples_leaf: int = 1, feature_weights: str | NDArray | None = None, verbose: int = 0) -> MondrianTree classmethod

Grow a Mondrian tree from the data bag X using rng.

The generator rng is advanced in place. String lifetime specifications ('sqrt_n' / 'density') grow a master tree to resolve the budget and then restore the RNG state, so the returned tree is a deterministic truncation of that master tree.

Parameters:

Name	Type	Description	Default
X	ndarray of shape (n, d)	Feature bag; the partition is a function of these values only.	required
rng	Generator	Random generator (advanced in place).	required
lifetime	float or {'sqrt_n', 'density'}	Mondrian time budget, or an unsupervised auto-tuning rule.	1.0
x_test	ndarray of shape (d,)	Reference point for the 'sqrt_n' rule (defaults to the last row of X, i.e. the augmented test point).	None
max_depth	int or None	Hard depth cap (None = no cap).	None
min_samples_leaf	int	Minimum number of points allowed in a leaf.	1
feature_weights	'variance'	Unsupervised per-dimension split weights.	'variance'
verbose	int	Verbosity level.	0

Returns:

Type	Description
MondrianTree

Source code in src/online_cp/mondrian/tree.py

@classmethod
def grow(
    cls,
    X: NDArray,
    rng: np.random.Generator,
    *,
    lifetime: float | str = 1.0,
    x_test: NDArray | None = None,
    max_depth: int | None = None,
    min_samples_leaf: int = 1,
    feature_weights: str | NDArray | None = None,
    verbose: int = 0,
) -> MondrianTree:
    """Grow a Mondrian tree from the data bag ``X`` using ``rng``.

    The generator ``rng`` is advanced in place. String ``lifetime``
    specifications (``'sqrt_n'`` / ``'density'``) grow a master tree to
    resolve the budget and then restore the RNG state, so the returned tree
    is a deterministic truncation of that master tree.

    Parameters
    ----------
    X : ndarray of shape (n, d)
        Feature bag; the partition is a function of these values only.
    rng : numpy.random.Generator
        Random generator (advanced in place).
    lifetime : float or {'sqrt_n', 'density'}
        Mondrian time budget, or an unsupervised auto-tuning rule.
    x_test : ndarray of shape (d,), optional
        Reference point for the ``'sqrt_n'`` rule (defaults to the last row
        of ``X``, i.e. the augmented test point).
    max_depth : int or None
        Hard depth cap (``None`` = no cap).
    min_samples_leaf : int
        Minimum number of points allowed in a leaf.
    feature_weights : {'variance'} or array-like of shape (d,) or None
        Unsupervised per-dimension split weights.
    verbose : int
        Verbosity level.

    Returns
    -------
    MondrianTree
    """
    X = np.asarray(X)
    fw_arr = _resolve_feature_weights(X, feature_weights)
    ref = X[-1] if x_test is None else np.asarray(x_test).ravel()
    lt_val = _resolve_lifetime(X, ref, lifetime, rng, fw_arr)
    root = _sample_mondrian_tree(
        rng,
        X,
        np.arange(X.shape[0]),
        parent_time=0.0,
        lifetime=lt_val,
        verbose=verbose,
        max_depth=max_depth,
        min_samples_leaf=min_samples_leaf,
        feature_weights=fw_arr,
    )
    return cls(root, X, lifetime=lt_val, feature_weights=fw_arr)

extend(x: NDArray, rng: np.random.Generator) -> MondrianTree

Return a new tree with x projected in (online ExtendMondrianBlock).

The result has the same law as :meth:grow on the augmented point set but costs only O(depth) (Roy & Teh 2009). This tree is left unchanged; x is appended as the last row of the returned tree's X.

Requires a fixed float lifetime (the projective regime); string lifetimes are not supported online.

Source code in src/online_cp/mondrian/tree.py

def extend(self, x: NDArray, rng: np.random.Generator) -> MondrianTree:
    """Return a new tree with ``x`` projected in (online ``ExtendMondrianBlock``).

    The result has the same law as :meth:`grow` on the augmented point set
    but costs only ``O(depth)`` (Roy & Teh 2009). This tree is left unchanged;
    ``x`` is appended as the last row of the returned tree's ``X``.

    Requires a fixed float ``lifetime`` (the projective regime); string
    lifetimes are not supported online.
    """
    if not isinstance(self.lifetime, (int, float)):
        raise ValueError(
            "extend() requires a fixed float lifetime; got "
            f"{self.lifetime!r}. String/auto-tuned lifetimes are batch-only."
        )
    x = np.asarray(x, dtype=float).ravel()
    x_idx = self.X.shape[0]
    new_root = _extend_mondrian(
        self.root, x, x_idx, rng, float(self.lifetime), self.feature_weights
    )
    return MondrianTree(
        new_root,
        np.vstack([self.X, x]),
        lifetime=self.lifetime,
        feature_weights=self.feature_weights,
    )

find_leaf(x: NDArray) -> _MondrianNode

Return the leaf node whose cell contains x.

Source code in src/online_cp/mondrian/tree.py

def find_leaf(self, x: NDArray) -> _MondrianNode:
    """Return the leaf node whose cell contains ``x``."""
    return _find_leaf(self.root, np.asarray(x).ravel())

collect_leaves() -> list[_MondrianNode]

Return all leaf nodes, left-to-right.

Source code in src/online_cp/mondrian/tree.py

def collect_leaves(self) -> list[_MondrianNode]:
    """Return all leaf nodes, left-to-right."""
    return _collect_leaves(self.root)

iter_nodes() -> list[tuple]

Depth-first pre-order over (node, depth, parent_id, node_id) tuples.

Source code in src/online_cp/mondrian/tree.py

def iter_nodes(self) -> list[tuple]:
    """Depth-first pre-order over ``(node, depth, parent_id, node_id)`` tuples."""
    return _iter_nodes(self.root)

struct_stats() -> dict

Return structural stats: n_nodes, n_leaves, n_branches, height.

Source code in src/online_cp/mondrian/tree.py

def struct_stats(self) -> dict:
    """Return structural stats: ``n_nodes``, ``n_leaves``, ``n_branches``, ``height``."""
    return _tree_struct_stats(self.root)

assign_counts(y: NDArray, label_to_idx: dict, K: int, n_train: int | None = None) -> None

Attach per-class count vectors to every leaf (modifies the tree in place).

Source code in src/online_cp/mondrian/tree.py

def assign_counts(
    self,
    y: NDArray,
    label_to_idx: dict,
    K: int,
    n_train: int | None = None,
) -> None:
    """Attach per-class count vectors to every leaf (modifies the tree in place)."""
    _assign_counts(self.root, y, label_to_idx, K, n_train=n_train)

to_dataframe()

Export the tree structure as a :class:pandas.DataFrame (one row per node).

Source code in src/online_cp/mondrian/tree.py

def to_dataframe(self):
    """Export the tree structure as a :class:`pandas.DataFrame` (one row per node)."""
    import pandas as pd

    rows = []
    node_id_map: dict[str, int] = {}
    next_id = 0
    for node, depth, parent_id_str, node_id_str in _iter_nodes(self.root):
        if node_id_str not in node_id_map:
            node_id_map[node_id_str] = next_id
            next_id += 1
        node_id_int = node_id_map[node_id_str]
        if parent_id_str is None:
            parent_id_int = None
        else:
            if parent_id_str not in node_id_map:
                node_id_map[parent_id_str] = next_id
                next_id += 1
            parent_id_int = node_id_map[parent_id_str]
        rows.append({
            "node_id": node_id_int,
            "parent_id": parent_id_int,
            "is_leaf": node.is_leaf(),
            "depth": depth,
            "split_dim": node.split_dim,
            "split_loc": float(node.split_loc) if not node.is_leaf() else None,
            "split_time": float(node.split_time) if not node.is_leaf() else None,
            "parent_time": float(node.parent_time),
            "n_points": node.n_points(),
        })
    return pd.DataFrame(rows)

draw(ax=None, max_depth=None, backend='auto')

Draw a node-link diagram of the tree.

See :meth:ConformalMondrianTreeClassifier.draw <online_cp.classifiers.ConformalMondrianTreeClassifier.draw> for the backend semantics. Uses a structural label (split test / leaf size).

Notes

For publication-quality vector output:

• Graphviz (when ax=None): Returns a :class:graphviz.Digraph which can be saved as SVG/PDF without quality loss::

  dg = tree.draw(backend="graphviz") # doctest: +SKIP dg.render("mytree", format="svg", cleanup=True) # doctest: +SKIP

• Matplotlib: Use backend="matplotlib" and save with bbox_inches='tight'::

  ax = tree.draw(backend="matplotlib") # doctest: +SKIP ax.figure.savefig("tree.svg", dpi=300, bbox_inches="tight") # doctest: +SKIP

Source code in src/online_cp/mondrian/tree.py

def draw(self, ax=None, max_depth=None, backend="auto"):
    """Draw a node-link diagram of the tree.

    See :meth:`ConformalMondrianTreeClassifier.draw
    <online_cp.classifiers.ConformalMondrianTreeClassifier.draw>` for the
    backend semantics. Uses a structural label (split test / leaf size).

    Notes
    -----
    For publication-quality vector output:

    * **Graphviz** (when ``ax=None``): Returns a :class:`graphviz.Digraph`
      which can be saved as SVG/PDF without quality loss::

          >>> dg = tree.draw(backend="graphviz")  # doctest: +SKIP
          >>> dg.render("mytree", format="svg", cleanup=True)  # doctest: +SKIP

    * **Matplotlib**: Use ``backend="matplotlib"`` and save with
      ``bbox_inches='tight'``::

          >>> ax = tree.draw(backend="matplotlib")  # doctest: +SKIP
          >>> ax.figure.savefig("tree.svg", dpi=300, bbox_inches="tight")  # doctest: +SKIP
    """
    def label_fn(node, depth, *, collapsed=False):
        if node.is_leaf() or collapsed:
            n = node.n_points()
            return f"leaf\nn={n}" if n else "empty"
        return f"x[{node.split_dim}] ≤\n{node.split_loc:.3g}"

    return _render_tree(
        self.root, ax, label_fn, max_depth, backend, title="Mondrian Tree",
    )

draw_partition(ax=None, scatter=True)

Draw the 2-D box partition (requires d == 2), one colour per leaf.

Source code in src/online_cp/mondrian/tree.py

def draw_partition(self, ax=None, scatter=True):
    """Draw the 2-D box partition (requires ``d == 2``), one colour per leaf."""
    if self.X.shape[1] != 2:
        raise ValueError(
            f"draw_partition() requires exactly 2 features, got {self.X.shape[1]}."
        )
    from matplotlib import colormaps

    ax = _get_ax(ax)
    leaves = self.collect_leaves()
    cmap = colormaps["tab20"].resampled(max(len(leaves), 1))
    leaf_order = {id(leaf): i for i, leaf in enumerate(leaves)}

    def leaf_color_fn(leaf):
        return cmap(leaf_order[id(leaf)] % cmap.N)

    _draw_partition(
        self.root, self.X, ax, leaf_color_fn,
        scatter_X=self.X if scatter else None,
    )
    return ax

Conformal tree & forest predictors

online_cp.classifiers.ConformalMondrianTreeClassifier

Bases: _MondrianClassifierInspection, ConformalClassifier

Conformal predictor using Mondrian trees.

A full (transductive) conformal predictor. At each prediction step the tree is rebuilt from the augmented bag {train ∪ (x_test, y_cand)}, guaranteeing exact validity under exchangeability (ALRW2 §2.2.9).

The nonconformity measure is (Leaf-Local Laplace Smoothing):

α_i = 1 - (C_{leaf(x_i), y_i} + λ) / (N_leaf(x_i) + K * λ)

where C_{leaf, k} is the count of class k in the leaf, N_leaf is the number of training points in the leaf, K is the number of classes, and λ (lambda_) is the Laplace smoothing parameter (default 1.0). This eliminates the discrete ties of raw fractions while confining the transductive delta strictly to leaf_star (O(1) blast radius).

Parameters:

Name	Type	Description	Default
lifetime	float or {'sqrt_n', 'density'}	Depth budget L for the Mondrian tree. Splits occur at times sampled from Exp(sum of feature ranges). A split is created only if its cumulative time is below L; otherwise the node becomes a leaf. Can also be a string for unsupervised automatic tuning (only X is used; labels are never accessed): 'sqrt_n': descend the master tree toward the test point and halt when the leaf count drops to ≤ √n. Coarsens the partition in sparse regions. 'density': sweep all split events of a master tree and choose the τ that maximises the histogram log-likelihood Σ_leaf n_leaf · log(n_leaf / V_leaf). Adapts to feature density. Effect on prediction quality: - Too small (e.g. 0.01): almost no splits → one big leaf → NCM is just the global class frequency → very weak discrimination. - Too large (e.g. np.inf): splits until every leaf is a singleton → NCM is 0 for every point → p-values are all equal (τ) → prediction sets are either full or empty, with no discrimination. - Well-chosen (default 1.0): leaves contain a handful of points, NCMs vary meaningfully, prediction sets are informative. Scaling guidance: lifetime is in the same units as the sum of feature ranges. For standardised data (zero mean, unit variance) a value of 1–5 is a reasonable starting range. For raw (unstandardised) features with large ranges (e.g. pixel values 0–255) you may need lifetime in the hundreds. When in doubt, standardise your features first and use the default.	1.0
lambda_	float	Laplace smoothing parameter λ. Controls the strength of the uniform prior added to each leaf's class counts. Default 1.0 (standard Laplace). Larger values push NCMs towards 1 - 1/K (uniform), reducing the influence of small leaves. Set to 0.0 to recover raw fraction NCMs.	1.0
label_space	array - like	Set of possible labels. If None, inferred from training data.	None
epsilon	float	Default significance level for predictions.	0.1
rnd_state	int	Seed for reproducible tree construction.	None
verbose	int	Verbosity level (0 = silent).	0
max_depth	int or None	Hard depth cap. Nodes at depth >= max_depth become leaves regardless of remaining lifetime. None = no cap.	None
feature_weights	'variance'	Unsupervised feature importance for split dimension sampling. Labels are never accessed. None (default): uniform — standard isotropic Mondrian. 'variance': split probabilities ∝ population variance per feature. High-variance dimensions are sampled more often. array of shape (d,): explicit non-negative weights (normalised internally; need not sum to 1).	'variance'

Attributes:

Name	Type	Description
X	ndarray of shape (n, d)	Training feature array.
y	ndarray of shape (n,)	Training labels.
label_space	ndarray	Sorted unique labels.
label_to_idx	dict	Mapping from label to integer index 0..K-1.
rnd_gen	Generator	Random generator instance.

Source code in src/online_cp/classifiers.py

class ConformalMondrianTreeClassifier(_MondrianClassifierInspection, ConformalClassifier):
    """Conformal predictor using Mondrian trees.

    A full (transductive) conformal predictor. At each prediction step the tree
    is rebuilt from the augmented bag {train ∪ (x_test, y_cand)}, guaranteeing
    exact validity under exchangeability (ALRW2 §2.2.9).

    The nonconformity measure is (Leaf-Local Laplace Smoothing):

        α_i = 1 - (C_{leaf(x_i), y_i} + λ) / (N_leaf(x_i) + K * λ)

    where C_{leaf, k} is the count of class k in the leaf, N_leaf is the
    number of training points in the leaf, K is the number of classes, and
    λ (`lambda_`) is the Laplace smoothing parameter (default 1.0). This
    eliminates the discrete ties of raw fractions while confining the
    transductive delta strictly to leaf_star (O(1) blast radius).

    Parameters
    ----------
    lifetime : float or {'sqrt_n', 'density'}
        Depth budget L for the Mondrian tree. Splits occur at times sampled
        from Exp(sum of feature ranges). A split is created only if its
        cumulative time is below L; otherwise the node becomes a leaf.

        Can also be a string for unsupervised automatic tuning (only X is
        used; labels are never accessed):

        - ``'sqrt_n'``: descend the master tree toward the test point and
          halt when the leaf count drops to ≤ √n. Coarsens the partition in
          sparse regions.
        - ``'density'``: sweep all split events of a master tree and choose
          the τ that maximises the histogram log-likelihood
          Σ_leaf n_leaf · log(n_leaf / V_leaf). Adapts to feature density.

        **Effect on prediction quality:**
        - Too small (e.g. 0.01): almost no splits → one big leaf → NCM is
          just the global class frequency → very weak discrimination.
        - Too large (e.g. np.inf): splits until every leaf is a singleton →
          NCM is 0 for every point → p-values are all equal (τ) → prediction
          sets are either full or empty, with no discrimination.
        - Well-chosen (default 1.0): leaves contain a handful of points,
          NCMs vary meaningfully, prediction sets are informative.

        **Scaling guidance:** `lifetime` is in the same units as the sum of
        feature ranges. For standardised data (zero mean, unit variance)
        a value of 1–5 is a reasonable starting range. For raw (unstandardised)
        features with large ranges (e.g. pixel values 0–255) you may need
        lifetime in the hundreds. When in doubt, standardise your features
        first and use the default.
    lambda_ : float
        Laplace smoothing parameter λ. Controls the strength of the uniform
        prior added to each leaf's class counts. Default 1.0 (standard Laplace).
        Larger values push NCMs towards 1 - 1/K (uniform), reducing the
        influence of small leaves. Set to 0.0 to recover raw fraction NCMs.
    label_space : array-like, optional
        Set of possible labels. If None, inferred from training data.
    epsilon : float
        Default significance level for predictions.
    rnd_state : int, optional
        Seed for reproducible tree construction.
    verbose : int
        Verbosity level (0 = silent).
    max_depth : int or None
        Hard depth cap. Nodes at depth >= max_depth become leaves regardless
        of remaining lifetime. None = no cap.
    feature_weights : {'variance'} or array-like of shape (d,) or None
        Unsupervised feature importance for split dimension sampling. Labels
        are never accessed.

        - ``None`` (default): uniform — standard isotropic Mondrian.
        - ``'variance'``: split probabilities ∝ population variance per
          feature. High-variance dimensions are sampled more often.
        - array of shape (d,): explicit non-negative weights (normalised
          internally; need not sum to 1).

    Attributes
    ----------
    X : ndarray of shape (n, d)
        Training feature array.
    y : ndarray of shape (n,)
        Training labels.
    label_space : ndarray
        Sorted unique labels.
    label_to_idx : dict
        Mapping from label to integer index 0..K-1.
    rnd_gen : numpy.random.Generator
        Random generator instance.
    """

    _SAVE_PARAMS: tuple = (
        "lifetime", "lambda_", "label_space", "epsilon", "rnd_state",
        "verbose", "max_depth", "feature_weights", "online",
    )
    _SAVE_STATE: tuple = ("X", "y", "label_to_idx", "label_space", "_online_tree")

    def __init__(
        self,
        lifetime: float | str = 1.0,
        lambda_: float = 1.0,
        label_space: NDArray | None = None,
        epsilon: float = 0.1,
        rnd_state: int | None = None,
        verbose: int = 0,
        max_depth: int | None = None,
        feature_weights: str | NDArray | None = None,
        online: bool = False,
    ):
        if max_depth is not None and (not isinstance(max_depth, int) or max_depth < 0):
            raise ValueError("max_depth must be a non-negative integer or None")
        if isinstance(lifetime, str) and lifetime not in ("sqrt_n", "density"):
            raise ValueError(
                f"Unknown lifetime string {lifetime!r}. "
                "Recognised values: 'sqrt_n', 'density'."
            )
        if isinstance(feature_weights, str) and feature_weights not in ("variance",):
            raise ValueError(
                f"Unknown feature_weights string {feature_weights!r}. "
                "Recognised value: 'variance'."
            )
        if online:
            if not isinstance(lifetime, (int, float)):
                raise ValueError(
                    "online=True requires a fixed float lifetime; string/auto-tuned "
                    "lifetimes recompute a master tree per step and are batch-only."
                )
            if max_depth is not None:
                raise ValueError(
                    "online=True requires max_depth=None (the projective Mondrian regime)."
                )
        super().__init__(epsilon=epsilon)
        self.lifetime = lifetime
        self.feature_weights = feature_weights
        self.lambda_ = lambda_
        self.label_space = label_space
        self.rnd_state = rnd_state
        self.verbose = verbose
        self.max_depth = max_depth
        self.online = online

        self.X = None
        self.y = None
        self.label_to_idx = None
        self.rnd_gen = np.random.default_rng(rnd_state)
        self._last_tree = None
        self._online_tree = None
        self._pending_tree = None  # extension scored by the last predict (online)

    def learn_initial_training_set(self, X: NDArray, y: NDArray) -> None:
        """Batch training phase.

        Args:
            X: Training features (n, d)
            y: Training labels (n,) — any hashable values
        """
        X = np.asarray(X)
        y = np.asarray(y)

        if X.shape[0] == 0:
            raise ValueError("Training set cannot be empty")
        if X.shape[0] != y.shape[0]:
            raise ValueError("X and y must have same length")

        self.X = X.copy()
        self.y = y.copy()

        # Infer or validate label space
        if self.label_space is None:
            self.label_space = np.unique(y)
        else:
            self.label_space = np.asarray(self.label_space)

        # Build label-to-index mapping
        self.label_to_idx = {label: idx for idx, label in enumerate(self.label_space)}

        if self.online:
            self._online_tree = MondrianTree.grow(
                self.X, self.rnd_gen,
                lifetime=self.lifetime, feature_weights=self.feature_weights,
            )

        if self.verbose >= 1:
            print(
                f"[MondrianTree] Initialized with {X.shape[0]} training points, "
                f"{X.shape[1]} dimensions, {len(self.label_space)} classes"
            )

    def learn_one(self, x: NDArray, y: Any, precomputed: dict | None = None) -> None:
        """Online learning step.

        Args:
            x: New feature vector (d,) or (1, d)
            y: New label (any hashable value in label_space)
            precomputed: Update dict from ``predict(..., return_update=True)`` (or
                ``compute_p_value(..., return_update=True)``) for this same ``x``.
                In online mode, if it carries the extension that scored ``x`` the
                persisted tree becomes that exact extension, so the realised
                stream is a single consistent Mondrian process (exact online
                validity). Without it — or if it is stale — a fresh ``extend`` is
                drawn. The last ``predict`` is also cached, so the plain
                ``predict(x)`` → ``learn_one(x, y)`` loop reuses it automatically.
        """
        x = np.asarray(x).ravel()

        if self.X is None:
            raise ValueError("Must call learn_initial_training_set first")

        if x.shape[0] != self.X.shape[1]:
            raise ValueError(f"Feature dimension mismatch: got {x.shape[0]}, expected {self.X.shape[1]}")

        if y not in self.label_to_idx:
            # New label encountered
            new_idx = len(self.label_space)
            self.label_space = np.append(self.label_space, y)
            self.label_to_idx[y] = new_idx

        self.X = np.vstack([self.X, x])
        self.y = np.append(self.y, y)
        if self.online:
            cand = precomputed.get("online_tree") if precomputed else None
            if not self._is_consistent_extension(cand, x):
                cand = getattr(self, "_pending_tree", None)
            self._pending_tree = None
            self._online_tree = (
                cand
                if self._is_consistent_extension(cand, x)
                else self._online_tree.extend(x, self.rnd_gen)
            )

    def _is_consistent_extension(self, cand, x) -> bool:
        """True iff ``cand`` is the persisted online tree extended by exactly ``x``.

        Guards reuse of a scored extension: a stale one (the base has advanced
        since it was scored) fails the row-count/last-row check, so a fresh
        ``extend`` is drawn instead.
        """
        return (
            cand is not None
            and self._online_tree is not None
            and cand.X.shape[0] == self._online_tree.X.shape[0] + 1
            and np.array_equal(cand.X[-1], np.asarray(x, dtype=float).ravel())
        )

    def predict(
        self,
        x: NDArray,
        epsilon: float | NDArray | None = None,
        return_p_values: bool = False,
        return_update: bool = False,
    ) -> ConformalPredictionSet | MultiLevelPredictionSet | tuple:
        """Make conformal prediction.

        Rebuilds the tree from the augmented bag [X, y, x, ?] where ? ranges
        over label_space. For each candidate label, computes nonconformity
        scores and p-values.

        Args:
            x: Test feature vector (d,) or (1, d)
            epsilon: Significance level(s). If None, uses self.epsilon.
                     Can be scalar or array-like.
            return_p_values: If True, also return dict of p-values keyed by label.
            return_update: If True, also return dict with tree info for efficiency.

        Returns:
            ConformalPredictionSet: If epsilon is scalar and neither flag is True
            MultiLevelPredictionSet: If epsilon is array-like
            Tuple: If return_p_values or return_update flag is set
        """
        x = np.asarray(x).ravel()

        if self.X is None:
            raise ValueError("Must call learn_initial_training_set first")

        if x.shape[0] != self.X.shape[1]:
            raise ValueError(f"Feature dimension mismatch: got {x.shape[0]}, expected {self.X.shape[1]}")

        if epsilon is None:
            epsilon = self.epsilon

        # Form augmented feature matrix (same for all candidates)
        X_aug = np.vstack([self.X, x])
        n = self.X.shape[0]
        n_total = n + 1
        K = len(self.label_space)

        # Build the Mondrian partition from the augmented bag (unsupervised in X).
        if self.online:
            ext_tree = self._online_tree.extend(x, self.rnd_gen)
            self._pending_tree = ext_tree
            tree = ext_tree.root
        else:
            self._pending_tree = None
            tree = MondrianTree.grow(
                X_aug,
                self.rnd_gen,
                lifetime=self.lifetime,
                x_test=x,
                max_depth=self.max_depth,
                feature_weights=self.feature_weights,
                verbose=self.verbose,
            ).root

        # Find leaf containing test point (index n in X_aug)
        leaf_star = _find_leaf(tree, X_aug[n])

        # Base state: assign training-only counts (test point excluded from counts)
        _assign_counts(tree, self.y, self.label_to_idx, K, n_train=n)

        n_star = leaf_star.n_points()  # leaf size including test point's structural slot
        n_star_train = n_star - 1  # training-only count in leaf_star
        c = leaf_star.counts.copy()  # training-only class counts per class, shape (K,)

        # Leaf-Local Laplace Smoothing
        LAMBDA = self.lambda_
        base_ncm = np.empty(n)
        for leaf in _collect_leaves(tree):
            n_leaf = leaf.n_points()
            has_test = int(n) in leaf.indices
            n_train_leaf = n_leaf - 1 if has_test else n_leaf
            denom = n_train_leaf + K * LAMBDA
            for idx in leaf.indices:
                if idx >= n:
                    continue  # skip the test point's structural slot
                yi = self.label_to_idx[self.y[idx]]
                base_ncm[idx] = 1.0 - (leaf.counts[yi] + LAMBDA) / denom

        # ncm_test[k]: test point's NCM under candidate label k (Laplace)
        denom_test = n_star_train + 1.0 + K * LAMBDA
        ncm_test = 1.0 - (c + 1.0 + LAMBDA) / denom_test  # shape (K,)

        # Vectorised (K, n) broadcast: compare all training NCMs against all ncm_test[k]
        _TOL = 1e-9
        diff = base_ncm[np.newaxis, :] - ncm_test[:, np.newaxis]  # (K, n)
        gt = np.sum(diff > _TOL, axis=1).astype(np.int64)  # (K,)
        eq = np.sum(np.abs(diff) <= _TOL, axis=1).astype(np.int64)  # (K,)

        # Delta correction: O(K²) — only leaf_star training points are affected.
        # Under candidate k, both numerator (if j==k) and denominator change.
        # All c[j] points of class j have identical base and actual NCMs.
        denom_base_star = n_star_train + K * LAMBDA
        for j in range(K):
            if c[j] == 0:
                continue
            base_j = 1.0 - (c[j] + LAMBDA) / denom_base_star
            count_j = int(c[j])
            for k in range(K):
                actual_j = 1.0 - (c[j] + LAMBDA + (1.0 if j == k else 0.0)) / denom_test
                ncm_test_k = ncm_test[k]
                # Classify base_j vs ncm_test_k
                if base_j > ncm_test_k + _TOL:
                    before = 0  # GT
                elif abs(base_j - ncm_test_k) <= _TOL:
                    before = 1  # EQ
                else:
                    before = 2  # LT
                # Classify actual_j vs ncm_test_k
                if actual_j > ncm_test_k + _TOL:
                    after = 0
                elif abs(actual_j - ncm_test_k) <= _TOL:
                    after = 1
                else:
                    after = 2
                if before == after:
                    continue
                if before == 0 and after == 1:
                    gt[k] -= count_j
                    eq[k] += count_j
                elif before == 0 and after == 2:
                    gt[k] -= count_j
                elif before == 1 and after == 2:
                    eq[k] -= count_j
                elif before == 1 and after == 0:
                    eq[k] -= count_j
                    gt[k] += count_j
                elif before == 2 and after == 1:
                    eq[k] += count_j
                elif before == 2 and after == 0:
                    gt[k] += count_j

        eq += 1  # test point always EQ to itself
        tau = self.rnd_gen.uniform(0, 1)
        p_values_arr = (gt + tau * eq) / n_total
        p_values = {label: float(p_values_arr[ki]) for ki, label in enumerate(self.label_space)}

        # Construct prediction set(s)
        result = self._compute_Gamma(p_values, epsilon)

        # Cache the tree for inspection utilities
        self._last_tree = tree

        # Handle return flags
        if return_update and return_p_values:
            return result, p_values, {"tree": tree, "leaf_star": leaf_star, "online_tree": self._pending_tree}
        elif return_p_values:
            return result, p_values
        elif return_update:
            return result, {"tree": tree, "leaf_star": leaf_star, "online_tree": self._pending_tree}
        else:
            return result

    def compute_p_value(self, x: NDArray, y: Any, return_update: bool = False) -> float | tuple[float, dict]:
        """Compute p-value for a single candidate label.

        Args:
            x: Test feature vector (d,)
            y: Candidate label
            return_update: If True, also return dict with tree info

        Returns:
            P-value (float) or (p-value, dict) if return_update=True
        """
        x = np.asarray(x).ravel()

        if y not in self.label_to_idx:
            raise ValueError(f"Label {y} not in label_space")

        # Build augmented dataset
        X_aug = np.vstack([self.X, x])
        n = self.X.shape[0]
        n_total = n + 1
        K = len(self.label_space)
        label_idx = self.label_to_idx[y]

        # Build the Mondrian partition from the augmented bag (unsupervised in X).
        if self.online:
            ext_tree = self._online_tree.extend(x, self.rnd_gen)
            self._pending_tree = ext_tree
            tree = ext_tree.root
        else:
            self._pending_tree = None
            tree = MondrianTree.grow(
                X_aug,
                self.rnd_gen,
                lifetime=self.lifetime,
                x_test=X_aug[n],
                max_depth=self.max_depth,
                feature_weights=self.feature_weights,
                verbose=self.verbose,
            ).root

        # Find leaf and assign training-only counts
        leaf_star = _find_leaf(tree, X_aug[n])
        _assign_counts(tree, self.y, self.label_to_idx, K, n_train=n)

        n_star = leaf_star.n_points()
        n_star_train = n_star - 1
        c = leaf_star.counts.copy()

        # Leaf-Local Laplace Smoothing
        LAMBDA = self.lambda_
        denom_test = n_star_train + 1.0 + K * LAMBDA
        ncm_test_k = 1.0 - (float(c[label_idx]) + 1.0 + LAMBDA) / denom_test

        # Build base_ncm for all training points (Laplace)
        base_ncm = np.empty(n)
        for leaf in _collect_leaves(tree):
            n_leaf = leaf.n_points()
            has_test = int(n) in leaf.indices
            n_train_leaf = n_leaf - 1 if has_test else n_leaf
            denom = n_train_leaf + K * LAMBDA
            for idx in leaf.indices:
                if idx >= n:
                    continue
                yi = self.label_to_idx[self.y[idx]]
                base_ncm[idx] = 1.0 - (leaf.counts[yi] + LAMBDA) / denom

        _TOL = 1e-9
        diff = base_ncm - ncm_test_k
        gt = int(np.sum(diff > _TOL))
        eq = int(np.sum(np.abs(diff) <= _TOL))

        # Delta correction: O(K) — for each class j in leaf_star
        denom_base_star = n_star_train + K * LAMBDA
        for j in range(K):
            if c[j] == 0:
                continue
            base_j = 1.0 - (c[j] + LAMBDA) / denom_base_star
            actual_j = 1.0 - (c[j] + LAMBDA + (1.0 if j == label_idx else 0.0)) / denom_test
            if base_j > ncm_test_k + _TOL:
                before = 0
            elif abs(base_j - ncm_test_k) <= _TOL:
                before = 1
            else:
                before = 2
            if actual_j > ncm_test_k + _TOL:
                after = 0
            elif abs(actual_j - ncm_test_k) <= _TOL:
                after = 1
            else:
                after = 2
            if before == after:
                continue
            count_j = int(c[j])
            if before == 0 and after == 1:
                gt -= count_j
                eq += count_j
            elif before == 0 and after == 2:
                gt -= count_j
            elif before == 1 and after == 2:
                eq -= count_j
            elif before == 1 and after == 0:
                eq -= count_j
                gt += count_j
            elif before == 2 and after == 1:
                eq += count_j
            elif before == 2 and after == 0:
                gt += count_j

        eq += 1  # test point always EQ to itself
        tau = self.rnd_gen.uniform(0, 1)
        p_val = float((gt + tau * eq) / n_total)

        # Cache the tree for inspection utilities
        self._last_tree = tree

        if return_update:
            return p_val, {"tree": tree, "leaf_star": leaf_star, "online_tree": self._pending_tree}
        else:
            return p_val

learn_initial_training_set(X: NDArray, y: NDArray) -> None

Batch training phase.

Args: X: Training features (n, d) y: Training labels (n,) — any hashable values

Source code in src/online_cp/classifiers.py

def learn_initial_training_set(self, X: NDArray, y: NDArray) -> None:
    """Batch training phase.

    Args:
        X: Training features (n, d)
        y: Training labels (n,) — any hashable values
    """
    X = np.asarray(X)
    y = np.asarray(y)

    if X.shape[0] == 0:
        raise ValueError("Training set cannot be empty")
    if X.shape[0] != y.shape[0]:
        raise ValueError("X and y must have same length")

    self.X = X.copy()
    self.y = y.copy()

    # Infer or validate label space
    if self.label_space is None:
        self.label_space = np.unique(y)
    else:
        self.label_space = np.asarray(self.label_space)

    # Build label-to-index mapping
    self.label_to_idx = {label: idx for idx, label in enumerate(self.label_space)}

    if self.online:
        self._online_tree = MondrianTree.grow(
            self.X, self.rnd_gen,
            lifetime=self.lifetime, feature_weights=self.feature_weights,
        )

    if self.verbose >= 1:
        print(
            f"[MondrianTree] Initialized with {X.shape[0]} training points, "
            f"{X.shape[1]} dimensions, {len(self.label_space)} classes"
        )

learn_one(x: NDArray, y: Any, precomputed: dict | None = None) -> None

Online learning step.

Args: x: New feature vector (d,) or (1, d) y: New label (any hashable value in label_space) precomputed: Update dict from predict(..., return_update=True) (or compute_p_value(..., return_update=True)) for this same x. In online mode, if it carries the extension that scored x the persisted tree becomes that exact extension, so the realised stream is a single consistent Mondrian process (exact online validity). Without it — or if it is stale — a fresh extend is drawn. The last predict is also cached, so the plain predict(x) → learn_one(x, y) loop reuses it automatically.

Source code in src/online_cp/classifiers.py

def learn_one(self, x: NDArray, y: Any, precomputed: dict | None = None) -> None:
    """Online learning step.

    Args:
        x: New feature vector (d,) or (1, d)
        y: New label (any hashable value in label_space)
        precomputed: Update dict from ``predict(..., return_update=True)`` (or
            ``compute_p_value(..., return_update=True)``) for this same ``x``.
            In online mode, if it carries the extension that scored ``x`` the
            persisted tree becomes that exact extension, so the realised
            stream is a single consistent Mondrian process (exact online
            validity). Without it — or if it is stale — a fresh ``extend`` is
            drawn. The last ``predict`` is also cached, so the plain
            ``predict(x)`` → ``learn_one(x, y)`` loop reuses it automatically.
    """
    x = np.asarray(x).ravel()

    if self.X is None:
        raise ValueError("Must call learn_initial_training_set first")

    if x.shape[0] != self.X.shape[1]:
        raise ValueError(f"Feature dimension mismatch: got {x.shape[0]}, expected {self.X.shape[1]}")

    if y not in self.label_to_idx:
        # New label encountered
        new_idx = len(self.label_space)
        self.label_space = np.append(self.label_space, y)
        self.label_to_idx[y] = new_idx

    self.X = np.vstack([self.X, x])
    self.y = np.append(self.y, y)
    if self.online:
        cand = precomputed.get("online_tree") if precomputed else None
        if not self._is_consistent_extension(cand, x):
            cand = getattr(self, "_pending_tree", None)
        self._pending_tree = None
        self._online_tree = (
            cand
            if self._is_consistent_extension(cand, x)
            else self._online_tree.extend(x, self.rnd_gen)
        )

predict(x: NDArray, epsilon: float | NDArray | None = None, return_p_values: bool = False, return_update: bool = False) -> ConformalPredictionSet | MultiLevelPredictionSet | tuple

Make conformal prediction.

Rebuilds the tree from the augmented bag [X, y, x, ?] where ? ranges over label_space. For each candidate label, computes nonconformity scores and p-values.

Args: x: Test feature vector (d,) or (1, d) epsilon: Significance level(s). If None, uses self.epsilon. Can be scalar or array-like. return_p_values: If True, also return dict of p-values keyed by label. return_update: If True, also return dict with tree info for efficiency.

Returns: ConformalPredictionSet: If epsilon is scalar and neither flag is True MultiLevelPredictionSet: If epsilon is array-like Tuple: If return_p_values or return_update flag is set

Source code in src/online_cp/classifiers.py

def predict(
    self,
    x: NDArray,
    epsilon: float | NDArray | None = None,
    return_p_values: bool = False,
    return_update: bool = False,
) -> ConformalPredictionSet | MultiLevelPredictionSet | tuple:
    """Make conformal prediction.

    Rebuilds the tree from the augmented bag [X, y, x, ?] where ? ranges
    over label_space. For each candidate label, computes nonconformity
    scores and p-values.

    Args:
        x: Test feature vector (d,) or (1, d)
        epsilon: Significance level(s). If None, uses self.epsilon.
                 Can be scalar or array-like.
        return_p_values: If True, also return dict of p-values keyed by label.
        return_update: If True, also return dict with tree info for efficiency.

    Returns:
        ConformalPredictionSet: If epsilon is scalar and neither flag is True
        MultiLevelPredictionSet: If epsilon is array-like
        Tuple: If return_p_values or return_update flag is set
    """
    x = np.asarray(x).ravel()

    if self.X is None:
        raise ValueError("Must call learn_initial_training_set first")

    if x.shape[0] != self.X.shape[1]:
        raise ValueError(f"Feature dimension mismatch: got {x.shape[0]}, expected {self.X.shape[1]}")

    if epsilon is None:
        epsilon = self.epsilon

    # Form augmented feature matrix (same for all candidates)
    X_aug = np.vstack([self.X, x])
    n = self.X.shape[0]
    n_total = n + 1
    K = len(self.label_space)

    # Build the Mondrian partition from the augmented bag (unsupervised in X).
    if self.online:
        ext_tree = self._online_tree.extend(x, self.rnd_gen)
        self._pending_tree = ext_tree
        tree = ext_tree.root
    else:
        self._pending_tree = None
        tree = MondrianTree.grow(
            X_aug,
            self.rnd_gen,
            lifetime=self.lifetime,
            x_test=x,
            max_depth=self.max_depth,
            feature_weights=self.feature_weights,
            verbose=self.verbose,
        ).root

    # Find leaf containing test point (index n in X_aug)
    leaf_star = _find_leaf(tree, X_aug[n])

    # Base state: assign training-only counts (test point excluded from counts)
    _assign_counts(tree, self.y, self.label_to_idx, K, n_train=n)

    n_star = leaf_star.n_points()  # leaf size including test point's structural slot
    n_star_train = n_star - 1  # training-only count in leaf_star
    c = leaf_star.counts.copy()  # training-only class counts per class, shape (K,)

    # Leaf-Local Laplace Smoothing
    LAMBDA = self.lambda_
    base_ncm = np.empty(n)
    for leaf in _collect_leaves(tree):
        n_leaf = leaf.n_points()
        has_test = int(n) in leaf.indices
        n_train_leaf = n_leaf - 1 if has_test else n_leaf
        denom = n_train_leaf + K * LAMBDA
        for idx in leaf.indices:
            if idx >= n:
                continue  # skip the test point's structural slot
            yi = self.label_to_idx[self.y[idx]]
            base_ncm[idx] = 1.0 - (leaf.counts[yi] + LAMBDA) / denom

    # ncm_test[k]: test point's NCM under candidate label k (Laplace)
    denom_test = n_star_train + 1.0 + K * LAMBDA
    ncm_test = 1.0 - (c + 1.0 + LAMBDA) / denom_test  # shape (K,)

    # Vectorised (K, n) broadcast: compare all training NCMs against all ncm_test[k]
    _TOL = 1e-9
    diff = base_ncm[np.newaxis, :] - ncm_test[:, np.newaxis]  # (K, n)
    gt = np.sum(diff > _TOL, axis=1).astype(np.int64)  # (K,)
    eq = np.sum(np.abs(diff) <= _TOL, axis=1).astype(np.int64)  # (K,)

    # Delta correction: O(K²) — only leaf_star training points are affected.
    # Under candidate k, both numerator (if j==k) and denominator change.
    # All c[j] points of class j have identical base and actual NCMs.
    denom_base_star = n_star_train + K * LAMBDA
    for j in range(K):
        if c[j] == 0:
            continue
        base_j = 1.0 - (c[j] + LAMBDA) / denom_base_star
        count_j = int(c[j])
        for k in range(K):
            actual_j = 1.0 - (c[j] + LAMBDA + (1.0 if j == k else 0.0)) / denom_test
            ncm_test_k = ncm_test[k]
            # Classify base_j vs ncm_test_k
            if base_j > ncm_test_k + _TOL:
                before = 0  # GT
            elif abs(base_j - ncm_test_k) <= _TOL:
                before = 1  # EQ
            else:
                before = 2  # LT
            # Classify actual_j vs ncm_test_k
            if actual_j > ncm_test_k + _TOL:
                after = 0
            elif abs(actual_j - ncm_test_k) <= _TOL:
                after = 1
            else:
                after = 2
            if before == after:
                continue
            if before == 0 and after == 1:
                gt[k] -= count_j
                eq[k] += count_j
            elif before == 0 and after == 2:
                gt[k] -= count_j
            elif before == 1 and after == 2:
                eq[k] -= count_j
            elif before == 1 and after == 0:
                eq[k] -= count_j
                gt[k] += count_j
            elif before == 2 and after == 1:
                eq[k] += count_j
            elif before == 2 and after == 0:
                gt[k] += count_j

    eq += 1  # test point always EQ to itself
    tau = self.rnd_gen.uniform(0, 1)
    p_values_arr = (gt + tau * eq) / n_total
    p_values = {label: float(p_values_arr[ki]) for ki, label in enumerate(self.label_space)}

    # Construct prediction set(s)
    result = self._compute_Gamma(p_values, epsilon)

    # Cache the tree for inspection utilities
    self._last_tree = tree

    # Handle return flags
    if return_update and return_p_values:
        return result, p_values, {"tree": tree, "leaf_star": leaf_star, "online_tree": self._pending_tree}
    elif return_p_values:
        return result, p_values
    elif return_update:
        return result, {"tree": tree, "leaf_star": leaf_star, "online_tree": self._pending_tree}
    else:
        return result

compute_p_value(x: NDArray, y: Any, return_update: bool = False) -> float | tuple[float, dict]

Compute p-value for a single candidate label.

Args: x: Test feature vector (d,) y: Candidate label return_update: If True, also return dict with tree info

Returns: P-value (float) or (p-value, dict) if return_update=True

Source code in src/online_cp/classifiers.py

def compute_p_value(self, x: NDArray, y: Any, return_update: bool = False) -> float | tuple[float, dict]:
    """Compute p-value for a single candidate label.

    Args:
        x: Test feature vector (d,)
        y: Candidate label
        return_update: If True, also return dict with tree info

    Returns:
        P-value (float) or (p-value, dict) if return_update=True
    """
    x = np.asarray(x).ravel()

    if y not in self.label_to_idx:
        raise ValueError(f"Label {y} not in label_space")

    # Build augmented dataset
    X_aug = np.vstack([self.X, x])
    n = self.X.shape[0]
    n_total = n + 1
    K = len(self.label_space)
    label_idx = self.label_to_idx[y]

    # Build the Mondrian partition from the augmented bag (unsupervised in X).
    if self.online:
        ext_tree = self._online_tree.extend(x, self.rnd_gen)
        self._pending_tree = ext_tree
        tree = ext_tree.root
    else:
        self._pending_tree = None
        tree = MondrianTree.grow(
            X_aug,
            self.rnd_gen,
            lifetime=self.lifetime,
            x_test=X_aug[n],
            max_depth=self.max_depth,
            feature_weights=self.feature_weights,
            verbose=self.verbose,
        ).root

    # Find leaf and assign training-only counts
    leaf_star = _find_leaf(tree, X_aug[n])
    _assign_counts(tree, self.y, self.label_to_idx, K, n_train=n)

    n_star = leaf_star.n_points()
    n_star_train = n_star - 1
    c = leaf_star.counts.copy()

    # Leaf-Local Laplace Smoothing
    LAMBDA = self.lambda_
    denom_test = n_star_train + 1.0 + K * LAMBDA
    ncm_test_k = 1.0 - (float(c[label_idx]) + 1.0 + LAMBDA) / denom_test

    # Build base_ncm for all training points (Laplace)
    base_ncm = np.empty(n)
    for leaf in _collect_leaves(tree):
        n_leaf = leaf.n_points()
        has_test = int(n) in leaf.indices
        n_train_leaf = n_leaf - 1 if has_test else n_leaf
        denom = n_train_leaf + K * LAMBDA
        for idx in leaf.indices:
            if idx >= n:
                continue
            yi = self.label_to_idx[self.y[idx]]
            base_ncm[idx] = 1.0 - (leaf.counts[yi] + LAMBDA) / denom

    _TOL = 1e-9
    diff = base_ncm - ncm_test_k
    gt = int(np.sum(diff > _TOL))
    eq = int(np.sum(np.abs(diff) <= _TOL))

    # Delta correction: O(K) — for each class j in leaf_star
    denom_base_star = n_star_train + K * LAMBDA
    for j in range(K):
        if c[j] == 0:
            continue
        base_j = 1.0 - (c[j] + LAMBDA) / denom_base_star
        actual_j = 1.0 - (c[j] + LAMBDA + (1.0 if j == label_idx else 0.0)) / denom_test
        if base_j > ncm_test_k + _TOL:
            before = 0
        elif abs(base_j - ncm_test_k) <= _TOL:
            before = 1
        else:
            before = 2
        if actual_j > ncm_test_k + _TOL:
            after = 0
        elif abs(actual_j - ncm_test_k) <= _TOL:
            after = 1
        else:
            after = 2
        if before == after:
            continue
        count_j = int(c[j])
        if before == 0 and after == 1:
            gt -= count_j
            eq += count_j
        elif before == 0 and after == 2:
            gt -= count_j
        elif before == 1 and after == 2:
            eq -= count_j
        elif before == 1 and after == 0:
            eq -= count_j
            gt += count_j
        elif before == 2 and after == 1:
            eq += count_j
        elif before == 2 and after == 0:
            gt += count_j

    eq += 1  # test point always EQ to itself
    tau = self.rnd_gen.uniform(0, 1)
    p_val = float((gt + tau * eq) / n_total)

    # Cache the tree for inspection utilities
    self._last_tree = tree

    if return_update:
        return p_val, {"tree": tree, "leaf_star": leaf_star, "online_tree": self._pending_tree}
    else:
        return p_val

online_cp.classifiers.ConformalMondrianForestClassifier

Bases: _MondrianClassifierInspection, ConformalClassifier

Conformal predictor using Mondrian forests (ensembles of Mondrian trees).

A full (transductive) conformal predictor that builds T Mondrian trees per prediction step from the same augmented bag but with different random seeds. The ensemble averages per-tree leaf posteriors, reducing variance while maintaining exact validity (average of invariant predictors is invariant).

The nonconformity measure is (Leaf-Local Laplace, averaged):

α_i = 1 - (1/T) Σ_t (C_{leaf_t(x_i), y_i} + λ) / (N_leaf_t(x_i) + K * λ)

where the sum averages the Laplace-smoothed class posterior for label y_i across all T trees. This smooths predictions, reduces discrete ties, and produces smaller (more efficient) prediction sets than a single tree.

Parameters:

Name	Type	Description	Default
n_trees	int	Number of Mondrian trees in the forest. Default: 10. Larger values reduce variance but increase computation per predict.	10
lifetime	float	Depth budget L for each tree (see ConformalMondrianTreeClassifier). Default: 1.0.	1.0
lambda_	float	Laplace smoothing parameter λ (see ConformalMondrianTreeClassifier). Default: 1.0.	1.0
label_space	array - like	Set of possible labels. If None, inferred from training data.	None
epsilon	float	Default significance level for predictions.	0.1
rnd_state	int	Seed for reproducible forest construction.	None
verbose	int	Verbosity level (0 = silent).	0

Attributes:

Name	Type	Description
X	ndarray of shape (n, d)	Training feature array.
y	ndarray of shape (n,)	Training labels.
label_space	ndarray	Sorted unique labels.
label_to_idx	dict	Mapping from label to integer index 0..K-1.
rnd_gen	Generator	Root random generator instance.

Source code in src/online_cp/classifiers.py

class ConformalMondrianForestClassifier(_MondrianClassifierInspection, ConformalClassifier):
    """Conformal predictor using Mondrian forests (ensembles of Mondrian trees).

    A full (transductive) conformal predictor that builds T Mondrian trees per
    prediction step from the same augmented bag but with different random seeds.
    The ensemble averages per-tree leaf posteriors, reducing variance while
    maintaining exact validity (average of invariant predictors is invariant).

    The nonconformity measure is (Leaf-Local Laplace, averaged):

        α_i = 1 - (1/T) Σ_t (C_{leaf_t(x_i), y_i} + λ) / (N_leaf_t(x_i) + K * λ)

    where the sum averages the Laplace-smoothed class posterior for label y_i
    across all T trees. This smooths predictions, reduces discrete ties, and
    produces smaller (more efficient) prediction sets than a single tree.

    Parameters
    ----------
    n_trees : int
        Number of Mondrian trees in the forest. Default: 10. Larger values
        reduce variance but increase computation per predict.
    lifetime : float
        Depth budget L for each tree (see `ConformalMondrianTreeClassifier`).
        Default: 1.0.
    lambda_ : float
        Laplace smoothing parameter λ (see `ConformalMondrianTreeClassifier`).
        Default: 1.0.
    label_space : array-like, optional
        Set of possible labels. If None, inferred from training data.
    epsilon : float
        Default significance level for predictions.
    rnd_state : int, optional
        Seed for reproducible forest construction.
    verbose : int
        Verbosity level (0 = silent).

    Attributes
    ----------
    X : ndarray of shape (n, d)
        Training feature array.
    y : ndarray of shape (n,)
        Training labels.
    label_space : ndarray
        Sorted unique labels.
    label_to_idx : dict
        Mapping from label to integer index 0..K-1.
    rnd_gen : numpy.random.Generator
        Root random generator instance.
    """

    _SAVE_PARAMS: tuple = (
        "n_trees", "lifetime", "lambda_", "label_space", "epsilon",
        "rnd_state", "verbose", "n_jobs", "max_depth", "online",
    )
    _SAVE_STATE: tuple = ("X", "y", "label_to_idx", "label_space", "_forest")

    def __init__(
        self,
        n_trees: int = 10,
        lifetime: float = 1.0,
        lambda_: float = 1.0,
        label_space: NDArray | None = None,
        epsilon: float = 0.1,
        rnd_state: int | None = None,
        verbose: int = 0,
        n_jobs: int = 1,
        max_depth: int | None = None,
        online: bool = False,
    ):
        if max_depth is not None and (not isinstance(max_depth, int) or max_depth < 0):
            raise ValueError("max_depth must be a non-negative integer or None")
        if online:
            if not isinstance(lifetime, (int, float)):
                raise ValueError(
                    "online=True requires a fixed float lifetime; string/auto-tuned "
                    "lifetimes recompute a master tree per step and are batch-only."
                )
            if max_depth is not None:
                raise ValueError(
                    "online=True requires max_depth=None (the projective Mondrian regime)."
                )
        super().__init__(epsilon=epsilon)
        self.n_trees = n_trees
        self.lifetime = lifetime
        self.lambda_ = lambda_
        self.label_space = label_space
        self.rnd_state = rnd_state
        self.verbose = verbose
        self.n_jobs = n_jobs
        self.max_depth = max_depth
        self.online = online

        self.X = None
        self.y = None
        self.label_to_idx = None
        self.rnd_gen = np.random.default_rng(rnd_state)
        self._last_tree = None
        self._last_seeds = None
        self._last_x = None
        self._forest = None
        self._pending_forest = None

    def learn_initial_training_set(self, X: NDArray, y: NDArray) -> None:
        """Batch training phase.

        Args:
            X: Training features (n, d)
            y: Training labels (n,) — any hashable values
        """
        X = np.asarray(X)
        y = np.asarray(y)

        if X.shape[0] == 0:
            raise ValueError("Training set cannot be empty")
        if X.shape[0] != y.shape[0]:
            raise ValueError("X and y must have same length")

        self.X = X.copy()
        self.y = y.copy()

        # Infer or validate label space
        if self.label_space is None:
            self.label_space = np.unique(y)
        else:
            self.label_space = np.asarray(self.label_space)

        # Build label-to-index mapping
        self.label_to_idx = {label: idx for idx, label in enumerate(self.label_space)}

        if self.online:
            seeds = self.rnd_gen.integers(0, 2**31, self.n_trees)
            self._forest = [
                MondrianTree.grow(self.X, np.random.default_rng(int(s)), lifetime=self.lifetime)
                for s in seeds
            ]

        if self.verbose >= 1:
            print(
                f"[MondrianForest] Initialized with {X.shape[0]} training points, "
                f"{X.shape[1]} dimensions, {len(self.label_space)} classes, {self.n_trees} trees"
            )

    def learn_one(self, x: NDArray, y: Any, precomputed: dict | None = None) -> None:
        """Online learning step.

        Args:
            x: New feature vector (d,) or (1, d)
            y: New label (any hashable value in label_space)
            precomputed: Update dict from ``predict(..., return_update=True)`` for
                this same ``x``. In online mode, if it carries the forest
                extensions that scored ``x`` they become the persisted forest, so
                each tree's stream is a single consistent Mondrian process (exact
                online validity). The last ``predict`` is also cached, so the plain
                ``predict(x)`` → ``learn_one(x, y)`` loop reuses it automatically.
        """
        x = np.asarray(x).ravel()

        if self.X is None:
            raise ValueError("Must call learn_initial_training_set first")

        if x.shape[0] != self.X.shape[1]:
            raise ValueError(f"Feature dimension mismatch: got {x.shape[0]}, expected {self.X.shape[1]}")

        if y not in self.label_to_idx:
            # New label encountered
            new_idx = len(self.label_space)
            self.label_space = np.append(self.label_space, y)
            self.label_to_idx[y] = new_idx

        self.X = np.vstack([self.X, x])
        self.y = np.append(self.y, y)
        if self.online:
            cand = precomputed.get("online_forest") if precomputed else None
            if not self._is_consistent_forest(cand, x):
                cand = getattr(self, "_pending_forest", None)
            self._pending_forest = None
            if self._is_consistent_forest(cand, x):
                self._forest = cand
            else:
                self._forest = [t.extend(x, self.rnd_gen) for t in self._forest]

    def _is_consistent_forest(self, cand, x) -> bool:
        """True iff ``cand`` is the persisted forest with every tree extended by exactly ``x``."""
        if cand is None or self._forest is None or len(cand) != len(self._forest):
            return False
        xr = np.asarray(x, dtype=float).ravel()
        base_n = self._forest[0].X.shape[0]
        return all(
            t.X.shape[0] == base_n + 1 and np.array_equal(t.X[-1], xr) for t in cand
        )

    def _forest_summaries(self, X_aug, x, n, K):
        """Per-tree ``(n_leaves, counts, ls_train_idx)`` summaries + viz/reuse caches.

        Batch: T fresh trees from seeds (joblib). Online: extend each persistent
        tree and summarise, caching the extensions for ``learn_one`` reuse.
        """
        if self.online:
            pending = [t.extend(x, self.rnd_gen) for t in self._forest]
            self._pending_forest = pending
            self._last_seeds = None
            self._last_x = x.copy()
            summaries = [
                _summarize_tree(t.root, self.y, self.label_to_idx, K, n) for t in pending
            ]
            self._last_tree = pending[0].root
            return summaries
        seeds = self.rnd_gen.integers(0, 2**31, self.n_trees)
        self._last_seeds = seeds.copy()
        self._last_x = x.copy()
        self._pending_forest = None
        _rng0 = np.random.default_rng(int(seeds[0]))
        _tree0 = MondrianTree.grow(X_aug, _rng0, lifetime=self.lifetime, max_depth=self.max_depth).root
        _assign_counts(_tree0, self.y, self.label_to_idx, K, n_train=n)
        self._last_tree = _tree0
        return Parallel(n_jobs=self.n_jobs)(
            delayed(_build_tree_summary)(
                int(seed), X_aug, self.y, self.label_to_idx, K, n, self.lifetime, self.max_depth
            )
            for seed in seeds
        )

    # ------------------------------------------------------------------
    # Forest indexing
    # ------------------------------------------------------------------

    def __len__(self) -> int:
        """Number of trees in the forest."""
        return self.n_trees

    def __iter__(self):
        """Iterate over individual tree views (calls predict first if needed)."""
        for i in range(self.n_trees):
            yield self[i]

    def __getitem__(self, i: int) -> ConformalMondrianTreeClassifier:
        """Return a single-tree view for tree ``i`` from the last predict/compute_p_value call.

        The returned object is a :class:`ConformalMondrianTreeClassifier` whose
        ``_last_tree`` is the exact i-th tree from the last call, enabling all
        T1/T2 inspection and visualisation methods (".summary", ".to_dataframe()",
        ".draw()", etc.). Its own ``predict()`` would rebuild a fresh tree (not
        tree i) — this object is intended for *inspection only*.

        Parameters
        ----------
        i : int
            Tree index, ``0 <= i < n_trees``.

        Returns
        -------
        ConformalMondrianTreeClassifier
            Single-tree view with ``_last_tree`` set to tree i.

        Raises
        ------
        RuntimeError
            If ``predict()`` or ``compute_p_value()`` has not been called yet.
        IndexError
            If ``i`` is out of range.
        """
        if self.online:
            pending = getattr(self, "_pending_forest", None)
            if pending is None:
                raise RuntimeError(
                    "No trees cached. Call predict() or compute_p_value() first."
                )
            if not (0 <= i < self.n_trees):
                raise IndexError(f"Tree index {i} out of range [0, {self.n_trees}).")
            n = self.X.shape[0]
            K = len(self.label_space)
            tree_i = pending[i].root
            _assign_counts(tree_i, self.y, self.label_to_idx, K, n_train=n)
            view = ConformalMondrianTreeClassifier(
                lifetime=self.lifetime,
                lambda_=self.lambda_,
                label_space=self.label_space.copy(),
                epsilon=self.epsilon,
                max_depth=self.max_depth,
            )
            view.X = self.X
            view.y = self.y
            view.label_to_idx = self.label_to_idx
            view._last_tree = tree_i
            return view

        if self._last_seeds is None:
            raise RuntimeError(
                "No trees cached. Call predict() or compute_p_value() first."
            )
        if not (0 <= i < self.n_trees):
            raise IndexError(f"Tree index {i} out of range [0, {self.n_trees}).")

        n = self.X.shape[0]
        X_aug = np.vstack([self.X, self._last_x])
        K = len(self.label_space)
        rng_i = np.random.default_rng(int(self._last_seeds[i]))
        tree_i = MondrianTree.grow(X_aug, rng_i, lifetime=self.lifetime, max_depth=self.max_depth).root
        _assign_counts(tree_i, self.y, self.label_to_idx, K, n_train=n)

        view = ConformalMondrianTreeClassifier(
            lifetime=self.lifetime,
            lambda_=self.lambda_,
            label_space=self.label_space.copy(),
            epsilon=self.epsilon,
            max_depth=self.max_depth,
        )
        view.X = self.X
        view.y = self.y
        view.label_to_idx = self.label_to_idx
        view._last_tree = tree_i
        return view

    def predict(
        self,
        x: NDArray,
        epsilon: float | NDArray | None = None,
        return_p_values: bool = False,
        return_update: bool = False,
    ) -> ConformalPredictionSet | MultiLevelPredictionSet | tuple:
        """Make conformal prediction using ensemble of Mondrian trees.

        For each candidate label, builds T trees from the augmented bag and
        computes the average leaf posterior across all trees. The NCM is then
        1 - avg_posterior, which is used to rank all n+1 points and compute
        the p-value.

        Args:
            x: Test feature vector (d,) or (1, d)
            epsilon: Significance level(s). If None, uses self.epsilon.
                     Can be scalar or array-like.
            return_p_values: If True, also return dict of p-values keyed by label.
            return_update: If True, also return dict with forest info.

        Returns:
            ConformalPredictionSet: If epsilon is scalar and neither flag is True
            MultiLevelPredictionSet: If epsilon is array-like
            Tuple: If return_p_values or return_update flag is set
        """
        x = np.asarray(x).ravel()

        if self.X is None:
            raise ValueError("Must call learn_initial_training_set first")

        if x.shape[0] != self.X.shape[1]:
            raise ValueError(f"Feature dimension mismatch: got {x.shape[0]}, expected {self.X.shape[1]}")

        if epsilon is None:
            epsilon = self.epsilon

        # Form augmented feature matrix (same for all candidates and all trees)
        X_aug = np.vstack([self.X, x])
        n = self.X.shape[0]
        n_total = n + 1
        K = len(self.label_space)

        # Build T tree summaries (batch: fresh trees from seeds; online: extend
        # the persistent forest). Caches viz/reuse state as a side effect.
        tree_summaries = self._forest_summaries(X_aug, x, n, K)
        n_leaves_all = [ts[0] for ts in tree_summaries]  # T x (n_total,)
        counts_all = [ts[1] for ts in tree_summaries]  # T x (n_total, K)
        ls_train_idx = [ts[2] for ts in tree_summaries]  # T x variable

        # Integer label indices for all training points
        LAMBDA = self.lambda_
        y_idx = np.array([self.label_to_idx[self.y[i]] for i in range(n)], dtype=np.int64)
        arange_n = np.arange(n)

        # base_ncm[i] = 1 - (1/T) * sum_t (counts_t[i, y_idx[i]] + λ) / (n_train_leaf_t[i] + K*λ)
        base_ncm_sum = np.zeros(n)
        for t in range(self.n_trees):
            nl = n_leaves_all[t][:n].astype(np.float64)  # (n,) structural sizes
            ct = counts_all[t][:n]  # (n, K)
            # leaf_star points have structural size including test slot
            n_train_leaf = nl.copy()
            for idx in ls_train_idx[t]:
                n_train_leaf[int(idx)] -= 1.0
            denom = n_train_leaf + K * LAMBDA
            base_ncm_sum += (ct[arange_n, y_idx].astype(np.float64) + LAMBDA) / denom
        base_ncm = 1.0 - base_ncm_sum / self.n_trees

        # ncm_test[k] = 1 - (1/T) * sum_t (counts_t[n, k] + 1 + λ) / (n_star_train_t + 1 + K*λ)
        ncm_test = np.zeros(K)
        for t in range(self.n_trees):
            n_star_train_t = float(n_leaves_all[t][n]) - 1.0
            denom_test_t = n_star_train_t + 1.0 + K * LAMBDA
            ncm_test += (counts_all[t][n].astype(np.float64) + 1.0 + LAMBDA) / denom_test_t
        ncm_test = 1.0 - ncm_test / self.n_trees  # shape (K,)

        # Vectorised (K, n) broadcast
        _TOL = 1e-9
        diff = base_ncm[np.newaxis, :] - ncm_test[:, np.newaxis]  # (K, n)
        gt = np.sum(diff > _TOL, axis=1).astype(np.int64)
        eq = np.sum(np.abs(diff) <= _TOL, axis=1).astype(np.int64)

        # Delta correction (Laplace): for each tree t, compute per-class delta
        # in leaf_star_t, then accumulate per-point adjustments.
        # Build per-tree delta matrices: delta_jk[j, k] = actual_prob - base_prob
        tree_deltas = []
        for t in range(self.n_trees):
            n_star_t = float(n_leaves_all[t][n])
            c_t = counts_all[t][n].astype(np.float64)
            denom_base_t = (n_star_t - 1.0) + K * LAMBDA
            denom_actual_t = n_star_t + K * LAMBDA
            delta_jk = np.empty((K, K))
            for j in range(K):
                base_p = (c_t[j] + LAMBDA) / denom_base_t
                for k in range(K):
                    actual_p = (c_t[j] + LAMBDA + (1.0 if j == k else 0.0)) / denom_actual_t
                    delta_jk[j, k] = actual_p - base_p
            tree_deltas.append(delta_jk)

        # Build point -> trees membership
        point_trees: dict[int, list] = {}
        for t in range(self.n_trees):
            for idx in ls_train_idx[t]:
                i_int = int(idx)
                if i_int not in point_trees:
                    point_trees[i_int] = []
                point_trees[i_int].append(t)

        # Reclassify affected points
        for i_int, trees_list in point_trees.items():
            j = int(y_idx[i_int])
            delta_sum_k = np.zeros(K)
            for t in trees_list:
                delta_sum_k += tree_deltas[t][j, :]
            delta_avg_k = delta_sum_k / self.n_trees

            base_val = base_ncm[i_int]
            for k in range(K):
                actual_val = base_val - delta_avg_k[k]
                ncm_test_k = ncm_test[k]
                # Before bucket
                if base_val > ncm_test_k + _TOL:
                    before = 0
                elif abs(base_val - ncm_test_k) <= _TOL:
                    before = 1
                else:
                    before = 2
                # After bucket
                if actual_val > ncm_test_k + _TOL:
                    after = 0
                elif abs(actual_val - ncm_test_k) <= _TOL:
                    after = 1
                else:
                    after = 2
                if before == after:
                    continue
                if before == 0 and after == 1:
                    gt[k] -= 1
                    eq[k] += 1
                elif before == 0 and after == 2:
                    gt[k] -= 1
                elif before == 1 and after == 2:
                    eq[k] -= 1
                elif before == 1 and after == 0:
                    eq[k] -= 1
                    gt[k] += 1
                elif before == 2 and after == 1:
                    eq[k] += 1
                elif before == 2 and after == 0:
                    gt[k] += 1

        eq += 1  # test point always EQ to itself
        tau = self.rnd_gen.uniform(0, 1)
        p_values_arr = (gt + tau * eq) / n_total
        p_values = {label: float(p_values_arr[ki]) for ki, label in enumerate(self.label_space)}

        result = self._compute_Gamma(p_values, epsilon)

        # Handle return flags
        _update = {"n_leaves_all": n_leaves_all, "counts_all": counts_all, "leaf_star_train_indices_all": ls_train_idx, "online_forest": self._pending_forest}
        if return_update and return_p_values:
            return result, p_values, _update
        elif return_p_values:
            return result, p_values
        elif return_update:
            return result, _update
        else:
            return result

    def compute_p_value(self, x: NDArray, y: Any, return_update: bool = False) -> float | tuple[float, dict]:
        """Compute p-value for a single candidate label.

        Args:
            x: Test feature vector (d,)
            y: Candidate label
            return_update: If True, also return dict with forest info

        Returns:
            P-value (float) or (p-value, dict) if return_update=True
        """
        x = np.asarray(x).ravel()

        if y not in self.label_to_idx:
            raise ValueError(f"Label {y} not in label_space")

        # Build augmented dataset
        X_aug = np.vstack([self.X, x])
        n = self.X.shape[0]
        n_total = n + 1
        K = len(self.label_space)
        label_idx = self.label_to_idx[y]

        # Build T tree summaries (possibly in parallel)
        tree_summaries = self._forest_summaries(X_aug, x, n, K)
        n_leaves_all = [ts[0] for ts in tree_summaries]
        counts_all = [ts[1] for ts in tree_summaries]
        ls_train_idx = [ts[2] for ts in tree_summaries]

        LAMBDA = self.lambda_
        y_idx = np.array([self.label_to_idx[self.y[i]] for i in range(n)], dtype=np.int64)
        arange_n = np.arange(n)

        # base_ncm for training points (Laplace)
        base_ncm_sum = np.zeros(n)
        for t in range(self.n_trees):
            nl = n_leaves_all[t][:n].astype(np.float64)
            ct = counts_all[t][:n]
            n_train_leaf = nl.copy()
            for idx in ls_train_idx[t]:
                n_train_leaf[int(idx)] -= 1.0
            denom = n_train_leaf + K * LAMBDA
            base_ncm_sum += (ct[arange_n, y_idx].astype(np.float64) + LAMBDA) / denom
        base_ncm = 1.0 - base_ncm_sum / self.n_trees

        # ncm_test for this single candidate (Laplace)
        ncm_test_k = 0.0
        for t in range(self.n_trees):
            n_star_train_t = float(n_leaves_all[t][n]) - 1.0
            denom_test_t = n_star_train_t + 1.0 + K * LAMBDA
            ncm_test_k += (float(counts_all[t][n, label_idx]) + 1.0 + LAMBDA) / denom_test_t
        ncm_test_k = 1.0 - ncm_test_k / self.n_trees

        _TOL = 1e-9
        diff = base_ncm - ncm_test_k
        gt = int(np.sum(diff > _TOL))
        eq = int(np.sum(np.abs(diff) <= _TOL))

        # Delta correction (Laplace): per-tree delta for leaf_star points
        tree_deltas = []
        for t in range(self.n_trees):
            n_star_t = float(n_leaves_all[t][n])
            c_t = counts_all[t][n].astype(np.float64)
            denom_base_t = (n_star_t - 1.0) + K * LAMBDA
            denom_actual_t = n_star_t + K * LAMBDA
            # delta[j] = actual_prob_j - base_prob_j for candidate label_idx
            delta_j = np.empty(K)
            for j in range(K):
                base_p = (c_t[j] + LAMBDA) / denom_base_t
                actual_p = (c_t[j] + LAMBDA + (1.0 if j == label_idx else 0.0)) / denom_actual_t
                delta_j[j] = actual_p - base_p
            tree_deltas.append(delta_j)

        # Build point -> trees membership
        point_trees: dict[int, list] = {}
        for t in range(self.n_trees):
            for idx in ls_train_idx[t]:
                i_int = int(idx)
                if i_int not in point_trees:
                    point_trees[i_int] = []
                point_trees[i_int].append(t)

        # Reclassify affected points
        for i_int, trees_list in point_trees.items():
            j = int(y_idx[i_int])
            delta_sum = 0.0
            for t in trees_list:
                delta_sum += tree_deltas[t][j]
            delta_avg = delta_sum / self.n_trees

            base_val = base_ncm[i_int]
            actual_val = base_val - delta_avg
            # Before bucket
            if base_val > ncm_test_k + _TOL:
                before = 0
            elif abs(base_val - ncm_test_k) <= _TOL:
                before = 1
            else:
                before = 2
            # After bucket
            if actual_val > ncm_test_k + _TOL:
                after = 0
            elif abs(actual_val - ncm_test_k) <= _TOL:
                after = 1
            else:
                after = 2
            if before == after:
                continue
            if before == 0 and after == 1:
                gt -= 1
                eq += 1
            elif before == 0 and after == 2:
                gt -= 1
            elif before == 1 and after == 2:
                eq -= 1
            elif before == 1 and after == 0:
                eq -= 1
                gt += 1
            elif before == 2 and after == 1:
                eq += 1
            elif before == 2 and after == 0:
                gt += 1

        eq += 1  # test point always EQ to itself
        tau = self.rnd_gen.uniform(0, 1)
        p_val = float((gt + tau * eq) / n_total)

        if return_update:
            return p_val, {
                "n_leaves_all": n_leaves_all,
                "counts_all": counts_all,
                "leaf_star_train_indices_all": ls_train_idx,
                "online_forest": self._pending_forest,
            }
        else:
            return p_val

learn_initial_training_set(X: NDArray, y: NDArray) -> None

Batch training phase.

Args: X: Training features (n, d) y: Training labels (n,) — any hashable values

Source code in src/online_cp/classifiers.py

def learn_initial_training_set(self, X: NDArray, y: NDArray) -> None:
    """Batch training phase.

    Args:
        X: Training features (n, d)
        y: Training labels (n,) — any hashable values
    """
    X = np.asarray(X)
    y = np.asarray(y)

    if X.shape[0] == 0:
        raise ValueError("Training set cannot be empty")
    if X.shape[0] != y.shape[0]:
        raise ValueError("X and y must have same length")

    self.X = X.copy()
    self.y = y.copy()

    # Infer or validate label space
    if self.label_space is None:
        self.label_space = np.unique(y)
    else:
        self.label_space = np.asarray(self.label_space)

    # Build label-to-index mapping
    self.label_to_idx = {label: idx for idx, label in enumerate(self.label_space)}

    if self.online:
        seeds = self.rnd_gen.integers(0, 2**31, self.n_trees)
        self._forest = [
            MondrianTree.grow(self.X, np.random.default_rng(int(s)), lifetime=self.lifetime)
            for s in seeds
        ]

    if self.verbose >= 1:
        print(
            f"[MondrianForest] Initialized with {X.shape[0]} training points, "
            f"{X.shape[1]} dimensions, {len(self.label_space)} classes, {self.n_trees} trees"
        )

learn_one(x: NDArray, y: Any, precomputed: dict | None = None) -> None

Online learning step.

Args: x: New feature vector (d,) or (1, d) y: New label (any hashable value in label_space) precomputed: Update dict from predict(..., return_update=True) for this same x. In online mode, if it carries the forest extensions that scored x they become the persisted forest, so each tree's stream is a single consistent Mondrian process (exact online validity). The last predict is also cached, so the plain predict(x) → learn_one(x, y) loop reuses it automatically.

Source code in src/online_cp/classifiers.py

def learn_one(self, x: NDArray, y: Any, precomputed: dict | None = None) -> None:
    """Online learning step.

    Args:
        x: New feature vector (d,) or (1, d)
        y: New label (any hashable value in label_space)
        precomputed: Update dict from ``predict(..., return_update=True)`` for
            this same ``x``. In online mode, if it carries the forest
            extensions that scored ``x`` they become the persisted forest, so
            each tree's stream is a single consistent Mondrian process (exact
            online validity). The last ``predict`` is also cached, so the plain
            ``predict(x)`` → ``learn_one(x, y)`` loop reuses it automatically.
    """
    x = np.asarray(x).ravel()

    if self.X is None:
        raise ValueError("Must call learn_initial_training_set first")

    if x.shape[0] != self.X.shape[1]:
        raise ValueError(f"Feature dimension mismatch: got {x.shape[0]}, expected {self.X.shape[1]}")

    if y not in self.label_to_idx:
        # New label encountered
        new_idx = len(self.label_space)
        self.label_space = np.append(self.label_space, y)
        self.label_to_idx[y] = new_idx

    self.X = np.vstack([self.X, x])
    self.y = np.append(self.y, y)
    if self.online:
        cand = precomputed.get("online_forest") if precomputed else None
        if not self._is_consistent_forest(cand, x):
            cand = getattr(self, "_pending_forest", None)
        self._pending_forest = None
        if self._is_consistent_forest(cand, x):
            self._forest = cand
        else:
            self._forest = [t.extend(x, self.rnd_gen) for t in self._forest]

__len__() -> int

Number of trees in the forest.

Source code in src/online_cp/classifiers.py

def __len__(self) -> int:
    """Number of trees in the forest."""
    return self.n_trees

__iter__()

Iterate over individual tree views (calls predict first if needed).

Source code in src/online_cp/classifiers.py

def __iter__(self):
    """Iterate over individual tree views (calls predict first if needed)."""
    for i in range(self.n_trees):
        yield self[i]

__getitem__(i: int) -> ConformalMondrianTreeClassifier

Return a single-tree view for tree i from the last predict/compute_p_value call.

The returned object is a :class:ConformalMondrianTreeClassifier whose _last_tree is the exact i-th tree from the last call, enabling all T1/T2 inspection and visualisation methods (".summary", ".to_dataframe()", ".draw()", etc.). Its own predict() would rebuild a fresh tree (not tree i) — this object is intended for inspection only.

Parameters:

Name	Type	Description	Default
i	int	Tree index, 0 <= i < n_trees.	required

Returns:

Type	Description
ConformalMondrianTreeClassifier	Single-tree view with _last_tree set to tree i.

Raises:

Type	Description
RuntimeError	If predict() or compute_p_value() has not been called yet.
IndexError	If i is out of range.

Source code in src/online_cp/classifiers.py

def __getitem__(self, i: int) -> ConformalMondrianTreeClassifier:
    """Return a single-tree view for tree ``i`` from the last predict/compute_p_value call.

    The returned object is a :class:`ConformalMondrianTreeClassifier` whose
    ``_last_tree`` is the exact i-th tree from the last call, enabling all
    T1/T2 inspection and visualisation methods (".summary", ".to_dataframe()",
    ".draw()", etc.). Its own ``predict()`` would rebuild a fresh tree (not
    tree i) — this object is intended for *inspection only*.

    Parameters
    ----------
    i : int
        Tree index, ``0 <= i < n_trees``.

    Returns
    -------
    ConformalMondrianTreeClassifier
        Single-tree view with ``_last_tree`` set to tree i.

    Raises
    ------
    RuntimeError
        If ``predict()`` or ``compute_p_value()`` has not been called yet.
    IndexError
        If ``i`` is out of range.
    """
    if self.online:
        pending = getattr(self, "_pending_forest", None)
        if pending is None:
            raise RuntimeError(
                "No trees cached. Call predict() or compute_p_value() first."
            )
        if not (0 <= i < self.n_trees):
            raise IndexError(f"Tree index {i} out of range [0, {self.n_trees}).")
        n = self.X.shape[0]
        K = len(self.label_space)
        tree_i = pending[i].root
        _assign_counts(tree_i, self.y, self.label_to_idx, K, n_train=n)
        view = ConformalMondrianTreeClassifier(
            lifetime=self.lifetime,
            lambda_=self.lambda_,
            label_space=self.label_space.copy(),
            epsilon=self.epsilon,
            max_depth=self.max_depth,
        )
        view.X = self.X
        view.y = self.y
        view.label_to_idx = self.label_to_idx
        view._last_tree = tree_i
        return view

    if self._last_seeds is None:
        raise RuntimeError(
            "No trees cached. Call predict() or compute_p_value() first."
        )
    if not (0 <= i < self.n_trees):
        raise IndexError(f"Tree index {i} out of range [0, {self.n_trees}).")

    n = self.X.shape[0]
    X_aug = np.vstack([self.X, self._last_x])
    K = len(self.label_space)
    rng_i = np.random.default_rng(int(self._last_seeds[i]))
    tree_i = MondrianTree.grow(X_aug, rng_i, lifetime=self.lifetime, max_depth=self.max_depth).root
    _assign_counts(tree_i, self.y, self.label_to_idx, K, n_train=n)

    view = ConformalMondrianTreeClassifier(
        lifetime=self.lifetime,
        lambda_=self.lambda_,
        label_space=self.label_space.copy(),
        epsilon=self.epsilon,
        max_depth=self.max_depth,
    )
    view.X = self.X
    view.y = self.y
    view.label_to_idx = self.label_to_idx
    view._last_tree = tree_i
    return view

predict(x: NDArray, epsilon: float | NDArray | None = None, return_p_values: bool = False, return_update: bool = False) -> ConformalPredictionSet | MultiLevelPredictionSet | tuple

Make conformal prediction using ensemble of Mondrian trees.

For each candidate label, builds T trees from the augmented bag and computes the average leaf posterior across all trees. The NCM is then 1 - avg_posterior, which is used to rank all n+1 points and compute the p-value.

Args: x: Test feature vector (d,) or (1, d) epsilon: Significance level(s). If None, uses self.epsilon. Can be scalar or array-like. return_p_values: If True, also return dict of p-values keyed by label. return_update: If True, also return dict with forest info.

Returns: ConformalPredictionSet: If epsilon is scalar and neither flag is True MultiLevelPredictionSet: If epsilon is array-like Tuple: If return_p_values or return_update flag is set

Source code in src/online_cp/classifiers.py

def predict(
    self,
    x: NDArray,
    epsilon: float | NDArray | None = None,
    return_p_values: bool = False,
    return_update: bool = False,
) -> ConformalPredictionSet | MultiLevelPredictionSet | tuple:
    """Make conformal prediction using ensemble of Mondrian trees.

    For each candidate label, builds T trees from the augmented bag and
    computes the average leaf posterior across all trees. The NCM is then
    1 - avg_posterior, which is used to rank all n+1 points and compute
    the p-value.

    Args:
        x: Test feature vector (d,) or (1, d)
        epsilon: Significance level(s). If None, uses self.epsilon.
                 Can be scalar or array-like.
        return_p_values: If True, also return dict of p-values keyed by label.
        return_update: If True, also return dict with forest info.

    Returns:
        ConformalPredictionSet: If epsilon is scalar and neither flag is True
        MultiLevelPredictionSet: If epsilon is array-like
        Tuple: If return_p_values or return_update flag is set
    """
    x = np.asarray(x).ravel()

    if self.X is None:
        raise ValueError("Must call learn_initial_training_set first")

    if x.shape[0] != self.X.shape[1]:
        raise ValueError(f"Feature dimension mismatch: got {x.shape[0]}, expected {self.X.shape[1]}")

    if epsilon is None:
        epsilon = self.epsilon

    # Form augmented feature matrix (same for all candidates and all trees)
    X_aug = np.vstack([self.X, x])
    n = self.X.shape[0]
    n_total = n + 1
    K = len(self.label_space)

    # Build T tree summaries (batch: fresh trees from seeds; online: extend
    # the persistent forest). Caches viz/reuse state as a side effect.
    tree_summaries = self._forest_summaries(X_aug, x, n, K)
    n_leaves_all = [ts[0] for ts in tree_summaries]  # T x (n_total,)
    counts_all = [ts[1] for ts in tree_summaries]  # T x (n_total, K)
    ls_train_idx = [ts[2] for ts in tree_summaries]  # T x variable

    # Integer label indices for all training points
    LAMBDA = self.lambda_
    y_idx = np.array([self.label_to_idx[self.y[i]] for i in range(n)], dtype=np.int64)
    arange_n = np.arange(n)

    # base_ncm[i] = 1 - (1/T) * sum_t (counts_t[i, y_idx[i]] + λ) / (n_train_leaf_t[i] + K*λ)
    base_ncm_sum = np.zeros(n)
    for t in range(self.n_trees):
        nl = n_leaves_all[t][:n].astype(np.float64)  # (n,) structural sizes
        ct = counts_all[t][:n]  # (n, K)
        # leaf_star points have structural size including test slot
        n_train_leaf = nl.copy()
        for idx in ls_train_idx[t]:
            n_train_leaf[int(idx)] -= 1.0
        denom = n_train_leaf + K * LAMBDA
        base_ncm_sum += (ct[arange_n, y_idx].astype(np.float64) + LAMBDA) / denom
    base_ncm = 1.0 - base_ncm_sum / self.n_trees

    # ncm_test[k] = 1 - (1/T) * sum_t (counts_t[n, k] + 1 + λ) / (n_star_train_t + 1 + K*λ)
    ncm_test = np.zeros(K)
    for t in range(self.n_trees):
        n_star_train_t = float(n_leaves_all[t][n]) - 1.0
        denom_test_t = n_star_train_t + 1.0 + K * LAMBDA
        ncm_test += (counts_all[t][n].astype(np.float64) + 1.0 + LAMBDA) / denom_test_t
    ncm_test = 1.0 - ncm_test / self.n_trees  # shape (K,)

    # Vectorised (K, n) broadcast
    _TOL = 1e-9
    diff = base_ncm[np.newaxis, :] - ncm_test[:, np.newaxis]  # (K, n)
    gt = np.sum(diff > _TOL, axis=1).astype(np.int64)
    eq = np.sum(np.abs(diff) <= _TOL, axis=1).astype(np.int64)

    # Delta correction (Laplace): for each tree t, compute per-class delta
    # in leaf_star_t, then accumulate per-point adjustments.
    # Build per-tree delta matrices: delta_jk[j, k] = actual_prob - base_prob
    tree_deltas = []
    for t in range(self.n_trees):
        n_star_t = float(n_leaves_all[t][n])
        c_t = counts_all[t][n].astype(np.float64)
        denom_base_t = (n_star_t - 1.0) + K * LAMBDA
        denom_actual_t = n_star_t + K * LAMBDA
        delta_jk = np.empty((K, K))
        for j in range(K):
            base_p = (c_t[j] + LAMBDA) / denom_base_t
            for k in range(K):
                actual_p = (c_t[j] + LAMBDA + (1.0 if j == k else 0.0)) / denom_actual_t
                delta_jk[j, k] = actual_p - base_p
        tree_deltas.append(delta_jk)

    # Build point -> trees membership
    point_trees: dict[int, list] = {}
    for t in range(self.n_trees):
        for idx in ls_train_idx[t]:
            i_int = int(idx)
            if i_int not in point_trees:
                point_trees[i_int] = []
            point_trees[i_int].append(t)

    # Reclassify affected points
    for i_int, trees_list in point_trees.items():
        j = int(y_idx[i_int])
        delta_sum_k = np.zeros(K)
        for t in trees_list:
            delta_sum_k += tree_deltas[t][j, :]
        delta_avg_k = delta_sum_k / self.n_trees

        base_val = base_ncm[i_int]
        for k in range(K):
            actual_val = base_val - delta_avg_k[k]
            ncm_test_k = ncm_test[k]
            # Before bucket
            if base_val > ncm_test_k + _TOL:
                before = 0
            elif abs(base_val - ncm_test_k) <= _TOL:
                before = 1
            else:
                before = 2
            # After bucket
            if actual_val > ncm_test_k + _TOL:
                after = 0
            elif abs(actual_val - ncm_test_k) <= _TOL:
                after = 1
            else:
                after = 2
            if before == after:
                continue
            if before == 0 and after == 1:
                gt[k] -= 1
                eq[k] += 1
            elif before == 0 and after == 2:
                gt[k] -= 1
            elif before == 1 and after == 2:
                eq[k] -= 1
            elif before == 1 and after == 0:
                eq[k] -= 1
                gt[k] += 1
            elif before == 2 and after == 1:
                eq[k] += 1
            elif before == 2 and after == 0:
                gt[k] += 1

    eq += 1  # test point always EQ to itself
    tau = self.rnd_gen.uniform(0, 1)
    p_values_arr = (gt + tau * eq) / n_total
    p_values = {label: float(p_values_arr[ki]) for ki, label in enumerate(self.label_space)}

    result = self._compute_Gamma(p_values, epsilon)

    # Handle return flags
    _update = {"n_leaves_all": n_leaves_all, "counts_all": counts_all, "leaf_star_train_indices_all": ls_train_idx, "online_forest": self._pending_forest}
    if return_update and return_p_values:
        return result, p_values, _update
    elif return_p_values:
        return result, p_values
    elif return_update:
        return result, _update
    else:
        return result

compute_p_value(x: NDArray, y: Any, return_update: bool = False) -> float | tuple[float, dict]

Compute p-value for a single candidate label.

Args: x: Test feature vector (d,) y: Candidate label return_update: If True, also return dict with forest info

Returns: P-value (float) or (p-value, dict) if return_update=True

Source code in src/online_cp/classifiers.py

def compute_p_value(self, x: NDArray, y: Any, return_update: bool = False) -> float | tuple[float, dict]:
    """Compute p-value for a single candidate label.

    Args:
        x: Test feature vector (d,)
        y: Candidate label
        return_update: If True, also return dict with forest info

    Returns:
        P-value (float) or (p-value, dict) if return_update=True
    """
    x = np.asarray(x).ravel()

    if y not in self.label_to_idx:
        raise ValueError(f"Label {y} not in label_space")

    # Build augmented dataset
    X_aug = np.vstack([self.X, x])
    n = self.X.shape[0]
    n_total = n + 1
    K = len(self.label_space)
    label_idx = self.label_to_idx[y]

    # Build T tree summaries (possibly in parallel)
    tree_summaries = self._forest_summaries(X_aug, x, n, K)
    n_leaves_all = [ts[0] for ts in tree_summaries]
    counts_all = [ts[1] for ts in tree_summaries]
    ls_train_idx = [ts[2] for ts in tree_summaries]

    LAMBDA = self.lambda_
    y_idx = np.array([self.label_to_idx[self.y[i]] for i in range(n)], dtype=np.int64)
    arange_n = np.arange(n)

    # base_ncm for training points (Laplace)
    base_ncm_sum = np.zeros(n)
    for t in range(self.n_trees):
        nl = n_leaves_all[t][:n].astype(np.float64)
        ct = counts_all[t][:n]
        n_train_leaf = nl.copy()
        for idx in ls_train_idx[t]:
            n_train_leaf[int(idx)] -= 1.0
        denom = n_train_leaf + K * LAMBDA
        base_ncm_sum += (ct[arange_n, y_idx].astype(np.float64) + LAMBDA) / denom
    base_ncm = 1.0 - base_ncm_sum / self.n_trees

    # ncm_test for this single candidate (Laplace)
    ncm_test_k = 0.0
    for t in range(self.n_trees):
        n_star_train_t = float(n_leaves_all[t][n]) - 1.0
        denom_test_t = n_star_train_t + 1.0 + K * LAMBDA
        ncm_test_k += (float(counts_all[t][n, label_idx]) + 1.0 + LAMBDA) / denom_test_t
    ncm_test_k = 1.0 - ncm_test_k / self.n_trees

    _TOL = 1e-9
    diff = base_ncm - ncm_test_k
    gt = int(np.sum(diff > _TOL))
    eq = int(np.sum(np.abs(diff) <= _TOL))

    # Delta correction (Laplace): per-tree delta for leaf_star points
    tree_deltas = []
    for t in range(self.n_trees):
        n_star_t = float(n_leaves_all[t][n])
        c_t = counts_all[t][n].astype(np.float64)
        denom_base_t = (n_star_t - 1.0) + K * LAMBDA
        denom_actual_t = n_star_t + K * LAMBDA
        # delta[j] = actual_prob_j - base_prob_j for candidate label_idx
        delta_j = np.empty(K)
        for j in range(K):
            base_p = (c_t[j] + LAMBDA) / denom_base_t
            actual_p = (c_t[j] + LAMBDA + (1.0 if j == label_idx else 0.0)) / denom_actual_t
            delta_j[j] = actual_p - base_p
        tree_deltas.append(delta_j)

    # Build point -> trees membership
    point_trees: dict[int, list] = {}
    for t in range(self.n_trees):
        for idx in ls_train_idx[t]:
            i_int = int(idx)
            if i_int not in point_trees:
                point_trees[i_int] = []
            point_trees[i_int].append(t)

    # Reclassify affected points
    for i_int, trees_list in point_trees.items():
        j = int(y_idx[i_int])
        delta_sum = 0.0
        for t in trees_list:
            delta_sum += tree_deltas[t][j]
        delta_avg = delta_sum / self.n_trees

        base_val = base_ncm[i_int]
        actual_val = base_val - delta_avg
        # Before bucket
        if base_val > ncm_test_k + _TOL:
            before = 0
        elif abs(base_val - ncm_test_k) <= _TOL:
            before = 1
        else:
            before = 2
        # After bucket
        if actual_val > ncm_test_k + _TOL:
            after = 0
        elif abs(actual_val - ncm_test_k) <= _TOL:
            after = 1
        else:
            after = 2
        if before == after:
            continue
        if before == 0 and after == 1:
            gt -= 1
            eq += 1
        elif before == 0 and after == 2:
            gt -= 1
        elif before == 1 and after == 2:
            eq -= 1
        elif before == 1 and after == 0:
            eq -= 1
            gt += 1
        elif before == 2 and after == 1:
            eq += 1
        elif before == 2 and after == 0:
            gt += 1

    eq += 1  # test point always EQ to itself
    tau = self.rnd_gen.uniform(0, 1)
    p_val = float((gt + tau * eq) / n_total)

    if return_update:
        return p_val, {
            "n_leaves_all": n_leaves_all,
            "counts_all": counts_all,
            "leaf_star_train_indices_all": ls_train_idx,
            "online_forest": self._pending_forest,
        }
    else:
        return p_val

online_cp.regressors.ConformalMondrianTreeRegressor

Bases: _MondrianRegressorInspection, ConformalRegressor

Conformal predictor using a single Mondrian tree for regression.

A full (transductive) conformal predictor that rebuilds one Mondrian tree at each prediction step from the augmented bag [X_train; x_test]. The nonconformity measure (NCM) is the absolute residual from the leaf mean:

α_i = |y_i − μ_leaf(x_i)|

where μ_leaf is the mean of all y-values (including the candidate) in the leaf. Because the mean is computed over the full augmented multiset this is a valid bag function and the resulting conformal predictor has exact, finite-sample validity (ALRW2 §2.2.9).

The prediction interval is computed algebraically via an O(n) knot-point solver: the p-value p(y) is piecewise-constant between O(n) critical y-values where NCM ranks change. No grid search is needed.

Mathematical note

Let A = number of training points in leaf_star and B = their y-sum. Under candidate y:

• Test NCM: |A·y − B| / (A+1) slope ±A/(A+1)
• Inside NCM_j: |(A+1)·y_j − B − y| / (A+1) slope ∓1/(A+1)

Since A ≥ 1, the test NCM grows strictly faster than any inside NCM as y → ±∞, so p(y) is unimodal (guaranteed contiguous interval).

Parameters:

Name	Type	Description	Default
lifetime	float or {'sqrt_n', 'density'}	Mondrian tree depth budget L. Larger → more splits → smaller leaves. Default 1.0. Can also be a string for unsupervised automatic tuning (only X is used; labels are never accessed): 'sqrt_n': halt when the leaf count drops to ≤ √n. 'density': choose τ maximising Σ_leaf n_leaf · log(n_leaf / V_leaf).	1.0
epsilon	float	Default significance level. Default 0.1.	0.1
rnd_state	int	Seed for reproducibility.	None
verbose	int	Verbosity level (0 = silent).	0
max_depth	int or None	Hard depth cap. Nodes at depth >= max_depth become leaves regardless of remaining lifetime. None = no cap.	None
feature_weights	'variance'	Unsupervised feature importance for split dimension sampling. Labels are never accessed. None (default): uniform — standard isotropic Mondrian. 'variance': split probabilities ∝ population variance per feature. High-variance dimensions are sampled more often. array of shape (d,): explicit non-negative weights (normalised internally; need not sum to 1).	'variance'

Source code in src/online_cp/regressors.py

class ConformalMondrianTreeRegressor(_MondrianRegressorInspection, ConformalRegressor):
    """Conformal predictor using a single Mondrian tree for regression.

    A full (transductive) conformal predictor that rebuilds one Mondrian tree
    at each prediction step from the augmented bag ``[X_train; x_test]``.
    The nonconformity measure (NCM) is the absolute residual from the leaf mean:

        α_i = |y_i − μ_leaf(x_i)|

    where μ_leaf is the mean of **all** y-values (including the candidate) in
    the leaf.  Because the mean is computed over the full augmented multiset
    this is a valid bag function and the resulting conformal predictor has
    exact, finite-sample validity (ALRW2 §2.2.9).

    The prediction interval is computed **algebraically** via an O(n) knot-point
    solver: the p-value p(y) is piecewise-constant between O(n) critical y-values
    where NCM ranks change.  No grid search is needed.

    Mathematical note
    -----------------
    Let A = number of training points in ``leaf_star`` and B = their y-sum.
    Under candidate y:

    * Test NCM:      ``|A·y − B| / (A+1)``        slope ±A/(A+1)
    * Inside NCM_j:  ``|(A+1)·y_j − B − y| / (A+1)``  slope ∓1/(A+1)

    Since A ≥ 1, the test NCM grows strictly faster than any inside NCM as
    y → ±∞, so p(y) is unimodal (guaranteed contiguous interval).

    Parameters
    ----------
    lifetime : float or {'sqrt_n', 'density'}
        Mondrian tree depth budget L. Larger → more splits → smaller leaves.
        Default 1.0.

        Can also be a string for unsupervised automatic tuning (only X is
        used; labels are never accessed):

        - ``'sqrt_n'``: halt when the leaf count drops to ≤ √n.
        - ``'density'``: choose τ maximising Σ_leaf n_leaf · log(n_leaf / V_leaf).
    epsilon : float
        Default significance level. Default 0.1.
    rnd_state : int, optional
        Seed for reproducibility.
    verbose : int
        Verbosity level (0 = silent).
    max_depth : int or None
        Hard depth cap. Nodes at depth >= max_depth become leaves regardless
        of remaining lifetime. None = no cap.
    feature_weights : {'variance'} or array-like of shape (d,) or None
        Unsupervised feature importance for split dimension sampling. Labels
        are never accessed.

        - ``None`` (default): uniform — standard isotropic Mondrian.
        - ``'variance'``: split probabilities ∝ population variance per
          feature. High-variance dimensions are sampled more often.
        - array of shape (d,): explicit non-negative weights (normalised
          internally; need not sum to 1).
    """

    _SAVE_PARAMS: tuple = (
        "lifetime", "epsilon", "rnd_state", "verbose", "max_depth",
        "feature_weights", "online",
    )
    _SAVE_STATE: tuple = ("X", "y", "_online_tree")

    def __init__(
        self,
        lifetime: float | str = 1.0,
        epsilon: float = 0.1,
        rnd_state: int | None = None,
        verbose: int = 0,
        max_depth: int | None = None,
        feature_weights: str | NDArray | None = None,
        online: bool = False,
    ) -> None:
        if max_depth is not None and (not isinstance(max_depth, int) or max_depth < 0):
            raise ValueError("max_depth must be a non-negative integer or None")
        if isinstance(lifetime, str) and lifetime not in ("sqrt_n", "density"):
            raise ValueError(
                f"Unknown lifetime string {lifetime!r}. "
                "Recognised values: 'sqrt_n', 'density'."
            )
        if isinstance(feature_weights, str) and feature_weights not in ("variance",):
            raise ValueError(
                f"Unknown feature_weights string {feature_weights!r}. "
                "Recognised value: 'variance'."
            )
        if online:
            if not isinstance(lifetime, (int, float)):
                raise ValueError(
                    "online=True requires a fixed float lifetime; string/auto-tuned "
                    "lifetimes recompute a master tree per step and are batch-only."
                )
            if max_depth is not None:
                raise ValueError(
                    "online=True requires max_depth=None (the projective Mondrian regime)."
                )
        super().__init__(epsilon=epsilon)
        self.lifetime = lifetime
        self.feature_weights = feature_weights
        self.rnd_state = rnd_state
        self.verbose = verbose
        self.max_depth = max_depth
        self.online = online
        self.X: NDArray | None = None
        self.y: NDArray | None = None
        self.rnd_gen = np.random.default_rng(rnd_state)
        self._last_tree = None
        self._online_tree = None
        self._pending_tree = None  # extension scored by the last predict (online)

    # ------------------------------------------------------------------
    # Training interface
    # ------------------------------------------------------------------

    def learn_initial_training_set(self, X: NDArray, y: NDArray) -> None:
        """Batch-initialize with training data.

        Parameters
        ----------
        X : ndarray, shape (n, d)
        y : ndarray, shape (n,)  — real-valued labels
        """
        X = np.atleast_2d(np.asarray(X, dtype=float))
        y = np.asarray(y, dtype=float).ravel()
        if X.shape[0] == 0:
            raise ValueError("Training set cannot be empty")
        if X.shape[0] != y.shape[0]:
            raise ValueError("X and y must have the same length")
        self.X = X.copy()
        self.y = y.copy()
        if self.online:
            self._online_tree = MondrianTree.grow(
                self.X, self.rnd_gen,
                lifetime=self.lifetime, feature_weights=self.feature_weights,
            )

    def learn_one(
        self,
        x: NDArray,
        y: float,
        precomputed: dict | None = None,
    ) -> None:
        """Incrementally add one observation.

        Parameters
        ----------
        x : array-like, shape (d,)
        y : float
        precomputed : dict, optional
            The update dict from ``predict(..., return_update=True)`` (or
            ``compute_p_value(..., return_update=True)``) for this same ``x``.
            In online mode, if it carries the extension that scored ``x`` the
            persisted tree becomes that exact extension, so the realised stream
            is a single consistent Mondrian process (exact online validity).
            Without it — or if it is stale — a fresh ``extend`` is drawn (still
            valid). The last ``predict`` is also cached, so the plain
            ``predict(x)`` → ``learn_one(x, y)`` loop reuses it automatically.
        """
        x = np.asarray(x, dtype=float).ravel()
        if self.X is None:
            # First training example: initialize with a single-point tree
            self.X = x.reshape(1, -1)
            self.y = np.array([float(y)])
            if self.online:
                self._online_tree = MondrianTree.grow(
                    self.X, self.rnd_gen,
                    lifetime=self.lifetime, feature_weights=self.feature_weights,
                )
            return
        if x.shape[0] != self.X.shape[1]:
            raise ValueError(f"Feature dimension mismatch: got {x.shape[0]}, expected {self.X.shape[1]}")
        self.X = np.vstack([self.X, x])
        self.y = np.append(self.y, float(y))
        if self.online:
            cand = precomputed.get("online_tree") if precomputed else None
            if not self._is_consistent_extension(cand, x):
                cand = getattr(self, "_pending_tree", None)
            self._pending_tree = None
            self._online_tree = (
                cand
                if self._is_consistent_extension(cand, x)
                else self._online_tree.extend(x, self.rnd_gen)
            )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _build_augmented_tree(self, x: NDArray):
        """Build (batch) or extend (online) the partition on [X_train; x].

        In online mode the extension is cached (``_pending_tree``) so a
        subsequent ``learn_one`` on the same ``x`` can persist the very tree
        that scored it — making the realised stream a single consistent
        Mondrian process (see :meth:`learn_one`).

        Returns ``(root, leaf_star, n)``.
        """
        x = np.asarray(x, dtype=float).ravel()
        n = self.X.shape[0]
        if self.online:
            ext_tree = self._online_tree.extend(x, self.rnd_gen)
            self._pending_tree = ext_tree
            tree = ext_tree.root
        else:
            self._pending_tree = None
            X_aug = np.vstack([self.X, x])
            tree = MondrianTree.grow(
                X_aug,
                self.rnd_gen,
                lifetime=self.lifetime,
                x_test=x,
                max_depth=self.max_depth,
                feature_weights=self.feature_weights,
                verbose=self.verbose,
            ).root
        leaf_star = _find_leaf(tree, x)
        return tree, leaf_star, n

    def _build_augmented_tree_with_rng(self, x: NDArray, rnd_gen: np.random.RandomState):
        """Build (batch) or extend (online) the partition on [X_train; x] with a given random generator.

        This is a variant of _build_augmented_tree that takes a random generator as input,
        allowing the tree to be built deterministically without advancing self.rnd_gen.
        This is essential for compute_p_value, where the tree must be the same for all
        candidate y values.

        Returns ``(root, leaf_star, n)``.
        """
        x = np.asarray(x, dtype=float).ravel()
        n = self.X.shape[0]
        if self.online:
            # For online mode, we can't easily support a custom rng without
            # modifying the online tree structure. Fall back to using self.rnd_gen.
            ext_tree = self._online_tree.extend(x, rnd_gen)
            self._pending_tree = ext_tree
            tree = ext_tree.root
        else:
            self._pending_tree = None
            X_aug = np.vstack([self.X, x])
            tree = MondrianTree.grow(
                X_aug,
                rnd_gen,
                lifetime=self.lifetime,
                x_test=x,
                max_depth=self.max_depth,
                feature_weights=self.feature_weights,
                verbose=self.verbose,
            ).root
        leaf_star = _find_leaf(tree, x)
        return tree, leaf_star, n

    def _is_consistent_extension(self, cand, x) -> bool:
        """True iff ``cand`` is the persisted online tree extended by exactly ``x``.

        Guards reuse of a scored extension: it must have one more row than the
        persisted tree and that row must be ``x`` — so a stale extension (the
        base has advanced since it was scored) falls back to a fresh ``extend``.
        """
        return (
            cand is not None
            and self._online_tree is not None
            and cand.X.shape[0] == self._online_tree.X.shape[0] + 1
            and np.array_equal(cand.X[-1], np.asarray(x, dtype=float).ravel())
        )

    def _outside_ncms_sorted(self, tree: _MondrianNode, leaf_star: _MondrianNode, n_train: int) -> NDArray:
        """Sorted nonconformity scores of all training points *outside* leaf_star.

        These scores are independent of the candidate ``y`` (the test point only
        affects its own leaf), so they are computed once per prediction — a
        single O(n) tree traversal + O(n log n) sort — and then reused for every
        candidate via binary search (see :meth:`_pvalues_at`).
        """
        outside: list[NDArray] = []
        for leaf in _collect_leaves(tree):
            if leaf is leaf_star:
                continue
            train_idx = leaf.indices[leaf.indices < n_train]
            if len(train_idx) == 0:
                continue
            mu = self.y[train_idx].mean()
            outside.append(np.abs(self.y[train_idx] - mu))
        ext = np.concatenate(outside) if outside else np.empty(0)
        ext.sort()
        return ext

    def _pvalues_at(
        self,
        y_grid: NDArray,
        A: int,
        B: float,
        ext_sorted: NDArray,
        y_star: NDArray,
        n_train: int,
        tau: float,
    ) -> NDArray:
        """Smoothed conformal p-values at candidate values ``y_grid`` (vectorised).

        The candidate-independent ``ext_sorted`` scores are counted by binary
        search (O(log n) per candidate); only the O(A) leaf-local scores are
        recomputed per candidate. This makes the exact solver O(n log n) instead
        of O(n^2). Numerically identical to evaluating the per-candidate p-value
        directly.
        """
        y_grid = np.atleast_1d(np.asarray(y_grid, dtype=float))
        ncm_test = np.abs(A * y_grid - B) / (A + 1)
        _TOL = 1e-9
        n_out = ext_sorted.size

        right = np.searchsorted(ext_sorted, ncm_test + _TOL, side="right")
        left = np.searchsorted(ext_sorted, ncm_test - _TOL, side="left")
        gt = (n_out - right).astype(np.int64)
        eq = (right - left).astype(np.int64)

        if A > 0 and y_star.size > 0:
            inside = np.abs((A + 1) * y_star[None, :] - B - y_grid[:, None]) / (A + 1)
            d = inside - ncm_test[:, None]
            gt = gt + np.sum(d > _TOL, axis=1)
            eq = eq + np.sum(np.abs(d) <= _TOL, axis=1)

        # eq now counts all points with alpha_i == alpha_n (excluding test point itself)
        # Apply tau to the ties, and add 1 for the test point comparing to itself
        p_values = (gt + tau * eq + tau) / (n_train + 1)
        return p_values

    # ------------------------------------------------------------------
    # Event-driven O(n log n) knot sweep
    # ------------------------------------------------------------------
    def _build_knot_events(
        self,
        A: int,
        B: float,
        ext_sorted: NDArray,
        y_star: NDArray,
    ) -> tuple[NDArray, int]:
        """Build knot events with crossing-direction flags for the O(n log n) sweep.

        Returns:
            knots_events: sorted array of (knot_value, flag) tuples, where flag is +1
                         when a point *enters* the {α_i > α_n} set as y increases through
                         the knot, and -1 when it *exits*.
            permanent_ties: count of training points j with α_j(y) ≡ α_n(y) for all y
                           (only happens when A = 1, in which case the single inside
                           point is permanently tied with the test point).

        The algorithm matches mondrian_trees.tex Appendix reg-knots and Prop. reg-exact.
        """
        knots: list[tuple[float, int]] = []
        permanent_ties = 0

        # ------------------------------------------------------------------
        # Outside-leaf points: α_i is constant (independent of y)
        # Each outside point i generates two knots where α_i crosses α_n:
        #   α_n(y) = |A·y - B| / (A+1)      (slope ±A/(A+1), vertex at B/A)
        #   α_i is fixed
        # Crossings at y = B/A ± ((A+1)/A)·α_i
        # Lower root (y_lower = B/A - ((A+1)/A)·α_i): as y increases, α_n overtakes α_i
        #   → for y < y_lower: α_n > α_i (test score higher)
        #   → for y > y_lower: α_i > α_n (training score higher) → *enter* G_i
        #   → tag y_lower with +1
        # Upper root (y_upper = B/A + ((A+1)/A)·α_i): as y increases, α_n moves away
        #   → for y < y_upper: α_i > α_n
        #   → for y > y_upper: α_n > α_i → *exit* G_i
        #   → tag y_upper with -1
        # ------------------------------------------------------------------
        n_out = ext_sorted.size
        if n_out > 0:
            fac = (A + 1) / A
            lower_knots = (B / A) - fac * ext_sorted  # +1 flag
            upper_knots = (B / A) + fac * ext_sorted  # -1 flag
            for y_knot in lower_knots:
                knots.append((float(y_knot), 1))  # entering
            for y_knot in upper_knots:
                knots.append((float(y_knot), -1))  # exiting

        # ------------------------------------------------------------------
        # Inside-leaf points: α_j(y) = |(A+1)·y_j - B - y| / (A+1)
        # This is a V-shape with vertex at y = y_j (slope ∓1/(A+1))
        # The test score α_n has vertex at y = B/A.
        # ------------------------------------------------------------------
        if A == 0:
            # Empty leaf: no inside points. Should not happen here (handled earlier).
            pass
        elif A == 1:
            # Singleton test leaf: the single inside point j has y_j = B/A (since A=1).
            # Corollary cor:slope: α_j(y) ≡ α_n(y) for all y (permanent tie).
            # No crossing knots — the inside point is always tied with the test point.
            permanent_ties = 1
        else:
            # A > 1: each inside point generates two knots:
            #   (i) y = y_j: test score vertex crosses inside score vertex
            #       For y < y_j: α_n > α_j (if y_j > B/A) or α_j > α_n (if y_j < B/A)
            #       At y = y_j: both scores have a kink
            #       As y increases through y_j, α_n overtakes α_j → *enter* G_j → +1
            #   (ii) y = (2B - (A+1)·y_j)/(A-1): reflection root where α_j = α_n again
            #       As y increases through this, α_n moves away → *exit* G_j → -1
            for y_j in y_star:
                # The valid crossing region G_j is bounded by the min and max roots
                refl = (2.0 * B - (A + 1) * y_j) / (A - 1)
                left_root = min(y_j, refl)
                right_root = max(y_j, refl)
                # Enter G_j at left root, exit at right root
                knots.append((float(left_root), 1))
                knots.append((float(right_root), -1))

        if not knots:
            # Degenerate: no outside points and A <= 1 (empty leaf or singleton)
            # Singleton handled via permanent_ties, so this must be empty leaf.
            # But empty leaf should be caught earlier. Return empty events.
            return np.array([], dtype=[("value", float), ("flag", int)]), permanent_ties

        # Sort by knot value, then by flag (process exits before entries to be safe)
        # Actually, we want to aggregate flags at coincident knots, so sort only by value
        knots.sort(key=lambda x: x[0])
        knots_arr = np.array(knots, dtype=[("value", float), ("flag", int)])
        return knots_arr, permanent_ties

    def _evaluate_pvalue_at_point(self, y, A, B, ext_sorted, y_star, n_train, tau, permanent_ties):
        """Evaluate smoothed conformal p-value at an arbitrary y (not just cell midpoints).

        This is used to check knot locations exactly, where p(y) may spike due to
        tie-breaking with tau, potentially making a singleton valid even if
        neighboring cells are invalid.

        Parameters
        ----------
        y : float
            Candidate label at which to evaluate p-value.
        A : int
            Number of training points in test leaf.
        B : float
            Sum of y-values in test leaf.
        ext_sorted : ndarray
            Sorted nonconformity scores of training points outside test leaf.
        y_star : ndarray
            y-values of training points inside test leaf.
        n_train : int
            Number of training points.
        tau : float
            Randomization variable in [0, 1] for smoothed p-value.
        permanent_ties : int
            Count of points permanently tied with test point (A == 1 case).

        Returns
        -------
        float
            Smoothed conformal p-value at y.
        """
        # Test nonconformity score: alpha_n(y) = |A*y - B| / (A+1)
        ncm_test = abs(A * y - B) / (A + 1)

        # Count outside points with alpha_i > ncm_test
        gt = int(np.sum(ext_sorted > ncm_test))

        # Count outside points with alpha_i == ncm_test (for tie-breaking)
        _TOL = 1e-9
        eq = int(np.sum(np.abs(ext_sorted - ncm_test) <= _TOL))

        # Add contributions from inside points (in leaf_star)
        if A > 0 and y_star.size > 0:
            # Inside scores: alpha_j(y) = |(A+1)*y_j - B - y| / (A+1)
            inside_ncm = np.abs((A + 1) * y_star - B - y) / (A + 1)
            gt += int(np.sum(inside_ncm > ncm_test))
            eq += int(np.sum(np.abs(inside_ncm - ncm_test) <= _TOL))

        # p-value = (gt + tau * (eq + 1 + permanent_ties)) / (n_train + 1)
        # The +1 is for the test point itself comparing to itself
        p_val = (gt + tau * (eq + 1 + permanent_ties)) / (n_train + 1)
        return p_val

    def _pvalue_from_sweep(
        self,
        knots_events: NDArray,
        A: int,
        B: float,
        ext_sorted: NDArray,
        y_star: NDArray,
        n_train: int,
        permanent_ties: int,
        epsilon: float,
        tau: float,
    ) -> tuple[float, float, list]:
        """Left-to-right sweep to find p(y) > epsilon interval.

        Returns (lower_bound, upper_bound, valid_knots) where valid_knots is a
        list of knot values where p(y) > epsilon (singleton valid regions).
        Uses O(1) amortised update per knot → O(n) total for all p-value evals.
        """
        # Edge case: no knots (A == 0 with no outside points, or A == 1 with no outside points)
        if knots_events.size == 0:
            # No knots → p(y) is constant everywhere
            # For A == 0: p = tau / (n+1) (no points to compare)
            # For A == 1 with permanent_ties: p = (0 + tau * (1 + 1)) / (n+1) = 2*tau/(n+1)
            p_const = tau * (1 + permanent_ties) / (n_train + 1)
            if p_const > epsilon:
                return -np.inf, np.inf, []
            else:
                # No valid region — return degenerate interval at peak
                peak = B / A if A > 0 else 0.0
                return float(peak), float(peak), [float(peak)]

        # ------------------------------------------------------------------
        # Baseline tally at y_start (far left of all knots)
        # Pick y_start strictly less than first knot
        y_start = knots_events[0]["value"] - 1.0
        # Evaluate all scores at y_start
        ncm_test_start = abs(A * y_start - B) / (A + 1)
        # Outside scores (fixed)
        greater_count = int(np.sum(ext_sorted > ncm_test_start))
        # Inside scores
        if A > 0 and y_star.size > 0:
            for y_j in y_star:
                ncm_j = abs((A + 1) * y_j - B - y_start) / (A + 1)
                if ncm_j > ncm_test_start:
                    greater_count += 1
        # ------------------------------------------------------------------

        # ------------------------------------------------------------------
        # Sweep left-to-right, updating greater_count at each knot
        # ------------------------------------------------------------------
        # We'll evaluate p(y) at the midpoint of each open cell
        # and track the first valid cell (for lower bound) and the last valid cell (for upper)
        # Also track valid knots (singletons at boundary points)
        first_valid_idx = -1
        last_valid_idx = -1
        valid_knots: list[float] = []
        prev_knot = None

        for i in range(len(knots_events) + 1):
            # Determine the cell's right boundary (current knot) and left boundary (prev knot)
            if i < len(knots_events):
                knot_val = knots_events[i]["value"]
                flag = knots_events[i]["flag"]
            else:
                knot_val = None  # rightmost end-cell

            # Compute midpoint of current cell (for potential debugging/logging)
            if prev_knot is None:
                # Leftmost end-cell: extend marginally beyond first knot
                cell_right = knots_events[0]["value"]
                _ = cell_right - max(abs(cell_right), 1.0) * 10.0
            elif knot_val is None:
                # Rightmost end-cell
                cell_left = prev_knot
                _ = cell_left + max(abs(cell_left), 1.0) * 10.0
            else:
                cell_left = prev_knot
                cell_right = knot_val
                _ = (cell_left + cell_right) / 2.0

            # Compute p-value for this cell
            # Apply tau to the ties (test point + any permanent ties)
            p_val = (greater_count + tau * (1 + permanent_ties)) / (n_train + 1)

            # Check validity of cell interior
            if p_val > epsilon:
                if first_valid_idx == -1:
                    first_valid_idx = i  # first valid cell index
                last_valid_idx = i     # update last valid cell index

            # Check validity of the knot itself (if not the first knot)
            # Evaluate p(y) exactly at the knot location
            if knot_val is not None and prev_knot is not None:
                # We're at knot i, which is the boundary between cell i-1 and cell i
                # Check if this knot is valid
                knot_p_val = self._evaluate_pvalue_at_point(
                    knot_val, A, B, ext_sorted, y_star, n_train, tau, permanent_ties
                )
                if knot_p_val > epsilon:
                    valid_knots.append(float(knot_val))

            # Update for next iteration: apply flags at this knot (if any)
            if knot_val is not None:
                greater_count += flag
                prev_knot = knot_val
            elif prev_knot is not None:
                # Already processed all knots, we're done
                break
            else:
                prev_knot = knots_events[0]["value"]

        # ------------------------------------------------------------------
        # Extract interval from valid cells, incorporating valid knots
        # ------------------------------------------------------------------
        if first_valid_idx == -1 and not valid_knots:
            # No valid region at all — return empty interval
            return np.nan, np.nan, []

        # Build the exact prediction set as union of valid cells and valid knots
        # Valid cells: cell i corresponds to open interval (knots[i-1], knots[i]) for i=1..len
        # Cell 0 is (-∞, knots[0]), cell len(knots) is (knots[-1], +∞)

        # If we have valid knots but no valid cells, return singleton(s)
        if first_valid_idx == -1:
            # All cells invalid, only knots valid
            if len(valid_knots) == 1:
                return valid_knots[0], valid_knots[0], valid_knots
            else:
                # Multiple disjoint singletons — return the smallest enclosing interval
                # (conservative but valid)
                return min(valid_knots), max(valid_knots), valid_knots

        # Determine bounds from first and last valid cells, incorporating valid knots
        lo = None
        hi = None

        if first_valid_idx == 0:
            # Cell 0 is (-∞, knots[0])
            lo = -np.inf
            # Check if knot[0] is valid (was already added to valid_knots above)
        else:
            lo = float(knots_events[first_valid_idx - 1]["value"])

        if last_valid_idx == len(knots_events):
            # Rightmost end-cell is valid
            hi = np.inf
        else:
            hi = float(knots_events[last_valid_idx]["value"])

        # If lo is a valid knot but not part of a valid cell, adjust to singleton
        if valid_knots and lo is not None and hi is not None:
            lo_knot = float(knots_events[first_valid_idx - 1]["value"]) if first_valid_idx > 0 else None

            # If lo_knot is valid and the cell to its left is invalid, lo should be singleton
            if lo_knot in valid_knots and first_valid_idx > 0:
                # The knot is valid but the cell to its left is not
                # This is a singleton at the left boundary
                pass  # Keep lo as is, valid_knots will record it

        return float(lo) if lo is not None else float(knots_events[0]["value"]), \
               float(hi) if hi is not None else float(knots_events[-1]["value"]), \
               valid_knots

    # ------------------------------------------------------------------
    # compute_p_value
    # ------------------------------------------------------------------

    def compute_p_value(
        self,
        x: NDArray,
        y_cand: float,
        return_update: bool = False,
        tau: float | None = None,
    ) -> float | tuple[float, dict]:
        """Compute conformal p-value for (x, y_cand).

        Parameters
        ----------
        x : array-like, shape (d,)
        y_cand : float
        return_update : bool
            If True, also return a dict with tree info.
        tau : float or None
            Randomization variable for smoothed p-value. If None, draws from U(0,1).
            For property testing and deterministic behavior, pass a fixed tau.

        Returns
        -------
        float or (float, dict)
        """
        x = np.asarray(x, dtype=float).ravel()
        n = self._safe_size_check(self.X)
        if n == 0:
            return (1.0, {}) if return_update else 1.0

        y_cand = float(y_cand)

        # FIX: Use a deterministic seed based on the test point x for tree building
        # This ensures the same tree is built for all candidate y values
        x_seed = int(hash(tuple(np.round(x, 10).flatten())) % (2**31))
        tree_rnd_gen = np.random.default_rng(x_seed)

        # tau must be INDEPENDENT of the tree, drawn from a separate source
        # If tau is not provided, draw it from a deterministic but SEPARATE source
        if tau is None:
            # Use a different deterministic seed for tau (independent of tree)
            tau_seed = (x_seed * 2654435761) % (2**32)  # Prime multiplier for decorrelation
            tau_rnd_gen = np.random.default_rng(tau_seed)
            tau = tau_rnd_gen.uniform(0.0, 1.0)

        tree, leaf_star, _ = self._build_augmented_tree_with_rng(x, tree_rnd_gen)
        leaf_star_train_idx = leaf_star.indices[leaf_star.indices < n]
        A = len(leaf_star_train_idx)
        B = float(self.y[leaf_star_train_idx].sum()) if A > 0 else 0.0
        ext_sorted = self._outside_ncms_sorted(tree, leaf_star, n)
        # tau was already set above if None
        p_val = float(
            self._pvalues_at(y_cand, A, B, ext_sorted, self.y[leaf_star_train_idx], n, tau)[0]
        )

        # Cache the tree for inspection utilities
        self._last_tree = tree

        if return_update:
            return p_val, {"tree": tree, "leaf_star": leaf_star, "online_tree": self._pending_tree}
        return p_val

    # ------------------------------------------------------------------
    # predict  (exact knot-point solver)
    # ------------------------------------------------------------------

    def predict(
        self,
        x: NDArray,
        epsilon: float | NDArray | None = None,
        return_update: bool = False,
    ) -> ContinuousPredictionSet | MultiLevelPredictionSet | tuple:
        """Predict a conformal interval for test object x.

        Uses an exact O(n) algebraic knot-point solver.  No grid search.

        Parameters
        ----------
        x : array-like, shape (d,)
        epsilon : float or array-like, optional
            Significance level(s).  Defaults to ``self.epsilon``.
        return_update : bool
            If True, return ``(interval, update_dict)`` for ``learn_one``.

        Returns
        -------
        ContinuousPredictionSet or MultiLevelPredictionSet
        """
        x = np.asarray(x, dtype=float).ravel()
        if epsilon is None:
            epsilon = self.epsilon

        n = self._safe_size_check(self.X)
        _inf_interval = self._construct_Gamma((-np.inf, np.inf), epsilon)

        if n == 0:
            if return_update:
                return _inf_interval, {}
            return _inf_interval

        tree, leaf_star, _ = self._build_augmented_tree(x)
        # Cache the tree for inspection utilities (draw_partition etc.),
        # regardless of which prediction branch is taken below.
        self._last_tree = tree

        leaf_star_train_idx = leaf_star.indices[leaf_star.indices < n]
        A = len(leaf_star_train_idx)
        B = self.y[leaf_star_train_idx].sum() if A > 0 else 0.0

        # ---- Draw ONE tau for the entire predict call to ensure deterministic behavior.
        # Use it for smoothed p-value calculation. This must be done BEFORE any A==0 checks.
        tau = self.rnd_gen.uniform(0.0, 1.0)

        # ---- A = 0: empty leaf → total epistemic uncertainty -----------
        if A == 0:
            ext_sorted = self._outside_ncms_sorted(tree, leaf_star, n)

            # Since test score is 0, rank depends entirely on outside scores >= 0
            gt = int(np.sum(ext_sorted > 0))
            eq = int(np.sum(ext_sorted == 0))
            p_const = (gt + tau * (eq + 1)) / (n + 1)

            eps_arr = np.atleast_1d(np.asarray(epsilon, dtype=float))
            predictions = {}
            for eps_val in eps_arr:
                if p_const > eps_val:
                    predictions[float(eps_val)] = self._construct_Gamma((-np.inf, np.inf), float(eps_val))
                else:
                    predictions[float(eps_val)] = self._construct_Gamma((np.nan, np.nan), float(eps_val))

            if hasattr(epsilon, "__iter__"):
                result = MultiLevelPredictionSet(predictions)
            else:
                result = predictions[float(eps_arr[0])]

            if return_update:
                return result, {"tree": tree, "leaf_star": leaf_star, "online_tree": self._pending_tree}
            return result

        # ---- Collect knot points ----------------------------------------
        # The candidate-independent outside-leaf scores are computed once here
        # and reused for the knots and for every p-value evaluation below.
        ext_sorted = self._outside_ncms_sorted(tree, leaf_star, n)

        # Outside-leaf_star points: ncm_i is fixed → rank changes when
        #   ncm_test(y) = ncm_i  ↔  y = (B ± ncm_i*(A+1)) / A
        knots: list[float] = []
        if ext_sorted.size:
            fac = (A + 1) / A
            knots.extend(((B / A) + fac * ext_sorted).tolist())
            knots.extend(((B / A) - fac * ext_sorted).tolist())

        # Inside-leaf_star points:
        #   always: y_j
        #   if A > 1: (2B − (A+1)*y_j) / (A−1)
        y_star = self.y[leaf_star_train_idx]
        knots.extend(y_star.tolist())
        if A > 1:
            knots.extend(((2.0 * B - (A + 1) * y_star) / (A - 1)).tolist())

        # Sort and de-duplicate (for potential future use)

        # Build knot events with +1/-1 flags
        knots_events, permanent_ties = self._build_knot_events(A, B, ext_sorted, y_star)

        # Perform left-to-right sweep to find valid region
        lo, hi, valid_knots = self._pvalue_from_sweep(
            knots_events, A, B, ext_sorted, y_star, n, permanent_ties,
            epsilon if not hasattr(epsilon, "__iter__") else epsilon[0], tau
        )

        # Build result with correct topology
        eps_arr = np.atleast_1d(np.asarray(epsilon, dtype=float))
        predictions = {}

        for eps_val in eps_arr:
            lo, hi, valid_knots = self._pvalue_from_sweep(
                knots_events, A, B, ext_sorted, y_star, n, permanent_ties, float(eps_val), tau
            )

            if len(valid_knots) == 1 and lo == hi == valid_knots[0]:
                predictions[float(eps_val)] = self._construct_Gamma((lo, hi), float(eps_val))
            else:
                predictions[float(eps_val)] = self._construct_Gamma((lo, hi), float(eps_val))

        if hasattr(epsilon, "__iter__"):
            result = MultiLevelPredictionSet(predictions)
        else:
            result = predictions[float(eps_arr[0])]


        if return_update:
            return result, {"tree": tree, "leaf_star": leaf_star, "online_tree": self._pending_tree}

        return result

    @staticmethod
    def _extract_interval(
        knots: NDArray,
        midpoints: NDArray,
        p_vals: NDArray,
        epsilon: float,
    ) -> tuple[float, float]:
        """Return (lower, upper) convex hull of valid region.

        The prediction set is {y : p(y) > epsilon} (ALRW2 convention).
        """
        threshold = epsilon
        valid = p_vals > threshold
        if not np.any(valid):
            # No candidate passes — return empty interval
            return np.nan, np.nan
        # Knot boundaries: the valid midpoints index into intervals between knots.
        # Interval i spans [knots[i-1], knots[i]] with midpoint midpoints[i].
        # midpoints[0] < knots[0] → left-open interval, midpoints[-1] > knots[-1] → right-open.
        valid_idx = np.where(valid)[0]
        lo_idx = valid_idx[0]
        hi_idx = valid_idx[-1]
        # Lower bound: left edge of first valid segment
        lo = -np.inf if lo_idx == 0 else float(knots[lo_idx - 1])
        # Upper bound: right edge of last valid segment
        hi = np.inf if hi_idx == len(knots) else float(knots[hi_idx])
        return lo, hi

    def _make_result(self, epsilon, lo: float, hi: float):
        """Construct a ContinuousPredictionSet or MultiLevelPredictionSet."""
        if hasattr(epsilon, "__iter__"):
            return MultiLevelPredictionSet(
                {float(eps): self._construct_Gamma((lo, hi), float(eps)) for eps in epsilon}
            )
        return self._construct_Gamma((lo, hi), float(epsilon))

    def _make_multi_result(
        self,
        epsilon,
        knots: NDArray,
        midpoints: NDArray,
        p_vals: NDArray,
    ):
        predictions = {}
        for eps in epsilon:
            lo, hi = self._extract_interval(knots, midpoints, p_vals, float(eps))
            predictions[float(eps)] = self._construct_Gamma((lo, hi), float(eps))
        return MultiLevelPredictionSet(predictions)

learn_initial_training_set(X: NDArray, y: NDArray) -> None

Batch-initialize with training data.

Parameters:

Name	Type	Description	Default
X	(ndarray, shape(n, d))		required
y	ndarray, shape (n,) — real-valued labels		required

Source code in src/online_cp/regressors.py

def learn_initial_training_set(self, X: NDArray, y: NDArray) -> None:
    """Batch-initialize with training data.

    Parameters
    ----------
    X : ndarray, shape (n, d)
    y : ndarray, shape (n,)  — real-valued labels
    """
    X = np.atleast_2d(np.asarray(X, dtype=float))
    y = np.asarray(y, dtype=float).ravel()
    if X.shape[0] == 0:
        raise ValueError("Training set cannot be empty")
    if X.shape[0] != y.shape[0]:
        raise ValueError("X and y must have the same length")
    self.X = X.copy()
    self.y = y.copy()
    if self.online:
        self._online_tree = MondrianTree.grow(
            self.X, self.rnd_gen,
            lifetime=self.lifetime, feature_weights=self.feature_weights,
        )

learn_one(x: NDArray, y: float, precomputed: dict | None = None) -> None

Incrementally add one observation.

Parameters:

Name	Type	Description	Default
x	(array - like, shape(d))		required
y	float		required
precomputed	dict	The update dict from predict(..., return_update=True) (or compute_p_value(..., return_update=True)) for this same x. In online mode, if it carries the extension that scored x the persisted tree becomes that exact extension, so the realised stream is a single consistent Mondrian process (exact online validity). Without it — or if it is stale — a fresh extend is drawn (still valid). The last predict is also cached, so the plain predict(x) → learn_one(x, y) loop reuses it automatically.	None

Source code in src/online_cp/regressors.py

def learn_one(
    self,
    x: NDArray,
    y: float,
    precomputed: dict | None = None,
) -> None:
    """Incrementally add one observation.

    Parameters
    ----------
    x : array-like, shape (d,)
    y : float
    precomputed : dict, optional
        The update dict from ``predict(..., return_update=True)`` (or
        ``compute_p_value(..., return_update=True)``) for this same ``x``.
        In online mode, if it carries the extension that scored ``x`` the
        persisted tree becomes that exact extension, so the realised stream
        is a single consistent Mondrian process (exact online validity).
        Without it — or if it is stale — a fresh ``extend`` is drawn (still
        valid). The last ``predict`` is also cached, so the plain
        ``predict(x)`` → ``learn_one(x, y)`` loop reuses it automatically.
    """
    x = np.asarray(x, dtype=float).ravel()
    if self.X is None:
        # First training example: initialize with a single-point tree
        self.X = x.reshape(1, -1)
        self.y = np.array([float(y)])
        if self.online:
            self._online_tree = MondrianTree.grow(
                self.X, self.rnd_gen,
                lifetime=self.lifetime, feature_weights=self.feature_weights,
            )
        return
    if x.shape[0] != self.X.shape[1]:
        raise ValueError(f"Feature dimension mismatch: got {x.shape[0]}, expected {self.X.shape[1]}")
    self.X = np.vstack([self.X, x])
    self.y = np.append(self.y, float(y))
    if self.online:
        cand = precomputed.get("online_tree") if precomputed else None
        if not self._is_consistent_extension(cand, x):
            cand = getattr(self, "_pending_tree", None)
        self._pending_tree = None
        self._online_tree = (
            cand
            if self._is_consistent_extension(cand, x)
            else self._online_tree.extend(x, self.rnd_gen)
        )

compute_p_value(x: NDArray, y_cand: float, return_update: bool = False, tau: float | None = None) -> float | tuple[float, dict]

Compute conformal p-value for (x, y_cand).

Parameters:

Name	Type	Description	Default
x	(array - like, shape(d))		required
y_cand	float		required
return_update	bool	If True, also return a dict with tree info.	False
tau	float or None	Randomization variable for smoothed p-value. If None, draws from U(0,1). For property testing and deterministic behavior, pass a fixed tau.	None

Returns:

Type	Description
float or (float, dict)

Source code in src/online_cp/regressors.py

def compute_p_value(
    self,
    x: NDArray,
    y_cand: float,
    return_update: bool = False,
    tau: float | None = None,
) -> float | tuple[float, dict]:
    """Compute conformal p-value for (x, y_cand).

    Parameters
    ----------
    x : array-like, shape (d,)
    y_cand : float
    return_update : bool
        If True, also return a dict with tree info.
    tau : float or None
        Randomization variable for smoothed p-value. If None, draws from U(0,1).
        For property testing and deterministic behavior, pass a fixed tau.

    Returns
    -------
    float or (float, dict)
    """
    x = np.asarray(x, dtype=float).ravel()
    n = self._safe_size_check(self.X)
    if n == 0:
        return (1.0, {}) if return_update else 1.0

    y_cand = float(y_cand)

    # FIX: Use a deterministic seed based on the test point x for tree building
    # This ensures the same tree is built for all candidate y values
    x_seed = int(hash(tuple(np.round(x, 10).flatten())) % (2**31))
    tree_rnd_gen = np.random.default_rng(x_seed)

    # tau must be INDEPENDENT of the tree, drawn from a separate source
    # If tau is not provided, draw it from a deterministic but SEPARATE source
    if tau is None:
        # Use a different deterministic seed for tau (independent of tree)
        tau_seed = (x_seed * 2654435761) % (2**32)  # Prime multiplier for decorrelation
        tau_rnd_gen = np.random.default_rng(tau_seed)
        tau = tau_rnd_gen.uniform(0.0, 1.0)

    tree, leaf_star, _ = self._build_augmented_tree_with_rng(x, tree_rnd_gen)
    leaf_star_train_idx = leaf_star.indices[leaf_star.indices < n]
    A = len(leaf_star_train_idx)
    B = float(self.y[leaf_star_train_idx].sum()) if A > 0 else 0.0
    ext_sorted = self._outside_ncms_sorted(tree, leaf_star, n)
    # tau was already set above if None
    p_val = float(
        self._pvalues_at(y_cand, A, B, ext_sorted, self.y[leaf_star_train_idx], n, tau)[0]
    )

    # Cache the tree for inspection utilities
    self._last_tree = tree

    if return_update:
        return p_val, {"tree": tree, "leaf_star": leaf_star, "online_tree": self._pending_tree}
    return p_val

predict(x: NDArray, epsilon: float | NDArray | None = None, return_update: bool = False) -> ContinuousPredictionSet | MultiLevelPredictionSet | tuple

Predict a conformal interval for test object x.

Uses an exact O(n) algebraic knot-point solver. No grid search.

Parameters:

Name	Type	Description	Default
x	(array - like, shape(d))		required
epsilon	float or array - like	Significance level(s). Defaults to self.epsilon.	None
return_update	bool	If True, return (interval, update_dict) for learn_one.	False

Returns:

Type	Description
ContinuousPredictionSet or MultiLevelPredictionSet

Source code in src/online_cp/regressors.py

def predict(
    self,
    x: NDArray,
    epsilon: float | NDArray | None = None,
    return_update: bool = False,
) -> ContinuousPredictionSet | MultiLevelPredictionSet | tuple:
    """Predict a conformal interval for test object x.

    Uses an exact O(n) algebraic knot-point solver.  No grid search.

    Parameters
    ----------
    x : array-like, shape (d,)
    epsilon : float or array-like, optional
        Significance level(s).  Defaults to ``self.epsilon``.
    return_update : bool
        If True, return ``(interval, update_dict)`` for ``learn_one``.

    Returns
    -------
    ContinuousPredictionSet or MultiLevelPredictionSet
    """
    x = np.asarray(x, dtype=float).ravel()
    if epsilon is None:
        epsilon = self.epsilon

    n = self._safe_size_check(self.X)
    _inf_interval = self._construct_Gamma((-np.inf, np.inf), epsilon)

    if n == 0:
        if return_update:
            return _inf_interval, {}
        return _inf_interval

    tree, leaf_star, _ = self._build_augmented_tree(x)
    # Cache the tree for inspection utilities (draw_partition etc.),
    # regardless of which prediction branch is taken below.
    self._last_tree = tree

    leaf_star_train_idx = leaf_star.indices[leaf_star.indices < n]
    A = len(leaf_star_train_idx)
    B = self.y[leaf_star_train_idx].sum() if A > 0 else 0.0

    # ---- Draw ONE tau for the entire predict call to ensure deterministic behavior.
    # Use it for smoothed p-value calculation. This must be done BEFORE any A==0 checks.
    tau = self.rnd_gen.uniform(0.0, 1.0)

    # ---- A = 0: empty leaf → total epistemic uncertainty -----------
    if A == 0:
        ext_sorted = self._outside_ncms_sorted(tree, leaf_star, n)

        # Since test score is 0, rank depends entirely on outside scores >= 0
        gt = int(np.sum(ext_sorted > 0))
        eq = int(np.sum(ext_sorted == 0))
        p_const = (gt + tau * (eq + 1)) / (n + 1)

        eps_arr = np.atleast_1d(np.asarray(epsilon, dtype=float))
        predictions = {}
        for eps_val in eps_arr:
            if p_const > eps_val:
                predictions[float(eps_val)] = self._construct_Gamma((-np.inf, np.inf), float(eps_val))
            else:
                predictions[float(eps_val)] = self._construct_Gamma((np.nan, np.nan), float(eps_val))

        if hasattr(epsilon, "__iter__"):
            result = MultiLevelPredictionSet(predictions)
        else:
            result = predictions[float(eps_arr[0])]

        if return_update:
            return result, {"tree": tree, "leaf_star": leaf_star, "online_tree": self._pending_tree}
        return result

    # ---- Collect knot points ----------------------------------------
    # The candidate-independent outside-leaf scores are computed once here
    # and reused for the knots and for every p-value evaluation below.
    ext_sorted = self._outside_ncms_sorted(tree, leaf_star, n)

    # Outside-leaf_star points: ncm_i is fixed → rank changes when
    #   ncm_test(y) = ncm_i  ↔  y = (B ± ncm_i*(A+1)) / A
    knots: list[float] = []
    if ext_sorted.size:
        fac = (A + 1) / A
        knots.extend(((B / A) + fac * ext_sorted).tolist())
        knots.extend(((B / A) - fac * ext_sorted).tolist())

    # Inside-leaf_star points:
    #   always: y_j
    #   if A > 1: (2B − (A+1)*y_j) / (A−1)
    y_star = self.y[leaf_star_train_idx]
    knots.extend(y_star.tolist())
    if A > 1:
        knots.extend(((2.0 * B - (A + 1) * y_star) / (A - 1)).tolist())

    # Sort and de-duplicate (for potential future use)

    # Build knot events with +1/-1 flags
    knots_events, permanent_ties = self._build_knot_events(A, B, ext_sorted, y_star)

    # Perform left-to-right sweep to find valid region
    lo, hi, valid_knots = self._pvalue_from_sweep(
        knots_events, A, B, ext_sorted, y_star, n, permanent_ties,
        epsilon if not hasattr(epsilon, "__iter__") else epsilon[0], tau
    )

    # Build result with correct topology
    eps_arr = np.atleast_1d(np.asarray(epsilon, dtype=float))
    predictions = {}

    for eps_val in eps_arr:
        lo, hi, valid_knots = self._pvalue_from_sweep(
            knots_events, A, B, ext_sorted, y_star, n, permanent_ties, float(eps_val), tau
        )

        if len(valid_knots) == 1 and lo == hi == valid_knots[0]:
            predictions[float(eps_val)] = self._construct_Gamma((lo, hi), float(eps_val))
        else:
            predictions[float(eps_val)] = self._construct_Gamma((lo, hi), float(eps_val))

    if hasattr(epsilon, "__iter__"):
        result = MultiLevelPredictionSet(predictions)
    else:
        result = predictions[float(eps_arr[0])]


    if return_update:
        return result, {"tree": tree, "leaf_star": leaf_star, "online_tree": self._pending_tree}

    return result

online_cp.regressors.ConformalMondrianForestRegressor

Bases: _MondrianRegressorInspection, ConformalRegressor

Conformal predictor using an ensemble of Mondrian trees for regression.

Builds T Mondrian trees per prediction step from the same augmented bag with different seeds. The per-tree NCMs are averaged to give a single smoothed NCM per point, reducing variance while maintaining exact validity.

α_i = (1/T) Σ_t |y_i − μ_{leaf_t(x_i)}|

The prediction set is found via an exact O(Mn log(Mn)) sweepline solver. For forests, the prediction set may be a union of disjoint intervals (the p-value is not guaranteed to be quasi-concave). By default, the convex hull is returned for compatibility with downstream tasks; set return_interval to False to get the exact union of intervals.

Parameters:

Name	Type	Description	Default
n_trees	int	Number of Mondrian trees. Default 10.	10
lifetime	float	Mondrian depth budget L. Default 1.0.	1.0
epsilon	float	Default significance level. Default 0.1.	0.1
rnd_state	int	Seed for the root RNG (each tree gets a derived seed).	None
verbose	int	Verbosity level (0 = silent).	0
n_jobs	int	Number of parallel workers for tree construction. Default 1.	1
grid_resolution	int	Number of grid points for the initial interval search. Default 200.	200
bisection_tol	float	Convergence tolerance for bisection refinement. Default 1e-6.	1e-06
return_interval	bool	If True, return the convex hull (single interval) of the exact prediction set. If False, return the exact prediction set as a union of intervals. Default True for backward compatibility.	True
max_depth	int or None	Hard depth cap for Mondrian trees. Default None.	None
online	bool	If True, use the online Mondrian forest mode. Default False.	False

Source code in src/online_cp/regressors.py

class ConformalMondrianForestRegressor(_MondrianRegressorInspection, ConformalRegressor):
    """Conformal predictor using an ensemble of Mondrian trees for regression.

    Builds T Mondrian trees per prediction step from the same augmented bag
    with different seeds. The per-tree NCMs are averaged to give a single
    smoothed NCM per point, reducing variance while maintaining exact validity.

        α_i = (1/T) Σ_t |y_i − μ_{leaf_t(x_i)}|

    The prediction set is found via an exact O(Mn log(Mn)) sweepline solver.
    For forests, the prediction set may be a union of disjoint intervals (the
    p-value is not guaranteed to be quasi-concave). By default, the convex hull
    is returned for compatibility with downstream tasks; set `return_interval`
    to `False` to get the exact union of intervals.

    Parameters
    ----------
    n_trees : int
        Number of Mondrian trees. Default 10.
    lifetime : float
        Mondrian depth budget L. Default 1.0.
    epsilon : float
        Default significance level. Default 0.1.
    rnd_state : int, optional
        Seed for the root RNG (each tree gets a derived seed).
    verbose : int
        Verbosity level (0 = silent).
    n_jobs : int
        Number of parallel workers for tree construction. Default 1.
    grid_resolution : int
        Number of grid points for the initial interval search. Default 200.
    bisection_tol : float
        Convergence tolerance for bisection refinement. Default 1e-6.
    return_interval : bool, optional
        If True, return the convex hull (single interval) of the exact prediction
        set. If False, return the exact prediction set as a union of intervals.
        Default True for backward compatibility.
    max_depth : int or None, optional
        Hard depth cap for Mondrian trees. Default None.
    online : bool, optional
        If True, use the online Mondrian forest mode. Default False.
    """

    _SAVE_PARAMS: tuple = (
        "n_trees", "lifetime", "epsilon", "rnd_state", "verbose",
        "n_jobs", "grid_resolution", "bisection_tol", "max_depth", "online",
        "return_interval",
    )
    _SAVE_STATE: tuple = ("X", "y", "_forest")

    def __init__(
        self,
        n_trees: int = 10,
        lifetime: float = 1.0,
        epsilon: float = 0.1,
        rnd_state: int | None = None,
        verbose: int = 0,
        n_jobs: int = 1,
        grid_resolution: int = 200,
        bisection_tol: float = 1e-6,
        max_depth: int | None = None,
        online: bool = False,
        return_interval: bool = True,
    ) -> None:
        if max_depth is not None and (not isinstance(max_depth, int) or max_depth < 0):
            raise ValueError("max_depth must be a non-negative integer or None")
        if online:
            if not isinstance(lifetime, (int, float)):
                raise ValueError(
                    "online=True requires a fixed float lifetime; string/auto-tuned "
                    "lifetimes recompute a master tree per step and are batch-only."
                )
            if max_depth is not None:
                raise ValueError(
                    "online=True requires max_depth=None (the projective Mondrian regime)."
                )
        super().__init__(epsilon=epsilon)
        self.n_trees = n_trees
        self.lifetime = lifetime
        self.rnd_state = rnd_state
        self.verbose = verbose
        self.n_jobs = n_jobs
        self.grid_resolution = grid_resolution
        self.bisection_tol = bisection_tol
        self.max_depth = max_depth
        self.online = online
        self.return_interval = return_interval
        self.X: NDArray | None = None
        self.y: NDArray | None = None
        self.rnd_gen = np.random.default_rng(rnd_state)
        self._last_tree = None
        self._last_seeds = None
        self._last_x = None
        self._forest = None
        self._pending_forest = None

    # ------------------------------------------------------------------
    # Training interface
    # ------------------------------------------------------------------

    def learn_initial_training_set(self, X: NDArray, y: NDArray) -> None:
        """Batch-initialize with training data."""
        X = np.atleast_2d(np.asarray(X, dtype=float))
        y = np.asarray(y, dtype=float).ravel()
        if X.shape[0] == 0:
            raise ValueError("Training set cannot be empty")
        if X.shape[0] != y.shape[0]:
            raise ValueError("X and y must have the same length")
        self.X = X.copy()
        self.y = y.copy()
        if self.online:
            seeds = self.rnd_gen.integers(0, 2**31, self.n_trees)
            self._forest = [
                MondrianTree.grow(self.X, np.random.default_rng(int(s)), lifetime=self.lifetime)
                for s in seeds
            ]

    def learn_one(
        self,
        x: NDArray,
        y: float,
        precomputed: dict | None = None,
    ) -> None:
        """Incrementally add one observation.

        In online mode the forest extensions that scored ``x`` (from
        ``precomputed`` or the last ``predict``) are persisted, so each tree's
        stream is a single consistent Mondrian process; otherwise fresh
        extensions are drawn.
        """
        x = np.asarray(x, dtype=float).ravel()
        if self.X is None:
            raise ValueError("Must call learn_initial_training_set first")
        if x.shape[0] != self.X.shape[1]:
            raise ValueError(f"Feature dimension mismatch: got {x.shape[0]}, expected {self.X.shape[1]}")
        self.X = np.vstack([self.X, x])
        self.y = np.append(self.y, float(y))
        if self.online:
            cand = precomputed.get("online_forest") if precomputed else None
            if not self._is_consistent_forest(cand, x):
                cand = getattr(self, "_pending_forest", None)
            self._pending_forest = None
            if self._is_consistent_forest(cand, x):
                self._forest = cand
            else:
                self._forest = [t.extend(x, self.rnd_gen) for t in self._forest]

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _is_consistent_forest(self, cand, x) -> bool:
        """True iff ``cand`` is the persisted forest with every tree extended by exactly ``x``."""
        if cand is None or self._forest is None or len(cand) != len(self._forest):
            return False
        xr = np.asarray(x, dtype=float).ravel()
        base_n = self._forest[0].X.shape[0]
        return all(
            t.X.shape[0] == base_n + 1 and np.array_equal(t.X[-1], xr) for t in cand
        )

    def _forest_summaries(self, X_aug, x, n, seeds: NDArray | None = None):
        """Per-tree ``(base_ncm, ls_idx, ls_mu)`` summaries + viz/reuse caches.

        Batch: T fresh trees from seeds (joblib). Online: extend each persistent
        tree and summarise, caching the extensions for ``learn_one`` reuse.

        Parameters
        ----------
        X_aug : array-like, shape (n+1, d)
            Augmented feature matrix with test point appended.
        x : array-like, shape (d,)
            Test point features.
        n : int
            Number of training points.
        seeds : array-like or None
            Seeds for each tree. If None, uses _tree_seeds().
        """
        if self.online:
            # In online mode, use deterministic seeds if provided
            if seeds is None:
                pending = [t.extend(x, self.rnd_gen) for t in self._forest]
            else:
                pending = [t.extend(x, np.random.default_rng(int(s))) for t, s in zip(self._forest, seeds)]
            self._pending_forest = pending
            self._last_seeds = None
            self._last_x = x.copy()
            self._last_tree = pending[0].root
            return [_summarize_tree_reg(t.root, self.y, n, x) for t in pending]

        # Batch mode
        if seeds is None:
            seeds = self._tree_seeds()
        self._last_seeds = seeds.copy()
        self._last_x = x.copy()
        self._pending_forest = None
        _rng0 = np.random.default_rng(int(seeds[0]))
        self._last_tree = MondrianTree.grow(X_aug, _rng0, lifetime=self.lifetime, max_depth=self.max_depth).root
        return Parallel(n_jobs=self.n_jobs)(
            delayed(_build_tree_summary_reg)(s, X_aug, self.y, n, self.lifetime, self.max_depth) for s in seeds
        )

    def _tree_seeds(self) -> NDArray:
        """Generate n_trees integer seeds from the root RNG."""
        return self.rnd_gen.integers(0, 2**31 - 1, size=self.n_trees)

    def _compute_avg_ncms(self, x: NDArray, y_cand: float) -> tuple[NDArray, float]:
        """Compute averaged NCMs for all training points + test point.

        Returns
        -------
        avg_train_ncms : ndarray, shape (n,)
            Averaged NCM for each training point.
        avg_test_ncm : float
            Averaged NCM for the test point.
        """
        n = self.X.shape[0]
        X_aug = np.vstack([self.X, x.reshape(1, -1)])
        results = self._forest_summaries(X_aug, x, n)

        # Accumulate per-tree NCMs
        # base_ncm[i] = |y_i - mu_leaf_i| (training-only mean)
        # For test point: ncm_test = |A*y_cand - B| / (A+1)
        # For inside points j: ncm_j = |(A+1)*y_j - B - y_cand| / (A+1)
        sum_train_ncms = np.zeros(n, dtype=float)
        sum_test_ncm = 0.0

        for base_ncm, ls_train_idx, ls_mu_train in results:
            A = len(ls_train_idx)
            B = ls_mu_train * A if A > 0 else 0.0

            # Test NCM for this tree
            test_ncm_t = abs(A * y_cand - B) / (A + 1)
            sum_test_ncm += test_ncm_t

            # Training NCMs: start from base, override leaf_star
            ncm_t = base_ncm.copy()
            if A > 0:
                y_star = self.y[ls_train_idx]
                ncm_t[ls_train_idx] = np.abs((A + 1) * y_star - B - y_cand) / (A + 1)

            sum_train_ncms += ncm_t

        T = float(self.n_trees)
        return sum_train_ncms / T, sum_test_ncm / T

    def _compute_avg_ncms_with_seeds(self, x: NDArray, y_cand: float, seeds: NDArray) -> tuple[NDArray, float]:
        """Compute averaged NCMs with deterministic seeds.

        Parameters
        ----------
        x : array-like, shape (d,)
        y_cand : float
        seeds : array-like, shape (n_trees,)
            Deterministic seeds for each tree.

        Returns
        -------
        avg_train_ncms : ndarray, shape (n,)
        avg_test_ncm : float
        """
        n = self.X.shape[0]
        X_aug = np.vstack([self.X, x.reshape(1, -1)])
        results = self._forest_summaries(X_aug, x, n, seeds=seeds)

        sum_train_ncms = np.zeros(n, dtype=float)
        sum_test_ncm = 0.0

        for base_ncm, ls_train_idx, ls_mu_train in results:
            A = len(ls_train_idx)
            B = ls_mu_train * A if A > 0 else 0.0

            test_ncm_t = abs(A * y_cand - B) / (A + 1)
            sum_test_ncm += test_ncm_t

            ncm_t = base_ncm.copy()
            if A > 0:
                y_star = self.y[ls_train_idx]
                ncm_t[ls_train_idx] = np.abs((A + 1) * y_star - B - y_cand) / (A + 1)

            sum_train_ncms += ncm_t

        T = float(self.n_trees)
        return sum_train_ncms / T, sum_test_ncm / T

    # ------------------------------------------------------------------
    # compute_p_value
    # ------------------------------------------------------------------

    def compute_p_value(
        self,
        x: NDArray,
        y_cand: float,
        return_update: bool = False,
        tau: float | None = None,
    ) -> float | tuple[float, dict]:
        """Compute conformal p-value for (x, y_cand).

        Parameters
        ----------
        x : array-like, shape (d,)
        y_cand : float
        return_update : bool
        tau : float or None
            Randomization variable for smoothed p-value. If None, draws from U(0,1).
            For property testing and deterministic behavior, pass a fixed tau.

        Returns
        -------
        float or (float, dict)
        """
        x = np.asarray(x, dtype=float).ravel()
        n = self._safe_size_check(self.X)
        if n == 0:
            return (1.0, {}) if return_update else 1.0

        # FIX: Use deterministic seeds based on x to ensure the same forest for all y
        # The seed for each tree must be deterministic based on x, not on self.rnd_gen
        # Use a hash of x to create reproducible seeds for each tree
        x_seed = int(hash(tuple(np.round(x, 10).flatten())) % (2**31))
        seeds = np.array([x_seed + i for i in range(self.n_trees)])

        # tau is drawn ONCE in predict() and passed to compute_p_value
        # If tau is not provided, draw it from a SEPARATE deterministic source
        # Tau must be independent of the tree-building process
        if tau is None:
            # Use a different deterministic seed for tau (independent of tree seeds)
            tau_seed = (x_seed * 2654435761) % (2**32)  # Prime multiplier for decorrelation
            tau = np.random.default_rng(tau_seed).uniform(0.0, 1.0)

        # Compute NCMs with deterministic seeds
        avg_train_ncms, avg_test_ncm = self._compute_avg_ncms_with_seeds(x, float(y_cand), seeds)

        all_ncms = np.append(avg_train_ncms, avg_test_ncm)
        p_val = float(self._compute_p_value(all_ncms, tau))

        if return_update:
            return p_val, {"online_forest": getattr(self, "_pending_forest", None)}
        return p_val

    # ------------------------------------------------------------------
    # predict  (grid + bisection)
    # ------------------------------------------------------------------

    # ------------------------------------------------------------------
    # Forest indexing
    # ------------------------------------------------------------------

    def __len__(self) -> int:
        """Number of trees in the forest."""
        return self.n_trees

    def __iter__(self):
        """Iterate over individual tree views."""
        for i in range(self.n_trees):
            yield self[i]

    def __getitem__(self, i: int) -> ConformalMondrianTreeRegressor:
        """Return a single-tree view for tree ``i`` from the last predict/compute_p_value call.

        The returned object is a :class:`ConformalMondrianTreeRegressor` whose
        ``_last_tree`` is the exact i-th tree from the last call, enabling all
        T1/T2 inspection and visualisation methods. Its own ``predict()`` would
        rebuild a fresh tree — this object is intended for *inspection only*.

        Parameters
        ----------
        i : int
            Tree index, ``0 <= i < n_trees``.

        Raises
        ------
        RuntimeError
            If ``predict()`` or ``compute_p_value()`` has not been called yet.
        IndexError
            If ``i`` is out of range.
        """
        if self.online:
            pending = getattr(self, "_pending_forest", None)
            if pending is None:
                raise RuntimeError(
                    "No trees cached. Call predict() or compute_p_value() first."
                )
            if not (0 <= i < self.n_trees):
                raise IndexError(f"Tree index {i} out of range [0, {self.n_trees}).")
            view = ConformalMondrianTreeRegressor(
                lifetime=self.lifetime,
                epsilon=self.epsilon,
                max_depth=self.max_depth,
            )
            view.X = self.X
            view.y = self.y
            view._last_tree = pending[i].root
            return view

        if self._last_seeds is None:
            raise RuntimeError(
                "No trees cached. Call predict() or compute_p_value() first."
            )
        if not (0 <= i < self.n_trees):
            raise IndexError(f"Tree index {i} out of range [0, {self.n_trees}).")

        X_aug = np.vstack([self.X, self._last_x])
        rng_i = np.random.default_rng(int(self._last_seeds[i]))
        tree_i = MondrianTree.grow(X_aug, rng_i, lifetime=self.lifetime, max_depth=self.max_depth).root

        view = ConformalMondrianTreeRegressor(
            lifetime=self.lifetime,
            epsilon=self.epsilon,
            max_depth=self.max_depth,
        )
        view.X = self.X
        view.y = self.y
        view._last_tree = tree_i
        return view

    def predict(
        self,
        x: NDArray,
        epsilon: float | NDArray | None = None,
        return_update: bool = False,
        return_interval: bool | None = None,
    ) -> ContinuousPredictionSet | MultiLevelPredictionSet | tuple:
        """Predict a conformal interval for test object x.

        Uses an exact O(Mn log(Mn)) sweepline solver to locate p-value crossings.
        For forests, the prediction set may be a union of disjoint intervals.

        Parameters
        ----------
        x : array-like, shape (d,)
        epsilon : float or array-like, optional
        return_update : bool
        return_interval : bool or None, optional
            If True, return the convex hull (single interval) of the exact prediction
            set. If False, return the exact union of intervals as a DiscretePredictionSet.
            If None (default), use ``self.return_interval``.

        Returns
        -------
        ContinuousPredictionSet or MultiLevelPredictionSet
            Returns a DiscretePredictionSet with multiple intervals if return_interval=False
            and the exact prediction set is disconnected.
        """
        x = np.asarray(x, dtype=float).ravel()
        if epsilon is None:
            epsilon = self.epsilon

        n = self._safe_size_check(self.X)
        if n == 0:
            result = self._make_result(epsilon, -np.inf, np.inf)
            if return_update:
                return result, {}
            return result

        # Use return_interval parameter if provided, otherwise fall back to self.return_interval
        if return_interval is None:
            return_interval = self.return_interval

        # ---- Build all T trees ONCE (batch: seeds; online: extend forest) --
        X_aug = np.vstack([self.X, x.reshape(1, -1)])
        summaries = self._forest_summaries(X_aug, x, n)

        T = float(self.n_trees)
        _TOL = 1e-9
        # Draw ONE tau for the entire predict call to ensure deterministic save/load.
        # Use it for smoothed p-value calculation.
        tau = self.rnd_gen.uniform(0.0, 1.0)

        def p_value_at(y_cand: float) -> float:
            """Evaluate p(y_cand) using pre-built tree summaries and smoothed tau."""
            sum_train = np.zeros(n, dtype=float)
            sum_test = 0.0
            for base_ncm, ls_idx, ls_mu in summaries:
                A = len(ls_idx)
                B = ls_mu * A if A > 0 else 0.0
                test_t = abs(A * y_cand - B) / (A + 1)
                sum_test += test_t
                ncm_t = base_ncm.copy()
                if A > 0:
                    y_star = self.y[ls_idx]
                    ncm_t[ls_idx] = np.abs((A + 1) * y_star - B - y_cand) / (A + 1)
                sum_train += ncm_t
            avg_train = sum_train / T
            avg_test = sum_test / T
            alpha_y = avg_test
            gt = int(np.sum(avg_train > alpha_y + _TOL))
            eq = int(np.sum(np.abs(avg_train - alpha_y) <= _TOL))  # exclude test point, apply tau below
            return float((gt + tau * eq) / (n + 1))

        # ---- Exact O(Mn log(Mn)) Sweepline Solver ----------------------
        # Use the Numba-accelerated solver from _forest_solver module
        M = len(summaries)
        eps_arr = np.atleast_1d(np.asarray(epsilon, dtype=float))
        predictions = {}

        for eps_val in eps_arr:
            if return_interval:
                lo, hi = forest_sweepline_solver(
                    summaries, self.y, float(eps_val), tau, n, M, return_exact=False
                )
                # Single interval mode: pass as (lo, hi) tuple
                if np.isnan(lo) and np.isnan(hi):
                    predictions[float(eps_val)] = self._construct_Gamma((np.nan, np.nan), float(eps_val))
                else:
                    predictions[float(eps_val)] = self._construct_Gamma((lo, hi), float(eps_val))
            else:
                intervals, Y_eval = forest_sweepline_solver(
                    summaries, self.y, float(eps_val), tau, n, M, return_exact=True
                )
                # intervals is a list of (lo, hi) tuples for the exact prediction set
                if not intervals:
                    predictions[float(eps_val)] = self._construct_Gamma([], float(eps_val))
                else:
                    predictions[float(eps_val)] = self._construct_Gamma(intervals, float(eps_val))

        if hasattr(epsilon, "__iter__"):
            result = MultiLevelPredictionSet(predictions)
        else:
            result = predictions[float(eps_arr[0])]

        if return_update:
            return result, {"online_forest": getattr(self, "_pending_forest", None)}
        return result

    @staticmethod
    def _make_result(epsilon, lo: float, hi: float):
        if hasattr(epsilon, "__iter__"):
            return MultiLevelPredictionSet(
                {float(eps): ContinuousPredictionSet([(lo, hi)], float(eps)) for eps in epsilon}
            )
        return ContinuousPredictionSet([(lo, hi)], float(epsilon))

learn_initial_training_set(X: NDArray, y: NDArray) -> None

Batch-initialize with training data.

Source code in src/online_cp/regressors.py

def learn_initial_training_set(self, X: NDArray, y: NDArray) -> None:
    """Batch-initialize with training data."""
    X = np.atleast_2d(np.asarray(X, dtype=float))
    y = np.asarray(y, dtype=float).ravel()
    if X.shape[0] == 0:
        raise ValueError("Training set cannot be empty")
    if X.shape[0] != y.shape[0]:
        raise ValueError("X and y must have the same length")
    self.X = X.copy()
    self.y = y.copy()
    if self.online:
        seeds = self.rnd_gen.integers(0, 2**31, self.n_trees)
        self._forest = [
            MondrianTree.grow(self.X, np.random.default_rng(int(s)), lifetime=self.lifetime)
            for s in seeds
        ]

learn_one(x: NDArray, y: float, precomputed: dict | None = None) -> None

Incrementally add one observation.

In online mode the forest extensions that scored x (from precomputed or the last predict) are persisted, so each tree's stream is a single consistent Mondrian process; otherwise fresh extensions are drawn.

Source code in src/online_cp/regressors.py

def learn_one(
    self,
    x: NDArray,
    y: float,
    precomputed: dict | None = None,
) -> None:
    """Incrementally add one observation.

    In online mode the forest extensions that scored ``x`` (from
    ``precomputed`` or the last ``predict``) are persisted, so each tree's
    stream is a single consistent Mondrian process; otherwise fresh
    extensions are drawn.
    """
    x = np.asarray(x, dtype=float).ravel()
    if self.X is None:
        raise ValueError("Must call learn_initial_training_set first")
    if x.shape[0] != self.X.shape[1]:
        raise ValueError(f"Feature dimension mismatch: got {x.shape[0]}, expected {self.X.shape[1]}")
    self.X = np.vstack([self.X, x])
    self.y = np.append(self.y, float(y))
    if self.online:
        cand = precomputed.get("online_forest") if precomputed else None
        if not self._is_consistent_forest(cand, x):
            cand = getattr(self, "_pending_forest", None)
        self._pending_forest = None
        if self._is_consistent_forest(cand, x):
            self._forest = cand
        else:
            self._forest = [t.extend(x, self.rnd_gen) for t in self._forest]

compute_p_value(x: NDArray, y_cand: float, return_update: bool = False, tau: float | None = None) -> float | tuple[float, dict]

Compute conformal p-value for (x, y_cand).

Parameters:

Name	Type	Description	Default
x	(array - like, shape(d))		required
y_cand	float		required
return_update	bool		False
tau	float or None	Randomization variable for smoothed p-value. If None, draws from U(0,1). For property testing and deterministic behavior, pass a fixed tau.	None

Returns:

Type	Description
float or (float, dict)

Source code in src/online_cp/regressors.py

def compute_p_value(
    self,
    x: NDArray,
    y_cand: float,
    return_update: bool = False,
    tau: float | None = None,
) -> float | tuple[float, dict]:
    """Compute conformal p-value for (x, y_cand).

    Parameters
    ----------
    x : array-like, shape (d,)
    y_cand : float
    return_update : bool
    tau : float or None
        Randomization variable for smoothed p-value. If None, draws from U(0,1).
        For property testing and deterministic behavior, pass a fixed tau.

    Returns
    -------
    float or (float, dict)
    """
    x = np.asarray(x, dtype=float).ravel()
    n = self._safe_size_check(self.X)
    if n == 0:
        return (1.0, {}) if return_update else 1.0

    # FIX: Use deterministic seeds based on x to ensure the same forest for all y
    # The seed for each tree must be deterministic based on x, not on self.rnd_gen
    # Use a hash of x to create reproducible seeds for each tree
    x_seed = int(hash(tuple(np.round(x, 10).flatten())) % (2**31))
    seeds = np.array([x_seed + i for i in range(self.n_trees)])

    # tau is drawn ONCE in predict() and passed to compute_p_value
    # If tau is not provided, draw it from a SEPARATE deterministic source
    # Tau must be independent of the tree-building process
    if tau is None:
        # Use a different deterministic seed for tau (independent of tree seeds)
        tau_seed = (x_seed * 2654435761) % (2**32)  # Prime multiplier for decorrelation
        tau = np.random.default_rng(tau_seed).uniform(0.0, 1.0)

    # Compute NCMs with deterministic seeds
    avg_train_ncms, avg_test_ncm = self._compute_avg_ncms_with_seeds(x, float(y_cand), seeds)

    all_ncms = np.append(avg_train_ncms, avg_test_ncm)
    p_val = float(self._compute_p_value(all_ncms, tau))

    if return_update:
        return p_val, {"online_forest": getattr(self, "_pending_forest", None)}
    return p_val

__len__() -> int

Number of trees in the forest.

Source code in src/online_cp/regressors.py

def __len__(self) -> int:
    """Number of trees in the forest."""
    return self.n_trees

__iter__()

Iterate over individual tree views.

Source code in src/online_cp/regressors.py

def __iter__(self):
    """Iterate over individual tree views."""
    for i in range(self.n_trees):
        yield self[i]

__getitem__(i: int) -> ConformalMondrianTreeRegressor

Return a single-tree view for tree i from the last predict/compute_p_value call.

The returned object is a :class:ConformalMondrianTreeRegressor whose _last_tree is the exact i-th tree from the last call, enabling all T1/T2 inspection and visualisation methods. Its own predict() would rebuild a fresh tree — this object is intended for inspection only.

Parameters:

Name	Type	Description	Default
i	int	Tree index, 0 <= i < n_trees.	required

Raises:

Type	Description
RuntimeError	If predict() or compute_p_value() has not been called yet.
IndexError	If i is out of range.

Source code in src/online_cp/regressors.py

def __getitem__(self, i: int) -> ConformalMondrianTreeRegressor:
    """Return a single-tree view for tree ``i`` from the last predict/compute_p_value call.

    The returned object is a :class:`ConformalMondrianTreeRegressor` whose
    ``_last_tree`` is the exact i-th tree from the last call, enabling all
    T1/T2 inspection and visualisation methods. Its own ``predict()`` would
    rebuild a fresh tree — this object is intended for *inspection only*.

    Parameters
    ----------
    i : int
        Tree index, ``0 <= i < n_trees``.

    Raises
    ------
    RuntimeError
        If ``predict()`` or ``compute_p_value()`` has not been called yet.
    IndexError
        If ``i`` is out of range.
    """
    if self.online:
        pending = getattr(self, "_pending_forest", None)
        if pending is None:
            raise RuntimeError(
                "No trees cached. Call predict() or compute_p_value() first."
            )
        if not (0 <= i < self.n_trees):
            raise IndexError(f"Tree index {i} out of range [0, {self.n_trees}).")
        view = ConformalMondrianTreeRegressor(
            lifetime=self.lifetime,
            epsilon=self.epsilon,
            max_depth=self.max_depth,
        )
        view.X = self.X
        view.y = self.y
        view._last_tree = pending[i].root
        return view

    if self._last_seeds is None:
        raise RuntimeError(
            "No trees cached. Call predict() or compute_p_value() first."
        )
    if not (0 <= i < self.n_trees):
        raise IndexError(f"Tree index {i} out of range [0, {self.n_trees}).")

    X_aug = np.vstack([self.X, self._last_x])
    rng_i = np.random.default_rng(int(self._last_seeds[i]))
    tree_i = MondrianTree.grow(X_aug, rng_i, lifetime=self.lifetime, max_depth=self.max_depth).root

    view = ConformalMondrianTreeRegressor(
        lifetime=self.lifetime,
        epsilon=self.epsilon,
        max_depth=self.max_depth,
    )
    view.X = self.X
    view.y = self.y
    view._last_tree = tree_i
    return view

predict(x: NDArray, epsilon: float | NDArray | None = None, return_update: bool = False, return_interval: bool | None = None) -> ContinuousPredictionSet | MultiLevelPredictionSet | tuple

Predict a conformal interval for test object x.

Uses an exact O(Mn log(Mn)) sweepline solver to locate p-value crossings. For forests, the prediction set may be a union of disjoint intervals.

Parameters:

Name	Type	Description	Default
x	(array - like, shape(d))		required
epsilon	float or array - like		None
return_update	bool		False
return_interval	bool or None	If True, return the convex hull (single interval) of the exact prediction set. If False, return the exact union of intervals as a DiscretePredictionSet. If None (default), use self.return_interval.	None

Returns:

Type	Description
ContinuousPredictionSet or MultiLevelPredictionSet	Returns a DiscretePredictionSet with multiple intervals if return_interval=False and the exact prediction set is disconnected.

Source code in src/online_cp/regressors.py

def predict(
    self,
    x: NDArray,
    epsilon: float | NDArray | None = None,
    return_update: bool = False,
    return_interval: bool | None = None,
) -> ContinuousPredictionSet | MultiLevelPredictionSet | tuple:
    """Predict a conformal interval for test object x.

    Uses an exact O(Mn log(Mn)) sweepline solver to locate p-value crossings.
    For forests, the prediction set may be a union of disjoint intervals.

    Parameters
    ----------
    x : array-like, shape (d,)
    epsilon : float or array-like, optional
    return_update : bool
    return_interval : bool or None, optional
        If True, return the convex hull (single interval) of the exact prediction
        set. If False, return the exact union of intervals as a DiscretePredictionSet.
        If None (default), use ``self.return_interval``.

    Returns
    -------
    ContinuousPredictionSet or MultiLevelPredictionSet
        Returns a DiscretePredictionSet with multiple intervals if return_interval=False
        and the exact prediction set is disconnected.
    """
    x = np.asarray(x, dtype=float).ravel()
    if epsilon is None:
        epsilon = self.epsilon

    n = self._safe_size_check(self.X)
    if n == 0:
        result = self._make_result(epsilon, -np.inf, np.inf)
        if return_update:
            return result, {}
        return result

    # Use return_interval parameter if provided, otherwise fall back to self.return_interval
    if return_interval is None:
        return_interval = self.return_interval

    # ---- Build all T trees ONCE (batch: seeds; online: extend forest) --
    X_aug = np.vstack([self.X, x.reshape(1, -1)])
    summaries = self._forest_summaries(X_aug, x, n)

    T = float(self.n_trees)
    _TOL = 1e-9
    # Draw ONE tau for the entire predict call to ensure deterministic save/load.
    # Use it for smoothed p-value calculation.
    tau = self.rnd_gen.uniform(0.0, 1.0)

    def p_value_at(y_cand: float) -> float:
        """Evaluate p(y_cand) using pre-built tree summaries and smoothed tau."""
        sum_train = np.zeros(n, dtype=float)
        sum_test = 0.0
        for base_ncm, ls_idx, ls_mu in summaries:
            A = len(ls_idx)
            B = ls_mu * A if A > 0 else 0.0
            test_t = abs(A * y_cand - B) / (A + 1)
            sum_test += test_t
            ncm_t = base_ncm.copy()
            if A > 0:
                y_star = self.y[ls_idx]
                ncm_t[ls_idx] = np.abs((A + 1) * y_star - B - y_cand) / (A + 1)
            sum_train += ncm_t
        avg_train = sum_train / T
        avg_test = sum_test / T
        alpha_y = avg_test
        gt = int(np.sum(avg_train > alpha_y + _TOL))
        eq = int(np.sum(np.abs(avg_train - alpha_y) <= _TOL))  # exclude test point, apply tau below
        return float((gt + tau * eq) / (n + 1))

    # ---- Exact O(Mn log(Mn)) Sweepline Solver ----------------------
    # Use the Numba-accelerated solver from _forest_solver module
    M = len(summaries)
    eps_arr = np.atleast_1d(np.asarray(epsilon, dtype=float))
    predictions = {}

    for eps_val in eps_arr:
        if return_interval:
            lo, hi = forest_sweepline_solver(
                summaries, self.y, float(eps_val), tau, n, M, return_exact=False
            )
            # Single interval mode: pass as (lo, hi) tuple
            if np.isnan(lo) and np.isnan(hi):
                predictions[float(eps_val)] = self._construct_Gamma((np.nan, np.nan), float(eps_val))
            else:
                predictions[float(eps_val)] = self._construct_Gamma((lo, hi), float(eps_val))
        else:
            intervals, Y_eval = forest_sweepline_solver(
                summaries, self.y, float(eps_val), tau, n, M, return_exact=True
            )
            # intervals is a list of (lo, hi) tuples for the exact prediction set
            if not intervals:
                predictions[float(eps_val)] = self._construct_Gamma([], float(eps_val))
            else:
                predictions[float(eps_val)] = self._construct_Gamma(intervals, float(eps_val))

    if hasattr(epsilon, "__iter__"):
        result = MultiLevelPredictionSet(predictions)
    else:
        result = predictions[float(eps_arr[0])]

    if return_update:
        return result, {"online_forest": getattr(self, "_pending_forest", None)}
    return result

Mondrian tree CPS

The conformal predictive system: a leaf-local signed-residual NCM that yields a full predictive distribution (step CDF of the category's training labels, with the ALRW2 §7.2 randomisation band). See the conformal predictive systems page for the API reference and .roadmap/archive/mondrian-cps-verification.md for the construction and its finite-sample validity boundary.

online_cp.CPS.MondrianTreePredictionMachine

Bases: ConformalPredictiveSystem

Conformal predictive system based on Mondrian trees with leaf-local signed residuals.

This CPS uses the Mondrian tree partition together with the leaf-local signed residual as the conformity measure. For a test point x, the category is the leaf of the partition containing x — or, when that leaf holds no training data, the deepest ancestor whose subtree does (the empty-leaf walk-up convention, see below). If the category holds A training points with y-sum B, then for a candidate response y the signed residuals (category mean including the candidate) are

• test point: (A·y − B) / (A+1)
• category point j: ((A+1)·y_j − B − y) / (A+1)

so that α_n^y − α_j^y = y − y_j is strictly increasing in y (slope 1). The critical values are therefore exactly the A category labels, and the conformal predictive distribution is the interval-valued step CDF of those labels (ALRW2 §7.2 band):

P̂(y) ∈ [ #{j : y_j < y}/A ,  #{j : y_j ≤ y}/A ].

The CPD is thus genuinely Mondrian: test points in different categories (different label sets) receive different CPDs, and the partition drives the prediction. Once the partition is known the CPD costs O(A log A) in the category size A.

Validity

Given the X-only partition, the test label is exchangeable with the category training labels (ALRW2 §2.2.9, Mondrian taxonomy) whenever the conditional label distribution is (approximately) constant over the category box, so the CPD is calibrated conditional on the category — asymptotically in the category size A:

P( y ∈ Γ^ε | category ) → 1 − ε   as A → ∞.

At finite A the two-sided quantile set under-covers by ≈ (1 − Aε/2)/(A+1) under exact exchangeability (see .roadmap/archive/mondrian-cps-verification.md).

Empty-leaf convention

A leaf may contain the test point but no training data. The category is then the deepest ancestor holding at least one training point. (Using the empty leaf itself would give a p-value that is identically the randomisation τ for every y — not a valid randomised predictive system, as it never moves from 0 to 1.)

Note

The partition is grown with the shared rnd_gen (advanced in place), as in the Mondrian predictors; within one prediction step use the return_update pattern if you evaluate the CPD more than once.

Parameters:

Name	Type	Description	Default
lifetime	float or {'sqrt_n', 'density'}	Mondrian tree depth budget L. Larger → more splits → smaller leaves. Default 1.0. Can also be a string for unsupervised automatic tuning (only X is used; labels are never accessed): 'sqrt_n': halt when the leaf count drops to ≤ √n. 'density': choose τ maximising Σ_leaf n_leaf · log(n_leaf / V_leaf).	1.0
epsilon	float	Default significance level. Default 0.1.	0.1
rnd_state	int	Seed for reproducibility.	None
verbose	int	Verbosity level (0 = silent).	0
max_depth	int or None	Hard depth cap. Nodes at depth >= max_depth become leaves regardless of remaining lifetime. None = no cap.	None
feature_weights	'variance'	Unsupervised feature importance for split dimension sampling. Labels are never accessed. None (default): uniform — standard isotropic Mondrian. 'variance': split probabilities ∝ population variance per feature. High-variance dimensions are sampled more often. array of shape (d,): explicit non-negative weights (normalised internally; need not sum to 1).	'variance'

Source code in src/online_cp/CPS.py

class MondrianTreePredictionMachine(ConformalPredictiveSystem):
    r"""Conformal predictive system based on Mondrian trees with leaf-local signed residuals.

    This CPS uses the Mondrian tree partition together with the
    **leaf-local signed residual** as the conformity measure.  For a
    test point ``x``, the *category* is the leaf of the partition
    containing ``x`` — or, when that leaf holds no training data, the
    deepest ancestor whose subtree does (the empty-leaf walk-up
    convention, see below).  If the category holds ``A`` training
    points with y-sum ``B``, then for a candidate response ``y`` the
    signed residuals (category mean **including** the candidate) are

    * test point:       ``(A·y − B) / (A+1)``
    * category point j: ``((A+1)·y_j − B − y) / (A+1)``

    so that ``α_n^y − α_j^y = y − y_j`` is strictly increasing in ``y``
    (slope 1).  The critical values are therefore exactly the ``A``
    category labels, and the conformal predictive distribution is the
    interval-valued step CDF of those labels (ALRW2 §7.2 band):

        P̂(y) ∈ [ #{j : y_j < y}/A ,  #{j : y_j ≤ y}/A ].

    The CPD is thus genuinely *Mondrian*: test points in different
    categories (different label sets) receive different CPDs, and the
    partition drives the prediction.  Once the partition is known the
    CPD costs O(A log A) in the category size A.

    Validity
    --------
    Given the X-only partition, the test label is exchangeable with the
    category training labels (ALRW2 §2.2.9, Mondrian taxonomy) whenever
    the conditional label distribution is (approximately) constant over
    the category box, so the CPD is calibrated conditional on the
    category — asymptotically in the category size A:

        P( y ∈ Γ^ε | category ) → 1 − ε   as A → ∞.

    At finite A the two-sided quantile set under-covers by ≈ (1 − Aε/2)/(A+1)
    under exact exchangeability (see `.roadmap/archive/mondrian-cps-verification.md`).

    Empty-leaf convention
    ---------------------
    A leaf may contain the test point but no training data.  The
    category is then the deepest ancestor holding at least one
    training point.  (Using the empty leaf itself would give a
    p-value that is identically the randomisation ``τ`` for every ``y``
    — not a valid randomised predictive system, as it never moves from
    0 to 1.)

    Note
    ----
    The partition is grown with the shared ``rnd_gen`` (advanced in
    place), as in the Mondrian predictors; within one prediction step
    use the ``return_update`` pattern if you evaluate the CPD more than
    once.

    Parameters
    ----------
    lifetime : float or {'sqrt_n', 'density'}
        Mondrian tree depth budget L. Larger → more splits → smaller leaves.
        Default 1.0.

        Can also be a string for unsupervised automatic tuning (only X is
        used; labels are never accessed):

        - ``'sqrt_n'``: halt when the leaf count drops to ≤ √n.
        - ``'density'``: choose τ maximising Σ_leaf n_leaf · log(n_leaf / V_leaf).
    epsilon : float
        Default significance level. Default 0.1.
    rnd_state : int, optional
        Seed for reproducibility.
    verbose : int
        Verbosity level (0 = silent).
    max_depth : int or None
        Hard depth cap. Nodes at depth >= max_depth become leaves regardless
        of remaining lifetime. None = no cap.
    feature_weights : {'variance'} or array-like of shape (d,) or None
        Unsupervised feature importance for split dimension sampling. Labels
        are never accessed.

        - ``None`` (default): uniform — standard isotropic Mondrian.
        - ``'variance'``: split probabilities ∝ population variance per
          feature. High-variance dimensions are sampled more often.
        - array of shape (d,): explicit non-negative weights (normalised
          internally; need not sum to 1).
    """

    _SAVE_PARAMS: tuple = (
        "lifetime", "epsilon", "rnd_state", "verbose", "max_depth",
        "feature_weights",
    )
    _SAVE_STATE: tuple = ("X", "y")

    def __init__(
        self,
        lifetime: float | str = 1.0,
        epsilon: float = 0.1,
        rnd_state: int | None = None,
        verbose: int = 0,
        max_depth: int | None = None,
        feature_weights: str | NDArray | None = None,
    ) -> None:
        if max_depth is not None and (not isinstance(max_depth, int) or max_depth < 0):
            raise ValueError("max_depth must be a non-negative integer or None")
        if isinstance(lifetime, str) and lifetime not in ("sqrt_n", "density"):
            raise ValueError(
                f"Unknown lifetime string {lifetime!r}. "
                "Recognised values: 'sqrt_n', 'density'."
            )
        if isinstance(feature_weights, str) and feature_weights not in ("variance",):
            raise ValueError(
                f"Unknown feature_weights string {feature_weights!r}. "
                "Recognised value: 'variance'."
            )
        super().__init__(epsilon=epsilon)
        self.lifetime = lifetime
        self.feature_weights = feature_weights
        self.rnd_state = rnd_state
        self.verbose = verbose
        self.max_depth = max_depth
        self.X: NDArray | None = None
        self.y: NDArray | None = None
        self.rnd_gen = np.random.default_rng(rnd_state)

    def learn_initial_training_set(self, X: NDArray, y: NDArray) -> None:
        """Batch-initialize with training data.

        Parameters
        ----------
        X : ndarray, shape (n, d)
        y : ndarray, shape (n,)  — real-valued labels
        """
        X = np.atleast_2d(np.asarray(X, dtype=float))
        y = np.asarray(y, dtype=float).ravel()
        if X.shape[0] == 0:
            raise ValueError("Training set cannot be empty")
        if X.shape[0] != y.shape[0]:
            raise ValueError("X and y must have the same length")
        self.X = X.copy()
        self.y = y.copy()

    def learn_one(self, x: NDArray, y: float, precomputed: dict | None = None) -> None:
        """Incrementally add one observation.

        Parameters
        ----------
        x : array-like, shape (d,)
        y : float
        precomputed : dict, optional
            Not used. Present for API consistency.
        """
        x = np.asarray(x, dtype=float).ravel()
        if self.X is None:
            raise ValueError("Must call learn_initial_training_set first")
        if x.shape[0] != self.X.shape[1]:
            raise ValueError(f"Feature dimension mismatch: got {x.shape[0]}, expected {self.X.shape[1]}")
        self.X = np.vstack([self.X, x])
        self.y = np.append(self.y, float(y))

    def predict_cpd(self, x, return_update=False):
        r"""Compute the conformal predictive distribution for test object x.

        The CPD is the interval-valued step CDF of the training labels
        in the *category* of ``x`` — the leaf of the Mondrian partition
        containing ``x``, or, when that leaf holds no training data,
        the deepest ancestor whose subtree does.

        Parameters
        ----------
        x : array-like, shape (d,)
        return_update : bool
            If True, return ``(cpd, precomputed)`` where ``precomputed`` is a
            dict that can be passed to ``learn_one``. For Mondrian trees this
            is currently empty but reserved for future incremental updates.

        Returns
        -------
        MondrianPredictiveDistributionFunction
            The conformal predictive distribution for x.
        """
        x = np.asarray(x, dtype=float).ravel()
        n = self._safe_size_check(self.X)

        if n == 0:
            # Empty training set → completely non-informative CPS
            Y = np.array([-np.inf, np.inf])
            cpd = MondrianPredictiveDistributionFunction(Y, np.array([0.0, 1.0]), np.array([0.0, 1.0]), epsilon=self.epsilon)
            if return_update:
                return cpd, {}
            return cpd

        # Lazy imports
        MondrianTree = _get_mondrian_tree()
        _select_category, _training_indices = _get_cps_helpers()

        # Build Mondrian tree on augmented data [X; x]
        X_aug = np.vstack([self.X, x])
        tree = MondrianTree.grow(
            X_aug,
            self.rnd_gen,
            lifetime=self.lifetime,
            x_test=x,
            max_depth=self.max_depth,
            feature_weights=self.feature_weights,
            verbose=self.verbose,
        ).root

        # Category: the leaf containing x, or (empty-leaf convention) the
        # deepest ancestor holding at least one training point.
        category = _select_category(tree, x, n)
        labels = self.y[_training_indices(category, n)]

        # Knots at the sorted category labels; signed-residual monotonicity
        # (α_n^y − α_j^y = y − y_j) makes the category labels the critical
        # values, so the category mean B does not appear in the final CPD.
        A = labels.size
        Y = np.empty(A + 2)
        Y[0] = -np.inf
        Y[-1] = np.inf
        Y[1:-1] = np.sort(labels)

        # Knot bands: L[j] = #{labels < Y[j]}/A, U[j] = #{labels <= Y[j]}/A
        L = np.empty(A + 2)
        U = np.empty(A + 2)
        L[0] = U[0] = 0.0
        L[-1] = U[-1] = 1.0
        L[1:-1] = np.searchsorted(Y[1:-1], Y[1:-1], side="left") / A
        U[1:-1] = np.searchsorted(Y[1:-1], Y[1:-1], side="right") / A

        cpd = MondrianPredictiveDistributionFunction(Y, L, U, epsilon=self.epsilon)

        if return_update:
            return cpd, {}
        else:
            return cpd

    def _safe_size_check(self, X):
        """Helper to check if X is None or empty."""
        if X is None:
            return 0
        return X.shape[0]

learn_initial_training_set(X: NDArray, y: NDArray) -> None

Batch-initialize with training data.

Parameters:

Name	Type	Description	Default
X	(ndarray, shape(n, d))		required
y	ndarray, shape (n,) — real-valued labels		required

Source code in src/online_cp/CPS.py

def learn_initial_training_set(self, X: NDArray, y: NDArray) -> None:
    """Batch-initialize with training data.

    Parameters
    ----------
    X : ndarray, shape (n, d)
    y : ndarray, shape (n,)  — real-valued labels
    """
    X = np.atleast_2d(np.asarray(X, dtype=float))
    y = np.asarray(y, dtype=float).ravel()
    if X.shape[0] == 0:
        raise ValueError("Training set cannot be empty")
    if X.shape[0] != y.shape[0]:
        raise ValueError("X and y must have the same length")
    self.X = X.copy()
    self.y = y.copy()

learn_one(x: NDArray, y: float, precomputed: dict | None = None) -> None

Incrementally add one observation.

Parameters:

Name	Type	Description	Default
x	(array - like, shape(d))		required
y	float		required
precomputed	dict	Not used. Present for API consistency.	None

Source code in src/online_cp/CPS.py

def learn_one(self, x: NDArray, y: float, precomputed: dict | None = None) -> None:
    """Incrementally add one observation.

    Parameters
    ----------
    x : array-like, shape (d,)
    y : float
    precomputed : dict, optional
        Not used. Present for API consistency.
    """
    x = np.asarray(x, dtype=float).ravel()
    if self.X is None:
        raise ValueError("Must call learn_initial_training_set first")
    if x.shape[0] != self.X.shape[1]:
        raise ValueError(f"Feature dimension mismatch: got {x.shape[0]}, expected {self.X.shape[1]}")
    self.X = np.vstack([self.X, x])
    self.y = np.append(self.y, float(y))

predict_cpd(x, return_update=False)

Compute the conformal predictive distribution for test object x.

The CPD is the interval-valued step CDF of the training labels in the category of x — the leaf of the Mondrian partition containing x, or, when that leaf holds no training data, the deepest ancestor whose subtree does.

Parameters:

Name	Type	Description	Default
x	(array - like, shape(d))		required
return_update	bool	If True, return (cpd, precomputed) where precomputed is a dict that can be passed to learn_one. For Mondrian trees this is currently empty but reserved for future incremental updates.	False

Returns:

Type	Description
MondrianPredictiveDistributionFunction	The conformal predictive distribution for x.

Source code in src/online_cp/CPS.py

def predict_cpd(self, x, return_update=False):
    r"""Compute the conformal predictive distribution for test object x.

    The CPD is the interval-valued step CDF of the training labels
    in the *category* of ``x`` — the leaf of the Mondrian partition
    containing ``x``, or, when that leaf holds no training data,
    the deepest ancestor whose subtree does.

    Parameters
    ----------
    x : array-like, shape (d,)
    return_update : bool
        If True, return ``(cpd, precomputed)`` where ``precomputed`` is a
        dict that can be passed to ``learn_one``. For Mondrian trees this
        is currently empty but reserved for future incremental updates.

    Returns
    -------
    MondrianPredictiveDistributionFunction
        The conformal predictive distribution for x.
    """
    x = np.asarray(x, dtype=float).ravel()
    n = self._safe_size_check(self.X)

    if n == 0:
        # Empty training set → completely non-informative CPS
        Y = np.array([-np.inf, np.inf])
        cpd = MondrianPredictiveDistributionFunction(Y, np.array([0.0, 1.0]), np.array([0.0, 1.0]), epsilon=self.epsilon)
        if return_update:
            return cpd, {}
        return cpd

    # Lazy imports
    MondrianTree = _get_mondrian_tree()
    _select_category, _training_indices = _get_cps_helpers()

    # Build Mondrian tree on augmented data [X; x]
    X_aug = np.vstack([self.X, x])
    tree = MondrianTree.grow(
        X_aug,
        self.rnd_gen,
        lifetime=self.lifetime,
        x_test=x,
        max_depth=self.max_depth,
        feature_weights=self.feature_weights,
        verbose=self.verbose,
    ).root

    # Category: the leaf containing x, or (empty-leaf convention) the
    # deepest ancestor holding at least one training point.
    category = _select_category(tree, x, n)
    labels = self.y[_training_indices(category, n)]

    # Knots at the sorted category labels; signed-residual monotonicity
    # (α_n^y − α_j^y = y − y_j) makes the category labels the critical
    # values, so the category mean B does not appear in the final CPD.
    A = labels.size
    Y = np.empty(A + 2)
    Y[0] = -np.inf
    Y[-1] = np.inf
    Y[1:-1] = np.sort(labels)

    # Knot bands: L[j] = #{labels < Y[j]}/A, U[j] = #{labels <= Y[j]}/A
    L = np.empty(A + 2)
    U = np.empty(A + 2)
    L[0] = U[0] = 0.0
    L[-1] = U[-1] = 1.0
    L[1:-1] = np.searchsorted(Y[1:-1], Y[1:-1], side="left") / A
    U[1:-1] = np.searchsorted(Y[1:-1], Y[1:-1], side="right") / A

    cpd = MondrianPredictiveDistributionFunction(Y, L, U, epsilon=self.epsilon)

    if return_update:
        return cpd, {}
    else:
        return cpd

Mondrian Venn predictor

Uses the Mondrian leaf as the Venn taxonomy category (a bag function → Venn validity, ALRW2 Thm 6.4).

online_cp.venn.MondrianVennPredictor

Bases: VennPredictor

Online Venn predictor using the Mondrian tree partition as taxonomy.

Uses a Mondrian tree leaf as the Venn taxonomy category (ALRW2 §6.2). At each call to predict(x) a single Mondrian tree is built from the augmented dataset {X_train ∪ {x_test}}. Each example is assigned to a leaf; the multiprobability prediction is the empirical label distribution within the test point's leaf.

The Mondrian partition depends only on X-values (bounding boxes), so the same tree is reused for all hypothesised test labels v ∈ Y, making each predict call build only one tree regardless of label-space size. The partition is a bag function (permutation-invariant in the augmented training set), guaranteeing exact Venn validity under exchangeability (ALRW2 Thm 6.4).

Parameters:

Name	Type	Description	Default
lifetime	float or {'sqrt_n', 'density'}	Mondrian time budget. Larger → more splits → smaller leaves. Leaves are created when the cumulative split time exceeds lifetime. Default 1.0. Can also be a string for unsupervised automatic tuning (only X is used; labels are never accessed): 'sqrt_n': halt when the leaf count drops to ≤ √n. Coarsens the partition in sparse regions. 'density': choose τ maximising the histogram log-likelihood Σ_leaf n_leaf · log(n_leaf / V_leaf). Adapts to feature density.	1.0
max_depth	int or None	Hard depth cap on the tree. Nodes at depth >= max_depth become leaves regardless of remaining lifetime. None = no cap.	None
min_samples_leaf	int	Minimum number of points in each leaf. A split that would produce a child with fewer than min_samples_leaf points is refused. Default 1 (no restriction; equivalent to the canonical Mondrian algorithm). Set to 2 to eliminate singleton leaves and avoid extreme p0/p1 outputs.	1
label_space	array - like or None	Explicit set of possible labels. If None, inferred from training data. When provided, the label space is fixed and labels outside it are rejected.	None
rnd_state	int or None	Seed for the random generator used to build Mondrian trees. Predictions are reproducible for a fixed seed and fixed call order.	None
feature_weights	'variance'	Unsupervised feature importance for split dimension sampling. Labels are never accessed. None (default): uniform — standard isotropic Mondrian. 'variance': split probabilities ∝ population variance per feature. High-variance dimensions are sampled more often. array of shape (d,): explicit non-negative weights (normalised internally; need not sum to 1).	'variance'

Examples:

>>> import numpy as np
>>> np.random.seed(0)
>>> X = np.random.randn(40, 2)
>>> y = (X[:, 0] + X[:, 1] > 0).astype(int)
>>> mvp = MondrianVennPredictor(lifetime=1.0, rnd_state=42)
>>> mvp.learn_initial_training_set(X[:30], y[:30])
>>> pred = mvp.predict(X[30])
>>> bool(0.0 <= pred.p0 <= 1.0 and 0.0 <= pred.p1 <= 1.0)
True

Source code in src/online_cp/venn.py

class MondrianVennPredictor(VennPredictor):
    """Online Venn predictor using the Mondrian tree partition as taxonomy.

    Uses a Mondrian tree leaf as the Venn taxonomy category (ALRW2 §6.2).
    At each call to ``predict(x)`` a single Mondrian tree is built from the
    augmented dataset {X_train ∪ {x_test}}.  Each example is assigned to a
    leaf; the multiprobability prediction is the empirical label distribution
    within the test point's leaf.

    The Mondrian partition depends only on X-values (bounding boxes), so the
    same tree is reused for all hypothesised test labels v ∈ Y, making each
    ``predict`` call build only one tree regardless of label-space size.
    The partition is a bag function (permutation-invariant in the augmented
    training set), guaranteeing exact Venn validity under exchangeability
    (ALRW2 Thm 6.4).

    Parameters
    ----------
    lifetime : float or {'sqrt_n', 'density'}
        Mondrian time budget. Larger → more splits → smaller leaves.
        Leaves are created when the cumulative split time exceeds ``lifetime``.
        Default 1.0.

        Can also be a string for unsupervised automatic tuning (only X is
        used; labels are never accessed):

        - ``'sqrt_n'``: halt when the leaf count drops to ≤ √n. Coarsens
          the partition in sparse regions.
        - ``'density'``: choose τ maximising the histogram log-likelihood
          Σ_leaf n_leaf · log(n_leaf / V_leaf). Adapts to feature density.
    max_depth : int or None
        Hard depth cap on the tree. Nodes at depth >= max_depth become leaves
        regardless of remaining ``lifetime``. None = no cap.
    min_samples_leaf : int
        Minimum number of points in each leaf. A split that would produce a
        child with fewer than ``min_samples_leaf`` points is refused. Default 1
        (no restriction; equivalent to the canonical Mondrian algorithm). Set
        to 2 to eliminate singleton leaves and avoid extreme p0/p1 outputs.
    label_space : array-like or None
        Explicit set of possible labels. If None, inferred from training data.
        When provided, the label space is fixed and labels outside it are
        rejected.
    rnd_state : int or None
        Seed for the random generator used to build Mondrian trees. Predictions
        are reproducible for a fixed seed and fixed call order.
    feature_weights : {'variance'} or array-like of shape (d,) or None
        Unsupervised feature importance for split dimension sampling. Labels
        are never accessed.

        - ``None`` (default): uniform — standard isotropic Mondrian.
        - ``'variance'``: split probabilities ∝ population variance per
          feature. High-variance dimensions are sampled more often.
        - array of shape (d,): explicit non-negative weights (normalised
          internally; need not sum to 1).

    Examples
    --------
    >>> import numpy as np
    >>> np.random.seed(0)
    >>> X = np.random.randn(40, 2)
    >>> y = (X[:, 0] + X[:, 1] > 0).astype(int)
    >>> mvp = MondrianVennPredictor(lifetime=1.0, rnd_state=42)
    >>> mvp.learn_initial_training_set(X[:30], y[:30])
    >>> pred = mvp.predict(X[30])
    >>> bool(0.0 <= pred.p0 <= 1.0 and 0.0 <= pred.p1 <= 1.0)
    True
    """

    _SAVE_PARAMS: tuple = (
        "lifetime", "max_depth", "min_samples_leaf", "label_space", "rnd_state",
        "feature_weights",
    )
    _SAVE_STATE: tuple = ("X", "y", "label_space", "_label_space_fixed")

    def __init__(
        self,
        lifetime: float | str = 1.0,
        max_depth: int | None = None,
        min_samples_leaf: int = 1,
        label_space=None,
        rnd_state: int | None = None,
        feature_weights: str | NDArray | None = None,
    ) -> None:
        if max_depth is not None and (not isinstance(max_depth, int) or max_depth < 0):
            raise ValueError("max_depth must be a non-negative integer or None")
        if not isinstance(min_samples_leaf, int) or min_samples_leaf < 1:
            raise ValueError("min_samples_leaf must be a positive integer")
        if isinstance(lifetime, str) and lifetime not in ("sqrt_n", "density"):
            raise ValueError(
                f"Unknown lifetime string {lifetime!r}. "
                "Recognised values: 'sqrt_n', 'density'."
            )
        if isinstance(feature_weights, str) and feature_weights not in ("variance",):
            raise ValueError(
                f"Unknown feature_weights string {feature_weights!r}. "
                "Recognised value: 'variance'."
            )

        super().__init__(label_space=label_space)

        self.lifetime = lifetime
        self.feature_weights = feature_weights
        self.max_depth = max_depth
        self.min_samples_leaf = min_samples_leaf
        self.rnd_state = rnd_state
        self.rnd_gen = np.random.default_rng(rnd_state)

        self.X: NDArray | None = None
        self.y: NDArray | None = None

    # ------------------------------------------------------------------
    # Online learning interface
    # ------------------------------------------------------------------

    def learn_initial_training_set(self, X: NDArray, y: NDArray) -> None:
        """Batch-initialise with a training set.

        Parameters
        ----------
        X : array-like, shape (n, d)
            Training feature vectors.
        y : array-like, shape (n,)
            Integer class labels.
        """
        X = np.atleast_2d(np.asarray(X, dtype=np.float64))
        y = np.asarray(y, dtype=int)
        self._update_label_space_batch(y)
        self.X = X
        self.y = y

    def learn_one(self, x: NDArray, y: int) -> None:
        """Incrementally add one labelled observation.

        Parameters
        ----------
        x : array-like, shape (d,)
            Feature vector.
        y : int
            True label.
        """
        x = np.asarray(x, dtype=np.float64).ravel()
        y = int(y)
        self._update_label_space_one(y)
        if self.X is None:
            self.X = x.reshape(1, -1)
            self.y = np.array([y], dtype=int)
        else:
            self.X = np.vstack([self.X, x.reshape(1, -1)])
            self.y = np.append(self.y, y)

    # ------------------------------------------------------------------
    # Mondrian taxonomy helpers
    # ------------------------------------------------------------------

    def _build_augmented_leaf_ids(self, x: NDArray) -> NDArray:
        """Build one Mondrian tree from X_aug = [X_train; x] and return leaf IDs.

        The leaf assignment is label-independent — the same tree is used for
        all hypothesised test labels v, so we build it only once per
        ``predict()`` call.

        Uses ``leaf.indices`` stored on each leaf node (O(n) total, avoids
        per-point tree traversals).

        Parameters
        ----------
        x : ndarray, shape (d,)
            Test feature vector (already flattened).

        Returns
        -------
        leaf_ids : ndarray of int, shape (n+1,)
            Integer leaf identifier for each point in the augmented bag.
            Points sharing the same id belong to the same Mondrian leaf.
        """
        n = len(self.y)
        X_aug = np.vstack([self.X, x.reshape(1, -1)])

        # Delegate partition construction to the standalone MondrianTree core.
        # grow() resolves feature weights and (string) lifetimes and advances
        # self.rnd_gen in the exact same order as the raw engine calls, so
        # seeded predictions are unchanged.
        tree = MondrianTree.grow(
            X_aug,
            self.rnd_gen,
            lifetime=self.lifetime,
            x_test=x,
            max_depth=self.max_depth,
            min_samples_leaf=self.min_samples_leaf,
            feature_weights=self.feature_weights,
        )

        # Assign leaf IDs from leaf.indices — O(n), no redundant tree walks.
        leaves = tree.collect_leaves()
        leaf_ids = np.empty(n + 1, dtype=np.intp)
        for leaf_id, leaf in enumerate(leaves):
            leaf_ids[leaf.indices] = leaf_id

        return leaf_ids

    def _categories_for_hypothesis(
        self, taxonomy_data: Any, v: int
    ) -> tuple:
        """Return (leaf_ids, labels_aug, test_idx) for hypothesis y_test = v.

        Parameters
        ----------
        taxonomy_data : ndarray, shape (n+1,)
            Leaf IDs precomputed by :meth:`_build_augmented_leaf_ids`.
        v : int
            Hypothesised test label.

        Returns
        -------
        leaf_ids : ndarray, shape (n+1,)
        labels_aug : ndarray, shape (n+1,)
        test_idx : int
        """
        leaf_ids = taxonomy_data
        test_idx = len(self.y)
        labels_aug = np.append(self.y, v)
        return leaf_ids, labels_aug, test_idx

    # ------------------------------------------------------------------
    # Prediction
    # ------------------------------------------------------------------

    def predict(self, x: NDArray) -> VennPrediction:
        """Produce a Venn multiprobability prediction.

        Builds one Mondrian tree from the augmented dataset {X_train ∪ {x}},
        assigns each point to a leaf, then computes empirical label frequencies
        within the test point's leaf under each hypothesis v ∈ Y.

        Parameters
        ----------
        x : array-like, shape (d,)
            Test feature vector.

        Returns
        -------
        VennPrediction
            Binary: contains p0, p1. Multiclass: |Y|×|Y| probs matrix.
        """
        x = np.asarray(x, dtype=np.float64).ravel()

        if self.X is None or len(self.y) == 0:
            return self._empty_prediction()

        leaf_ids = self._build_augmented_leaf_ids(x)
        return self._venn_predict_from_taxonomy(leaf_ids)

learn_initial_training_set(X: NDArray, y: NDArray) -> None

Batch-initialise with a training set.

Parameters:

Name	Type	Description	Default
X	(array - like, shape(n, d))	Training feature vectors.	required
y	(array - like, shape(n))	Integer class labels.	required

Source code in src/online_cp/venn.py

def learn_initial_training_set(self, X: NDArray, y: NDArray) -> None:
    """Batch-initialise with a training set.

    Parameters
    ----------
    X : array-like, shape (n, d)
        Training feature vectors.
    y : array-like, shape (n,)
        Integer class labels.
    """
    X = np.atleast_2d(np.asarray(X, dtype=np.float64))
    y = np.asarray(y, dtype=int)
    self._update_label_space_batch(y)
    self.X = X
    self.y = y

learn_one(x: NDArray, y: int) -> None

Incrementally add one labelled observation.

Parameters:

Name	Type	Description	Default
x	(array - like, shape(d))	Feature vector.	required
y	int	True label.	required

Source code in src/online_cp/venn.py

def learn_one(self, x: NDArray, y: int) -> None:
    """Incrementally add one labelled observation.

    Parameters
    ----------
    x : array-like, shape (d,)
        Feature vector.
    y : int
        True label.
    """
    x = np.asarray(x, dtype=np.float64).ravel()
    y = int(y)
    self._update_label_space_one(y)
    if self.X is None:
        self.X = x.reshape(1, -1)
        self.y = np.array([y], dtype=int)
    else:
        self.X = np.vstack([self.X, x.reshape(1, -1)])
        self.y = np.append(self.y, y)

predict(x: NDArray) -> VennPrediction

Produce a Venn multiprobability prediction.

Builds one Mondrian tree from the augmented dataset {X_train ∪ {x}}, assigns each point to a leaf, then computes empirical label frequencies within the test point's leaf under each hypothesis v ∈ Y.

Parameters:

Name	Type	Description	Default
x	(array - like, shape(d))	Test feature vector.	required

Returns:

Type	Description
VennPrediction	Binary: contains p0, p1. Multiclass: |Y|×|Y| probs matrix.

Source code in src/online_cp/venn.py

def predict(self, x: NDArray) -> VennPrediction:
    """Produce a Venn multiprobability prediction.

    Builds one Mondrian tree from the augmented dataset {X_train ∪ {x}},
    assigns each point to a leaf, then computes empirical label frequencies
    within the test point's leaf under each hypothesis v ∈ Y.

    Parameters
    ----------
    x : array-like, shape (d,)
        Test feature vector.

    Returns
    -------
    VennPrediction
        Binary: contains p0, p1. Multiclass: |Y|×|Y| probs matrix.
    """
    x = np.asarray(x, dtype=np.float64).ravel()

    if self.X is None or len(self.y) == 0:
        return self._empty_prediction()

    leaf_ids = self._build_augmented_leaf_ids(x)
    return self._venn_predict_from_taxonomy(leaf_ids)

Group-conditional wrappers

The general Mondrian taxonomy framework: a single pooled model with the calibration step restricted to same-category examples (Vovk et al., 2005).

online_cp.mondrian.framework.MondrianConformalRegressor

Bases: SerializableMixin

Mondrian conformal regressor: group-conditional coverage.

Wraps a conformal regressor (ridge, kernel ridge, or lasso). A single pooled model is trained on ALL data; the Mondrian construction lives in the calibration step, not in the model. At prediction time the nonconformity scores are computed from the pooled model, but the interval is calibrated only against scores from training examples in the SAME Mondrian category as the test object. This is the Mondrian taxonomy construction of [ALRW2 §4.6] (Vovk et al.).

The result is valid group-conditional (object-conditional) coverage: for every category \(k\),

\[ \mathbb{P}\bigl(y \in \Gamma^\epsilon(x) \mid \kappa(x) = k\bigr) \geq 1 - \epsilon , \]

where \(\kappa\) is the category function. Unconditional conformal predictors only guarantee this marginally; conditioning on a Mondrian taxonomy restores the guarantee within each group, at the cost of a smaller effective calibration set per category.

Parameters:

Name	Type	Description	Default
base_model	ConformalRidgeRegressor, KernelConformalRidgeRegressor, or ConformalLassoRegressor	The underlying conformal regressor. Trained on all data (pooled).	required
category_fn	callable	A function (x) -> hashable mapping each object to its Mondrian category \(\kappa(x)\).	required

Source code in src/online_cp/mondrian/framework.py

class MondrianConformalRegressor(SerializableMixin):
    r"""Mondrian conformal regressor: group-conditional coverage.

    Wraps a conformal regressor (ridge, kernel ridge, or lasso). A *single*
    pooled model is trained on ALL data; the Mondrian construction lives in the
    calibration step, not in the model. At prediction time the nonconformity
    scores are computed from the pooled model, but the interval is calibrated
    only against scores from training examples in the SAME Mondrian category as
    the test object. This is the Mondrian taxonomy construction of [ALRW2 §4.6]
    (Vovk et al.).

    The result is valid **group-conditional** (object-conditional) coverage: for
    every category $k$,

    $$
    \mathbb{P}\bigl(y \in \Gamma^\epsilon(x) \mid \kappa(x) = k\bigr)
        \geq 1 - \epsilon ,
    $$

    where $\kappa$ is the category function. Unconditional conformal predictors
    only guarantee this marginally; conditioning on a Mondrian taxonomy restores
    the guarantee within each group, at the cost of a smaller effective
    calibration set per category.

    Parameters
    ----------
    base_model : ConformalRidgeRegressor, KernelConformalRidgeRegressor, or ConformalLassoRegressor
        The underlying conformal regressor. Trained on all data (pooled).
    category_fn : callable
        A function ``(x) -> hashable`` mapping each object to its Mondrian
        category $\kappa(x)$.
    """

    _SAVE_PARAMS: tuple = ("base_model", "category_fn")
    _SAVE_STATE: tuple = ("categories_",)
    _SAVE_CALLABLES: tuple = ("category_fn",)

    def __init__(
        self, base_model: ConformalRegressor, category_fn: Callable[[NDArray[np.floating[Any]]], Hashable]
    ) -> None:
        if not isinstance(
            base_model,
            (ConformalRidgeRegressor, KernelConformalRidgeRegressor, ConformalLassoRegressor),
        ):
            raise TypeError(
                f"base_model must be a ConformalRidgeRegressor, "
                f"KernelConformalRidgeRegressor, or ConformalLassoRegressor, "
                f"got {type(base_model).__name__}"
            )
        self.base_model = base_model
        self.category_fn = category_fn
        self.categories_ = []

    def learn_initial_training_set(self, X, y):
        r"""Batch-train the pooled model and record each object's category.

        Fits ``base_model`` on the full training set ``(X, y)`` and stores the
        Mondrian category $\kappa(x_i)$ of every training object, so later
        predictions can be calibrated within categories.

        Parameters
        ----------
        X : array-like of shape (n, d)
            Training objects.
        y : array-like of shape (n,)
            Training responses.
        """
        X = np.asarray(X)
        y = np.asarray(y)
        self.base_model.learn_initial_training_set(X, y)
        self.categories_ = [self.category_fn(x_i) for x_i in X]

    def learn_one(self, x, y, **kwargs):
        """Update the pooled model with one example and record its category.

        Parameters
        ----------
        x : array-like of shape (d,)
            New object.
        y : float
            Observed response.
        **kwargs
            Forwarded to the underlying model's ``learn_one`` (e.g. a
            ``precomputed`` cache returned by :meth:`predict`).
        """
        self.base_model.learn_one(x, y, **kwargs)
        self.categories_.append(self.category_fn(x))

    def predict(
        self,
        x: NDArray[np.floating[Any]],
        epsilon: float | NDArray[np.floating[Any]] | None = None,
        bounds: str = "both",
    ) -> ContinuousPredictionSet | MultiLevelPredictionSet:
        """Compute the Mondrian (group-conditional) prediction interval.

        Computes the pooled model's $A$/$B$ nonconformity decomposition for the
        test object, then builds the interval using only the training scores
        whose category matches the test object's Mondrian category (a boolean
        ``cat_mask`` over the pooled scores). This restricts calibration to the
        test object's group while reusing the single pooled fit.

        Parameters
        ----------
        x : array-like of shape (d,)
            Test object.
        epsilon : float, array-like, or None
            Significance level(s). If None, uses ``base_model.epsilon``. An
            iterable yields a :class:`MultiLevelPredictionSet`.
        bounds : {"both", "lower", "upper"}, default "both"
            Which side(s) of the interval to compute.

        Returns
        -------
        ContinuousPredictionSet or MultiLevelPredictionSet
            The group-conditional prediction set at ``epsilon``.
        """
        if epsilon is None:
            epsilon = self.base_model.epsilon

        cat = self.category_fn(x)

        if isinstance(self.base_model, ConformalRidgeRegressor):
            return self._predict_ridge(x, epsilon, bounds, cat)
        elif isinstance(self.base_model, KernelConformalRidgeRegressor):
            return self._predict_kernel_ridge(x, epsilon, bounds, cat)
        else:
            return self._predict_lasso(x, epsilon, cat)

    def _predict_ridge(self, x, epsilon, bounds, cat):
        model = self.base_model
        if model._safe_size_check(model.X) == 0:
            return self._inf_interval(epsilon)

        X = np.append(model.X, x.reshape(1, -1), axis=0)
        XTXinv = model.XTXinv - (model.XTXinv @ np.outer(x, x) @ model.XTXinv) / (1 + x.T @ model.XTXinv @ x)
        A, B = model.compute_A_and_B(X, XTXinv, model.y)

        cat_mask = np.array([c == cat for c in self.categories_] + [True])
        n_cat = int(cat_mask.sum())

        eps_check = max(epsilon) if hasattr(epsilon, "__iter__") else epsilon
        if bounds == "both" and not (eps_check >= 2 / n_cat):
            return self._inf_interval(epsilon)
        if bounds != "both" and not (eps_check >= 1 / n_cat):
            return self._inf_interval(epsilon)

        A_cat = A[cat_mask]
        B_cat = B[cat_mask]
        l_dic, u_dic = ConformalRegressor._vectorised_l_and_u(A_cat, B_cat)
        return self._build_interval(l_dic, u_dic, epsilon, n_cat, bounds)

    def _predict_kernel_ridge(self, x, epsilon, bounds, cat):
        model = self.base_model
        if model.X is None:
            return self._inf_interval(epsilon)

        k = model.kernel(model.X, x).reshape(-1, 1)
        kappa = model.kernel(x, x)
        K = KernelConformalRidgeRegressor._update_K(model.K, k, kappa)
        Kinv = KernelConformalRidgeRegressor._update_Kinv(model.Kinv, k, kappa + model.a)
        X = np.append(model.X, x.reshape(1, -1), axis=0)
        A, B = model.compute_A_and_B(X, K, Kinv, model.y)

        cat_mask = np.array([c == cat for c in self.categories_] + [True])
        n_cat = int(cat_mask.sum())

        eps_check = max(epsilon) if hasattr(epsilon, "__iter__") else epsilon
        if bounds == "both" and not (eps_check >= 2 / n_cat):
            return self._inf_interval(epsilon)
        if bounds != "both" and not (eps_check >= 1 / n_cat):
            return self._inf_interval(epsilon)

        A_cat = A[cat_mask]
        B_cat = B[cat_mask]
        l_dic, u_dic = ConformalRegressor._vectorised_l_and_u(A_cat, B_cat)
        return self._build_interval(l_dic, u_dic, epsilon, n_cat, bounds)

    def _predict_lasso(self, x, epsilon, cat):
        model = self.base_model
        x = np.atleast_1d(x).ravel()
        if model.X is None or model.X.shape[0] < 2:
            return self._inf_interval(epsilon)
        if hasattr(epsilon, "__iter__"):
            predictions = {}
            for eps in epsilon:
                predictions[eps] = self._predict_lasso_single(x, eps, cat)
            return MultiLevelPredictionSet(predictions)
        return self._predict_lasso_single(x, epsilon, cat)

    def _predict_lasso_single(self, x, epsilon, cat):
        model = self.base_model
        cat_mask_train = np.array([c == cat for c in self.categories_])
        n_cat = int(cat_mask_train.sum()) + 1  # +1 for test point
        if n_cat < 2:
            return ContinuousPredictionSet([(-np.inf, np.inf)], epsilon)

        threshold = int(np.ceil(n_cat * (1 - epsilon)))
        y0 = x @ model.beta
        y_range = model.y.max() - model.y.min()
        if y_range == 0:
            y_range = 1.0
        t_max = (model.y.max() - y0) + model.search_range_factor * y_range
        t_min = (model.y.min() - y0) - model.search_range_factor * y_range

        intervals_pos = self._run_homotopy_mondrian(x, +1, t_max, cat_mask_train, threshold)
        intervals_neg = self._run_homotopy_mondrian(x, -1, t_min, cat_mask_train, threshold)

        all_intervals = []
        for a, b in intervals_neg:
            all_intervals.append((y0 + a, y0 + b))
        if self._t_in_set(x, 0.0, cat_mask_train, threshold):
            all_intervals.append((y0, y0))
        for a, b in intervals_pos:
            all_intervals.append((y0 + a, y0 + b))

        merged = ConformalLassoRegressor._merge_intervals(all_intervals)
        if not merged:
            return ContinuousPredictionSet([], epsilon)
        elif len(merged) == 1:
            return ContinuousPredictionSet([merged[0]], epsilon)
        else:
            return ContinuousPredictionSet(merged, epsilon)

    def _t_in_set(self, x_new, t, cat_mask_train, threshold):
        model = self.base_model
        X_aug = np.vstack([model.X, x_new.reshape(1, -1)])
        y_aug = np.append(model.y, x_new @ model.beta + t)
        beta_t = _solve_lasso(X_aug, y_aug, model.lam, rho=model.rho, warm_start=model.beta)
        abs_res = np.abs(y_aug - X_aug @ beta_t)
        # Ascending rank of |r_{n+1}(t)| among the pooled same-category fitted
        # residuals, including the test point itself.  Group-conditional
        # conformal criterion: the candidate is in the set iff this rank is NO
        # HIGHER than the threshold (test residual among the smallest in the
        # category), matching the base ConformalLassoRegressor.
        rank = int(np.sum(abs_res[:-1][cat_mask_train] <= abs_res[-1])) + 1
        return rank <= threshold

    def _run_homotopy_mondrian(self, x_new, direction, t_bound, cat_mask_train, threshold):

        model = self.base_model
        sign = 1 if direction > 0 else -1
        t_bound_abs = abs(t_bound)
        n = model.X.shape[0]
        p = model.X.shape[1]
        lam = model.lam
        beta_k = model.beta.copy()
        J_k = np.where(np.abs(beta_k) > 1e-12)[0]
        residuals = model.y - model.X @ beta_k
        v_full = model.X.T @ residuals
        J_c_k = np.setdiff1d(np.arange(p), J_k)
        v_inactive = v_full[J_c_k]
        r_train = residuals.copy()
        r_test = 0.0
        XtX_aug = model.X.T @ model.X + np.outer(x_new, x_new)
        t_accumulated = 0.0
        intervals_in_set = []

        for _ in range(model.max_homotopy_steps):
            if t_accumulated >= t_bound_abs:
                break
            if len(J_k) == 0:
                if len(J_c_k) == 0:
                    break
                gamma_k = sign * x_new[J_c_k]
                dt_dual = np.full(len(J_c_k), np.inf)
                for idx in range(len(J_c_k)):
                    g = gamma_k[idx]
                    if g > 1e-15:
                        dt_dual[idx] = (lam - sign * v_inactive[idx]) / g
                    elif g < -1e-15:
                        dt_dual[idx] = (-lam - sign * v_inactive[idx]) / g
                dt_dual = np.where(dt_dual > 1e-12, dt_dual, np.inf)
                dt_k = min(float(np.min(dt_dual)), t_bound_abs - t_accumulated)
                sub = self._find_intervals_mondrian(
                    r_train,
                    r_test,
                    np.zeros(n),
                    sign * 1.0,
                    dt_k,
                    t_accumulated,
                    sign,
                    threshold,
                    cat_mask_train,
                )
                intervals_in_set.extend(sub)
                t_accumulated += dt_k
                r_test += sign * 1.0 * dt_k
                v_inactive += gamma_k * dt_k
                if dt_k < t_bound_abs - t_accumulated + 1e-12:
                    entering = J_c_k[np.argmin(dt_dual)]
                    J_k = np.append(J_k, entering)
                    J_c_k = np.setdiff1d(np.arange(p), J_k)
                    v_inactive = v_full[J_c_k]
                continue

            Sigma_J = XtX_aug[np.ix_(J_k, J_k)] + model.rho * np.eye(len(J_k))
            x_J = x_new[J_k]
            try:
                Sigma_J_inv = np.linalg.inv(Sigma_J)
            except np.linalg.LinAlgError:
                break
            Sigma_J_inv_x = Sigma_J_inv @ x_J
            eta_k = sign * Sigma_J_inv_x
            if len(J_c_k) > 0:
                gamma_k = sign * (x_new[J_c_k] - XtX_aug[np.ix_(J_c_k, J_k)] @ Sigma_J_inv_x)
            else:
                gamma_k = np.array([])
            slopes_train = -(model.X[:, J_k] @ eta_k)
            slope_test = sign * (1.0 - x_J @ Sigma_J_inv_x)
            beta_J = beta_k[J_k]
            dt_primal = np.full(len(J_k), np.inf)
            for idx in range(len(J_k)):
                if abs(eta_k[idx]) > 1e-15:
                    dt = -beta_J[idx] / eta_k[idx]
                    if dt > 1e-12:
                        dt_primal[idx] = dt
            dt_dual = np.full(len(J_c_k), np.inf)
            for idx in range(len(J_c_k)):
                g = gamma_k[idx]
                v_j = v_inactive[idx]
                if abs(g) > 1e-15:
                    dt1 = (lam - v_j) / g
                    dt2 = (-lam - v_j) / g
                    cands = []
                    if dt1 > 1e-12:
                        cands.append(dt1)
                    if dt2 > 1e-12:
                        cands.append(dt2)
                    if cands:
                        dt_dual[idx] = min(cands)
            dt_k = min(
                float(np.min(dt_primal)) if len(dt_primal) > 0 else np.inf,
                float(np.min(dt_dual)) if len(dt_dual) > 0 else np.inf,
            )
            dt_k = min(dt_k, t_bound_abs - t_accumulated)
            if dt_k <= 0 or not np.isfinite(dt_k):
                break
            sub = self._find_intervals_mondrian(
                r_train,
                r_test,
                slopes_train,
                slope_test,
                dt_k,
                t_accumulated,
                sign,
                threshold,
                cat_mask_train,
            )
            intervals_in_set.extend(sub)
            beta_k[J_k] += eta_k * dt_k
            r_train += slopes_train * dt_k
            r_test += slope_test * dt_k
            if len(J_c_k) > 0:
                v_inactive += gamma_k * dt_k
            t_accumulated += dt_k
            min_primal = float(np.min(dt_primal)) if len(dt_primal) > 0 else np.inf
            min_dual = float(np.min(dt_dual)) if len(dt_dual) > 0 else np.inf
            if dt_k >= t_bound_abs - (t_accumulated - dt_k):
                break
            if min_primal <= min_dual:
                leaving_idx = np.argmin(dt_primal)
                beta_k[J_k[leaving_idx]] = 0.0
                J_k = np.delete(J_k, leaving_idx)
            else:
                entering_var = J_c_k[np.argmin(dt_dual)]
                J_k = np.append(J_k, entering_var)
            J_c_k = np.setdiff1d(np.arange(p), J_k)
            v_full = model.X.T @ r_train + x_new * r_test
            v_inactive = v_full[J_c_k]

        return intervals_in_set

    @staticmethod
    def _find_intervals_mondrian(
        r_train,
        r_test,
        slopes_train,
        slope_test,
        dt_k,
        t_accumulated,
        sign,
        threshold,
        cat_mask_train,
    ):
        from online_cp.regressors import _compute_crossings

        n = len(r_train)
        raw_crossings = _compute_crossings(r_train, slopes_train, r_test, slope_test, n, dt_k)
        crossings = sorted(set([0.0] + list(raw_crossings) + [dt_k]))
        result_intervals = []
        for idx in range(len(crossings) - 1):
            d_start = crossings[idx]
            d_end = crossings[idx + 1]
            if d_end - d_start < 1e-14:
                continue
            d_mid = (d_start + d_end) / 2
            r_i_mid = r_train + slopes_train * d_mid
            r_test_mid = r_test + slope_test * d_mid
            # Ascending rank of the test residual among the pooled
            # same-category fitted residuals (including itself).  Group-
            # conditional conformal criterion: the candidate is kept iff this
            # rank is NO HIGHER than the threshold (test residual among the
            # smallest in the category), matching the base
            # ConformalLassoRegressor.
            abs_r_cat = np.abs(r_i_mid[cat_mask_train])
            abs_r_test = np.abs(r_test_mid)
            rank = int(np.sum(abs_r_cat <= abs_r_test)) + 1
            if rank <= threshold:
                t_start = sign * (t_accumulated + d_start)
                t_end = sign * (t_accumulated + d_end)
                if t_start > t_end:
                    t_start, t_end = t_end, t_start
                result_intervals.append((t_start, t_end))
        return result_intervals

    def compute_p_value(self, x, y, **kwargs):
        r"""Mondrian conformal p-value for the candidate pair $(x, y)$.

        Computes the (optionally smoothed) conformal p-value using only the
        pooled nonconformity scores from the same Mondrian category as $x$, so
        the p-value is calibrated group-conditionally. Dispatches to the
        appropriate backend (ridge, kernel ridge, or lasso).

        Parameters
        ----------
        x : array-like of shape (d,)
            Test object.
        y : float
            Candidate response.
        **kwargs
            Backend options such as ``bounds``, ``smoothed`` and ``tau``.

        Returns
        -------
        float
            The group-conditional p-value.
        """
        cat = self.category_fn(x)
        model = self.base_model
        if isinstance(model, ConformalRidgeRegressor):
            return self._p_value_ridge(x, y, cat, **kwargs)
        elif isinstance(model, KernelConformalRidgeRegressor):
            return self._p_value_kernel_ridge(x, y, cat, **kwargs)
        else:
            return self._p_value_lasso(x, y, cat, **kwargs)

    def _p_value_ridge(self, x, y, cat, bounds="both", smoothed=True, tau=None):
        model = self.base_model
        if tau is None and smoothed:
            tau = model.rnd_gen.uniform(0, 1)
        if model.XTXinv is None:
            return tau if smoothed else 1
        X = np.append(model.X, x.reshape(1, -1), axis=0)
        XTXinv = model.XTXinv - (model.XTXinv @ np.outer(x, x) @ model.XTXinv) / (1 + x.T @ model.XTXinv @ x)
        A, B = model.compute_A_and_B(X, XTXinv, model.y)
        cat_mask = np.array([c == cat for c in self.categories_] + [True])
        A_cat, B_cat = A[cat_mask], B[cat_mask]
        if bounds == "both":
            E = A_cat + y * B_cat
            Alpha = np.array([min((E >= e).sum(), (E <= e).sum()) for e in E])
            c_type = "conformity"
        elif bounds == "lower":
            Alpha = -(A_cat + y * B_cat)
            c_type = "nonconformity"
        else:
            Alpha = A_cat + y * B_cat
            c_type = "nonconformity"
        return ConformalRegressor._compute_p_value(Alpha, tau if smoothed else None, c_type=c_type)

    def _p_value_kernel_ridge(self, x, y, cat, bounds="both", smoothed=True, tau=None):
        model = self.base_model
        if tau is None and smoothed:
            tau = model.rnd_gen.uniform(0, 1)
        if model.Kinv is None:
            return tau if smoothed else 1
        k = model.kernel(model.X, x).reshape(-1, 1)
        kappa = model.kernel(x, x)
        K = KernelConformalRidgeRegressor._update_K(model.K, k, kappa)
        Kinv = KernelConformalRidgeRegressor._update_Kinv(model.Kinv, k, kappa + model.a)
        X = np.append(model.X, x.reshape(1, -1), axis=0)
        A, B = model.compute_A_and_B(X, K, Kinv, model.y)
        cat_mask = np.array([c == cat for c in self.categories_] + [True])
        A_cat, B_cat = A[cat_mask], B[cat_mask]
        if bounds == "both":
            E = A_cat + y * B_cat
            Alpha = np.array([min((E >= e).sum(), (E <= e).sum()) for e in E])
            c_type = "conformity"
        elif bounds == "lower":
            Alpha = -(A_cat + y * B_cat)
            c_type = "nonconformity"
        else:
            Alpha = A_cat + y * B_cat
            c_type = "nonconformity"
        return ConformalRegressor._compute_p_value(Alpha, tau if smoothed else None, c_type=c_type)

    def _p_value_lasso(self, x, y, cat, smoothed=True, tau=None):
        model = self.base_model
        x = np.atleast_1d(x).ravel()
        if tau is None and smoothed:
            tau = model.rnd_gen.uniform()
        if model.X is None:
            return tau if smoothed else 1
        X_aug = np.vstack([model.X, x.reshape(1, -1)])
        y_aug = np.append(model.y, y)
        beta_aug = _solve_lasso(X_aug, y_aug, model.lam, rho=model.rho, warm_start=model.beta)
        residuals = np.abs(y_aug - X_aug @ beta_aug)
        cat_mask_train = np.array([c == cat for c in self.categories_])
        r_cat = residuals[:-1][cat_mask_train]
        r_test = residuals[-1]
        if smoothed and tau is not None:
            gt = np.sum(r_cat > r_test)
            eq = np.sum(r_cat == r_test) + 1
            p = (gt + tau * eq) / (len(r_cat) + 1)
        else:
            geq = np.sum(r_cat >= r_test) + 1
            p = geq / (len(r_cat) + 1)
        return float(p)

    @staticmethod
    def _build_interval(l_dic, u_dic, epsilon, n, bounds):
        if hasattr(epsilon, "__iter__"):
            predictions = {}
            for eps in epsilon:
                if bounds == "both":
                    lo = ConformalRegressor._get_lower(l_dic=l_dic, epsilon=eps / 2, n=n)
                    up = ConformalRegressor._get_upper(u_dic=u_dic, epsilon=eps / 2, n=n)
                elif bounds == "lower":
                    lo = ConformalRegressor._get_lower(l_dic=l_dic, epsilon=eps, n=n)
                    up = np.inf
                else:
                    lo = -np.inf
                    up = ConformalRegressor._get_upper(u_dic=u_dic, epsilon=eps, n=n)
                predictions[eps] = ContinuousPredictionSet([(lo, up)], eps)
            return MultiLevelPredictionSet(predictions)
        if bounds == "both":
            lo = ConformalRegressor._get_lower(l_dic=l_dic, epsilon=epsilon / 2, n=n)
            up = ConformalRegressor._get_upper(u_dic=u_dic, epsilon=epsilon / 2, n=n)
        elif bounds == "lower":
            lo = ConformalRegressor._get_lower(l_dic=l_dic, epsilon=epsilon, n=n)
            up = np.inf
        else:
            lo = -np.inf
            up = ConformalRegressor._get_upper(u_dic=u_dic, epsilon=epsilon, n=n)
        return ContinuousPredictionSet([(lo, up)], epsilon)

    @staticmethod
    def _inf_interval(epsilon):
        if hasattr(epsilon, "__iter__"):
            return MultiLevelPredictionSet(
                {eps: ContinuousPredictionSet([(-np.inf, np.inf)], eps) for eps in epsilon}
            )
        return ContinuousPredictionSet([(-np.inf, np.inf)], epsilon)

    @property
    def categories(self):
        """Set of Mondrian categories discovered in the training stream so far."""
        return set(self.categories_)

    def __repr__(self):
        return (
            f"MondrianConformalRegressor(n_categories={len(self.categories)}, "
            f"n_total={len(self.categories_)}, "
            f"base={self.base_model.__class__.__name__})"
        )

categories property

Set of Mondrian categories discovered in the training stream so far.

learn_initial_training_set(X, y)

Batch-train the pooled model and record each object's category.

Fits base_model on the full training set (X, y) and stores the Mondrian category \(\kappa(x_i)\) of every training object, so later predictions can be calibrated within categories.

Parameters:

Name	Type	Description	Default
X	array-like of shape (n, d)	Training objects.	required
y	array-like of shape (n,)	Training responses.	required

Source code in src/online_cp/mondrian/framework.py

def learn_initial_training_set(self, X, y):
    r"""Batch-train the pooled model and record each object's category.

    Fits ``base_model`` on the full training set ``(X, y)`` and stores the
    Mondrian category $\kappa(x_i)$ of every training object, so later
    predictions can be calibrated within categories.

    Parameters
    ----------
    X : array-like of shape (n, d)
        Training objects.
    y : array-like of shape (n,)
        Training responses.
    """
    X = np.asarray(X)
    y = np.asarray(y)
    self.base_model.learn_initial_training_set(X, y)
    self.categories_ = [self.category_fn(x_i) for x_i in X]

learn_one(x, y, **kwargs)

Update the pooled model with one example and record its category.

Parameters:

Name	Type	Description	Default
x	array-like of shape (d,)	New object.	required
y	float	Observed response.	required
**kwargs		Forwarded to the underlying model's learn_one (e.g. a precomputed cache returned by :meth:predict).	{}

Source code in src/online_cp/mondrian/framework.py

def learn_one(self, x, y, **kwargs):
    """Update the pooled model with one example and record its category.

    Parameters
    ----------
    x : array-like of shape (d,)
        New object.
    y : float
        Observed response.
    **kwargs
        Forwarded to the underlying model's ``learn_one`` (e.g. a
        ``precomputed`` cache returned by :meth:`predict`).
    """
    self.base_model.learn_one(x, y, **kwargs)
    self.categories_.append(self.category_fn(x))

predict(x: NDArray[np.floating[Any]], epsilon: float | NDArray[np.floating[Any]] | None = None, bounds: str = 'both') -> ContinuousPredictionSet | MultiLevelPredictionSet

Compute the Mondrian (group-conditional) prediction interval.

Computes the pooled model's \(A\)/\(B\) nonconformity decomposition for the test object, then builds the interval using only the training scores whose category matches the test object's Mondrian category (a boolean cat_mask over the pooled scores). This restricts calibration to the test object's group while reusing the single pooled fit.

Parameters:

Name	Type	Description	Default
x	array-like of shape (d,)	Test object.	required
epsilon	float, array-like, or None	Significance level(s). If None, uses base_model.epsilon. An iterable yields a :class:MultiLevelPredictionSet.	None
bounds	('both', 'lower', 'upper')	Which side(s) of the interval to compute.	"both"

Returns:

Type	Description
ContinuousPredictionSet or MultiLevelPredictionSet	The group-conditional prediction set at epsilon.

Source code in src/online_cp/mondrian/framework.py

def predict(
    self,
    x: NDArray[np.floating[Any]],
    epsilon: float | NDArray[np.floating[Any]] | None = None,
    bounds: str = "both",
) -> ContinuousPredictionSet | MultiLevelPredictionSet:
    """Compute the Mondrian (group-conditional) prediction interval.

    Computes the pooled model's $A$/$B$ nonconformity decomposition for the
    test object, then builds the interval using only the training scores
    whose category matches the test object's Mondrian category (a boolean
    ``cat_mask`` over the pooled scores). This restricts calibration to the
    test object's group while reusing the single pooled fit.

    Parameters
    ----------
    x : array-like of shape (d,)
        Test object.
    epsilon : float, array-like, or None
        Significance level(s). If None, uses ``base_model.epsilon``. An
        iterable yields a :class:`MultiLevelPredictionSet`.
    bounds : {"both", "lower", "upper"}, default "both"
        Which side(s) of the interval to compute.

    Returns
    -------
    ContinuousPredictionSet or MultiLevelPredictionSet
        The group-conditional prediction set at ``epsilon``.
    """
    if epsilon is None:
        epsilon = self.base_model.epsilon

    cat = self.category_fn(x)

    if isinstance(self.base_model, ConformalRidgeRegressor):
        return self._predict_ridge(x, epsilon, bounds, cat)
    elif isinstance(self.base_model, KernelConformalRidgeRegressor):
        return self._predict_kernel_ridge(x, epsilon, bounds, cat)
    else:
        return self._predict_lasso(x, epsilon, cat)

compute_p_value(x, y, **kwargs)

Mondrian conformal p-value for the candidate pair \((x, y)\).

Computes the (optionally smoothed) conformal p-value using only the pooled nonconformity scores from the same Mondrian category as \(x\), so the p-value is calibrated group-conditionally. Dispatches to the appropriate backend (ridge, kernel ridge, or lasso).

Parameters:

Name	Type	Description	Default
x	array-like of shape (d,)	Test object.	required
y	float	Candidate response.	required
**kwargs		Backend options such as bounds, smoothed and tau.	{}

Returns:

Type	Description
float	The group-conditional p-value.

Source code in src/online_cp/mondrian/framework.py

def compute_p_value(self, x, y, **kwargs):
    r"""Mondrian conformal p-value for the candidate pair $(x, y)$.

    Computes the (optionally smoothed) conformal p-value using only the
    pooled nonconformity scores from the same Mondrian category as $x$, so
    the p-value is calibrated group-conditionally. Dispatches to the
    appropriate backend (ridge, kernel ridge, or lasso).

    Parameters
    ----------
    x : array-like of shape (d,)
        Test object.
    y : float
        Candidate response.
    **kwargs
        Backend options such as ``bounds``, ``smoothed`` and ``tau``.

    Returns
    -------
    float
        The group-conditional p-value.
    """
    cat = self.category_fn(x)
    model = self.base_model
    if isinstance(model, ConformalRidgeRegressor):
        return self._p_value_ridge(x, y, cat, **kwargs)
    elif isinstance(model, KernelConformalRidgeRegressor):
        return self._p_value_kernel_ridge(x, y, cat, **kwargs)
    else:
        return self._p_value_lasso(x, y, cat, **kwargs)

online_cp.mondrian.framework.MondrianConformalClassifier

Bases: SerializableMixin

Mondrian conformal classifier: group-conditional coverage.

Wraps a conformal classifier (KNN or SVM). A single pooled model is trained on ALL data; only the p-value computation is restricted to training examples in the SAME Mondrian category as the (hypothesised) test example. This yields valid group-conditional coverage: for every category \(k\),

\[ \mathbb{P}\bigl(y \in \Gamma^\epsilon(x) \mid \kappa(x, y) = k\bigr) \geq 1 - \epsilon . \]

The most common special case is label-conditional conformal prediction [ALRW2 §4.6.7], where the category of an example is its label, \(\kappa(n, (x, y)) = y\). Use category_fn="label" for this; it guarantees per-class coverage and is the standard remedy for class imbalance.

Parameters:

Name	Type	Description	Default
base_model	ConformalNearestNeighboursClassifier or ConformalSupportVectorMachine	The underlying conformal classifier. Trained on all data (pooled).	required
category_fn	str or callable	Determines the Mondrian taxonomy: "label" — label-conditional: category = label (most common). A callable (x, y) -> hashable — general Mondrian taxonomy depending on both features and label. A callable (x) -> hashable — object-conditional taxonomy depending only on features (legacy interface).	required

Source code in src/online_cp/mondrian/framework.py

class MondrianConformalClassifier(SerializableMixin):
    r"""Mondrian conformal classifier: group-conditional coverage.

    Wraps a conformal classifier (KNN or SVM). A single pooled model is trained
    on ALL data; only the p-value computation is restricted to training examples
    in the SAME Mondrian category as the (hypothesised) test example. This yields
    valid **group-conditional** coverage: for every category $k$,

    $$
    \mathbb{P}\bigl(y \in \Gamma^\epsilon(x) \mid \kappa(x, y) = k\bigr)
        \geq 1 - \epsilon .
    $$

    The most common special case is **label-conditional** conformal prediction
    [ALRW2 §4.6.7], where the category of an example is its label,
    $\kappa(n, (x, y)) = y$. Use ``category_fn="label"`` for this; it guarantees
    per-class coverage and is the standard remedy for class imbalance.

    Parameters
    ----------
    base_model : ConformalNearestNeighboursClassifier or ConformalSupportVectorMachine
        The underlying conformal classifier. Trained on all data (pooled).
    category_fn : str or callable
        Determines the Mondrian taxonomy:

        - ``"label"`` — label-conditional: category = label (most common).
        - A callable ``(x, y) -> hashable`` — general Mondrian taxonomy
          depending on both features and label.
        - A callable ``(x) -> hashable`` — object-conditional taxonomy
          depending only on features (legacy interface).
    """

    _SAVE_PARAMS: tuple = ("base_model", "category_fn")
    _SAVE_STATE: tuple = ("categories_",)
    _SAVE_CALLABLES: tuple = ("category_fn",)

    def __init__(self, base_model: Any, category_fn: str | Callable[..., Hashable]) -> None:
        self.base_model = base_model

        if category_fn == "label":
            self._category_fn = lambda x, y: y
            self._label_aware = True
        elif callable(category_fn):
            import inspect

            sig = inspect.signature(category_fn)
            n_params = sum(
                1
                for p in sig.parameters.values()
                if p.default is inspect.Parameter.empty and p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD)
            )
            if n_params >= 2:
                self._category_fn = category_fn
                self._label_aware = True
            else:
                self._category_fn = category_fn
                self._label_aware = False
        else:
            raise TypeError(f"category_fn must be 'label' or a callable, got {type(category_fn).__name__}")

        self.category_fn = category_fn
        self.categories_ = []

    def learn_initial_training_set(self, X: NDArray[np.floating[Any]], y: NDArray[Any]) -> None:
        r"""Batch-train the pooled model and record each example's category.

        Fits ``base_model`` on the full training set and stores the Mondrian
        category of every training example. For a label-aware taxonomy (e.g.
        ``category_fn="label"``) the category depends on both $x_i$ and $y_i$;
        for an object-conditional taxonomy it depends on $x_i$ only.

        Parameters
        ----------
        X : array-like of shape (n, d)
            Training objects.
        y : array-like of shape (n,)
            Training labels.
        """
        X = np.asarray(X)
        y = np.asarray(y)
        self.base_model.learn_initial_training_set(X, y)
        if self._label_aware:
            self.categories_ = [self._category_fn(x_i, y_i) for x_i, y_i in zip(X, y)]
        else:
            self.categories_ = [self._category_fn(x_i) for x_i in X]

    def learn_one(self, x: NDArray[np.floating[Any]], y: Any, **kwargs: Any) -> None:
        """Update the pooled model with one example and record its category.

        Parameters
        ----------
        x : array-like of shape (d,)
            New object.
        y : hashable
            Observed label.
        **kwargs
            Forwarded to the underlying model's ``learn_one``.
        """
        self.base_model.learn_one(x, y, **kwargs)
        if self._label_aware:
            self.categories_.append(self._category_fn(x, y))
        else:
            self.categories_.append(self._category_fn(x))

    def predict(
        self,
        x: NDArray[np.floating[Any]],
        epsilon: float | NDArray[np.floating[Any]] | None = None,
        return_p_values: bool = False,
    ) -> Any:
        r"""Compute the Mondrian (group-conditional) prediction set.

        For each candidate label the conformal p-value is computed from the
        pooled scores, restricted to the matching Mondrian category:

        - **Label-aware taxonomy** (e.g. ``category_fn="label"``): each
          hypothesised label $y$ defines its own category mask $\kappa(x, y)$,
          so the candidate is calibrated against training examples sharing that
          category. This is what makes label-conditional coverage possible.
        - **Object-conditional taxonomy**: a single mask $\kappa(x)$ is shared
          across all candidate labels.

        For a multiclass SVM the scores are first restricted to the one-vs-rest
        positive class before the category mask is applied, matching the SVM's
        per-class nonconformity construction.

        Parameters
        ----------
        x : array-like of shape (d,)
            Test object.
        epsilon : float, array-like, or None
            Significance level(s). If None, uses ``base_model.epsilon``.
        return_p_values : bool, default False
            If True, also return the ``{label: p_value}`` dict.

        Returns
        -------
        prediction set, or ``(prediction set, p_values)`` if ``return_p_values``.
        """
        from online_cp.classifiers import (
            ConformalNearestNeighboursClassifier,
            ConformalSupportVectorMachine,
        )

        if epsilon is None:
            epsilon = self.base_model.epsilon
        model = self.base_model
        tau = model.rnd_gen.uniform(0, 1)

        if self._label_aware:
            # Label-aware: per-hypothesis category mask
            cat_masks = {}
            for label in model.label_space:
                cat = self._category_fn(x, label)
                cat_masks[label] = np.array([c == cat for c in self.categories_])
        else:
            # Object-conditional: one shared mask for all hypotheses
            cat = self._category_fn(x)
            shared_mask = np.array([c == cat for c in self.categories_])
            cat_masks = {label: shared_mask for label in model.label_space}

        if isinstance(model, ConformalNearestNeighboursClassifier):
            p_values = self._predict_knn(x, cat_masks, tau)
        elif isinstance(model, ConformalSupportVectorMachine):
            p_values = self._predict_svm(x, cat_masks, tau)
        else:
            raise TypeError(f"Unsupported classifier: {type(model).__name__}")

        Gamma = model._compute_Gamma(p_values, epsilon)
        if return_p_values:
            return Gamma, p_values
        return Gamma

    def _predict_knn(self, x, cat_masks, tau):
        model = self.base_model
        p_values = {}
        if model.y.shape[0] < 1:
            for label in model.label_space:
                p_values[label] = 1.0
            return p_values

        d = model.distance_func(model.X, x)
        D = model.update_distance_matrix(model.D, d)
        for label in model.label_space:
            y = np.append(model.y, label)
            same_d, diff_d = model._find_nearest_distances(D, y)
            Alpha_full = np.nan_to_num(same_d / diff_d, nan=np.inf)
            # Filter to same-category + test point (last element always included)
            cat_mask_aug = np.append(cat_masks[label], True)
            Alpha_cat = Alpha_full[cat_mask_aug]
            p_values[label] = self._compute_mondrian_p(Alpha_cat, tau)
        return p_values

    def _predict_svm(self, x, cat_masks, tau):
        from online_cp.classifiers import _smo_solve

        model = self.base_model
        p_values = {}
        x = np.atleast_1d(x).ravel()
        if model.X is None or model.y.shape[0] == 0:
            for label in model.label_space:
                p_values[label] = 1.0
            return p_values

        k_row = model._compute_kernel_row(model.X, x)
        kappa = model._kernel(x.reshape(1, -1))
        if np.ndim(kappa) > 0:
            kappa = kappa.item()
        n = model.K.shape[0]
        K_aug = np.empty((n + 1, n + 1))
        K_aug[:n, :n] = model.K
        K_aug[:n, n] = k_row
        K_aug[n, :n] = k_row
        K_aug[n, n] = kappa

        for label in model.label_space:
            cat_mask = cat_masks[label]
            y_aug = np.append(model.y, float(label))
            y_binary = np.where(y_aug == label, 1.0, -1.0)
            alpha, _ = _smo_solve(K_aug, y_binary, model.C, tol=model.smo_tol, max_iter=model.smo_max_iter)
            if len(model.label_space) > 2:
                # Multiclass: filter to same-class first, then apply cat_mask
                same_class_mask = y_binary == 1.0
                alpha_class = alpha[same_class_mask]
                cat_indices = np.where(same_class_mask)[0]
                cat_mask_subset = np.array([(cat_mask[i] if i < n else True) for i in cat_indices])
                Alpha_cat = alpha_class[cat_mask_subset]
            else:
                Alpha_cat = alpha[np.append(cat_mask, True)]
            p_values[label] = self._compute_mondrian_p(Alpha_cat, tau)
        return p_values

    @staticmethod
    def _compute_mondrian_p(Alpha, tau):
        """Smoothed p-value from category-filtered scores (last = test)."""
        alpha_n = Alpha[-1]
        gt = np.sum(Alpha > alpha_n)
        eq = np.sum(Alpha == alpha_n)
        return float((gt + tau * eq) / Alpha.shape[0])

    @property
    def categories(self):
        """Set of Mondrian categories discovered in the training stream so far."""
        return set(self.categories_)

    def __repr__(self):
        return (
            f"MondrianConformalClassifier(n_categories={len(self.categories)}, "
            f"n_total={len(self.categories_)}, "
            f"base={self.base_model.__class__.__name__})"
        )

categories property

Set of Mondrian categories discovered in the training stream so far.

learn_initial_training_set(X: NDArray[np.floating[Any]], y: NDArray[Any]) -> None

Batch-train the pooled model and record each example's category.

Fits base_model on the full training set and stores the Mondrian category of every training example. For a label-aware taxonomy (e.g. category_fn="label") the category depends on both \(x_i\) and \(y_i\); for an object-conditional taxonomy it depends on \(x_i\) only.

Parameters:

Name	Type	Description	Default
X	array-like of shape (n, d)	Training objects.	required
y	array-like of shape (n,)	Training labels.	required

Source code in src/online_cp/mondrian/framework.py

def learn_initial_training_set(self, X: NDArray[np.floating[Any]], y: NDArray[Any]) -> None:
    r"""Batch-train the pooled model and record each example's category.

    Fits ``base_model`` on the full training set and stores the Mondrian
    category of every training example. For a label-aware taxonomy (e.g.
    ``category_fn="label"``) the category depends on both $x_i$ and $y_i$;
    for an object-conditional taxonomy it depends on $x_i$ only.

    Parameters
    ----------
    X : array-like of shape (n, d)
        Training objects.
    y : array-like of shape (n,)
        Training labels.
    """
    X = np.asarray(X)
    y = np.asarray(y)
    self.base_model.learn_initial_training_set(X, y)
    if self._label_aware:
        self.categories_ = [self._category_fn(x_i, y_i) for x_i, y_i in zip(X, y)]
    else:
        self.categories_ = [self._category_fn(x_i) for x_i in X]

learn_one(x: NDArray[np.floating[Any]], y: Any, **kwargs: Any) -> None

Update the pooled model with one example and record its category.

Parameters:

Name	Type	Description	Default
x	array-like of shape (d,)	New object.	required
y	hashable	Observed label.	required
**kwargs	Any	Forwarded to the underlying model's learn_one.	{}

Source code in src/online_cp/mondrian/framework.py

def learn_one(self, x: NDArray[np.floating[Any]], y: Any, **kwargs: Any) -> None:
    """Update the pooled model with one example and record its category.

    Parameters
    ----------
    x : array-like of shape (d,)
        New object.
    y : hashable
        Observed label.
    **kwargs
        Forwarded to the underlying model's ``learn_one``.
    """
    self.base_model.learn_one(x, y, **kwargs)
    if self._label_aware:
        self.categories_.append(self._category_fn(x, y))
    else:
        self.categories_.append(self._category_fn(x))

predict(x: NDArray[np.floating[Any]], epsilon: float | NDArray[np.floating[Any]] | None = None, return_p_values: bool = False) -> Any

Compute the Mondrian (group-conditional) prediction set.

For each candidate label the conformal p-value is computed from the pooled scores, restricted to the matching Mondrian category:

• Label-aware taxonomy (e.g. category_fn="label"): each hypothesised label \(y\) defines its own category mask \(\kappa(x, y)\), so the candidate is calibrated against training examples sharing that category. This is what makes label-conditional coverage possible.
• Object-conditional taxonomy: a single mask \(\kappa(x)\) is shared across all candidate labels.

For a multiclass SVM the scores are first restricted to the one-vs-rest positive class before the category mask is applied, matching the SVM's per-class nonconformity construction.

Parameters:

Name	Type	Description	Default
x	array-like of shape (d,)	Test object.	required
epsilon	float, array-like, or None	Significance level(s). If None, uses base_model.epsilon.	None
return_p_values	bool	If True, also return the {label: p_value} dict.	False

Returns:

Type	Description
prediction set, or ``(prediction set, p_values)`` if ``return_p_values``.

Source code in src/online_cp/mondrian/framework.py

def predict(
    self,
    x: NDArray[np.floating[Any]],
    epsilon: float | NDArray[np.floating[Any]] | None = None,
    return_p_values: bool = False,
) -> Any:
    r"""Compute the Mondrian (group-conditional) prediction set.

    For each candidate label the conformal p-value is computed from the
    pooled scores, restricted to the matching Mondrian category:

    - **Label-aware taxonomy** (e.g. ``category_fn="label"``): each
      hypothesised label $y$ defines its own category mask $\kappa(x, y)$,
      so the candidate is calibrated against training examples sharing that
      category. This is what makes label-conditional coverage possible.
    - **Object-conditional taxonomy**: a single mask $\kappa(x)$ is shared
      across all candidate labels.

    For a multiclass SVM the scores are first restricted to the one-vs-rest
    positive class before the category mask is applied, matching the SVM's
    per-class nonconformity construction.

    Parameters
    ----------
    x : array-like of shape (d,)
        Test object.
    epsilon : float, array-like, or None
        Significance level(s). If None, uses ``base_model.epsilon``.
    return_p_values : bool, default False
        If True, also return the ``{label: p_value}`` dict.

    Returns
    -------
    prediction set, or ``(prediction set, p_values)`` if ``return_p_values``.
    """
    from online_cp.classifiers import (
        ConformalNearestNeighboursClassifier,
        ConformalSupportVectorMachine,
    )

    if epsilon is None:
        epsilon = self.base_model.epsilon
    model = self.base_model
    tau = model.rnd_gen.uniform(0, 1)

    if self._label_aware:
        # Label-aware: per-hypothesis category mask
        cat_masks = {}
        for label in model.label_space:
            cat = self._category_fn(x, label)
            cat_masks[label] = np.array([c == cat for c in self.categories_])
    else:
        # Object-conditional: one shared mask for all hypotheses
        cat = self._category_fn(x)
        shared_mask = np.array([c == cat for c in self.categories_])
        cat_masks = {label: shared_mask for label in model.label_space}

    if isinstance(model, ConformalNearestNeighboursClassifier):
        p_values = self._predict_knn(x, cat_masks, tau)
    elif isinstance(model, ConformalSupportVectorMachine):
        p_values = self._predict_svm(x, cat_masks, tau)
    else:
        raise TypeError(f"Unsupported classifier: {type(model).__name__}")

    Gamma = model._compute_Gamma(p_values, epsilon)
    if return_p_values:
        return Gamma, p_values
    return Gamma