fix docstrings (#28)

Author: aloctavodia

pymc_experimental/bart/__init__.py

@@ -20,6 +20,6 @@
 __all__ = ["BART", "PGBART"]
 
 
-
 import pymc as pm
+
 pm.STEP_METHODS = list(pm.STEP_METHODS) + [PGBART]

pymc_experimental/bart/bart.py

@@ -25,9 +25,7 @@
 
 
 class BARTRV(RandomVariable):
-    """
-    Base class for BART
-    """
+    """Base class for BART."""
 
     name = "BART"
     ndim_supp = 1

pymc_experimental/bart/pgbart.py

@@ -32,7 +32,7 @@
 
 class PGBART(ArrayStepShared):
     """
-    Particle Gibss BART sampling step
+    Particle Gibss BART sampling step.
 
     Parameters
     ----------
@@ -208,9 +208,7 @@ def astep(self, _):
         return self.sum_trees, [stats]
 
     def normalize(self, particles):
-        """
-        Use logsumexp trick to get W_t and softmax to get normalized_weights
-        """
+        """Use logsumexp trick to get W_t and softmax to get normalized_weights."""
         log_w = np.array([p.log_weight for p in particles])
         log_w_max = log_w.max()
         log_w_ = log_w - log_w_max
@@ -224,9 +222,7 @@ def normalize(self, particles):
         return W_t, normalized_weights
 
     def init_particles(self, tree_id: int) -> np.ndarray:
-        """
-        Initialize particles
-        """
+        """Initialize particles."""
         p = self.all_particles[tree_id]
         particles = [p]
         particles.append(copy(p))
@@ -238,7 +234,7 @@ def init_particles(self, tree_id: int) -> np.ndarray:
 
     def update_weight(self, particle, old=False):
         """
-        Update the weight of a particle
+        Update the weight of a particle.
 
         Since the prior is used as the proposal,the weights are updated additively as the ratio of
         the new and old log-likelihoods.
@@ -253,19 +249,15 @@ def update_weight(self, particle, old=False):
 
     @staticmethod
     def competence(var, has_grad):
-        """
-        PGBART is only suitable for BART distributions
-        """
+        """PGBART is only suitable for BART distributions."""
         dist = getattr(var.owner, "op", None)
         if isinstance(dist, BARTRV):
             return Competence.IDEAL
         return Competence.INCOMPATIBLE
 
 
 class ParticleTree:
-    """
-    Particle tree
-    """
+    """Particle tree."""
 
     def __init__(self, tree):
         self.tree = tree.copy()  # keeps the tree that we care at the moment
@@ -340,6 +332,7 @@ def rvs(self):
 def compute_prior_probability(alpha):
     """
     Calculate the probability of the node being a LeafNode (1 - p(being SplitNode)).
+
     Taken from equation 19 in [Rockova2018].
 
     Parameters
@@ -463,7 +456,7 @@ def get_new_idx_data_points(split_value, idx_data_points, selected_predictor, X)
 
 
 def draw_leaf_value(Y_mu_pred, X_mu, mean, m, normal, mu_std):
-    """Draw Gaussian distributed leaf values"""
+    """Draw Gaussian distributed leaf values."""
     if Y_mu_pred.size == 0:
         return 0
     else:
@@ -504,9 +497,7 @@ def discrete_uniform_sampler(upper_value):
 
 
 class NormalSampler:
-    """
-    Cache samples from a standard normal distribution
-    """
+    """Cache samples from a standard normal distribution."""
 
     def __init__(self):
         self.size = 1000

pymc_experimental/bart/tree.py

@@ -21,7 +21,7 @@
 
 
 class Tree:
-    """Full binary tree
+    """Full binary tree.
 
     A full binary tree is a tree where each node has exactly zero or two children.
     This structure is used as the basic component of the Bayesian Additive Regression Tree (BART)
@@ -135,17 +135,16 @@ def _traverse_tree(self, x, node_index=0):
     @staticmethod
     def init_tree(leaf_node_value, idx_data_points):
         """
+        Initialize tree.
 
         Parameters
         ----------
         leaf_node_value
         idx_data_points
-        m : int
-            number of trees in BART
 
         Returns
         -------
-
+        tree
         """
         new_tree = Tree(len(idx_data_points))
         new_tree[0] = LeafNode(index=0, value=leaf_node_value, idx_data_points=idx_data_points)

pymc_experimental/bart/utils.py

@@ -1,3 +1,5 @@
+"""Utility function for variable selection and bart interpretability."""
+
 import arviz as az
 import matplotlib.pyplot as plt
 import numpy as np
@@ -10,7 +12,7 @@
 
 def predict(idata, rng, X_new=None, size=None, excluded=None):
     """
-    Generate samples from the BART-posterior
+    Generate samples from the BART-posterior.
 
     Parameters
     ----------
@@ -75,7 +77,7 @@ def plot_dependence(
     ax=None,
 ):
     """
-    Partial dependence or individual conditional expectation plot
+    Partial dependence or individual conditional expectation plot.
 
     Parameters
     ----------
@@ -107,12 +109,12 @@ def plot_dependence(
     instances : int
         Number of instances of X to plot. Only relevant if ice ``kind="ice"`` plots.
     random_seed : int
-        random_seed used to sample from the posterior. Defaults to None.
+        Seed used to sample from the posterior. Defaults to None.
     sharey : bool
         Controls sharing of properties among y-axes. Defaults to True.
     rug : bool
         Whether to include a rugplot. Defaults to True.
-    smooth=True,
+    smooth : bool
         If True the result will be smoothed by first computing a linear interpolation of the data
         over a regular grid and then applying the Savitzky-Golay filter to the interpolated data.
         Defaults to True.
@@ -302,7 +304,7 @@ def plot_dependence(
 
 def plot_variable_importance(idata, labels=None, figsize=None, samples=100, random_seed=None):
     """
-    Estimates variable importance from the BART-posterior
+    Estimates variable importance from the BART-posterior.
 
     Parameters
     ----------