Merge pull request #3600 from aloctavodia/smc_lh

reduce number of logp evaluations

Author: lucianopaz

RELEASE-NOTES.md

@@ -18,6 +18,7 @@
 - SMC: stabilize covariance matrix [3573](https://github.com/pymc-devs/pymc3/pull/3573)
 - SMC is no longer a step method of `pm.sample` now it should be called using `pm.sample_smc` [3579](https://github.com/pymc-devs/pymc3/pull/3579)
 - SMC: improve computation of the proposal scaling factor [3594](https://github.com/pymc-devs/pymc3/pull/3594)
+- SMC: reduce number of logp evaluations [3600](https://github.com/pymc-devs/pymc3/pull/3600)
 - Now uses `multiprocessong` rather than `psutil` to count CPUs, which results in reliable core counts on Chromebooks.
 - `sample_posterior_predictive` now preallocates the memory required for its output to improve memory usage. Addresses problems raised in this [discourse thread](https://discourse.pymc.io/t/memory-error-with-posterior-predictive-sample/2891/4).
 - Fixed a bug in `Categorical.logp`. In the case of multidimensional `p`'s, the indexing was done wrong leading to incorrectly shaped tensors that consumed `O(n**2)` memory instead of `O(n)`. This fixes issue [#3535](https://github.com/pymc-devs/pymc3/issues/3535)

pymc3/smc/smc.py

@@ -191,13 +191,12 @@ def sample_smc(
 
     if parallel and cores > 1:
         pool = mp.Pool(processes=cores)
+        results = pool.starmap(likelihood_logp, [(sample,) for sample in posterior])
+    else:
+        results = [likelihood_logp(sample) for sample in posterior]
+    likelihoods = np.array(results).squeeze()
 
     while beta < 1:
-        if parallel and cores > 1:
-            results = pool.starmap(likelihood_logp, [(sample,) for sample in posterior])
-        else:
-            results = [likelihood_logp(sample) for sample in posterior]
-        likelihoods = np.array(results).squeeze()
         beta, old_beta, weights, sj = calc_beta(beta, likelihoods, threshold)
 
         model.marginal_likelihood *= sj
@@ -238,16 +237,20 @@ def sample_smc(
         if parallel and cores > 1:
             results = pool.starmap(
                 metrop_kernel,
-                [(posterior[draw], tempered_logp[draw], *parameters) for draw in range(draws)],
+                [
+                    (posterior[draw], tempered_logp[draw], likelihoods[draw], *parameters)
+                    for draw in range(draws)
+                ],
             )
         else:
             results = [
-                metrop_kernel(posterior[draw], tempered_logp[draw], *parameters)
+                metrop_kernel(posterior[draw], tempered_logp[draw], likelihoods[draw], *parameters)
                 for draw in tqdm(range(draws), disable=not progressbar)
             ]
 
-        posterior, acc_list = zip(*results)
+        posterior, acc_list, likelihoods = zip(*results)
         posterior = np.array(posterior)
+        likelihoods = np.array(likelihoods)
         acc_rate = sum(acc_list) / proposed
         stage += 1
 

pymc3/smc/smc_utils.py

@@ -110,6 +110,7 @@ def _posterior_to_trace(posterior, variables, model, var_info):
 def metrop_kernel(
     q_old,
     old_tempered_logp,
+    old_likelihood,
     proposal,
     scaling,
     accepted,
@@ -147,9 +148,10 @@ def metrop_kernel(
         q_old, accept = metrop_select(new_tempered_logp - old_tempered_logp, q_new, q_old)
         if accept:
             accepted += 1
+            old_likelihood = ll
             old_tempered_logp = new_tempered_logp
 
-    return q_old, accepted
+    return q_old, accepted, old_likelihood
 
 
 def calc_beta(beta, likelihoods, threshold=0.5, psis=True):