Adapt prior and param-init to forward map

colibri/analytic_fit.py

@@ -130,14 +130,14 @@ def analytic_fit(
     """
     # 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(pdf_model, forward_map, FIT_XGRID, data)
+    check_pdf_model_is_linear(forward_map, FIT_XGRID, data)
 
     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
+    parameters = forward_map.param_names
 
     # Precompute predictions for the basis of the model
     bases = jnp.identity(len(parameters))

colibri/bayes_prior.py

@@ -20,6 +20,10 @@ def bayesian_prior(prior_settings, forward_map):
     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

colibri/checks.py

@@ -7,7 +7,7 @@
 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.theory_predictions import fast_kernel_arrays
 
 from colibri.utils import get_fit_path, get_pdf_model
 
@@ -42,16 +42,15 @@ def check_pdf_models_equal(prior_settings, forward_map, theoryid):
             )
 
 
-def check_pdf_model_is_linear(pdf_model, forward_map, FIT_XGRID, data):
+def check_pdf_model_is_linear(forward_map, FIT_XGRID, data):
     """
     Decorator that can be added to functions to check that the
     PDF model is linear.
     """
 
-    pred_data = make_pred_data(data, FIT_XGRID)
     fk = fast_kernel_arrays(data, FIT_XGRID)
 
-    parameters = pdf_model.param_names
+    parameters = forward_map.param_names
     intercept, _ = forward_map(fk, jnp.zeros(len(parameters)))
 
     # Run the check for 10 random points in the parameter space

colibri/hessian_fit.py

@@ -15,7 +15,7 @@
 
 
 def hessian_fit(
-    pdf_model,
+    forward_map,
     log_likelihood,
     optimizer_provider,
     max_epochs,
@@ -27,8 +27,9 @@ def hessian_fit(
 
     Parameters
     ----------
-    pdf_model: pdf_model.PDFModel
-        The PDF model to be fitted.
+    forward_map: forward_map.ForwardMap
+        The forward map whose .param_names (PDF params + extra params) are used
+        for initialization and result labeling.
     log_likelihood: callable
         The log-likelihood function to be maximized.
     optimizer_provider: optax.GradientTransformation
@@ -72,7 +73,7 @@ def valid_chi2(params):
         log.info(f"Hessian fit initialization iteration {i+1}")
         # Generate random initial parameters
         initial_parameters = pdf_initial_parameters(
-            pdf_model, param_initialiser_settings, rng_seed + i
+            forward_map, param_initialiser_settings, rng_seed + i
         )
 
         # Delegate to generic gradient descent
@@ -191,7 +192,7 @@ def valid_chi2(params):
         hessian=hessian,
         cov_params=cov_params,
         resampled_posterior=hessian_param_set,
-        param_names=pdf_model.param_names,
+        param_names=forward_map.param_names,
     )
 
 

colibri/monte_carlo_fit.py

@@ -108,7 +108,7 @@ def loss_validation(parameters):
     )
 
 
-def run_monte_carlo_fit(monte_carlo_fit, pdf_model, output_path, replica_index, Q0):
+def run_monte_carlo_fit(monte_carlo_fit, forward_map, output_path, replica_index, Q0):
     """
     Runs the Monte Carlo fit and writes the output to the output directory.
 
@@ -117,8 +117,10 @@ def run_monte_carlo_fit(monte_carlo_fit, pdf_model, output_path, replica_index,
     monte_carlo_fit: MonteCarloFit
         The results of the Monte Carlo fit.
 
-    pdf_model: pdf_model.PDFModel
-        The PDF model used in the fit.
+    forward_map: forward_map.ForwardMap
+        The forward map used in the fit. Its .param_names (PDF + extra params)
+        label the optimized parameters; its .pdf_model is used to write the
+        export grid.
 
     output_path: pathlib.PosixPath
         Path to the output folder.
@@ -131,7 +133,7 @@ def run_monte_carlo_fit(monte_carlo_fit, pdf_model, output_path, replica_index,
     """
     mc_fit = monte_carlo_fit
 
-    df = pd.DataFrame(mc_fit.optimized_parameters, index=pdf_model.param_names).T
+    df = pd.DataFrame(mc_fit.optimized_parameters, index=forward_map.param_names).T
 
     # In a Monte Carlo fit, replicas are written to the fit_replicas
     # directory, and mc_postfit must then be applied to select valid ones
@@ -144,7 +146,7 @@ def run_monte_carlo_fit(monte_carlo_fit, pdf_model, output_path, replica_index,
     log.info(f"Writing exportgrid for replica {replica_index}")
     write_exportgrid_mc(
         jnp.array(df.iloc[0, :].tolist()),
-        pdf_model,
+        forward_map.pdf_model,
         replica_index,
         output_path,
         Q0,

colibri/param_initialisation.py

@@ -5,14 +5,15 @@
 log = logging.getLogger(__name__)
 
 
-def pdf_initial_parameters(pdf_model, param_initialiser_settings, replica_index=-1):
+def pdf_initial_parameters(forward_map, param_initialiser_settings, replica_index=-1):
     """
-    This function provides initial parameters for the PDF model.
+    This function provides initial parameters for the forward map.
 
     Parameters
     ----------
-    pdf_model: pdf_mode.PDFModel
-        The PDF model to initialise the parameters for.
+    forward_map: forward_map.ForwardMap
+        The forward map whose .param_names (PDF parameters + any extra parameters)
+        are used to determine the size and ordering of the initialised array.
 
     param_initialiser_settings: dict
         The settings for the initialiser.
@@ -34,7 +35,7 @@ def pdf_initial_parameters(pdf_model, param_initialiser_settings, replica_index=
         param_initialiser_settings["type"] = "zeros"
 
     if param_initialiser_settings["type"] == "zeros":
-        return jnp.array([0.0] * len(pdf_model.param_names))
+        return jnp.array([0.0] * len(forward_map.param_names))
 
     if "random_seed" in param_initialiser_settings:
         random_seed = jax.random.PRNGKey(
@@ -43,7 +44,7 @@ def pdf_initial_parameters(pdf_model, param_initialiser_settings, replica_index=
     else:
         random_seed = jax.random.PRNGKey(replica_index)
 
-    param_names = pdf_model.param_names
+    param_names = forward_map.param_names
 
     if param_initialiser_settings["type"] == "normal":
         means_setting = param_initialiser_settings.get("means", 0.0)
@@ -127,7 +128,7 @@ def expand(setting, default, name):
 
         initial_values = jax.random.uniform(
             key=random_seed,
-            shape=(len(pdf_model.param_names),),
+            shape=(len(forward_map.param_names),),
             minval=min_val,
             maxval=max_val,
         )

colibri/tests/test_checks.py

@@ -111,9 +111,8 @@ def test_check_pdf_models_equal_false_param_names(
         check_pdf_models_equal.__wrapped__(prior_settings, forward_map, theoryid)
 
 
-@patch("colibri.checks.make_pred_data")
 @patch("colibri.checks.fast_kernel_arrays")
-def test_check_pdf_model_is_linear(mock_fast_kernel_arrays, mock_make_pred_data):
+def test_check_pdf_model_is_linear(mock_fast_kernel_arrays):
     # Create test data
     FIT_XGRID = jnp.array([1.0, 2.0, 3.0])
     data = jnp.array([0.1, 0.2, 0.3])
@@ -140,7 +139,7 @@ def pdf_linear_model(params):
     )
 
     # Test for linear model (should not raise an exception)
-    check_pdf_model_is_linear(mock_pdf_model, forward_map_lin, FIT_XGRID, data)
+    check_pdf_model_is_linear(forward_map_lin, FIT_XGRID, data)
 
     # Now mock a non-linear model to ensure the ValueError is raised
     non_linear_model = FKTableForwardMap(
@@ -152,4 +151,4 @@ def pdf_linear_model(params):
 
     # Ensure ValueError is raised for non-linear model
     with pytest.raises(ValueError):
-        check_pdf_model_is_linear(mock_pdf_model, non_linear_model, FIT_XGRID, data)
+        check_pdf_model_is_linear(non_linear_model, FIT_XGRID, data)

colibri/tests/test_hessian_fit.py

@@ -11,9 +11,9 @@
 from numpy.testing import assert_allclose
 
 from colibri.hessian_fit import HessianFit, hessian_fit, run_hessian_fit
-from colibri.tests.conftest import MOCK_PDF_MODEL
+from colibri.tests.conftest import MOCK_PDF_MODEL, TEST_FORWARD_MAP_DIS
 
-N_PARAMS = len(MOCK_PDF_MODEL.param_names)
+N_PARAMS = len(TEST_FORWARD_MAP_DIS.param_names)
 
 
 class MockOptimizerProvider:
@@ -53,7 +53,7 @@ def test_hessian_fit_runs_and_shapes():
     param_initialiser_settings = {"type": "zeros"}
 
     result = hessian_fit(
-        pdf_model=MOCK_PDF_MODEL,
+        forward_map=TEST_FORWARD_MAP_DIS,
         log_likelihood=log_likelihood,
         optimizer_provider=MockOptimizerProvider(),
         max_epochs=5,
@@ -62,7 +62,7 @@ def test_hessian_fit_runs_and_shapes():
     )
 
     assert isinstance(result, HessianFit)
-    assert result.param_names == MOCK_PDF_MODEL.param_names
+    assert result.param_names == TEST_FORWARD_MAP_DIS.param_names
 
     # Hessian of ||p||^2 with the internal 0.5 factor should be identity
     assert_allclose(result.hessian, jnp.eye(N_PARAMS), rtol=1e-12, atol=1e-12)
@@ -92,7 +92,7 @@ def test_run_hessian_fit_exports(mock_write_exportgrid, tmp_path):
     param_initialiser_settings = {"type": "zeros"}
 
     fit_result = hessian_fit(
-        pdf_model=MOCK_PDF_MODEL,
+        forward_map=TEST_FORWARD_MAP_DIS,
         log_likelihood=log_likelihood,
         optimizer_provider=MockOptimizerProvider(),
         max_epochs=1,
@@ -127,7 +127,7 @@ def test_hessian_fit_raises_when_require_local_min_fails():
 
     with pytest.raises(ValueError):
         hessian_fit(
-            pdf_model=MOCK_PDF_MODEL,
+            forward_map=TEST_FORWARD_MAP_DIS,
             log_likelihood=log_likelihood,
             optimizer_provider=MockOptimizerProvider(),
             max_epochs=1,
@@ -149,7 +149,7 @@ def test_hessian_fit_logs_warning_on_local_min_check_failed(caplog):
 
     with caplog.at_level("WARNING"):
         res = hessian_fit(
-            pdf_model=MOCK_PDF_MODEL,
+            forward_map=TEST_FORWARD_MAP_DIS,
             log_likelihood=log_likelihood,
             optimizer_provider=MockOptimizerProvider(),
             max_epochs=1,
@@ -179,7 +179,7 @@ def test_hessian_fit_logs_critical_on_non_pd_hessian(caplog):
 
     with caplog.at_level("CRITICAL"):
         hessian_fit(
-            pdf_model=MOCK_PDF_MODEL,
+            forward_map=TEST_FORWARD_MAP_DIS,
             log_likelihood=bad_loglike,
             optimizer_provider=MockOptimizerProvider(),
             max_epochs=1,
@@ -210,7 +210,7 @@ def test_hessian_fit_raises_on_nonfinite_min_chi2():
 
     with pytest.raises(ValueError):
         hessian_fit(
-            pdf_model=MOCK_PDF_MODEL,
+            forward_map=TEST_FORWARD_MAP_DIS,
             log_likelihood=nan_loglike,
             optimizer_provider=MockOptimizerProvider(),
             max_epochs=1,

colibri/tests/test_monte_carlo_fit.py

@@ -12,11 +12,10 @@
 from numpy.testing import assert_allclose
 
 from colibri.monte_carlo_fit import MonteCarloFit, monte_carlo_fit, run_monte_carlo_fit
-from colibri.tests.conftest import MOCK_PDF_MODEL
+from colibri.tests.conftest import MOCK_PDF_MODEL, TEST_FORWARD_MAP_DIS
 from colibri.data_batch import data_batches
 
-mock_pdf_model = MOCK_PDF_MODEL
-N_PARAMS = len(MOCK_PDF_MODEL.param_names)
+N_PARAMS = len(TEST_FORWARD_MAP_DIS.param_names)
 
 
 class MockOptimizerProvider:
@@ -86,7 +85,11 @@ def create_directory_side_effect(*args, **kwargs):
     output_path = str(tmp_path)
 
     run_monte_carlo_fit(
-        mock_monte_carlo_fit, mock_pdf_model, output_path, replica_index=1, Q0=1.65
+        mock_monte_carlo_fit,
+        TEST_FORWARD_MAP_DIS,
+        output_path,
+        replica_index=1,
+        Q0=1.65,
     )
 
     # Check if the write_exportgrid function was called once as expected