seqtrees 0.1.0__py3-none-any.whl

seqtrees/__init__.py

@@ -0,0 +1,6 @@
+"""Sequential tree models for tabular synthetic data generation."""
+
+from .synthesizer import SequentialTreeSynthesizer
+
+__all__ = ["SequentialTreeSynthesizer"]
+__version__ = "0.1.0"

seqtrees/_backends.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from typing import Any
+
+from ._tree import ConditionalTree
+
+
+SUPPORTED_TREE_BACKENDS = {"native", "sklearn", "lightgbm", "auto"}
+
+
+def create_tree_backend(
+    backend: str,
+    *,
+    max_depth: int,
+    min_samples_leaf: int,
+    min_impurity_decrease: float,
+    max_thresholds: int,
+    random_state: int | None = None,
+    n_jobs: int | None = None,
+) -> Any:
+    if backend not in SUPPORTED_TREE_BACKENDS:
+        raise ValueError(f"tree_backend must be one of {sorted(SUPPORTED_TREE_BACKENDS)}.")
+
+    if backend == "sklearn":
+        from ._sklearn_tree import SklearnConditionalTree
+
+        return SklearnConditionalTree(
+            max_depth=max_depth,
+            min_samples_leaf=min_samples_leaf,
+            min_impurity_decrease=min_impurity_decrease,
+            max_thresholds=max_thresholds,
+            random_state=random_state,
+        )
+
+    if backend == "lightgbm":
+        from ._lightgbm_tree import LightGBMConditionalTree
+
+        return LightGBMConditionalTree(
+            max_depth=max_depth,
+            min_samples_leaf=min_samples_leaf,
+            min_impurity_decrease=min_impurity_decrease,
+            max_thresholds=max_thresholds,
+            random_state=random_state,
+            n_jobs=n_jobs,
+        )
+
+    if backend == "auto":
+        try:
+            from ._lightgbm_tree import LightGBMConditionalTree
+            import lightgbm  # noqa: F401
+
+            return LightGBMConditionalTree(
+                max_depth=max_depth,
+                min_samples_leaf=min_samples_leaf,
+                min_impurity_decrease=min_impurity_decrease,
+                max_thresholds=max_thresholds,
+                random_state=random_state,
+                n_jobs=n_jobs,
+            )
+        except ImportError:
+            pass
+        try:
+            from ._sklearn_tree import SklearnConditionalTree
+            import sklearn  # noqa: F401
+
+            return SklearnConditionalTree(
+                max_depth=max_depth,
+                min_samples_leaf=min_samples_leaf,
+                min_impurity_decrease=min_impurity_decrease,
+                max_thresholds=max_thresholds,
+                random_state=random_state,
+            )
+        except ImportError:
+            pass
+
+    return ConditionalTree(
+        max_depth=max_depth,
+        min_samples_leaf=min_samples_leaf,
+        min_impurity_decrease=min_impurity_decrease,
+        max_thresholds=max_thresholds,
+    )

seqtrees/_distributions.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import math
+import random
+from collections import Counter
+from typing import Any
+
+from ._utils import is_number
+
+
+class EmpiricalDistribution:
+    """A finite empirical distribution over observed values."""
+
+    def __init__(self, values: list[Any]):
+        if not values:
+            raise ValueError("EmpiricalDistribution requires at least one value.")
+        counts = Counter(values)
+        self.values_ = list(counts.keys())
+        self.weights_ = [counts[value] for value in self.values_]
+        self.n_samples_ = len(values)
+
+    def entropy(self) -> float:
+        total = sum(self.weights_)
+        score = 0.0
+        for weight in self.weights_:
+            probability = weight / total
+            score -= probability * math.log2(probability)
+        return score
+
+    def sample(self, rng: random.Random) -> Any:
+        return rng.choices(self.values_, weights=self.weights_, k=1)[0]
+
+    def sample_interpolated(self, rng: random.Random) -> Any:
+        numeric_values = [value for value in self.values_ if is_number(value)]
+        if len(numeric_values) < 2:
+            return self.sample(rng)
+
+        left = self.sample(rng)
+        right = self.sample(rng)
+        if not is_number(left) or not is_number(right):
+            return self.sample(rng)
+        if left == right:
+            return left
+        weight = rng.random()
+        return left + weight * (right - left)
+
+    def probability(self, value: Any, alpha: float = 1.0) -> float:
+        total = sum(self.weights_)
+        size = len(self.values_)
+        for candidate, weight in zip(self.values_, self.weights_):
+            if candidate == value:
+                return (weight + alpha) / (total + alpha * size)
+        return alpha / (total + alpha * (size + 1))
+
+
+def entropy(values: list[Any]) -> float:
+    return EmpiricalDistribution(values).entropy()

seqtrees/_lightgbm_tree.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from typing import Any
+
+from ._distributions import EmpiricalDistribution
+from ._utils import require_numeric_values
+
+
+class LightGBMConditionalTree:
+    """Conditional sampler backed by LightGBM's histogram regression trees."""
+
+    def __init__(
+        self,
+        *,
+        max_depth: int = 5,
+        min_samples_leaf: int = 5,
+        min_impurity_decrease: float = 1e-9,
+        max_thresholds: int = 16,
+        random_state: int | None = None,
+        n_jobs: int | None = None,
+    ) -> None:
+        del max_thresholds
+        self.max_depth = max_depth
+        self.min_samples_leaf = min_samples_leaf
+        self.min_impurity_decrease = min_impurity_decrease
+        self.random_state = random_state
+        self.n_jobs = n_jobs
+
+    def fit(self, records: list[dict[str, Any]], features: list[str], target: str) -> "LightGBMConditionalTree":
+        try:
+            from lightgbm import LGBMRegressor
+        except ImportError as exc:
+            raise ImportError("Install seqtrees[lightgbm] to use tree_backend='lightgbm'.") from exc
+
+        self.features_ = list(features)
+        self.target_ = target
+        target_values = [record[target] for record in records]
+        self.marginal_ = EmpiricalDistribution(target_values)
+        require_numeric_values(records, self.features_ + [target], backend="lightgbm")
+
+        if self.max_depth == 0 or not self.features_ or len(set(target_values)) <= 1:
+            self.model_ = None
+            self.leaf_distributions_ = {}
+            return self
+
+        x_matrix = [self._encode_features(record) for record in records]
+
+        max_depth = -1 if self.max_depth is None else self.max_depth
+        num_leaves = 31 if max_depth <= 0 else max(2, min(31, 2**max_depth))
+        self.model_ = LGBMRegressor(
+            boosting_type="gbdt",
+            objective="regression",
+            n_estimators=32,
+            learning_rate=0.2,
+            num_leaves=num_leaves,
+            max_depth=max_depth,
+            min_child_samples=self.min_samples_leaf,
+            min_split_gain=self.min_impurity_decrease,
+            random_state=self.random_state,
+            n_jobs=self.n_jobs,
+            verbosity=-1,
+        )
+        self.model_.fit(x_matrix, target_values)
+
+        grouped_values: dict[tuple[int, ...], list[Any]] = {}
+        for leaf_id, value in zip(self._leaf_ids(x_matrix), target_values):
+            grouped_values.setdefault(leaf_id, []).append(value)
+        self.leaf_distributions_ = {
+            leaf_id: EmpiricalDistribution(values) for leaf_id, values in grouped_values.items()
+        }
+        return self
+
+    def sample(self, row: dict[str, Any], rng, *, interpolate: bool = False) -> Any:
+        distribution = self._distribution_for_row(row)
+        if interpolate:
+            return distribution.sample_interpolated(rng)
+        return distribution.sample(rng)
+
+    def leaf_probability(self, row: dict[str, Any], value: Any, alpha: float = 1.0) -> float:
+        return self._distribution_for_row(row).probability(value, alpha=alpha)
+
+    def conditional_entropy(self, records: list[dict[str, Any]]) -> float:
+        if not records:
+            return 0.0
+        total = 0.0
+        for record in records:
+            total += self._distribution_for_row(record).entropy()
+        return total / len(records)
+
+    def _distribution_for_row(self, row: dict[str, Any]) -> EmpiricalDistribution:
+        if self.model_ is None:
+            return self.marginal_
+        leaf_id = self._leaf_ids([self._encode_features(row)])[0]
+        return self.leaf_distributions_.get(leaf_id, self.marginal_)
+
+    def _encode_features(self, record: dict[str, Any]) -> list[int | float]:
+        return [float("nan") if record.get(feature) is None else record.get(feature) for feature in self.features_]
+
+    def _leaf_ids(self, x_matrix: list[list[int | float]]) -> list[tuple[int, ...]]:
+        leaves = self.model_.predict(x_matrix, pred_leaf=True)
+        result = []
+        for row in leaves:
+            if hasattr(row, "tolist"):
+                row = row.tolist()
+            if isinstance(row, list):
+                result.append(tuple(int(value) for value in row))
+            else:
+                result.append((int(row),))
+        return result

seqtrees/_parallel.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import os
+from concurrent.futures import ThreadPoolExecutor
+from typing import Callable, Iterable, TypeVar
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+def resolve_n_jobs(n_jobs: int | None) -> int:
+    if n_jobs is None:
+        return 1
+    if n_jobs == 0:
+        raise ValueError("n_jobs cannot be 0. Use None, 1, a positive integer, or -1 for all cores.")
+    if n_jobs < 0:
+        return max(1, os.cpu_count() or 1)
+    return n_jobs
+
+
+def parallel_map(func: Callable[[T], R], items: Iterable[T], n_jobs: int | None) -> list[R]:
+    values = list(items)
+    workers = resolve_n_jobs(n_jobs)
+    if workers == 1 or len(values) <= 1:
+        return [func(value) for value in values]
+    with ThreadPoolExecutor(max_workers=workers) as executor:
+        return list(executor.map(func, values))

seqtrees/_sklearn_tree.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+from typing import Any
+
+from ._distributions import EmpiricalDistribution
+from ._utils import require_numeric_values
+
+
+class SklearnConditionalTree:
+    """Conditional sampler backed by scikit-learn's compiled regression tree."""
+
+    def __init__(
+        self,
+        *,
+        max_depth: int = 5,
+        min_samples_leaf: int = 5,
+        min_impurity_decrease: float = 1e-9,
+        max_thresholds: int = 16,
+        random_state: int | None = None,
+    ) -> None:
+        del max_thresholds
+        self.max_depth = max_depth
+        self.min_samples_leaf = min_samples_leaf
+        self.min_impurity_decrease = min_impurity_decrease
+        self.random_state = random_state
+
+    def fit(self, records: list[dict[str, Any]], features: list[str], target: str) -> "SklearnConditionalTree":
+        try:
+            from sklearn.tree import DecisionTreeRegressor
+        except ImportError as exc:
+            raise ImportError("Install seqtrees[sklearn] to use tree_backend='sklearn'.") from exc
+
+        self.features_ = list(features)
+        self.target_ = target
+        target_values = [record[target] for record in records]
+        self.marginal_ = EmpiricalDistribution(target_values)
+        require_numeric_values(records, self.features_ + [target], backend="sklearn")
+
+        if self.max_depth == 0 or not self.features_ or len(set(target_values)) <= 1:
+            self.model_ = None
+            self.leaf_distributions_ = {}
+            return self
+
+        x_matrix = [self._encode_features(record) for record in records]
+
+        max_depth = None if self.max_depth is None else self.max_depth
+        self.model_ = DecisionTreeRegressor(
+            criterion="squared_error",
+            max_depth=max_depth,
+            min_samples_leaf=self.min_samples_leaf,
+            min_impurity_decrease=self.min_impurity_decrease,
+            random_state=self.random_state,
+        )
+        self.model_.fit(x_matrix, target_values)
+
+        leaves = self.model_.apply(x_matrix)
+        grouped_values: dict[int, list[Any]] = {}
+        for leaf_id, value in zip(leaves, target_values):
+            grouped_values.setdefault(int(leaf_id), []).append(value)
+        self.leaf_distributions_ = {
+            leaf_id: EmpiricalDistribution(values) for leaf_id, values in grouped_values.items()
+        }
+        return self
+
+    def sample(self, row: dict[str, Any], rng, *, interpolate: bool = False) -> Any:
+        distribution = self._distribution_for_row(row)
+        if interpolate:
+            return distribution.sample_interpolated(rng)
+        return distribution.sample(rng)
+
+    def leaf_probability(self, row: dict[str, Any], value: Any, alpha: float = 1.0) -> float:
+        return self._distribution_for_row(row).probability(value, alpha=alpha)
+
+    def conditional_entropy(self, records: list[dict[str, Any]]) -> float:
+        if not records:
+            return 0.0
+        total = 0.0
+        for record in records:
+            total += self._distribution_for_row(record).entropy()
+        return total / len(records)
+
+    def _distribution_for_row(self, row: dict[str, Any]) -> EmpiricalDistribution:
+        if self.model_ is None:
+            return self.marginal_
+        leaf_id = int(self.model_.apply([self._encode_features(row)])[0])
+        return self.leaf_distributions_.get(leaf_id, self.marginal_)
+
+    def _encode_features(self, record: dict[str, Any]) -> list[int]:
+        return [record.get(feature) for feature in self.features_]

seqtrees/_tree.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from ._distributions import EmpiricalDistribution, entropy
+from ._utils import is_number
+
+
+@dataclass
+class Split:
+    feature: str
+    threshold: float | None = None
+    category: Any | None = None
+
+    def go_left(self, row: dict[str, Any]) -> bool:
+        value = row.get(self.feature)
+        if self.threshold is not None:
+            return is_number(value) and value <= self.threshold
+        return value == self.category
+
+    @property
+    def label(self) -> str:
+        if self.threshold is not None:
+            return f"{self.feature} <= {self.threshold:g}"
+        return f"{self.feature} == {self.category!r}"
+
+
+@dataclass
+class TreeNode:
+    distribution: EmpiricalDistribution
+    split: Split | None = None
+    left: "TreeNode | None" = None
+    right: "TreeNode | None" = None
+
+    @property
+    def is_leaf(self) -> bool:
+        return self.split is None
+
+
+class ConditionalTree:
+    """Small CART-style conditional sampler used internally by SeqTrees."""
+
+    def __init__(
+        self,
+        *,
+        max_depth: int = 5,
+        min_samples_leaf: int = 5,
+        min_impurity_decrease: float = 1e-9,
+        max_thresholds: int = 16,
+    ) -> None:
+        self.max_depth = max_depth
+        self.min_samples_leaf = min_samples_leaf
+        self.min_impurity_decrease = min_impurity_decrease
+        self.max_thresholds = max_thresholds
+
+    def fit(self, records: list[dict[str, Any]], features: list[str], target: str) -> "ConditionalTree":
+        self.features_ = list(features)
+        self.target_ = target
+        self.root_ = self._fit_node(records, depth=0)
+        return self
+
+    def sample(self, row: dict[str, Any], rng, *, interpolate: bool = False) -> Any:
+        node = self.root_
+        while not node.is_leaf:
+            node = node.left if node.split.go_left(row) else node.right
+        if interpolate:
+            return node.distribution.sample_interpolated(rng)
+        return node.distribution.sample(rng)
+
+    def leaf_probability(self, row: dict[str, Any], value: Any, alpha: float = 1.0) -> float:
+        node = self.root_
+        while not node.is_leaf:
+            node = node.left if node.split.go_left(row) else node.right
+        return node.distribution.probability(value, alpha=alpha)
+
+    def conditional_entropy(self, records: list[dict[str, Any]]) -> float:
+        if not records:
+            return 0.0
+        total = 0.0
+        for record in records:
+            node = self.root_
+            while not node.is_leaf:
+                node = node.left if node.split.go_left(record) else node.right
+            total += node.distribution.entropy()
+        return total / len(records)
+
+    def _fit_node(self, records: list[dict[str, Any]], depth: int) -> TreeNode:
+        values = [record[self.target_] for record in records]
+        distribution = EmpiricalDistribution(values)
+        node = TreeNode(distribution=distribution)
+
+        if depth >= self.max_depth or len(records) < 2 * self.min_samples_leaf or distribution.entropy() == 0:
+            return node
+
+        split, gain = self._best_split(records)
+        if split is None or gain <= self.min_impurity_decrease:
+            return node
+
+        left_records, right_records = self._partition(records, split)
+        node.split = split
+        node.left = self._fit_node(left_records, depth + 1)
+        node.right = self._fit_node(right_records, depth + 1)
+        return node
+
+    def _best_split(self, records: list[dict[str, Any]]) -> tuple[Split | None, float]:
+        parent_entropy = entropy([record[self.target_] for record in records])
+        best_split = None
+        best_gain = 0.0
+        for feature in self.features_:
+            for split in self._candidate_splits(records, feature):
+                left_records, right_records = self._partition(records, split)
+                if len(left_records) < self.min_samples_leaf or len(right_records) < self.min_samples_leaf:
+                    continue
+                weighted_entropy = (
+                    len(left_records) * entropy([record[self.target_] for record in left_records])
+                    + len(right_records) * entropy([record[self.target_] for record in right_records])
+                ) / len(records)
+                gain = parent_entropy - weighted_entropy
+                if gain > best_gain:
+                    best_split = split
+                    best_gain = gain
+        return best_split, best_gain
+
+    def _candidate_splits(self, records: list[dict[str, Any]], feature: str) -> list[Split]:
+        values = [record.get(feature) for record in records if record.get(feature) is not None]
+        unique = sorted(set(values)) if all(is_number(value) for value in values) else list(dict.fromkeys(values))
+        if len(unique) <= 1:
+            return []
+
+        if all(is_number(value) for value in unique) and len(unique) > 2:
+            thresholds = [(unique[i] + unique[i + 1]) / 2 for i in range(len(unique) - 1)]
+            if len(thresholds) > self.max_thresholds:
+                step = len(thresholds) / self.max_thresholds
+                thresholds = [thresholds[int(i * step)] for i in range(self.max_thresholds)]
+            return [Split(feature=feature, threshold=threshold) for threshold in thresholds]
+
+        return [Split(feature=feature, category=value) for value in unique[:-1]]
+
+    @staticmethod
+    def _partition(records: list[dict[str, Any]], split: Split) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
+        left_records = []
+        right_records = []
+        for record in records:
+            if split.go_left(record):
+                left_records.append(record)
+            else:
+                right_records.append(record)
+        return left_records, right_records

seqtrees/_utils.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+
+def is_pandas_dataframe(value: Any) -> bool:
+    return hasattr(value, "columns") and hasattr(value, "to_dict") and hasattr(value, "__len__")
+
+
+def normalize_table(data: Any) -> tuple[list[dict[str, Any]], list[str], bool]:
+    """Return records, column names, and whether the input looked DataFrame-like."""
+    if is_pandas_dataframe(data):
+        columns = [str(column) for column in data.columns]
+        records = [{str(key): val for key, val in row.items()} for row in data.to_dict("records")]
+        return records, columns, True
+
+    if not isinstance(data, Sequence) or isinstance(data, (str, bytes)):
+        raise TypeError("X must be a pandas DataFrame, a list of dictionaries, or a list of rows.")
+
+    rows = list(data)
+    if not rows:
+        raise ValueError("X must contain at least one row.")
+
+    first = rows[0]
+    if isinstance(first, Mapping):
+        columns = [str(column) for column in first.keys()]
+        records = []
+        for row in rows:
+            if not isinstance(row, Mapping):
+                raise TypeError("All rows must use the same representation.")
+            records.append({str(key): value for key, value in row.items()})
+        return records, columns, False
+
+    if isinstance(first, Sequence) and not isinstance(first, (str, bytes)):
+        width = len(first)
+        columns = [f"x{i}" for i in range(width)]
+        records = []
+        for row in rows:
+            if not isinstance(row, Sequence) or isinstance(row, (str, bytes)) or len(row) != width:
+                raise TypeError("All rows must be sequences with the same length.")
+            records.append({columns[i]: value for i, value in enumerate(row)})
+        return records, columns, False
+
+    raise TypeError("Rows must be dictionaries or sequences.")
+
+
+def resolve_columns(columns: list[str], requested: Sequence[str | int] | None) -> list[str]:
+    if requested is None:
+        return list(columns)
+
+    resolved = []
+    for item in requested:
+        if isinstance(item, int):
+            try:
+                resolved.append(columns[item])
+            except IndexError as exc:
+                raise ValueError(f"Column index out of range: {item}") from exc
+        else:
+            name = str(item)
+            if name not in columns:
+                raise ValueError(f"Unknown column in variable_order: {name!r}")
+            resolved.append(name)
+
+    if len(set(resolved)) != len(resolved):
+        raise ValueError("variable_order cannot contain duplicate columns.")
+    if set(resolved) != set(columns):
+        missing = sorted(set(columns) - set(resolved))
+        extra = sorted(set(resolved) - set(columns))
+        raise ValueError(f"variable_order must include every column exactly once; missing={missing}, extra={extra}")
+    return resolved
+
+
+def resolve_column_subset(columns: list[str], requested: Sequence[str | int] | None, *, parameter: str) -> list[str]:
+    if requested is None:
+        return list(columns)
+
+    resolved = []
+    for item in requested:
+        if isinstance(item, int):
+            try:
+                resolved.append(columns[item])
+            except IndexError as exc:
+                raise ValueError(f"Column index out of range in {parameter}: {item}") from exc
+        else:
+            name = str(item)
+            if name not in columns:
+                raise ValueError(f"Unknown column in {parameter}: {name!r}")
+            resolved.append(name)
+
+    if len(set(resolved)) != len(resolved):
+        raise ValueError(f"{parameter} cannot contain duplicate columns.")
+    return resolved
+
+
+def is_number(value: Any) -> bool:
+    return isinstance(value, (int, float)) and not isinstance(value, bool)
+
+
+def is_integer_code(value: Any) -> bool:
+    return isinstance(value, int) and not isinstance(value, bool)
+
+
+def is_float_value(value: Any) -> bool:
+    return isinstance(value, float) and not isinstance(value, bool)
+
+
+def require_numeric_values(records: list[dict[str, Any]], columns: list[str], *, backend: str) -> None:
+    for column in columns:
+        for row_index, record in enumerate(records):
+            value = record.get(column)
+            if value is None:
+                continue
+            if not is_number(value):
+                raise TypeError(
+                    f"tree_backend={backend!r} requires preprocessed numeric data; "
+                    f"column {column!r} has non-numeric value {value!r} at row {row_index}."
+                )
+
+
+def validate_no_nulls(records: list[dict[str, Any]], columns: list[str]) -> None:
+    for row_index, record in enumerate(records):
+        for column in columns:
+            if record.get(column) is None:
+                raise ValueError(f"SeqTrees does not accept null values; found null in column {column!r} at row {row_index}.")
+
+
+def validate_variable_types(
+    records: list[dict[str, Any]],
+    *,
+    continuous_columns: set[str],
+    discrete_columns: set[str],
+) -> None:
+    for column in continuous_columns:
+        for row_index, record in enumerate(records):
+            value = record[column]
+            if not is_float_value(value):
+                raise TypeError(
+                    f"continuous column {column!r} must contain float values; "
+                    f"found {value!r} at row {row_index}."
+                )
+
+    for column in discrete_columns:
+        for row_index, record in enumerate(records):
+            value = record[column]
+            if not is_integer_code(value):
+                raise TypeError(
+                    f"discrete column {column!r} must contain integer category codes; "
+                    f"found {value!r} at row {row_index}."
+                )