wip return mean and std variable importance

Author: aloctavodia

pymc3/distributions/bart.py

@@ -144,10 +144,9 @@ def get_new_idx_data_points(self, current_split_node, idx_data_points):
         idx_split_variable = current_split_node.idx_split_variable
         split_value = current_split_node.split_value
 
-        left_idx = np.nonzero(self.X[idx_data_points, idx_split_variable] <= split_value)
+        left_idx = self.X[idx_data_points, idx_split_variable] <= split_value
         left_node_idx_data_points = idx_data_points[left_idx]
-        right_idx = np.nonzero(~(self.X[idx_data_points, idx_split_variable] <= split_value))
-        right_node_idx_data_points = idx_data_points[right_idx]
+        right_node_idx_data_points = idx_data_points[~left_idx]
 
         return left_node_idx_data_points, right_node_idx_data_points
 

pymc3/sampling.py

@@ -607,8 +607,11 @@ def sample(
     trace.report._t_sampling = t_sampling
 
     if "variable_inclusion" in trace.stat_names:
-        variable_inclusion = trace.get_sampler_stats("variable_inclusion")[::-draws]
-        trace.report.variable_importance = np.mean([vi / vi.sum() for vi in variable_inclusion], 0)
+        variable_inclusion = np.vstack(trace.get_sampler_stats("variable_inclusion"))
+        variable_inclusion = np.split(variable_inclusion, 50)
+        dada = np.vstack([v.sum(0) / v.sum() for v in variable_inclusion])
+        trace.report.variable_importance_m = dada.mean(0)
+        trace.report.variable_importance_s = dada.std(0)
 
     n_chains = len(trace.chains)
     _log.info(

pymc3/step_methods/pgbart.py

@@ -60,7 +60,7 @@ def __init__(self, vars=None, num_particles=10, max_stages=5000, chunk="auto", m
         self.idx = 0
         if chunk == "auto":
             self.chunk = max(1, int(self.bart.m * 0.1))
-        self.variable_inclusion = np.zeros(self.bart.num_variates)
+        self.variable_inclusion = np.zeros(self.bart.num_variates, dtype="int")
         self.num_particles = num_particles
         self.log_num_particles = np.log(num_particles)
         self.indices = list(range(1, num_particles))
@@ -77,6 +77,7 @@ def __init__(self, vars=None, num_particles=10, max_stages=5000, chunk="auto", m
     def astep(self, _):
         bart = self.bart
         num_observations = bart.num_observations
+        variable_inclusion = self.variable_inclusion
 
         # For the tunning phase we restrict max_stages to a low number, otherwise it is almost sure
         # we will reach max_stages given that our first set of m trees is not good at all.
@@ -141,10 +142,11 @@ def astep(self, _):
             bart.sum_trees_output = bart.Y - R_j + new_prediction
 
             if not self.tune:
+                variable_inclusion = self.variable_inclusion
                 for index in new_tree.used_variates:
-                    self.variable_inclusion[index] += 1
+                    variable_inclusion[index] += 1
 
-        stats = {"variable_inclusion": self.variable_inclusion}
+        stats = {"variable_inclusion": variable_inclusion}
 
         return bart.sum_trees_output, [stats]