reduce number of logp evaluations II (#3601)

* reduce number of logp evaluations II

* floatX

Author: aloctavodia

pymc3/smc/smc.py

@@ -191,10 +191,14 @@ def sample_smc(
 
     if parallel and cores > 1:
         pool = mp.Pool(processes=cores)
-        results = pool.starmap(likelihood_logp, [(sample,) for sample in posterior])
+        priors = pool.starmap(prior_logp, [(sample,) for sample in posterior])
+        likelihoods = pool.starmap(likelihood_logp, [(sample,) for sample in posterior])
     else:
-        results = [likelihood_logp(sample) for sample in posterior]
-    likelihoods = np.array(results).squeeze()
+        priors = [prior_logp(sample) for sample in posterior]
+        likelihoods = [likelihood_logp(sample) for sample in posterior]
+
+    priors = np.array(priors).squeeze()
+    likelihoods = np.array(likelihoods).squeeze()
 
     while beta < 1:
         beta, old_beta, weights, sj = calc_beta(beta, likelihoods, threshold)
@@ -203,6 +207,7 @@ def sample_smc(
         # resample based on plausibility weights (selection)
         resampling_indexes = np.random.choice(np.arange(draws), size=draws, p=weights)
         posterior = posterior[resampling_indexes]
+        priors = priors[resampling_indexes]
         likelihoods = likelihoods[resampling_indexes]
 
         # compute proposal distribution based on weights
@@ -219,7 +224,6 @@ def sample_smc(
         pm._log.info("Stage: {:3d} Beta: {:.3f} Steps: {:3d}".format(stage, beta, n_steps))
         # Apply Metropolis kernel (mutation)
         proposed = draws * n_steps
-        priors = np.array([prior_logp(sample) for sample in posterior]).squeeze()
         tempered_logp = priors + likelihoods * beta
 
         parameters = (
@@ -238,18 +242,31 @@ def sample_smc(
             results = pool.starmap(
                 metrop_kernel,
                 [
-                    (posterior[draw], tempered_logp[draw], likelihoods[draw], *parameters)
+                    (
+                        posterior[draw],
+                        tempered_logp[draw],
+                        priors[draw],
+                        likelihoods[draw],
+                        *parameters,
+                    )
                     for draw in range(draws)
                 ],
             )
         else:
             results = [
-                metrop_kernel(posterior[draw], tempered_logp[draw], likelihoods[draw], *parameters)
+                metrop_kernel(
+                    posterior[draw],
+                    tempered_logp[draw],
+                    priors[draw],
+                    likelihoods[draw],
+                    *parameters
+                )
                 for draw in tqdm(range(draws), disable=not progressbar)
             ]
 
-        posterior, acc_list, likelihoods = zip(*results)
+        posterior, acc_list, priors, likelihoods = zip(*results)
         posterior = np.array(posterior)
+        priors = np.array(priors)
         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_prior,
     old_likelihood,
     proposal,
     scaling,
@@ -142,16 +143,18 @@ def metrop_kernel(
             q_new = floatX(q_old + delta)
 
         ll = likelihood_logp(q_new)
+        pl = prior_logp(q_new)
 
-        new_tempered_logp = prior_logp(q_new) + ll * beta
+        new_tempered_logp = pl + ll * beta
 
         q_old, accept = metrop_select(new_tempered_logp - old_tempered_logp, q_new, q_old)
         if accept:
             accepted += 1
+            old_prior = pl
             old_likelihood = ll
             old_tempered_logp = new_tempered_logp
 
-    return q_old, accepted, old_likelihood
+    return q_old, accepted, old_prior, old_likelihood
 
 
 def calc_beta(beta, likelihoods, threshold=0.5, psis=True):