liesel-gam 0.0.4__py3-none-any.whl → 0.0.6a4__py3-none-any.whl

liesel_gam/builder/category_mapping.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import numpy as np
+import pandas as pd
+
+Array = Any
+
+
+class CategoryError(KeyError):
+    pass
+
+
+class UnknownLabelError(CategoryError):
+    pass
+
+
+class UnknownCodeError(CategoryError):
+    pass
+
+
+class CategoryMapping:
+    """Wraps a category mapping of labels to integers."""
+
+    def __init__(self, labels_to_integers_map: dict[Any, int]) -> None:
+        self._code_for_unknown_label = -1
+        self._label_for_unknown_code = None
+
+        self.labels_to_integers_map = labels_to_integers_map
+        self.integers_to_labels_map = {
+            code: label for label, code in self.labels_to_integers_map.items()
+        }
+
+    @classmethod
+    def from_series(cls, series: pd.Series | pd.Categorical) -> CategoryMapping:
+        """
+        When series is a pd.Categorical, the category sorting is kept.
+        When series is a series of dtype str or object, categories are sorted
+        alphabetically.
+        """
+        is_series = isinstance(series, pd.Series)
+        has_cat_dtype = isinstance(series.dtype, pd.CategoricalDtype)
+        is_cat = isinstance(series, pd.Categorical)
+        if is_cat:
+            unique_labels = np.asarray(series.categories)
+        elif is_series and has_cat_dtype:
+            unique_labels = np.asarray(series.cat.categories)
+        elif is_series:
+            cat = pd.Categorical(series)
+            unique_labels = np.sort(np.asarray(cat.categories))
+        else:
+            raise TypeError(
+                f"series must be a pd.Series or pd.Categorical, got {type(series)}."
+            )
+
+        mapping = {val: i for i, val in enumerate(unique_labels)}
+        return cls(mapping)
+
+    def to_integers(
+        self, labels_or_integers: np.typing.ArrayLike | Sequence[int] | Sequence[str]
+    ) -> np.typing.NDArray[np.int_]:
+        arr = np.asarray(labels_or_integers)
+
+        # Case 1: Already an integer array
+        if np.issubdtype(arr.dtype, np.integer):
+            valid_integers = np.array(list(self.integers_to_labels_map.keys()))
+            if not np.isin(arr, valid_integers).all():
+                invalid = arr[~np.isin(arr, valid_integers)]
+                raise ValueError(
+                    f"Unknown integer codes: {invalid.tolist()} "
+                    f"(valid integers: {valid_integers.tolist()})"
+                )
+            return arr.astype(int, copy=False)
+
+        # Case 2: Otherwise treat as labels
+        return self.labels_to_integers(arr)
+
+    def to_labels(
+        self, labels_or_integers: np.typing.ArrayLike | Sequence[int] | Sequence[str]
+    ) -> np.typing.NDArray[Any]:
+        arr = np.asarray(labels_or_integers)
+
+        # Case 1: It is an integer array
+        if np.issubdtype(arr.dtype, np.integer):
+            return self.integers_to_labels(arr)
+
+        # Case 2: Otherwise treat as labels
+        valid_labels = np.array(list(self.labels_to_integers_map.keys()))
+        if not np.isin(arr, valid_labels).all():
+            invalid = arr[~np.isin(arr, valid_labels)]
+            raise ValueError(
+                f"Unknown labels: {invalid.tolist()} "
+                f"(valid labels: {valid_labels.tolist()})"
+            )
+        return arr
+
+    def labels_to_integers(
+        self, labels: np.typing.ArrayLike | Sequence[str]
+    ) -> np.typing.NDArray[np.int_]:
+        """
+        A function of labels -> integers.
+
+        For unknown labels, returns -1.
+        """
+        labels = np.asarray(labels)
+        labels_flat = labels.flatten()
+        codes_flat = np.zeros_like(labels_flat, dtype=int)
+
+        for i, xi in enumerate(labels_flat):
+            codes_flat[i] = self.labels_to_integers_map.get(
+                xi, self._code_for_unknown_label
+            )
+            if codes_flat[i] == self._code_for_unknown_label:
+                raise UnknownLabelError(f"Category label {xi} is unknown.")
+
+        codes = np.reshape(codes_flat, shape=labels.shape)
+
+        return np.astype(codes, np.int_)
+
+    def integers_to_labels(
+        self, integers: np.typing.ArrayLike | Sequence[int]
+    ) -> np.typing.NDArray[Any]:
+        """
+        A function of integers -> labels.
+
+        For integers without labels, returns
+        """
+        integers = np.asarray(integers)
+        integers_flat = integers.flatten()
+        labels_flat_list = []
+
+        for xi in integers_flat:
+            label = self.integers_to_labels_map.get(xi, self._label_for_unknown_code)
+            if label == self._label_for_unknown_code:
+                raise UnknownCodeError(f"Category code {xi} is unknown.")
+            labels_flat_list.append(label)
+
+        labels_flat = np.asarray(labels_flat_list)
+        labels = np.reshape(labels_flat, shape=integers.shape)
+        return labels
+
+
+def series_is_categorical(series: pd.Series | pd.Categorical) -> bool:
+    """
+    Provides a liberal interpretation of when a series is categorical. The following
+    are treated as categorical:
+
+    - Series with dtype str
+    - Series with dtype object
+    - Series with dtype CategoricalDtype
+    """
+    # This corresponds to how formulaic determines categorical columns.
+    # See formulaic.materializers.pandas.PandasMaterializer._is_categorical
+    is_cat1 = series.dtype in ("str", "object")
+    is_cat2 = isinstance(series.dtype, pd.CategoricalDtype)
+    return is_cat1 or is_cat2

liesel_gam/builder/consolidate_bases.py

@@ -0,0 +1,105 @@
+"""
+Instances of :class:`.Basis` may use non-jittable basis functions.
+In batched optimization, this may lead to inefficient repeated basis evaluation.
+If the basis functions depend on ryp for interfacing to R, which many do, this will
+not only be inefficient but fail completely, because R is not thread-safe.
+
+To solve these issues, this module provides utility functions to create models that
+can be safely and efficiently used in batched operations:
+
+- :func:`.consolidate_bases` splits a model into the model, where
+   all bases are turned into strong, observed varibales, and a model for the bases.
+   The former can be used in batched optimization, the latter can be used to still
+   conveniently evaluate all relevant bases based on their original inputs.
+
+- :func:`.evaluate_bases` takes a position of input data, evaluate the corresponding
+  bases in the provided model, and returns a position of the evaluated bases.
+"""
+
+import liesel.model as lsl
+from liesel.goose.types import Position
+
+from ..var import Basis
+
+
+def _remove_singleton_vars(gb: lsl.GraphBuilder) -> lsl.GraphBuilder:
+    """
+    Removes all singleton variables from the provided GraphBuilder.
+    """
+    model = gb.build_model()
+
+    G = model.var_graph
+    singletons1 = [n for n, d in G.degree() if d == 0]
+    singletons2 = [n for n in G.nodes() if G.in_degree(n) == 0 and G.out_degree(n) == 0]
+    singleton_vars = set(singletons1 + singletons2)
+
+    G = model.node_graph
+    singletons1 = [n for n, d in G.degree() if d == 0]
+    singletons2 = [n for n in G.nodes() if G.in_degree(n) == 0 and G.out_degree(n) == 0]
+    singleton_nodes = set(singletons1 + singletons2)
+
+    nodes, vars_ = model.pop_nodes_and_vars()
+
+    for var in singleton_vars:
+        vars_.pop(var.name, None)
+
+    for node in singleton_nodes:
+        nodes.pop(node.name, None)
+
+    gb = lsl.GraphBuilder(to_float32=model._to_float32)
+    gb.add(*vars_.values())
+    return gb
+
+
+def consolidate_bases(
+    model: lsl.Model, copy: bool = True
+) -> tuple[lsl.Model, lsl.Model]:
+    """
+    Turns all :class:`.Basis` variables in the provided model into strong,
+    observed :class:`liesel.model.Var` variables.
+
+    Returns a new model that depends only on the strong bases, and a model that
+    holds the original bases and their input data.
+
+    If ``copy=False``, all data will be extracted from the original model, instead
+    of creating copies. This saves memory, but renders the original model empty.
+    """
+    if copy:
+        nodes, vars_ = model.copy_nodes_and_vars()
+    else:
+        nodes, vars_ = model.pop_nodes_and_vars()
+
+    gb = lsl.GraphBuilder(to_float32=model._to_float32)
+    gb.add(*nodes.values(), *vars_.values())
+
+    weak_bases = []
+
+    for var in gb.vars:
+        if not isinstance(var, Basis):
+            continue
+        weak_basis = var
+        strong_basis = lsl.Var.new_obs(weak_basis.update().value, name=weak_basis.name)
+        gb.replace_var(old=weak_basis, new=strong_basis)
+        weak_bases.append(weak_basis)
+
+    gb = _remove_singleton_vars(gb)
+
+    bases_model = lsl.GraphBuilder().add(*weak_bases).build_model()
+    model = gb.build_model()
+
+    return model, bases_model
+
+
+def evaluate_bases(newdata: Position, model: lsl.Model) -> Position:
+    """
+    Evaluates all :class:`.Basis` variables in the provided model at the provided
+    newdata position.
+    """
+    state = model.update_state(newdata)
+
+    basis_names = []
+    for var in model.vars.values():
+        if isinstance(var, Basis):
+            basis_names.append(var.name)
+
+    return Position(model.extract_position(basis_names, state))