update variable importance report

Author: aloctavodia

pymc3/sampling.py

@@ -608,10 +608,12 @@ def sample(
 
     if "variable_inclusion" in trace.stat_names:
         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)
+        variable_inclusion = np.vstack(
+            [v.sum(0) / v.sum() for v in np.array_split(variable_inclusion, 50)]
+        )
+        trace.report.variable_importance = np.empty((variable_inclusion.shape[1], 3))
+        trace.report.variable_importance[:, 0] = variable_inclusion.mean(0)
+        trace.report.variable_importance[:, 1:3] = arviz.hdi(variable_inclusion, hdi_prob=0.68)
 
     n_chains = len(trace.chains)
     _log.info(

pymc3/step_methods/pgbart.py

@@ -60,7 +60,6 @@ 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, dtype="int")
         self.num_particles = num_particles
         self.log_num_particles = np.log(num_particles)
         self.indices = list(range(1, num_particles))
@@ -77,7 +76,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
+        variable_inclusion = np.zeros(bart.num_variates, dtype="int")
 
         # 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.
@@ -142,7 +141,6 @@ 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:
                     variable_inclusion[index] += 1