renaming

bpd/pipelines/shear_inference.py

@@ -6,7 +6,7 @@
 from jax._src.prng import PRNGKeyArray
 from jax.scipy import stats
 
-from bpd.chains import do_inference_nuts
+from bpd.chains import run_inference_nuts
 from bpd.likelihood import shear_loglikelihood
 from bpd.prior import ellip_mag_prior
 
@@ -21,12 +21,14 @@ def logtarget_density(g: Array, *, data: Array, loglikelihood: Callable):
 def pipeline_shear_inference(
     rng_key: PRNGKeyArray,
     e_post: Array,
+    *,
     true_g: Array,
     sigma_e: float,
     sigma_e_int: float,
     n_samples: int,
+    initial_step_size: float,
     n_warmup_steps: int = 500,
-    initial_step_size: float = 1e-2,
+    max_num_doublings: int = 2,
 ):
     prior = partial(ellip_mag_prior, sigma=sigma_e)
     interim_prior = partial(ellip_mag_prior, sigma=sigma_e_int)
@@ -38,12 +40,12 @@ def pipeline_shear_inference(
     _logtarget = partial(logtarget_density, loglikelihood=_loglikelihood)
 
     _do_inference = partial(
-        do_inference_nuts,
+        run_inference_nuts,
         data=e_post,
         logtarget=_logtarget,
         n_samples=n_samples,
         n_warmup_steps=n_warmup_steps,
-        max_num_doublings=2,
+        max_num_doublings=max_num_doublings,
         intial_step_size=initial_step_size,
     )