Merge branch 'pymc-devs:main' into gpv4update

Author: bwengals

.pre-commit-config.yaml

@@ -14,7 +14,7 @@ repos:
         exclude: ^requirements-dev\.txt$
     -   id: trailing-whitespace
 - repo: https://github.com/PyCQA/isort
-  rev: 5.9.3
+  rev: 5.10.1
   hooks:
     - id: isort
       name: isort

CONTRIBUTING.md

@@ -146,7 +146,7 @@ docker exec -it pymc jupyter notebook list
 ## Style guide
 
 We have configured a pre-commit hook that checks for `black`-compliant code style.
-We encourage you to configure the pre-commit hook as described in the [PyMC Python Code Style Wiki Page](https://github.com/pymc-devs/pymc/wiki/PyMC-Python-Code-Style), because it will automatically enforce the code style on your commits.
+We encourage you to configure the pre-commit hook as described in the [PyMC Python Code Style Wiki Page](https://github.com/pymc-devs/pymc/wiki/Python-Code-Style), because it will automatically enforce the code style on your commits.
 
 Similarly, consult the [PyMC's Jupyter Notebook Style](https://github.com/pymc-devs/pymc/wiki/PyMC-Jupyter-Notebook-Style-Guide) guides for notebooks.
 

RELEASE-NOTES.md

@@ -93,8 +93,9 @@ This includes API changes we did not warn about since at least `3.11.0` (2021-01
 - `pm.DensityDist` can now accept an optional `logcdf` keyword argument to pass in a function to compute the cummulative density function of the distribution (see [5026](https://github.com/pymc-devs/pymc/pull/5026)).
 - `pm.DensityDist` can now accept an optional `get_moment` keyword argument to pass in a function to compute the moment of the distribution (see [5026](https://github.com/pymc-devs/pymc/pull/5026)).
 - New features for BART:
-  -  Added linear response, increased number of trees fitted per step [5044](https://github.com/pymc-devs/pymc3/pull/5044).
-  -  Added partial dependence plots and individual conditional expectation plots [5091](https://github.com/pymc-devs/pymc3/pull/5091).
+  - Added linear response, increased number of trees fitted per step [5044](https://github.com/pymc-devs/pymc3/pull/5044).
+  - Added partial dependence plots and individual conditional expectation plots [5091](https://github.com/pymc-devs/pymc3/pull/5091).
+  - Modify how particle weights are computed. This improves accuracy of the modeled function (see [5177](https://github.com/pymc-devs/pymc3/pull/5177)).
 - `pm.Data` now passes additional kwargs to `aesara.shared`. [#5098](https://github.com/pymc-devs/pymc/pull/5098)
 - ...
 

conda-envs/environment-dev-py37.yml

@@ -4,7 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
-- aeppl>=0.0.13
+- aeppl=0.0.15
 - aesara>=2.2.6
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/environment-dev-py38.yml

@@ -4,7 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
-- aeppl>=0.0.13
+- aeppl=0.0.15
 - aesara>=2.2.6
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/environment-dev-py39.yml

@@ -4,7 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
-- aeppl>=0.0.13
+- aeppl=0.0.15
 - aesara>=2.2.6
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/environment-test-py37.yml

@@ -4,7 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
-- aeppl>=0.0.13
+- aeppl=0.0.15
 - aesara>=2.2.6
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/environment-test-py38.yml

@@ -4,7 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
-- aeppl>=0.0.13
+- aeppl==0.0.15
 - aesara>=2.2.6
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/environment-test-py39.yml

@@ -4,7 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
-- aeppl>=0.0.13
+- aeppl=0.0.15
 - aesara>=2.2.6
 - arviz>=0.11.4
 - cachetools

conda-envs/windows-environment-dev-py38.yml

@@ -4,7 +4,7 @@ channels:
 - defaults
 dependencies:
  # base dependencies (see install guide for Windows)
-- aeppl>=0.0.13
+- aeppl=0.0.15
 - aesara>=2.2.6
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/windows-environment-test-py38.yml

@@ -4,7 +4,7 @@ channels:
 - defaults
 dependencies:
  # base dependencies (see install guide for Windows)
-- aeppl>=0.0.13
+- aeppl=0.0.15
 - aesara>=2.2.6
 - arviz>=0.11.2
 - cachetools

docs/source/conf.py

@@ -19,7 +19,6 @@
 import os
 import sys
 
-
 # If extensions (or modules to document with autodoc) are in another directory,
 # add these directories to sys.path here. If the directory is relative to the
 # documentation root, use os.path.abspath to make it absolute, like shown here.

pymc/aesaraf.py

@@ -335,7 +335,7 @@ def rvs_to_value_vars(
     initial_replacements: Optional[Dict[TensorVariable, TensorVariable]] = None,
     **kwargs,
 ) -> Tuple[TensorVariable, Dict[TensorVariable, TensorVariable]]:
-    """Replace random variables in graphs with their value variables.
+    """Clone and replace random variables in graphs with their value variables.
 
     This will *not* recompute test values in the resulting graphs.
 
@@ -383,6 +383,16 @@ def transform_replacements(var, replacements):
         # Walk the transformed variable and make replacements
         return [trans_rv_value]
 
+    # Clone original graphs
+    inputs = [i for i in graph_inputs(graphs) if not isinstance(i, Constant)]
+    equiv = clone_get_equiv(inputs, graphs, False, False, {})
+    graphs = [equiv[n] for n in graphs]
+
+    if initial_replacements:
+        initial_replacements = {
+            equiv.get(k, k): equiv.get(v, v) for k, v in initial_replacements.items()
+        }
+
     return replace_rvs_in_graphs(graphs, transform_replacements, initial_replacements, **kwargs)
 
 

pymc/backends/arviz.py

@@ -42,6 +42,26 @@
 Var = Any  # pylint: disable=invalid-name
 
 
+def find_observations(model: Optional["Model"]) -> Optional[Dict[str, Var]]:
+    """If there are observations available, return them as a dictionary."""
+    if model is None:
+        return None
+
+    observations = {}
+    for obs in model.observed_RVs:
+        aux_obs = getattr(obs.tag, "observations", None)
+        if aux_obs is not None:
+            try:
+                obs_data = extract_obs_data(aux_obs)
+                observations[obs.name] = obs_data
+            except TypeError:
+                warnings.warn(f"Could not extract data from symbolic observation {obs}")
+        else:
+            warnings.warn(f"No data for observation {obs}")
+
+    return observations
+
+
 class _DefaultTrace:
     """
     Utility for collecting samples into a dictionary.
@@ -196,25 +216,7 @@ def arbitrary_element(dct: Dict[Any, np.ndarray]) -> np.ndarray:
             self.dims = {**model_dims, **self.dims}
 
         self.density_dist_obs = density_dist_obs
-        self.observations = self.find_observations()
-
-    def find_observations(self) -> Optional[Dict[str, Var]]:
-        """If there are observations available, return them as a dictionary."""
-        if self.model is None:
-            return None
-        observations = {}
-        for obs in self.model.observed_RVs:
-            aux_obs = getattr(obs.tag, "observations", None)
-            if aux_obs is not None:
-                try:
-                    obs_data = extract_obs_data(aux_obs)
-                    observations[obs.name] = obs_data
-                except TypeError:
-                    warnings.warn(f"Could not extract data from symbolic observation {obs}")
-            else:
-                warnings.warn(f"No data for observation {obs}")
-
-        return observations
+        self.observations = find_observations(self.model)
 
     def split_trace(self) -> Tuple[Union[None, "MultiTrace"], Union[None, "MultiTrace"]]:
         """Split MultiTrace object into posterior and warmup.

pymc/bart/pgbart.py

@@ -101,9 +101,10 @@ class PGBART(ArrayStepShared):
         Number of particles for the conditional SMC sampler. Defaults to 10
     max_stages : int
         Maximum number of iterations of the conditional SMC sampler. Defaults to 100.
-    batch : int
+    batch : int or tuple
         Number of trees fitted per step. Defaults to  "auto", which is the 10% of the `m` trees
-        during tuning and 20% after tuning.
+        during tuning and 20% after tuning. If a tuple is passed the first element is the batch size
+        during tuning and the second the batch size after tuning.
     model: PyMC Model
         Optional model for sampling step. Defaults to None (taken from context).
 
@@ -138,9 +139,9 @@ def __init__(self, vars=None, num_particles=10, max_stages=100, batch="auto", mo
         self.alpha = self.bart.alpha
         self.k = self.bart.k
         self.response = self.bart.response
-        self.split_prior = self.bart.split_prior
-        if self.split_prior is None:
-            self.split_prior = np.ones(self.X.shape[1])
+        self.alpha_vec = self.bart.split_prior
+        if self.alpha_vec is None:
+            self.alpha_vec = np.ones(self.X.shape[1])
 
         self.init_mean = self.Y.mean()
         # if data is binary
@@ -149,7 +150,7 @@ def __init__(self, vars=None, num_particles=10, max_stages=100, batch="auto", mo
             self.mu_std = 6 / (self.k * self.m ** 0.5)
         # maybe we need to check for count data
         else:
-            self.mu_std = self.Y.std() / (self.k * self.m ** 0.5)
+            self.mu_std = (2 * self.Y.std()) / (self.k * self.m ** 0.5)
 
         self.num_observations = self.X.shape[0]
         self.num_variates = self.X.shape[1]
@@ -167,14 +168,18 @@ def __init__(self, vars=None, num_particles=10, max_stages=100, batch="auto", mo
 
         self.normal = NormalSampler()
         self.prior_prob_leaf_node = compute_prior_probability(self.alpha)
-        self.ssv = SampleSplittingVariable(self.split_prior)
+        self.ssv = SampleSplittingVariable(self.alpha_vec)
 
         self.tune = True
-        self.idx = 0
-        self.batch = batch
 
-        if self.batch == "auto":
-            self.batch = max(1, int(self.m * 0.1))
+        if batch == "auto":
+            self.batch = (max(1, int(self.m * 0.1)), max(1, int(self.m * 0.2)))
+        else:
+            if isinstance(batch, (tuple, list)):
+                self.batch = batch
+            else:
+                self.batch = (batch, batch)
+
         self.log_num_particles = np.log(num_particles)
         self.indices = list(range(1, num_particles))
         self.len_indices = len(self.indices)
@@ -187,6 +192,9 @@ def __init__(self, vars=None, num_particles=10, max_stages=100, batch="auto", mo
         self.all_particles = []
         for i in range(self.m):
             self.a_tree.tree_id = i
+            self.a_tree.leaf_node_value = (
+                self.init_mean / self.m + self.normal.random() * self.mu_std,
+            )
             p = ParticleTree(
                 self.a_tree,
                 self.init_log_weight,
@@ -201,20 +209,16 @@ def astep(self, q: RaveledVars) -> Tuple[RaveledVars, List[Dict[str, Any]]]:
         sum_trees_output = q.data
         variable_inclusion = np.zeros(self.num_variates, dtype="int")
 
-        if self.idx == self.m:
-            self.idx = 0
-
-        for tree_id in range(self.idx, self.idx + self.batch):
-            if tree_id >= self.m:
-                break
+        tree_ids = np.random.randint(0, self.m, size=self.batch[~self.tune])
+        for tree_id in tree_ids:
             # Generate an initial set of SMC particles
             # at the end of the algorithm we return one of these particles as the new tree
             particles = self.init_particles(tree_id)
             # Compute the sum of trees without the tree we are attempting to replace
             self.sum_trees_output_noi = sum_trees_output - particles[0].tree.predict_output()
 
             # The old tree is not growing so we update the weights only once.
-            self.update_weight(particles[0])
+            self.update_weight(particles[0], new=True)
             for t in range(self.max_stages):
                 # Sample each particle (try to grow each tree), except for the first one.
                 for p in particles[1:]:
@@ -235,15 +239,15 @@ def astep(self, q: RaveledVars) -> Tuple[RaveledVars, List[Dict[str, Any]]]:
                     if tree_grew:
                         self.update_weight(p)
                 # Normalize weights
-                W_t, normalized_weights = self.normalize(particles)
+                W_t, normalized_weights = self.normalize(particles[1:])
 
                 # Resample all but first particle
-                re_n_w = normalized_weights[1:] / normalized_weights[1:].sum()
+                re_n_w = normalized_weights
                 new_indices = np.random.choice(self.indices, size=self.len_indices, p=re_n_w)
                 particles[1:] = particles[new_indices]
 
                 # Set the new weights
-                for p in particles:
+                for p in particles[1:]:
                     p.log_weight = W_t
 
                 # Check if particles can keep growing, otherwise stop iterating
@@ -254,23 +258,25 @@ def astep(self, q: RaveledVars) -> Tuple[RaveledVars, List[Dict[str, Any]]]:
                 if all(non_available_nodes_for_expansion):
                     break
 
+            for p in particles[1:]:
+                p.log_weight = p.old_likelihood_logp
+
+            _, normalized_weights = self.normalize(particles)
             # Get the new tree and update
             new_particle = np.random.choice(particles, p=normalized_weights)
             new_tree = new_particle.tree
-            self.all_trees[self.idx] = new_tree
+            self.all_trees[tree_id] = new_tree
             new_particle.log_weight = new_particle.old_likelihood_logp - self.log_num_particles
             self.all_particles[tree_id] = new_particle
             sum_trees_output = self.sum_trees_output_noi + new_tree.predict_output()
 
             if self.tune:
+                self.ssv = SampleSplittingVariable(self.alpha_vec)
                 for index in new_particle.used_variates:
-                    self.split_prior[index] += 1
-                    self.ssv = SampleSplittingVariable(self.split_prior)
+                    self.alpha_vec[index] += 1
             else:
-                self.batch = max(1, int(self.m * 0.2))
                 for index in new_particle.used_variates:
                     variable_inclusion[index] += 1
-            self.idx += 1
 
         stats = {"variable_inclusion": variable_inclusion, "bart_trees": copy(self.all_trees)}
         sum_trees_output = RaveledVars(sum_trees_output, point_map_info)
@@ -323,7 +329,7 @@ def init_particles(self, tree_id: int) -> np.ndarray:
 
         return np.array(particles)
 
-    def update_weight(self, particle: List[ParticleTree]) -> None:
+    def update_weight(self, particle: List[ParticleTree], new=False) -> None:
         """
         Update the weight of a particle
 
@@ -333,20 +339,22 @@ def update_weight(self, particle: List[ParticleTree]) -> None:
         new_likelihood = self.likelihood_logp(
             self.sum_trees_output_noi + particle.tree.predict_output()
         )
-        particle.log_weight += new_likelihood - particle.old_likelihood_logp
-        particle.old_likelihood_logp = new_likelihood
+        if new:
+            particle.log_weight = new_likelihood
+        else:
+            particle.log_weight += new_likelihood - particle.old_likelihood_logp
+            particle.old_likelihood_logp = new_likelihood
 
 
 class SampleSplittingVariable:
-    def __init__(self, alpha_prior):
+    def __init__(self, alpha_vec):
         """
-        Sample splitting variables proportional to `alpha_prior`.
+        Sample splitting variables proportional to `alpha_vec`.
 
-        This is equivalent as sampling weights from a Dirichlet distribution with `alpha_prior`
-        parameter and then using those weights to sample from the available spliting variables.
+        This is equivalent to compute the posterior mean of a Dirichlet-Multinomial model.
         This enforce sparsity.
         """
-        self.enu = list(enumerate(np.cumsum(alpha_prior / alpha_prior.sum())))
+        self.enu = list(enumerate(np.cumsum(alpha_vec / alpha_vec.sum())))
 
     def rvs(self):
         r = np.random.random()

pymc/distributions/__init__.py

@@ -13,10 +13,8 @@
 #   limitations under the License.
 
 from pymc.distributions.logprob import (  # isort:skip
-    _logcdf,
     logcdf,
     logp,
-    logcdfpt,
     logp_transform,
     logpt,
     logpt_sum,
@@ -195,6 +193,5 @@
     "logp",
     "logp_transform",
     "logcdf",
-    "_logcdf",
     "logpt_sum",
 ]