Random Forest

"""Random forest - bootstrap, random feature subsets, average. Transcribed.

Two sources of randomness, and the second is the one that matters. The
bootstrap gives each tree a different sample; the per-node feature subset
forces trees to use different features, which is what drives the pairwise
correlation down and makes the averaging worth doing.
"""

import random


def gini(labels):
    if not labels:
        return 0.0
    counts = {}
    for label in labels:
        counts[label] = counts.get(label, 0) + 1
    impurity = 1.0
    for count in counts.values():
        proportion = count / len(labels)
        impurity -= proportion * proportion
    return impurity


def majority(labels):
    counts = {}
    for label in labels:
        counts[label] = counts.get(label, 0) + 1
    best_label, best_count = None, -1
    for label, count in counts.items():
        if count > best_count:
            best_label, best_count = label, count
    return best_label


def best_split(X, y, feature_subset):
    """Same criterion as CART, restricted to a random subset of features."""
    n = len(X)
    best = (None, None, gini(y))

    for feature in feature_subset:
        values = sorted({X[i][feature] for i in range(n)})
        for position in range(len(values) - 1):
            threshold = (values[position] + values[position + 1]) / 2.0

            left, right = [], []
            for i in range(n):
                if X[i][feature] <= threshold:
                    left.append(y[i])
                else:
                    right.append(y[i])

            if not left or not right:
                continue

            weighted = len(left) / n * gini(left) + len(right) / n * gini(right)
            if weighted < best[2]:
                best = (feature, threshold, weighted)

    return best


def build_tree(X, y, m_try, rng, depth=0, max_depth=20, min_samples_leaf=1):
    if depth >= max_depth or len(set(y)) == 1 or len(y) < 2 * min_samples_leaf:
        return {"leaf": True, "prediction": majority(y)}

    # Redrawn at EVERY node, not once per tree: drawing once would just give a
    # forest of trees each fitted on one fixed random subspace, which
    # decorrelates far less.
    feature_subset = rng.sample(range(len(X[0])), m_try)
    feature, threshold, _ = best_split(X, y, feature_subset)
    if feature is None:
        return {"leaf": True, "prediction": majority(y)}

    left_X, left_y, right_X, right_y = [], [], [], []
    for i in range(len(X)):
        if X[i][feature] <= threshold:
            left_X.append(X[i])
            left_y.append(y[i])
        else:
            right_X.append(X[i])
            right_y.append(y[i])

    return {
        "leaf": False,
        "feature": feature,
        "threshold": threshold,
        "left": build_tree(left_X, left_y, m_try, rng, depth + 1, max_depth, min_samples_leaf),
        "right": build_tree(right_X, right_y, m_try, rng, depth + 1, max_depth, min_samples_leaf),
    }


def fit(X, y, n_trees=100, m_try=None, seed=0):
    n = len(X)
    d = len(X[0])
    m_try = m_try or max(1, int(math_sqrt(d)))
    rng = random.Random(seed)

    forest = []
    out_of_bag = []          # rows NOT drawn for each tree - the free holdout

    for _ in range(n_trees):
        indices = [rng.randrange(n) for _ in range(n)]
        drawn = set(indices)
        # About 37% of rows fall outside any given bootstrap sample.
        out_of_bag.append([i for i in range(n) if i not in drawn])

        sample_X = [X[i] for i in indices]
        sample_y = [y[i] for i in indices]
        forest.append(build_tree(sample_X, sample_y, m_try, rng))

    return forest, out_of_bag


def math_sqrt(value):
    return value ** 0.5


def predict_one(forest, x):
    votes = {}
    for tree in forest:
        node = tree
        while not node["leaf"]:
            node = node["left"] if x[node["feature"]] <= node["threshold"] else node["right"]
        votes[node["prediction"]] = votes.get(node["prediction"], 0) + 1

    best_label, best_count = None, -1
    for label, count in votes.items():
        if count > best_count:
            best_label, best_count = label, count
    return best_label

"""Random forest - typed, index-partitioned, with an out-of-bag estimate."""

from dataclasses import dataclass, field

import numpy as np
from numpy.typing import NDArray

Matrix = NDArray[np.float64]
Labels = NDArray[np.int64]


@dataclass(frozen=True)
class TreeNode:
    prediction: int
    feature: int = -1
    threshold: float = 0.0
    left: int = -1
    right: int = -1

    @property
    def is_leaf(self) -> bool:
        return self.feature < 0


@dataclass
class Tree:
    """Nodes in one flat list; children are indices, so prediction is a walk."""

    nodes: list[TreeNode] = field(default_factory=list)

    def predict(self, X: Matrix) -> Labels:
        out = np.empty(X.shape[0], dtype=np.int64)
        for row_index, row in enumerate(X):
            cursor = 0
            while not self.nodes[cursor].is_leaf:
                node = self.nodes[cursor]
                cursor = node.left if row[node.feature] <= node.threshold else node.right
            out[row_index] = self.nodes[cursor].prediction
        return out


def _grow(
    X: Matrix,
    y: Labels,
    rows: NDArray,
    tree: Tree,
    n_classes: int,
    m_try: int,
    rng: np.random.Generator,
    depth: int,
    max_depth: int,
    min_samples_leaf: int,
) -> int:
    counts = np.bincount(y[rows], minlength=n_classes)
    index = len(tree.nodes)
    tree.nodes.append(TreeNode(prediction=int(counts.argmax())))

    if depth >= max_depth or counts.max() == rows.size or rows.size < 2 * min_samples_leaf:
        return index

    # Redrawn at every node. Drawing once per tree would give a forest of
    # fixed random subspaces, which decorrelates far less.
    subset = rng.choice(X.shape[1], size=m_try, replace=False)

    best_impurity = 1.0 - float(((counts / rows.size) ** 2).sum())
    best: tuple[int, float] | None = None

    for feature in subset:
        order = rows[np.argsort(X[rows, feature], kind="stable")]
        values = X[order, feature]
        left_counts = np.zeros(n_classes, dtype=np.int64)
        right_counts = counts.copy()

        for position in range(rows.size - 1):
            label = y[order[position]]
            left_counts[label] += 1
            right_counts[label] -= 1

            if values[position] == values[position + 1]:
                continue
            n_left, n_right = position + 1, rows.size - position - 1
            if n_left < min_samples_leaf or n_right < min_samples_leaf:
                continue

            weighted = (
                n_left * (1.0 - ((left_counts / n_left) ** 2).sum())
                + n_right * (1.0 - ((right_counts / n_right) ** 2).sum())
            ) / rows.size

            if weighted < best_impurity:
                best_impurity = weighted
                best = (int(feature), float((values[position] + values[position + 1]) / 2.0))

    if best is None:
        return index

    feature, threshold = best
    mask = X[rows, feature] <= threshold
    left = _grow(X, y, rows[mask], tree, n_classes, m_try, rng, depth + 1, max_depth, min_samples_leaf)
    right = _grow(X, y, rows[~mask], tree, n_classes, m_try, rng, depth + 1, max_depth, min_samples_leaf)

    tree.nodes[index] = TreeNode(
        prediction=tree.nodes[index].prediction,
        feature=feature,
        threshold=threshold,
        left=left,
        right=right,
    )
    return index


@dataclass(frozen=True)
class RandomForest:
    trees: list[Tree]
    n_classes: int
    oob_score: float

    def predict(self, X: Matrix) -> Labels:
        votes = np.zeros((X.shape[0], self.n_classes), dtype=np.int64)
        for tree in self.trees:
            votes[np.arange(X.shape[0]), tree.predict(X)] += 1
        return votes.argmax(axis=1)


def fit(
    X: Matrix,
    y: Labels,
    n_trees: int = 100,
    m_try: int | None = None,
    max_depth: int = 20,
    min_samples_leaf: int = 1,
    seed: int = 0,
) -> RandomForest:
    """Grow a forest. Raises ValueError on malformed input."""
    if X.ndim != 2:
        raise ValueError(f"X must be 2-D, got shape {X.shape}")
    if X.shape[0] != y.shape[0]:
        raise ValueError(f"X has {X.shape[0]} rows but y has {y.shape[0]}")
    if n_trees < 1:
        raise ValueError(f"n_trees must be at least 1, got {n_trees}")

    n, d = X.shape
    n_classes = int(y.max()) + 1
    features_per_split = m_try or max(1, int(np.sqrt(d)))
    rng = np.random.default_rng(seed)

    trees: list[Tree] = []
    # Out-of-bag votes: about 37% of rows sit outside each bootstrap sample,
    # which gives an honest holdout estimate without splitting the data.
    oob_votes = np.zeros((n, n_classes), dtype=np.int64)

    for _ in range(n_trees):
        indices = rng.integers(0, n, size=n)
        tree = Tree()
        _grow(X, y, indices, tree, n_classes, features_per_split, rng, 0, max_depth, min_samples_leaf)
        trees.append(tree)

        oob_mask = np.ones(n, dtype=bool)
        oob_mask[indices] = False
        if oob_mask.any():
            oob_votes[np.flatnonzero(oob_mask), tree.predict(X[oob_mask])] += 1

    scored = oob_votes.sum(axis=1) > 0
    oob_score = float((oob_votes[scored].argmax(axis=1) == y[scored]).mean()) if scored.any() else 0.0

    return RandomForest(trees=trees, n_classes=n_classes, oob_score=oob_score)

"""Random forest - shared binned features, flat trees, batched voting."""

import numpy as np
from numpy.typing import NDArray

Matrix = NDArray[np.float64]

N_BINS = 256


class BinnedForest:
    """The two changes that matter once B is large.

    First: the design matrix is quantized into 256 bins ONCE, before any tree
    exists, and every tree reads the same uint8 matrix. Binning costs one pass
    and is amortized across all B trees, which is where it genuinely pays -
    for a single tree it is close to a wash.

    Second: trees are stored as parallel flat arrays rather than objects, so
    prediction over a batch is an index walk through contiguous int32 arrays
    with no attribute lookup per node.

    Splits are chosen from bin edges rather than every midpoint, so they are
    approximate. That is the standing trade of histogram trees; averaging over
    a forest tends to absorb it, which is why production forests bin.
    """

    def __init__(self, X: Matrix, n_bins: int = N_BINS) -> None:
        design = np.ascontiguousarray(X, dtype=np.float64)
        self._edges = np.stack([
            np.quantile(design[:, feature], np.linspace(0.0, 1.0, n_bins + 1)[1:-1])
            for feature in range(design.shape[1])
        ])
        # Fortran order: split finding walks one feature at a time, so the
        # feature must be the contiguous axis. uint8 is 8x smaller than the
        # float64 it replaces, so far more stays in cache.
        self._binned = np.empty(design.shape, dtype=np.uint8, order="F")
        for feature in range(design.shape[1]):
            self._binned[:, feature] = np.searchsorted(
                self._edges[feature], design[:, feature], side="left"
            )

        # Trees as parallel flat arrays, appended to as the forest grows.
        self._feature: list[NDArray] = []
        self._bin: list[NDArray] = []
        self._left: list[NDArray] = []
        self._right: list[NDArray] = []
        self._value: list[NDArray] = []

    def bin_queries(self, X: Matrix) -> NDArray:
        """Queries must be binned with the SAME edges the forest was fitted on."""
        out = np.empty(X.shape, dtype=np.uint8, order="F")
        for feature in range(X.shape[1]):
            out[:, feature] = np.searchsorted(self._edges[feature], X[:, feature], side="left")
        return out

    def predict(self, X: Matrix, n_classes: int) -> NDArray:
        """Level-synchronous traversal: every row of the batch moves together."""
        binned = self.bin_queries(X)
        n = binned.shape[0]

        votes = np.zeros((n, n_classes), dtype=np.int32)      # allocated once
        cursor = np.empty(n, dtype=np.int32)
        rows = np.arange(n)

        for tree in range(len(self._feature)):
            feature = self._feature[tree]
            split_bin = self._bin[tree]
            left = self._left[tree]
            right = self._right[tree]
            value = self._value[tree]

            cursor.fill(0)
            active = rows

            # One iteration per level, not one per row: all rows still at an
            # internal node advance together, so the interpreter runs O(depth)
            # times instead of O(n * depth).
            while active.size:
                nodes = cursor[active]
                internal = feature[nodes] >= 0
                if not internal.any():
                    break

                moving = active[internal]
                at = cursor[moving]
                goes_left = binned[moving, feature[at]] <= split_bin[at]
                np.copyto(cursor[moving], np.where(goes_left, left[at], right[at]))
                active = moving

            np.add.at(votes, (rows, value[cursor]), 1)

        return votes.argmax(axis=1)