[MAINT] Separate unit tests in `test_tree.py` for pickling and min_impurity_decrease
@@ -981,6 +981,9 @@ def test_min_impurity_decrease():
actual_decrease, expected_decrease
)
+
+def test_pickle():
+ """Test pickling preserves Tree properties and performance."""
for name, TreeEstimator in ALL_TREES.items():
if "Classifier" in name:
X, y = iris.data, iris.target
@@ -990,23 +993,43 @@ def test_min_impurity_decrease():
est = TreeEstimator(random_state=0)
est.fit(X, y)
score = est.score(X, y)
- fitted_attribute = dict()
- for attribute in ["max_depth", "node_count", "capacity"]:
- fitted_attribute[attribute] = getattr(est.tree_, attribute)
+
+ # test that all class properties are maintained
+ attributes = [
+ "max_depth",
+ "node_count",
+ "capacity",
+ "n_classes",
+ "children_left",
+ "children_right",
+ "n_leaves",
+ "feature",
+ "threshold",
+ "impurity",
+ "n_node_samples",
+ "weighted_n_node_samples",
+ "value",
+ ]
+ fitted_attribute = {
+ attribute: getattr(est.tree_, attribute) for attribute in attributes
+ }
serialized_object = pickle.dumps(est)
est2 = pickle.loads(serialized_object)
assert type(est2) == est.__class__
+
score2 = est2.score(X, y)
assert (
score == score2
), "Failed to generate same score after pickling with {0}".format(name)
-
for attribute in fitted_attribute:
- assert (
- getattr(est2.tree_, attribute) == fitted_attribute[attribute]
- ), "Failed to generate same attribute {0} after pickling with {1}".format(
- attribute, name
+ assert_array_equal(
+ getattr(est2.tree_, attribute),
+ fitted_attribute[attribute],
+ err_msg=(
+ f"Failed to generate same attribute {attribute} after pickling with"
+ f" {name}"
+ ),
)
