Export more metadata from export_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

@@ -103,8 +103,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): 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:
@@ -175,8 +192,16 @@ def export_model(model: Model) -> Any:
     # 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.tree import (
+            DecisionTreeClassifier,
+            DecisionTreeRegressor,
+            ExtraTreeRegressor,
+        )
     except ImportError as e:
         raise TreeliteError("This function requires scikit-learn package") from e
 
@@ -225,6 +250,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

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

@@ -195,6 +195,44 @@ def test_skl_converter_iforest(dataset):
     np.testing.assert_almost_equal(out_pred, expected_pred)
 
 
+def test_iforest_round_trip():
+    """
+    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,
+    )
+    clf.fit(X)
+    tl_model = treelite.sklearn.import_model(clf)
+    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,
+        )
+
+
 @given(
     dataset=standard_regression_datasets(),
     max_feat=floats(min_value=0.2, max_value=0.8),