format

Author: malihass

papers/tutorial/calibration/post_cal.py

@@ -12,7 +12,7 @@ def post_process_cal(
     data_x,
     data_y,
     sigma,
-    args
+    args,
 ):
     # Post process
     ranges = []
@@ -45,8 +45,8 @@ def post_process_cal(
     else:
         filename += "_opt"
     filename += f"_{args.alpha}_{args.beta}_corner"
-    plt.savefig(filename+".png")
-    plt.savefig(filename+".eps")
+    plt.savefig(filename + ".png")
+    plt.savefig(filename + ".eps")
 
     for ax in fig.get_axes():
         ax.tick_params(
@@ -113,7 +113,7 @@ def post_process_cal(
         label="95% Model confidence interval",
     )
     plt.plot(rangex, std2_5_real, "--", color="k", linewidth=3)
-    pretty_labels("", "", 20, title=f"Missing phys. unc. = {sigma:.2g}") 
+    pretty_labels("", "", 20, title=f"Missing phys. unc. = {sigma:.2g}")
     filename = ""
     if args.useNN:
         filename += "Surr"
@@ -124,6 +124,6 @@ def post_process_cal(
     else:
         filename += "_opt"
     filename += f"_{args.alpha}_{args.beta}_prop"
-    plt.savefig(filename+".png")
-    plt.savefig(filename+".eps")
+    plt.savefig(filename + ".png")
+    plt.savefig(filename + ".eps")
     plt.show()

papers/tutorial/calibration/tut_calibration.py

@@ -155,6 +155,7 @@ def bayes_step_opt_err(y=None, y_err=0.1):
     std_obs = jnp.ones(y_model.shape[0]) * y_err
     numpyro.sample("obs", dist.Normal(y_model, std_obs), obs=y)
 
+
 def bayes_step_cal_err(y=None):
     mu_d = numpyro.sample("mu_d", dist.Uniform(0.01, 0.9))
     sigma_d = numpyro.sample("sigma_d", dist.Uniform(0.01, 0.9))
@@ -163,6 +164,7 @@ def bayes_step_cal_err(y=None):
     std_obs = jnp.ones(y_model.shape[0]) * err
     numpyro.sample("obs", dist.Normal(y_model, std_obs), obs=y)
 
+
 def mcmc_iter(y_err=0.1, mcmc_method="HMC", cal_err=False):
     rng_key = random.PRNGKey(0)
     rng_key, rng_key_ = random.split(rng_key)
@@ -176,7 +178,6 @@ def mcmc_iter(y_err=0.1, mcmc_method="HMC", cal_err=False):
         bayes_step = bayes_step_cal_err
     else:
         bayes_step = bayes_step_opt_err
-    
 
     # Hamiltonian Monte Carlo (HMC) with no u turn sampling (NUTS)
     if mcmc_method.lower() == "hmc":
@@ -195,9 +196,9 @@ def mcmc_iter(y_err=0.1, mcmc_method="HMC", cal_err=False):
         jit_model_args=True,
     )
     if cal_err:
-       mcmc.run(rng_key_, y=data_y)
+        mcmc.run(rng_key_, y=data_y)
     else:
-       mcmc.run(rng_key_, y=data_y, y_err=y_err)
+        mcmc.run(rng_key_, y=data_y, y_err=y_err)
     mcmc.print_summary()
 
     # Draw samples
@@ -246,7 +247,7 @@ def mcmc_iter(y_err=0.1, mcmc_method="HMC", cal_err=False):
     }
 
     if cal_err:
-        err = np.mean(np_mcmc_samples[:,-1])
+        err = np.mean(np_mcmc_samples[:, -1])
     else:
         err = y_err
     true_m95 = data_y - 2 * y_err
@@ -260,24 +261,29 @@ def mcmc_iter(y_err=0.1, mcmc_method="HMC", cal_err=False):
     else:
         return True, results
 
+
 if not args.cal_err:
     min_sigma = 1e-4
     max_sigma = 1
     guess_sigma = 0.07
     sigma = guess_sigma
-    
+
     for iteration_sigma in range(10):
         print(f"Doing sigma = {sigma:.3g}")
-        reduce_sigma, results = mcmc_iter(y_err=sigma, mcmc_method="hmc", cal_err=args.cal_err)
+        reduce_sigma, results = mcmc_iter(
+            y_err=sigma, mcmc_method="hmc", cal_err=args.cal_err
+        )
         if reduce_sigma:
             max_sigma = sigma
             sigma = sigma - (sigma - min_sigma) / 2
         else:
             min_sigma = sigma
             sigma = sigma + (max_sigma - sigma) / 2
-    
+
     if not reduce_sigma:
-        reduce_sigma, results = mcmc_iter(y_err=max_sigma, mcmc_method="hmc", cal_err=args.cal_err)
+        reduce_sigma, results = mcmc_iter(
+            y_err=max_sigma, mcmc_method="hmc", cal_err=args.cal_err
+        )
 else:
     _, results = mcmc_iter(mcmc_method="hmc", cal_err=args.cal_err)
 
@@ -293,7 +299,7 @@ def mcmc_iter(y_err=0.1, mcmc_method="HMC", cal_err=False):
 )
 
 if args.cal_err:
-    sigma = np.mean(np_mcmc_samples[:,-1])
+    sigma = np.mean(np_mcmc_samples[:, -1])
 
 
 post_process_cal(
@@ -305,5 +311,5 @@ def mcmc_iter(y_err=0.1, mcmc_method="HMC", cal_err=False):
     data_x,
     data_y,
     sigma,
-    args
+    args,
 )