Export more metadata from export_model()

python/treelite/model.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import ctypes
+import json
 import pathlib
 import platform
 from typing import Any, List, Optional, Union
@@ -72,6 +73,24 @@ def output_type(self) -> str:
         _check_call(_LIB.TreeliteGetOutputType(self.handle, ctypes.byref(out)))
         return py_str(out.value)
 
+    @property
+    def attributes(self) -> dict[Any, Any]:
+        """Optional model attributes (JSON string)"""
+        if self.handle is None:
+            raise AttributeError("Model not loaded yet")
+
+        obj = _TreelitePyBufferFrame()
+        _check_call(
+            _LIB.TreeliteGetHeaderField(
+                self.handle,
+                c_str("attributes"),
+                ctypes.byref(obj),
+            )
+        )
+        array = _pybuffer2numpy(obj)
+        attributes_str = array.tobytes().decode("utf-8")
+        return json.loads(attributes_str)
+
     @classmethod
     def concatenate(cls, model_objs: List[Model]) -> Model:
         """

python/treelite/model_builder.py

@@ -361,7 +361,7 @@ def data_count(self, data_count: int):
             Number of data points
         """
         _check_call(
-            _LIB.TreeliteModelBuilderGain(
+            _LIB.TreeliteModelBuilderDataCount(
                 self.handle,
                 ctypes.c_uint64(data_count),
             )

python/treelite/sklearn/exporter.py

@@ -1,5 +1,6 @@
 """Converter to export Treelite models as scikit-learn models (EXPERIMENTAL)"""
 
+import warnings
 from enum import IntEnum
 from typing import Any
 
@@ -103,8 +104,25 @@ def _export_tree(
     nodes["feature"] = tree_accessor.get_field("split_index")
     nodes["threshold"] = tree_accessor.get_field("threshold")
     nodes["impurity"] = np.nan
-    nodes["n_node_samples"] = -1
-    nodes["weighted_n_node_samples"] = np.nan
+    data_count = tree_accessor.get_field("data_count").astype(np.intp)
+    data_count_mask = tree_accessor.get_field("data_count_present").astype(np.bool_)
+    if data_count.size == 0:
+        nodes["n_node_samples"] = np.full((n_nodes,), fill_value=-1, dtype=np.intp)
+    else:
+        data_count[~data_count_mask] = -1
+        nodes["n_node_samples"] = data_count
+    # TODO(chyunsu3): In Treelite 5.0, rename field sum_hess -> weighted_data_count
+    weighted_data_count = tree_accessor.get_field("sum_hess").astype(np.float64)
+    weighted_data_count_mask = tree_accessor.get_field("sum_hess_present").astype(
+        np.bool_
+    )
+    if weighted_data_count.size == 0:
+        nodes["weighted_n_node_samples"] = np.full(
+            (n_nodes,), fill_value=np.nan, dtype=np.float64
+        )
+    else:
+        weighted_data_count[~weighted_data_count_mask] = np.nan
+        nodes["weighted_n_node_samples"] = weighted_data_count
     nodes["missing_go_to_left"] = tree_accessor.get_field("default_left")
 
     if n_targets == 1 and n_classes[0] == 1:
@@ -154,7 +172,8 @@ def export_model(model: Model) -> Any:
 
     Note
     ----
-    Currently only random forests can be exported as scikit-learn model objects.
+    Currently only random forests and isolation forests can be exported as
+    scikit-learn model objects.
     Support for gradient boosted trees and other kinds of tree models will be
     added in the future.
 
@@ -168,15 +187,23 @@ def export_model(model: Model) -> Any:
     sklearn_model : object of type \
                     :py:class:`~sklearn.ensemble.RandomForestRegressor` / \
                     :py:class:`~sklearn.ensemble.RandomForestClassifier` / \
-                    :py:class:`~sklearn.ensemble.GradientBoostingRegressor` / \
-                    :py:class:`~sklearn.ensemble.GradientBoostingClassifier`
+                    :py:class:`~sklearn.ensemble.IsolationForest`
         Scikit-learn model
     """
     # pylint: disable=too-many-locals
     try:
         from sklearn import __version__ as sklearn_version
-        from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
-        from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
+        from sklearn.ensemble import (
+            IsolationForest,
+            RandomForestClassifier,
+            RandomForestRegressor,
+        )
+        from sklearn.ensemble._iforest import _average_path_length
+        from sklearn.tree import (
+            DecisionTreeClassifier,
+            DecisionTreeRegressor,
+            ExtraTreeRegressor,
+        )
     except ImportError as e:
         raise TreeliteError("This function requires scikit-learn package") from e
 
@@ -225,6 +252,9 @@ def raise_not_rf_error(reason):
     if task_type in [_TaskType.kBinaryClf, _TaskType.kMultiClf]:
         estimator_class = RandomForestClassifier
         subestimator_class = DecisionTreeClassifier
+    elif task_type == _TaskType.kIsolationForest:
+        estimator_class = IsolationForest
+        subestimator_class = ExtraTreeRegressor
     else:
         estimator_class = RandomForestRegressor
         subestimator_class = DecisionTreeRegressor
@@ -266,6 +296,42 @@ def raise_not_rf_error(reason):
                     "classes_": [np.arange(n_classes[i]) for i in range(n_targets)],
                 }
             )
+    elif estimator_class is IsolationForest:
+        # Recover the `offset_` field; if missing, set to -0.5
+        try:
+            offset = model.attributes["sklearn_iforest_offset"]
+        except KeyError:
+            warnings.warn(
+                "Treelite model does not store attribute 'sklearn_iforest_offset'; "
+                "setting it to the default value of -0.5...",
+                UserWarning,
+            )
+            offset = -0.5
+
+        # Compute max_samples by taking the max over the weighted root counts
+        # (with bootstrap=True the unweighted root only counts distinct rows)
+        max_samples = int(
+            max(estimator.tree_.weighted_n_node_samples[0] for estimator in estimators)
+        )
+        state.update(
+            {
+                "_max_samples": max_samples,
+                "max_samples_": max_samples,
+                "offset_": offset,
+                "_average_path_length_per_tree": tuple(
+                    _average_path_length(est.tree_.n_node_samples) for est in estimators
+                ),
+                "_decision_path_lengths": tuple(
+                    est.tree_.compute_node_depths() for est in estimators
+                ),
+                # The exported trees reference features globally, so scoring uses
+                # the full feature set for every tree.
+                "_max_features": n_features,
+                "estimators_features_": [
+                    np.arange(n_features, dtype=np.int64) for _ in estimators
+                ],
+            }
+        )
     clf.__setstate__(state)
 
     return clf

python/treelite/sklearn/importer.py

@@ -1,14 +1,15 @@
 """Converter to ingest scikit-learn models into Treelite"""
 
 import ctypes
+import json
 from typing import Optional
 
 import numpy as np
 from packaging.version import parse as parse_version
 
 from ..core import _LIB, TreeliteError, _check_call
-from ..model import Model
-from ..util import c_array
+from ..model import Model, _numpy2pybuffer
+from ..util import c_array, c_str
 from .isolation_forest import calculate_depths, expected_depth
 
 
@@ -79,6 +80,18 @@ def import_model(sklearn_model) -> Model:
         # clf is an IsolationForest
         # tl_model is a Treelite representation of clf
 
+    To reproduce the output of :py:meth:`~sklearn.ensemble.IsolationForest.decision_function`,
+    retrieve the value of ``IsolationForest.offset_`` and apply it,
+    as follows:
+
+    .. code-block:: python
+
+        # Treelite model stores an optional list of attributes (as a JSON string).
+        # We can retrieve `offset_` from it.
+        offset = tl_model.attributes.get("sklearn_iforest_offset", -0.5)
+        # Apply offset_ to compute the decision function.
+        decision_function = -treelite.gtil.predict(tl_model, X) - offset
+
     Parameters
     ----------
     sklearn_model : object of type \
@@ -246,6 +259,7 @@ def import_model(sklearn_model) -> Model:
             )
         )
     elif isinstance(sklearn_model, IsolationForest):
+        # TODO(chyunsu3): In Treelite 5.0, pass offset_ field via TreeliteLoadSKLearnIsolationForest()
         _check_call(
             _LIB.TreeliteLoadSKLearnIsolationForest(
                 ctypes.c_int(sklearn_model.n_estimators),
@@ -263,6 +277,20 @@ def import_model(sklearn_model) -> Model:
                 ctypes.byref(handle),
             )
         )
+        # Store `offset_` field as a model attribute
+        attributes = {
+            "sklearn_iforest_offset": float(sklearn_model.offset_),
+        }
+        attributes_serialized = json.dumps(attributes)
+        _check_call(
+            _LIB.TreeliteSetHeaderField(
+                handle,
+                c_str("attributes"),
+                _numpy2pybuffer(
+                    np.frombuffer(attributes_serialized.encode("utf-8"), dtype="S1")
+                ),
+            )
+        )
     elif isinstance(sklearn_model, (RandomForestC, ExtraTreesC)):
         n_classes = np.array(sklearn_model.n_classes_, dtype=np.int32)
         _check_call(

tests/python/test_model_builder.py

@@ -139,3 +139,48 @@ def make_tree_stump(left_child_val, right_child_val):
         expected_pred = np.array([[2, 2], [1, 1]])
         pred = treelite.gtil.predict_leaf(model, dmat)
     np.testing.assert_almost_equal(pred, expected_pred, decimal=5)
+
+
+def test_data_count_setter():
+    """Test whether data count can be specified as part of model builder"""
+    # Tree stump with 3 nodes
+    builder = ModelBuilder(
+        threshold_type="float32",
+        leaf_output_type="float32",
+        metadata=Metadata(
+            num_feature=2,
+            task_type="kRegressor",
+            average_tree_output=False,
+            num_target=1,
+            num_class=[1],
+            leaf_vector_shape=(1, 1),
+        ),
+        tree_annotation=TreeAnnotation(num_tree=1, target_id=[0], class_id=[0]),
+        postprocessor=PostProcessorFunc(name="identity"),
+        base_scores=[0.0],
+    )
+    builder.start_tree()
+    builder.start_node(0)
+    builder.numerical_test(
+        feature_id=0,
+        threshold=0.0,
+        default_left=False,
+        opname="<=",
+        left_child_key=1,
+        right_child_key=2,
+    )
+    builder.data_count(100)
+    builder.end_node()
+    builder.start_node(1)
+    builder.leaf(-1.0)
+    builder.data_count(10)
+    builder.end_node()
+    builder.start_node(2)
+    builder.leaf(1.0)
+    builder.data_count(90)
+    builder.end_node()
+    builder.end_tree()
+
+    model = builder.commit()
+    data_count = model.get_tree_accessor(0).get_field("data_count")
+    np.testing.assert_array_equal(data_count, np.array([100, 10, 90], dtype=np.int32))

tests/python/test_sklearn_integration.py

@@ -187,11 +187,74 @@ def test_skl_converter_iforest(dataset):
         random_state=0,
     )
     clf.fit(X)
-    expected_pred = -clf.score_samples(X)
-    expected_pred = expected_pred.reshape((-1, 1, 1))
+    tl_model = treelite.sklearn.import_model(clf)
+
+    # 1. Compare raw anomaly scores
+    np.testing.assert_almost_equal(
+        -treelite.gtil.predict(tl_model, X),
+        clf.score_samples(X).reshape((-1, 1, 1)),
+    )
+
+    # 2. Compare decision_function
+    # (decision_function = score_samples - offset)
+    offset = tl_model.attributes["sklearn_iforest_offset"]
+    np.testing.assert_almost_equal(
+        -treelite.gtil.predict(tl_model, X) - offset,
+        clf.decision_function(X).reshape((-1, 1, 1)),
+    )
 
+
+@pytest.mark.parametrize("bootstrap", [True, False])
+@pytest.mark.parametrize("use_sample_weights", [True, False])
+def test_iforest_round_trip(bootstrap, use_sample_weights):
+    """
+    Ensure that Treelite preserve important attributes when importing
+    and exporting isolation forests.
+    """
+
+    n_samples, n_outliers = 120, 40
+    rng = np.random.RandomState(0)
+    covariance = np.array([[0.5, -0.1], [0.7, 0.4]])
+    cluster_1 = 0.4 * rng.randn(n_samples, 2) @ covariance + np.array([2, 2])
+    cluster_2 = 0.3 * rng.randn(n_samples, 2) + np.array([-2, -2])
+    outliers = rng.uniform(low=-4, high=4, size=(n_outliers, 2))
+
+    X = np.concatenate([cluster_1, cluster_2, outliers])
+
+    clf = IsolationForest(
+        max_samples=100,
+        n_estimators=100,
+        n_jobs=-1,
+        random_state=0,
+        bootstrap=bootstrap,
+    )
+    if use_sample_weights:
+        clf.fit(X, sample_weight=rng.uniform(low=0.2, high=0.8, size=(X.shape[0],)))
+    else:
+        clf.fit(X)
     tl_model = treelite.sklearn.import_model(clf)
-    out_pred = treelite.gtil.predict(tl_model, X)
+    exported_model = treelite.sklearn.export_model(tl_model)
+    assert type(exported_model) is type(clf)
+    assert len(clf.estimators_) == len(exported_model.estimators_)
+    for old_tree, new_tree in zip(clf.estimators_, exported_model.estimators_):
+        assert type(old_tree) is type(new_tree)
+        np.testing.assert_array_equal(
+            old_tree.tree_.n_node_samples, new_tree.tree_.n_node_samples
+        )
+        np.testing.assert_almost_equal(
+            old_tree.tree_.weighted_n_node_samples,
+            new_tree.tree_.weighted_n_node_samples,
+            decimal=5,
+        )
+    np.testing.assert_almost_equal(clf.offset_, exported_model.offset_)
+    np.testing.assert_almost_equal(clf.max_samples_, exported_model.max_samples_)
+
+    expected_pred = clf.score_samples(X)
+    out_pred = exported_model.score_samples(X)
+    np.testing.assert_almost_equal(out_pred, expected_pred)
+
+    expected_pred = clf.decision_function(X)
+    out_pred = exported_model.decision_function(X)
     np.testing.assert_almost_equal(out_pred, expected_pred)