Merge branch 'main' into best_epoch

Author: LucaMantani

.github/workflows/documentation.yml

@@ -30,7 +30,7 @@ jobs:
       - name: Build Sphinx HTML
         run: |
           cd colibri/doc/sphinx
-          make html
+          make html SPHINXOPTS="-W --keep-going"
       - name: Upload Pages artifact
         uses: actions/upload-pages-artifact@v3
         with:

README.md

@@ -39,7 +39,7 @@ Documentation regarding the usage and installation is available at <https://hep-
 - Command-line scripts for common workflows (`colibri`, `evolve_fit`, etc.)
 - Integration with external PDF model repositories
 
-<img src="./colibri-diagram-hessian.png" alt="colibri diagram" width="90%">
+<img src="./colibri-diagram-blackjax.png" alt="colibri diagram" width="90%">
 
 ---
 

colibri-diagram-blackjax.png

@@ -82,14 +82,11 @@ def analytic_evidence_uniform_prior(sol_covmat, sol_mean, max_logl, a_vec, b_vec
     return log_evidence, log_occam_factor
 
 
-@check_pdf_model_is_linear
 def analytic_fit(
     central_covmat_index,
-    _pred_data,
-    pdf_model,
+    forward_map,
     analytic_settings,
     prior_settings,
-    FIT_XGRID,
     fast_kernel_arrays,
 ):
     """
@@ -106,40 +103,35 @@ def analytic_fit(
     central_covmat_index: commondata_utils.CentralCovmatIndex
         dataclass containing central values and covariance matrix.
 
-    _pred_data: @jax.jit CompiledFunction
-        Prediction function for the fit.
-
-    pdf_model: pdf_model.PDFModel
-        PDF model to fit.
+    forward_map: @jax.jit CompiledFunction
+        Forward map function for the fit.
 
     analytic_settings: dict
         Settings for the analytic fit.
 
     prior_settings: PriorSettings
         Settings for the prior.
 
-    FIT_XGRID: np.ndarray
-        xgrid of the theory, computed by a production rule by taking
-        the sorted union of the xgrids of the datasets entering the fit.
-
     fast_kernel_arrays: tuple
         Tuple containing the fast kernel arrays.
     """
+    # Ensure that the PDF model is linear before running the fit.
+    log.info("Checking that the PDF model is linear...")
+    check_pdf_model_is_linear(forward_map, fast_kernel_arrays)
 
     log.warning("The prior is assumed to be flat in the parameters.")
     log.warning(
         "Assuming that the prior is wide enough to fully cover the gaussian likelihood."
     )
 
-    parameters = pdf_model.param_names
-    pred_and_pdf = pdf_model.pred_and_pdf_func(FIT_XGRID, forward_map=_pred_data)
+    parameters = forward_map.param_names
 
     # Precompute predictions for the basis of the model
     bases = jnp.identity(len(parameters))
     predictions = jnp.array(
-        [pred_and_pdf(basis, fast_kernel_arrays)[0] for basis in bases]
+        [forward_map(fast_kernel_arrays, basis)[0] for basis in bases]
     )
-    intercept = pred_and_pdf(jnp.zeros(len(parameters)), fast_kernel_arrays)[0]
+    intercept = forward_map(fast_kernel_arrays, jnp.zeros(len(parameters)))[0]
 
     # Construct the analytic solution
     central_values = central_covmat_index.central_values

colibri-diagram-hessian.png

@@ -32,6 +32,7 @@
     "colibri.param_initialisation",
     "colibri.export_results",
     "colibri.closure_test",
+    "colibri.forward_map",
     "reportengine.report",
 ]
 

colibri/analytic_fit.py

@@ -5,15 +5,13 @@
     cast_to_numpy,
     get_full_posterior,
 )
-from colibri.checks import check_pdf_models_equal
 from colibri.core import BayesianPrior
 import tensorflow_probability.substrates.jax as tfp
 
 tfd = tfp.distributions
 
 
-@check_pdf_models_equal
-def bayesian_prior(prior_settings, pdf_model):
+def bayesian_prior(prior_settings, forward_map):
     """
     Produces a prior transform function.
 
@@ -22,6 +20,10 @@ def bayesian_prior(prior_settings, pdf_model):
     prior_settings: dict
         The settings for the prior transform.
 
+    forward_map: ForwardMap
+        The forward map of the problem, used to determine parameter names and ordering.
+
+
     Returns
     -------
     prior_transform: @jax.jit CompiledFunction
@@ -31,8 +33,8 @@ def bayesian_prior(prior_settings, pdf_model):
         prior_specs = prior_settings.prior_distribution_specs
 
         if "bounds" in prior_specs:
-            # Use param names from the model to order bounds correctly
-            param_names = pdf_model.param_names
+            # Use param names from the forward map to order bounds correctly
+            param_names = forward_map.param_names
             bounds_dict = prior_specs["bounds"]
             missing = [p for p in param_names if p not in bounds_dict]
             if missing:
@@ -45,8 +47,9 @@ def bayesian_prior(prior_settings, pdf_model):
 
         elif "min_val" in prior_specs and "max_val" in prior_specs:
             # Global bounds for all parameters
-            mins = prior_specs["min_val"]
-            maxs = prior_specs["max_val"]
+            n_params = len(forward_map.param_names)
+            mins = jnp.array([float(prior_specs["min_val"])] * n_params)
+            maxs = jnp.array([float(prior_specs["max_val"])] * n_params)
 
         else:
             raise ValueError(

colibri/app.py

@@ -17,6 +17,7 @@
 from jax.scipy.special import logsumexp
 import tqdm
 import anesthetic
+import pandas as pd
 
 from colibri.core import BlackJAXFit
 from colibri.export_results import export_bayes_results, write_replicas
@@ -37,7 +38,7 @@
 
 
 def blackjax_fit(
-    pdf_model,
+    forward_map,
     bayesian_prior,
     blackjax_settings,
     log_likelihood,
@@ -47,8 +48,8 @@ def blackjax_fit(
 
     Parameters
     ----------
-    pdf_model: pdf_model.PDFModel
-        The PDF model to fit.
+    forward_map: ForwardMap
+        The forward map whose ``param_names`` enumerate all fit parameters.
 
     bayesian_prior: BayesianPrior, @jax.jit CompiledFunction
         The prior function for the model.
@@ -68,9 +69,9 @@ def blackjax_fit(
     log.info(f"Running fit with backend: {jax.default_backend()}")
 
     # set the BlackJAX seed
-    rng_key = jax.random.PRNGKey(blackjax_settings["seed"])
+    rng_key = jax.random.PRNGKey(blackjax_settings["blackjax_seed"])
     log.info(f"BlackJAX initialisation seed: {rng_key}")
-    n_dims = pdf_model.n_parameters
+    n_dims = len(forward_map.param_names)
     n_live = blackjax_settings["n_live"]
     n_delete = int(blackjax_settings["delete_fraction"] * n_live)
 
@@ -141,13 +142,17 @@ def one_step(carry, xs):
         data=final_states.particles,
         logL=final_states.loglikelihood,
         logL_birth=final_states.loglikelihood_birth,
-        columns=pdf_model.param_names,
+        columns=forward_map.param_names,
     )
     # write nested_samples.csv to blackjax_logs
     log_dir = blackjax_settings["log_dir"]
     os.makedirs(log_dir, exist_ok=True)  # Create directory if it doesn't exist
     nested_samples.to_csv(log_dir + "/nested_samples.csv")
 
+    # Export resampled posterior samples
+    posterior_df = pd.DataFrame(resampled_posterior, columns=forward_map.param_names)
+    posterior_df.to_csv(os.path.join(log_dir, "posterior_samples.csv"), index=False)
+
     # Compute bayesian metrics (similar to UltraNest)
     # Find maximum likelihood point
     max_ll_idx = jnp.argmax(final_states.loglikelihood)
@@ -167,7 +172,7 @@ def one_step(carry, xs):
             "logZ_err": logzs.std(),
             "ess": ess_value,
         },
-        param_names=pdf_model.param_names,
+        param_names=forward_map.param_names,
         resampled_posterior=resampled_posterior,
         full_posterior_samples=full_samples,
         bayesian_metrics={

colibri/bayes_prior.py

@@ -7,27 +7,26 @@
 from reportengine.checks import make_argcheck
 import jax.numpy as jnp
 import jax
-from colibri.theory_predictions import make_pred_data, fast_kernel_arrays
-
-from colibri.utils import get_fit_path, get_pdf_model, pdf_models_equal
+from colibri.utils import get_fit_path, get_pdf_model
 
 
 @make_argcheck
-def check_pdf_models_equal(prior_settings, pdf_model, theoryid):
+def check_pdf_models_equal(prior_settings, forward_map, theoryid):
     """
     Decorator that can be added to functions to check that the
     PDF model used as prior (eg when using prior_settings["type"] == "prior_from_gauss_posterior")
-    matches the PDF model used in the current fit (pdf_model).
+    matches the PDF model used in the current fit (via ``forward_map.pdf_param_names``).
     """
 
     if prior_settings.prior_distribution == "prior_from_gauss_posterior":
 
         prior_fit = prior_settings.prior_distribution_specs["prior_fit"]
         prior_pdf_model = get_pdf_model(prior_fit)
 
-        if not pdf_models_equal(prior_pdf_model, pdf_model):
+        if prior_pdf_model.param_names != list(forward_map.pdf_param_names):
             raise ValueError(
-                f"PDF model {pdf_model} does not match prior settings {prior_pdf_model}"
+                f"PDF param names from forward_map {list(forward_map.pdf_param_names)} "
+                f"do not match prior PDF model param names {prior_pdf_model.param_names}"
             )
 
         # load filter.yml runcard of the prior fit
@@ -41,19 +40,20 @@ def check_pdf_models_equal(prior_settings, pdf_model, theoryid):
             )
 
 
-@make_argcheck
-def check_pdf_model_is_linear(pdf_model, FIT_XGRID, data):
+def check_pdf_model_is_linear(forward_map, fast_kernel_arrays):
     """
     Decorator that can be added to functions to check that the
     PDF model is linear.
+
+    Note that the FK arrays are taken as an argument rather than rebuilt here,
+    so that they are guaranteed to be consistent with the
+    ``fill_fk_xgrid_with_zeros`` setting ``forward_map`` was built with.
     """
 
-    pred_data = make_pred_data(data, FIT_XGRID)
-    fk = fast_kernel_arrays(data, FIT_XGRID)
+    fk = fast_kernel_arrays
 
-    parameters = pdf_model.param_names
-    pred_and_pdf = pdf_model.pred_and_pdf_func(FIT_XGRID, forward_map=pred_data)
-    intercept = pred_and_pdf(jnp.zeros(len(parameters)), fk)[0]
+    parameters = forward_map.param_names
+    intercept, _ = forward_map(fk, jnp.zeros(len(parameters)))
 
     # Run the check for 10 random points in the parameter space
     for i in range(10):
@@ -65,16 +65,16 @@ def check_pdf_model_is_linear(pdf_model, FIT_XGRID, data):
 
         # Test additivity
         add_check = jnp.isclose(
-            pred_and_pdf(x1, fk)[0] + pred_and_pdf(x2, fk)[0],
-            pred_and_pdf(x1 + x2, fk)[0] + intercept,
+            forward_map(fk, x1)[0] + forward_map(fk, x2)[0],
+            forward_map(fk, x1 + x2)[0] + intercept,
         )
 
         # Test homogeneity
         c = jax.random.uniform(key, (1,))
 
         homogeneity_check = jnp.isclose(
-            c * (pred_and_pdf(x1, fk)[0] - intercept),
-            pred_and_pdf(c * x1, fk)[0] - intercept,
+            c * (forward_map(fk, x1)[0] - intercept),
+            forward_map(fk, c * x1)[0] - intercept,
         )
 
         if not add_check.all() or not homogeneity_check.all():

colibri/blackjax_fit.py

@@ -67,11 +67,7 @@ def level_0_commondata_tuple(
         Subset of flavour (evolution basis) indices to be used.
 
     fill_fk_xgrid_with_zeros: bool, default is False
-        If True, then the missing xgrid points in the FK table
-        will be filled with zeros. This is useful when the FK table
-        is needed as tensor of shape (Ndat, Nfl, Nfk_xgrid) with Nfk_xgrid and Nfl fixed
-        for all datasets.
-
+        Must match the value used to build ``fast_kernel_arrays``.
 
     Returns
     -------