fixed pickling of grf (#384)

* fixed pickling of grf

* changed it to using the self.max_values for even better consistency with sklearn logic for future easier extendibility.

Author: vsyrgkanis

econml/tests/test_grf_python.py

@@ -8,6 +8,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+import joblib
 from econml.grf import RegressionForest, CausalForest, CausalIVForest, MultiOutputGRF
 from econml.utilities import cross_product
 from copy import deepcopy
@@ -688,3 +689,20 @@ def test_multioutput(self,):
             np.testing.assert_allclose(imps[0, :], imps[1, :])
 
         return
+
+    def test_pickling(self,):
+
+        n_features = 2
+        n = 10
+        random_state = 123
+        X, y, _ = self._get_regression_data(n, n_features, random_state)
+
+        forest = RegressionForest(n_estimators=4, warm_start=True, random_state=123).fit(X, y)
+        forest.n_estimators = 8
+        forest.fit(X, y)
+        pred1 = forest.predict(X)
+
+        joblib.dump(forest, 'forest.jbl')
+        loaded_forest = joblib.load('forest.jbl')
+        np.testing.assert_equal(loaded_forest.n_estimators, forest.n_estimators)
+        np.testing.assert_allclose(loaded_forest.predict(X), pred1)

econml/tree/_tree.pyx

@@ -517,15 +517,17 @@ cdef class Tree:
 
         node_ndarray = d['nodes']
         value_ndarray = d['values']
-        
-        value_shape = (node_ndarray.shape[0], self.n_outputs)
+
         if (node_ndarray.ndim != 1 or
                 node_ndarray.dtype != NODE_DTYPE or
-                not node_ndarray.flags.c_contiguous or
-                value_ndarray.shape != value_shape or
+                not node_ndarray.flags.c_contiguous):
+            raise ValueError('Did not recognise loaded array layout for `node_ndarray`')
+
+        value_shape = (node_ndarray.shape[0], self.n_outputs, self.max_n_classes)
+        if (value_ndarray.shape != value_shape or
                 not value_ndarray.flags.c_contiguous or
                 value_ndarray.dtype != np.float64):
-            raise ValueError('Did not recognise loaded array layout')
+            raise ValueError('Did not recognise loaded array layout for `value_ndarray`')
 
         self.capacity = node_ndarray.shape[0]
         if self._resize_c(self.capacity) != 0:
@@ -541,15 +543,15 @@ cdef class Tree:
             if (jac_ndarray.shape != jac_shape or
                     not jac_ndarray.flags.c_contiguous or
                     jac_ndarray.dtype != np.float64):
-                raise ValueError('Did not recognise loaded array layout')
+                raise ValueError('Did not recognise loaded array layout for `jac_ndarray`')
             jac = memcpy(self.jac, (<np.ndarray> jac_ndarray).data,
                          self.capacity * self.jac_stride * sizeof(double))
             precond_ndarray = d['precond']
             precond_shape = (node_ndarray.shape[0], self.n_outputs)
             if (precond_ndarray.shape != precond_shape or
                     not precond_ndarray.flags.c_contiguous or
                     precond_ndarray.dtype != np.float64):
-                raise ValueError('Did not recognise loaded array layout')
+                raise ValueError('Did not recognise loaded array layout for `precond_ndarray`')
             precond = memcpy(self.precond, (<np.ndarray> precond_ndarray).data,
                              self.capacity * self.precond_stride * sizeof(double))
 
@@ -917,7 +919,7 @@ cdef class Tree:
         cdef np.npy_intp shape[3]
         shape[0] = <np.npy_intp> self.node_count
         shape[1] = <np.npy_intp> self.n_outputs
-        shape[2] = 1
+        shape[2] = <np.npy_intp> self.max_n_classes
         cdef np.ndarray arr
         arr = np.PyArray_SimpleNewFromData(3, shape, np.NPY_DOUBLE, self.value)
         Py_INCREF(self)