Merge branch 'main' into use_fitcommondata

Author: LucaMantani

colibri/analytic_fit.py

@@ -85,12 +85,9 @@ def analytic_evidence_uniform_prior(sol_covmat, sol_mean, max_logl, a_vec, b_vec
 def analytic_fit(
     central_covmat_index,
     forward_map,
-    pdf_model,
     analytic_settings,
     prior_settings,
-    FIT_XGRID,
     fast_kernel_arrays,
-    data,
 ):
     """
     Analytic fits, for any *linear* PDF model.
@@ -109,28 +106,18 @@ def analytic_fit(
     forward_map: @jax.jit CompiledFunction
         Forward map function for the fit.
 
-    pdf_model: pdf_model.PDFModel
-        PDF model to 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.
-
-    data: validphys.core.DataGroupSpec
-        The data group specification for the 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(forward_map, FIT_XGRID, data)
+    check_pdf_model_is_linear(forward_map, fast_kernel_arrays)
 
     log.warning("The prior is assumed to be flat in the parameters.")
     log.warning(

colibri/checks.py

@@ -7,8 +7,6 @@
 from reportengine.checks import make_argcheck
 import jax.numpy as jnp
 import jax
-from colibri.theory_predictions import fast_kernel_arrays
-
 from colibri.utils import get_fit_path, get_pdf_model
 
 
@@ -42,13 +40,17 @@ def check_pdf_models_equal(prior_settings, forward_map, theoryid):
             )
 
 
-def check_pdf_model_is_linear(forward_map, 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.
     """
 
-    fk = fast_kernel_arrays(data, FIT_XGRID)
+    fk = fast_kernel_arrays
 
     parameters = forward_map.param_names
     intercept, _ = forward_map(fk, jnp.zeros(len(parameters)))

colibri/commondata_utils.py

@@ -113,11 +113,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
     -------

colibri/tests/test_analytic_fit.py

@@ -11,7 +11,6 @@
 import jax.random
 import pytest
 
-from colibri.api import API as colibriAPI
 from colibri.analytic_fit import AnalyticFit, analytic_fit, run_analytic_fit
 from colibri.core import PriorSettings
 from colibri.forward_map import FKTableForwardMap
@@ -21,7 +20,6 @@
     TEST_FK_ARRAYS,
     TEST_PRIOR_SETTINGS_UNIFORM,
     TEST_XGRID,
-    TEST_DATASETS,
 )
 
 analytic_settings = {
@@ -30,8 +28,6 @@
     "n_posterior_samples": 10,
 }
 
-TEST_DATA = colibriAPI.data(**TEST_DATASETS)
-
 
 def test_analytic_fit_flat_direction():
     """
@@ -51,12 +47,9 @@ def test_analytic_fit_flat_direction():
         analytic_fit(
             MOCK_CENTRAL_COVMAT_INDEX,
             forward_map,
-            MOCK_PDF_MODEL,
             analytic_settings,
             TEST_PRIOR_SETTINGS_UNIFORM,
-            TEST_XGRID,
             TEST_FK_ARRAYS,
-            TEST_DATA,
         )
 
 
@@ -80,12 +73,9 @@ def test_analytic_fit(caplog, monkeypatch):
     result = analytic_fit(
         MOCK_CENTRAL_COVMAT_INDEX,
         forward_map,
-        MOCK_PDF_MODEL,
         analytic_settings,
         TEST_PRIOR_SETTINGS_UNIFORM,
-        TEST_XGRID,
         TEST_FK_ARRAYS,
-        TEST_DATA,
     )
 
     assert isinstance(result, AnalyticFit)
@@ -103,12 +93,9 @@ def test_analytic_fit(caplog, monkeypatch):
         result_2 = analytic_fit(
             MOCK_CENTRAL_COVMAT_INDEX,
             forward_map,
-            MOCK_PDF_MODEL,
             analytic_settings,
             TEST_PRIOR_SETTINGS_UNIFORM,
-            TEST_XGRID,
             TEST_FK_ARRAYS,
-            TEST_DATA,
         )
 
     # Check that an error message was logged, because the prior was not wide enough
@@ -147,12 +134,9 @@ def test_analytic_fit_different_priors(caplog, monkeypatch):
     result = analytic_fit(
         MOCK_CENTRAL_COVMAT_INDEX,
         forward_map,
-        MOCK_PDF_MODEL,
         analytic_settings,
         PRIOR_SETTINGS1,
-        TEST_XGRID,
         TEST_FK_ARRAYS,
-        TEST_DATA,
     )
 
     assert isinstance(result, AnalyticFit)
@@ -174,12 +158,9 @@ def test_analytic_fit_different_priors(caplog, monkeypatch):
     result = analytic_fit(
         MOCK_CENTRAL_COVMAT_INDEX,
         forward_map,
-        MOCK_PDF_MODEL,
         analytic_settings,
         PRIOR_SETTINGS2,
-        TEST_XGRID,
         TEST_FK_ARRAYS,
-        TEST_DATA,
     )
 
 

colibri/tests/test_checks.py

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

colibri/tests/test_theory_predictions.py

@@ -9,12 +9,18 @@
 from numpy.testing import assert_allclose
 from validphys.fkparser import load_fktable
 
+import numpy as np
+
+from validphys.api import API as vpAPI
+from validphys.convolution import central_predictions
+
 from colibri.api import API as colibriAPI
 from colibri.tests.conftest import (
     CLOSURE_TEST_PDFSET,
     TEST_DATASET,
     TEST_DATASET_HAD,
     TEST_DATASETS,
+    TEST_DATASETS_DIS_HAD,
     TEST_DATASETS_HAD,
 )
 from colibri.theory_predictions import (
@@ -30,26 +36,12 @@ def __init__(self, xgrid):
         self.xgrid = xgrid
 
 
-def test_fktable_xgrid_indices_fill_with_zeros():
-    # Case where fill_fk_xgrid_with_zeros is True
-    fktable = FKTableDataMock(xgrid=jnp.array([0.1, 0.2, 0.3]))
-    FIT_XGRID = jnp.array([0.05, 0.1, 0.15, 0.2, 0.25, 0.3])
-
-    expected_indices = jnp.arange(
-        len(FIT_XGRID)
-    )  # Should return indices for the entire FIT_XGRID
-    result = fktable_xgrid_indices(fktable, FIT_XGRID, fill_fk_xgrid_with_zeros=True)
-
-    assert jnp.array_equal(result, expected_indices)
-
-
-def test_fktable_xgrid_indices_no_fill():
-    # Case where fill_fk_xgrid_with_zeros is False
+def test_fktable_xgrid_indices():
     fktable = FKTableDataMock(xgrid=jnp.array([0.1, 0.2, 0.3]))
     FIT_XGRID = jnp.array([0.05, 0.1, 0.15, 0.2, 0.25, 0.3])
 
     expected_indices = jnp.array([1, 3, 5])  # Indices where fk_xgrid matches FIT_XGRID
-    result = fktable_xgrid_indices(fktable, FIT_XGRID, fill_fk_xgrid_with_zeros=False)
+    result = fktable_xgrid_indices(fktable, FIT_XGRID)
 
     assert jnp.array_equal(result, expected_indices)
 
@@ -61,7 +53,7 @@ def test_fktable_xgrid_indices_with_tolerance():
 
     # Due to tolerance, the indices should match as if they were the same
     expected_indices = jnp.array([1, 3, 5])
-    result = fktable_xgrid_indices(fktable, FIT_XGRID, fill_fk_xgrid_with_zeros=False)
+    result = fktable_xgrid_indices(fktable, FIT_XGRID)
 
     assert jnp.array_equal(result, expected_indices)
 
@@ -74,7 +66,7 @@ def test_fktable_xgrid_indices_no_matches():
     expected_indices = jnp.array(
         []
     )  # No matching indices, closest_indices returns empty array
-    result = fktable_xgrid_indices(fktable, FIT_XGRID, fill_fk_xgrid_with_zeros=False)
+    result = fktable_xgrid_indices(fktable, FIT_XGRID)
     assert jnp.array_equal(result, expected_indices)
 
 
@@ -122,6 +114,49 @@ def test_fast_kernel_arrays():
     assert jnp.any(fk_arrays_filled[0][0][:, :, non_zero_indices] != 0)
 
 
+def test_fast_kernel_arrays_hadronic_fill_with_zeros():
+    """
+    Test that fast_kernel_arrays correctly fills the x-grid with zeros for hadronic FK tables.
+    This is a regression test for the bug where the 4D hadronic array was assigned
+    into a 3D zeros array.
+    """
+    from colibri.utils import closest_indices
+    from validphys.fkparser import load_fktable
+
+    dataset = colibriAPI.data(**TEST_DATASETS_HAD)
+    ds = dataset.datasets[0]
+    FIT_XGRID = colibriAPI.FIT_XGRID(**TEST_DATASETS_HAD)
+
+    # This should not raise an error (regression check)
+    fk_arrays_filled = colibriAPI.fast_kernel_arrays(
+        **{**TEST_DATASETS_HAD, "fill_fk_xgrid_with_zeros": True}
+    )
+
+    fk_arr = fk_arrays_filled[0][0]
+
+    # Hadronic FK array should be 4D: (Ndat, Nfl, Nfit_x, Nfit_x)
+    assert fk_arr.ndim == 4
+    assert fk_arr.shape[2] == len(FIT_XGRID)
+    assert fk_arr.shape[3] == len(FIT_XGRID)
+
+    # Check that non-zero values are placed at the correct x-grid positions
+    fk_xgrid = load_fktable(ds.fkspecs[0]).xgrid
+    non_zero_indices = closest_indices(FIT_XGRID, fk_xgrid, atol=1e-8)
+    non_zero_indices = np.array(non_zero_indices)
+
+    # The non-zero block should contain non-zero values
+    assert jnp.any(
+        fk_arr[:, :, non_zero_indices[:, None], non_zero_indices[None, :]] != 0
+    )
+
+    # Entries outside the non-zero block should be zero
+    all_indices = np.arange(len(FIT_XGRID))
+    zero_indices = np.setdiff1d(all_indices, non_zero_indices)
+    if len(zero_indices) > 0:
+        assert jnp.all(fk_arr[:, :, zero_indices, :] == 0)
+        assert jnp.all(fk_arr[:, :, :, zero_indices] == 0)
+
+
 def test_make_dis_prediction():
     """
     Test make_dis_prediction function gives the same results
@@ -198,3 +233,54 @@ def test_make_pred_data():
 
     assert callable(eval_preds)
     assert pred_data.shape == (fk_arrs[0][0].shape[0],)
+
+
+def test_predictions_independent_of_fill_fk_xgrid_with_zeros():
+    """
+    Regression test: ``fill_fk_xgrid_with_zeros`` is a memory/layout option and
+    must not change the theory predictions.
+
+    Previously the prediction closures unconditionally indexed the FK array with
+    ``fktable_xgrid_indices``, which are indices into FIT_XGRID. That is only
+    valid for zero-padded FK arrays; with ``fill_fk_xgrid_with_zeros=False`` the
+    FK x-grid axis was silently re-shuffled and out-of-range indices were clamped
+    by jax, giving wrong predictions.
+
+    Uses a mixed DIS + hadronic setup so that both convolution paths are covered
+    and FIT_XGRID is a strict superset of at least one FK x-grid (otherwise the
+    index mapping is the identity and the bug is invisible).
+    """
+    data = colibriAPI.data(**TEST_DATASETS_DIS_HAD)
+    FIT_XGRID = colibriAPI.FIT_XGRID(**TEST_DATASETS_DIS_HAD)
+
+    # Guard: the setup must actually exercise a non-trivial x-grid mapping.
+    mappings = [
+        fktable_xgrid_indices(load_fktable(ds.fkspecs[0]).with_cuts(ds.cuts), FIT_XGRID)
+        for ds in data.datasets
+    ]
+    assert any(
+        not jnp.array_equal(idx, jnp.arange(len(idx))) for idx in mappings
+    ), "test setup is degenerate: every FK x-grid maps onto FIT_XGRID as the identity"
+
+    pdf_grid = colibriAPI.closure_test_central_pdf_grid(
+        **{**CLOSURE_TEST_PDFSET, **TEST_DATASETS_DIS_HAD}
+    )
+
+    preds = {}
+    for fill in (False, True):
+        inp = {**TEST_DATASETS_DIS_HAD, "fill_fk_xgrid_with_zeros": fill}
+        eval_preds = colibriAPI.make_pred_data(**inp)
+        preds[fill] = eval_preds(pdf_grid, colibriAPI.fast_kernel_arrays(**inp))
+
+    assert_allclose(preds[False], preds[True], rtol=1e-6)
+
+    # Cross-check both against validphys, which is the ground truth here.
+    # NOTE: use central_predictions rather than dataset_inputs_results: the latter
+    # convolves every replica of the PDF set (and makes LHAPDF load all 101 members)
+    # only for its central value to be taken, which is ~200x slower here.
+    pdf = vpAPI.pdf(pdf=CLOSURE_TEST_PDFSET["closure_test_pdf"])
+    vp_preds = np.concatenate(
+        [np.array(central_predictions(ds, pdf)).ravel() for ds in data.datasets]
+    )
+    assert_allclose(preds[False], vp_preds, rtol=1e-6)
+    assert_allclose(preds[True], vp_preds, rtol=1e-6)

colibri/theory_penalties.py

@@ -79,8 +79,13 @@ def make_penalty_posdataset(posdataset, FIT_XGRID, flavour_indices=None):
 
     """
 
+    # NOTE: positivity_fast_kernel_arrays never zero-pads the FK x-grid, so the
+    # prediction closures must not mask the FK arrays onto the fit x-grid.
     pred_funcs = pred_funcs_from_dataset(
-        posdataset, FIT_XGRID, flavour_indices, fill_fk_xgrid_with_zeros=False
+        posdataset,
+        FIT_XGRID,
+        flavour_indices,
+        fill_fk_xgrid_with_zeros=False,
     )
 
     def pos_penalty(pdf, alpha, lambda_positivity, fk_dataset):