Update mcmc.py

mcmc.py

@@ -145,32 +145,31 @@ if __name__ == '__main__':
     num_covariance_estimation_iterations = 50
     num_covariance_estimation_samples = 50
     num_final_samples = 10000
-    with tf.device("/CPU:0"):
-        start = time.perf_counter()
-        for i in range(num_covariance_estimation_iterations):
-            step_start = time.perf_counter()
-            samples, results = sample(num_covariance_estimation_samples,
-                                      initial_mcmc_state,
-                                      scale)
-            step_end = time.perf_counter()
-            print(f'{i} time {step_end - step_start}')
-            print("Acceptance: ", results.numpy().mean())
-            joint_posterior = tf.concat([joint_posterior, samples], axis=0)
-            cov = tfp.stats.covariance(tf.math.log(joint_posterior))
-            print(cov.numpy())
-            scale = cov * 2.38**2 / joint_posterior.shape[1]
-            initial_mcmc_state = joint_posterior[-1, :]
-
+    start = time.perf_counter()
+    for i in range(num_covariance_estimation_iterations):
         step_start = time.perf_counter()
-        samples, results = sample(num_final_samples,
-                                  init_state=joint_posterior[-1, :], scale=scale,)
-        joint_posterior = tf.concat([joint_posterior, samples], axis=0)
+        samples, results = sample(num_covariance_estimation_samples,
+                                  initial_mcmc_state,
+                                  scale)
         step_end = time.perf_counter()
-        print(f'Sampling step time {step_end - step_start}')
-        end = time.perf_counter()
-        print(f"Simulation complete in {end-start} seconds")
-        print("Acceptance: ", np.mean(results.numpy()))
-        print(tfp.stats.covariance(tf.math.log(joint_posterior)))
+        print(f'{i} time {step_end - step_start}')
+        print("Acceptance: ", results.numpy().mean())
+        joint_posterior = tf.concat([joint_posterior, samples], axis=0)
+        cov = tfp.stats.covariance(tf.math.log(joint_posterior))
+        print(cov.numpy())
+        scale = cov * 2.38**2 / joint_posterior.shape[1]
+        initial_mcmc_state = joint_posterior[-1, :]
+
+    step_start = time.perf_counter()
+    samples, results = sample(num_final_samples,
+                              init_state=joint_posterior[-1, :], scale=scale,)
+    joint_posterior = tf.concat([joint_posterior, samples], axis=0)
+    step_end = time.perf_counter()
+    print(f'Sampling step time {step_end - step_start}')
+    end = time.perf_counter()
+    print(f"Simulation complete in {end-start} seconds")
+    print("Acceptance: ", np.mean(results.numpy()))
+    print(tfp.stats.covariance(tf.math.log(joint_posterior)))
 
     fig, ax = plt.subplots(1, 3)
     ax[0].plot(joint_posterior[:, 0])