add numba, small tweaks

Author: aloctavodia

pymc_bart/pgbart.py

@@ -14,8 +14,8 @@
 
 import logging
 
-from copy import copy, deepcopy
-from numba import jit
+from copy import deepcopy
+from numba import njit
 
 import aesara
 import numpy as np
@@ -57,7 +57,7 @@ class PGBART(ArrayStepShared):
     def __init__(
         self,
         vars=None,
-        num_particles=40,
+        num_particles=20,
         batch="auto",
         model=None,
     ):
@@ -105,8 +105,6 @@ def __init__(
             idx_data_points=np.arange(self.num_observations, dtype="int32"),
             shape=self.shape,
         )
-        self.mean = fast_mean()
-
         self.normal = NormalSampler(mu_std, self.shape)
         self.uniform = UniformSampler(0.33, 0.75, self.shape)
         self.prior_prob_leaf_node = compute_prior_probability(self.alpha)
@@ -159,7 +157,6 @@ def astep(self, _):
                         self.X,
                         self.missing_data,
                         self.sum_trees,
-                        self.mean,
                         self.m,
                         self.normal,
                         self.shape,
@@ -204,7 +201,7 @@ def astep(self, _):
                 for index in used_variates:
                     variable_inclusion[index] += 1
 
-        stats = {"variable_inclusion": variable_inclusion, "bart_trees": copy(self.all_trees)}
+        stats = {"variable_inclusion": variable_inclusion, "bart_trees": self.all_trees}
         return self.sum_trees, [stats]
 
     def normalize(self, particles):
@@ -226,15 +223,20 @@ def normalize(self, particles):
         return w_t, normalized_weights
 
     def resample(self, particles, normalized_weights):
-        """Use systematic resample for all but first two particles"""
+        """
+        Use systematic resample for all but first two particles
+
+        Ensure particles are copied only if needed.
+        """
+        new_indices = systematic(normalized_weights)
         seen = []
         new_particles = []
         for idx in new_indices:
             if idx in seen:
                 new_particles.append(deepcopy(particles[idx]))
             else:
-                seen.append(idx)
                 new_particles.append(particles[idx])
+                seen.append(idx)
 
         particles[2:] = new_particles
 
@@ -243,15 +245,13 @@ def resample(self, particles, normalized_weights):
     def init_particles(self, tree_id: int) -> np.ndarray:
         """Initialize particles."""
         p0 = self.all_particles[tree_id]
-        p1 = copy(p0)
+        p1 = deepcopy(p0)
         p1.sample_leafs(
             self.sum_trees,
-            self.mean,
             self.m,
             self.normal,
             self.shape,
         )
-
         # The old tree and the one with new leafs do not grow so we update the weights only once
         self.update_weight(p0, old=True)
         self.update_weight(p1, old=True)
@@ -303,7 +303,6 @@ def __init__(self, tree):
         self.old_likelihood_logp = 0
         self.kf = 0.75
 
-
     def sample_tree(
         self,
         ssv,
@@ -312,7 +311,6 @@ def sample_tree(
         X,
         missing_data,
         sum_trees,
-        mean,
         m,
         normal,
         shape,
@@ -332,7 +330,6 @@ def sample_tree(
                     X,
                     missing_data,
                     sum_trees,
-                    mean,
                     m,
                     normal,
                     self.kf,
@@ -345,15 +342,14 @@ def sample_tree(
 
         return tree_grew
 
-    def sample_leafs(self, sum_trees, mean, m, normal, shape):
+    def sample_leafs(self, sum_trees, m, normal, shape):
 
         for idx in self.tree.idx_leaf_nodes:
             if idx > 0:
                 leaf = self.tree[idx]
                 idx_data_points = leaf.idx_data_points
                 node_value = draw_leaf_value(
                     sum_trees[:, idx_data_points],
-                    mean,
                     m,
                     normal,
                     self.kf,
@@ -414,7 +410,6 @@ def grow_tree(
     X,
     missing_data,
     sum_trees,
-    mean,
     m,
     normal,
     kf,
@@ -443,7 +438,6 @@ def grow_tree(
             idx_data_point = new_idx_data_points[idx]
             node_value = draw_leaf_value(
                 sum_trees[:, idx_data_point],
-                mean,
                 m,
                 normal,
                 kf,
@@ -496,7 +490,7 @@ def get_split_value(available_splitting_values, idx_data_points, missing_data):
         return split_value
 
 
-def draw_leaf_value(Y_mu_pred, mean, m, normal, kf, shape):
+def draw_leaf_value(Y_mu_pred, m, normal, kf, shape):
     """Draw Gaussian distributed leaf values."""
     if Y_mu_pred.size == 0:
         return np.zeros(shape)
@@ -505,40 +499,31 @@ def draw_leaf_value(Y_mu_pred, mean, m, normal, kf, shape):
         if Y_mu_pred.size == 1:
             mu_mean = np.full(shape, Y_mu_pred.item() / m)
         else:
-            mu_mean = mean(Y_mu_pred) / m
+            mu_mean = fast_mean(Y_mu_pred) / m
 
         draw = norm + mu_mean
         return draw
 
 
-def fast_mean():
-    """If available use Numba to speed up the computation of the mean."""
-    try:
-        from numba import jit
-    except ImportError:
-        from functools import partial
+@njit
+def fast_mean(a):
+    """Use Numba to speed up the computation of the mean."""
+
+    if a.ndim == 1:
+        count = a.shape[0]
+        suma = 0
+        for i in range(count):
+            suma += a[i]
+        return suma / count
+    elif a.ndim == 2:
+        res = np.zeros(a.shape[0])
+        count = a.shape[1]
+        for j in range(a.shape[0]):
+            for i in range(count):
+                res[j] += a[j, i]
+        return res / count
 
-        return partial(np.mean, axis=1)
 
-    @jit
-    def mean(a):
-        if a.ndim == 1:
-            count = a.shape[0]
-            suma = 0
-            for i in range(count):
-                suma += a[i]
-            return suma / count
-        elif a.ndim == 2:
-            res = np.zeros(a.shape[0])
-            count = a.shape[1]
-            for j in range(a.shape[0]):
-                for i in range(count):
-                    res[j] += a[j, i]
-            return res / count
-
-    return mean
-
-@jit()
 def discrete_uniform_sampler(upper_value):
     """Draw from the uniform distribution with bounds [0, upper_value).
 
@@ -555,15 +540,14 @@ def __init__(self, scale, shape):
         self.scale = scale
         self.shape = shape
         self.update()
-    
+
     def random(self):
         if self.idx == self.size:
             self.update()
         pop = self.cache[:, self.idx]
         self.idx += 1
         return pop
 
-
     def update(self):
         self.idx = 0
         self.cache = np.random.normal(loc=0.0, scale=self.scale, size=(self.shape, self.size))
@@ -586,44 +570,54 @@ def random(self):
         self.idx += 1
         return pop
 
-
     def update(self):
         self.idx = 0
         self.cache = np.random.uniform(
             self.lower_bound, self.upper_bound, size=(self.shape, self.size)
         )
 
-@jit()
-def systematic(W):
-    """Systematic resampling.
+
+def systematic(normalized_weights):
     """
-    M = len(W)
-    su = (np.random.rand(1) + np.arange(M)) / M
-    return inverse_cdf(su, W) + 2
-
-
-@jit(nopython=True)
-def inverse_cdf(su, W):
-    """Inverse CDF algorithm for a finite distribution.
-        Parameters
-        ----------
-        su: (M,) ndarray
-            M sorted uniform variates (i.e. M ordered points in [0,1]).
-        W: (N,) ndarray
-            a vector of N normalized weights (>=0 and sum to one)
-        Returns
-        -------
-        A: (M,) ndarray
-            a vector of M indices in range 0, ..., N-1
+    Systematic resampling.
+
+    Return indices in the range 2, ..., len(normalized_weights)+2
+
+    Note: adapted from https://github.com/nchopin/particles
+    """
+    lnw = len(normalized_weights)
+    single_uniform = (np.random.rand(1) + np.arange(lnw)) / lnw
+    return inverse_cdf(single_uniform, normalized_weights) + 2
+
+
+@njit
+def inverse_cdf(single_uniform, normalized_weights):
+    """
+    Inverse CDF algorithm for a finite distribution.
+
+    Parameters
+    ----------
+    single_uniform: ndarray
+        ordered points in [0,1]
+
+    normalized_weights: ndarray
+        normalized weights
+
+    Returns
+    -------
+    A: ndarray
+        a vector of indices in range 2, ..., len(normalized_weights)+2
+
+    Note: adapted from https://github.com/nchopin/particles
     """
     j = 0
-    s = W[0]
-    M = su.shape[0]
+    s = normalized_weights[0]
+    M = single_uniform.shape[0]
     A = np.empty(M, dtype=np.int64)
     for n in range(M):
-        while su[n] > s:
+        while single_uniform[n] > s:
             j += 1
-            s += W[j]
+            s += normalized_weights[j]
         A[n] = j
     return A
 

requirements.txt

@@ -1,2 +1,3 @@
 pymc>=4.1.7
 arviz>=0.12.1
+numba>=0.55.1