diff --git a/colibri/analytic_fit.py b/colibri/analytic_fit.py index 895befac..637ea0bd 100644 --- a/colibri/analytic_fit.py +++ b/colibri/analytic_fit.py @@ -82,15 +82,15 @@ 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, + forward_map, pdf_model, analytic_settings, prior_settings, FIT_XGRID, fast_kernel_arrays, + data, ): """ Analytic fits, for any *linear* PDF model. @@ -106,8 +106,8 @@ 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. + forward_map: @jax.jit CompiledFunction + Forward map function for the fit. pdf_model: pdf_model.PDFModel PDF model to fit. @@ -124,7 +124,13 @@ def analytic_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(pdf_model, forward_map, FIT_XGRID, data) log.warning("The prior is assumed to be flat in the parameters.") log.warning( @@ -132,14 +138,13 @@ def analytic_fit( ) parameters = pdf_model.param_names - pred_and_pdf = pdf_model.pred_and_pdf_func(FIT_XGRID, forward_map=_pred_data) # 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 diff --git a/colibri/app.py b/colibri/app.py index 0e19c397..bac72136 100644 --- a/colibri/app.py +++ b/colibri/app.py @@ -32,6 +32,7 @@ "colibri.param_initialisation", "colibri.export_results", "colibri.closure_test", + "colibri.forward_map", "reportengine.report", ] diff --git a/colibri/bayes_prior.py b/colibri/bayes_prior.py index b148924c..aa098ed5 100644 --- a/colibri/bayes_prior.py +++ b/colibri/bayes_prior.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. @@ -31,8 +29,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 +43,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 = jnp.array([float(prior_specs["min_val"])] * pdf_model.n_parameters) - maxs = jnp.array([float(prior_specs["max_val"])] * pdf_model.n_parameters) + 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( diff --git a/colibri/blackjax_fit.py b/colibri/blackjax_fit.py index 22ea86cc..bcb3bd49 100644 --- a/colibri/blackjax_fit.py +++ b/colibri/blackjax_fit.py @@ -38,7 +38,7 @@ def blackjax_fit( - pdf_model, + forward_map, bayesian_prior, blackjax_settings, log_likelihood, @@ -48,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. @@ -71,7 +71,7 @@ def blackjax_fit( # set the BlackJAX seed rng_key = jax.random.PRNGKey(blackjax_settings["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) @@ -142,7 +142,7 @@ 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"] @@ -150,7 +150,7 @@ def one_step(carry, xs): nested_samples.to_csv(log_dir + "/nested_samples.csv") # Export resampled posterior samples - posterior_df = pd.DataFrame(resampled_posterior, columns=pdf_model.param_names) + 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) @@ -172,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={ diff --git a/colibri/checks.py b/colibri/checks.py index 1ad34187..fc4aab4e 100644 --- a/colibri/checks.py +++ b/colibri/checks.py @@ -9,15 +9,15 @@ 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": @@ -25,9 +25,10 @@ def check_pdf_models_equal(prior_settings, pdf_model, theoryid): 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,8 +42,7 @@ 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(pdf_model, forward_map, FIT_XGRID, data): """ Decorator that can be added to functions to check that the PDF model is linear. @@ -52,8 +52,7 @@ def check_pdf_model_is_linear(pdf_model, FIT_XGRID, data): fk = fast_kernel_arrays(data, FIT_XGRID) 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] + 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 +64,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(): diff --git a/colibri/export_results.py b/colibri/export_results.py index d294874d..8df0bfce 100644 --- a/colibri/export_results.py +++ b/colibri/export_results.py @@ -246,10 +246,11 @@ def write_replicas( # Create the exportgrid lhapdf_interpolator = pdf_model.grid_values_func(xgrid) + n_pdf_params = len(pdf_model.param_names) # Finish by writing the replicas to export grids, ready for evolution for i in indices_per_process: - parameters = jnp.array(bayes_fit.resampled_posterior[i, :]) + parameters = jnp.array(bayes_fit.resampled_posterior[i, :n_pdf_params]) grid_for_writing = np.array(lhapdf_interpolator(parameters)) replica_index = i + 1 diff --git a/colibri/forward_map.py b/colibri/forward_map.py new file mode 100644 index 00000000..e7e1e111 --- /dev/null +++ b/colibri/forward_map.py @@ -0,0 +1,183 @@ +""" +colibri.forward_map.py + +Forward maps: parameters → theory predictions. + +A ``ForwardMap`` implements the final stage of the fit pipeline, turning the +fit parameter vector into theory predictions that can be compared with +data in the likelihood. It will also return the PDF values on the fit x-grid, +which is sometimes needed for computing penalties. + + +Design choice: fixed call signature +----------------------------------- +The log-likelihood calls every forward map with the same fixed signature:: + + (fk_tables, params) -> predictions, pdf + +The PDF grid function is bound at construction time (via ``pdf_grid_func`` +stored on the instance), so ``fk_tables`` remain a dynamic argument that JAX +traces as an abstract input — keeping them out of the compiled binary and +avoiding expensive recompilation when values change. + +Parameter convention +-------------------- +``params`` is a 1-D array containing *all* fit parameters. In colibri we allow +for "extra" fit parameters beyond the PDF model parameters (e.g. nuisance-like factors, +or parameters of a custom prediction function). + +By convention: + +``params[:self.n_pdf_params]`` are PDF parameters consumed by the bound +``pdf_grid_func``; any remaining entries are "extra" parameters interpreted by +the chosen ``ForwardMap`` implementation. + +Example - fitting a normalisation factor on top of the PDF +---------------------------------------------------------- +:: + + class NormForwardMap(ForwardMap): + def __init__(self, pred_func, pdf_model, pdf_grid_func): + super().__init__(pdf_model, extra_param_names=["norm"]) + self._pred_func = pred_func + self._pdf_grid_func = pdf_grid_func + + def __call__(self, fk_tables, params): + pdf = self._pdf_grid_func(params[: self.n_pdf_params]) + norm = params[self.n_pdf_params] # first extra parameter + return norm * self._pred_func(pdf, fk_tables), pdf + +Example - fixed PDF, fitting only extra parameters +-------------------------------------------------- +:: + + class FixedPDFForwardMap(ForwardMap): + def __init__(self, pred_func, fixed_pdf, fk_tables, pdf_model=None): + super().__init__(pdf_model) + self._pred_func = pred_func + self.fixed_pdf = fixed_pdf + self._fixed_pred = self._pred_func(fixed_pdf, fk_tables) + + def __call__(self, fk_tables, params): + scale = params[0] + return scale * self._fixed_pred, self.fixed_pdf +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any, Callable + +import jax.numpy as jnp + + +class ForwardMap(ABC): + """Abstract base class for forward maps. + + A forward map turns fit parameters into theory predictions that can be + compared with experimental data inside the likelihood. + + All forward maps share the same call signature: + + ``(fk_tables, params) -> predictions`` + + Notes + ----- + The split point between PDF parameters and "extra" parameters is owned + by the forward map via ``self.n_pdf_params``. + """ + + def __init__(self, pdf_model, extra_param_names: list[str] = []): + + self.pdf_model = pdf_model + if pdf_model is not None: + self.pdf_param_names = pdf_model.param_names + else: + self.pdf_param_names = [] + self.extra_param_names = extra_param_names + + @property + def n_pdf_params(self) -> int: + """Number of PDF parameters, derived from ``pdf_param_names``.""" + return len(self.pdf_param_names) + + @property + def param_names(self) -> list[str]: + """All fit parameter names: PDF parameters followed by extra parameters.""" + return list(self.pdf_param_names) + list(self.extra_param_names) + + @abstractmethod + def __call__( + self, + fk_tables: Any, + params: jnp.ndarray, + ) -> jnp.ndarray: + """Compute theory predictions from fit parameters. + + Parameters + ---------- + fk_tables : jnp.ndarray + Fast-kernel tables needed by the prediction function. + + params : jnp.ndarray + 1-D array containing all fit parameters. By convention: + * ``params[:self.n_pdf_params]`` are PDF parameters + * the remaining entries are extra parameters interpreted by the + specific ``ForwardMap`` implementation. + + Returns + ------- + jnp.ndarray + Theory predictions (1-D array with one entry per data point). + jnp.ndarray + The PDF values (2-D array with shape (N_fl, N_x)). + + """ + raise NotImplementedError + + +class FKTableForwardMap(ForwardMap): + """Default forward map: params → PDF → FK-table convolution. + + This is the standard pipeline used in colibri PDF fits. + """ + + def __init__( + self, + pred_func: Callable[[jnp.ndarray, Any], jnp.ndarray], + pdf_model, + pdf_grid_func: Callable[[jnp.ndarray], jnp.ndarray], + extra_param_names: list[str] = [], + ): + super().__init__(pdf_model, extra_param_names=extra_param_names) + self._pred_func = pred_func + self._pdf_grid_func = pdf_grid_func + + def __call__(self, fk_tables, params): + pdf_params = params[: self.n_pdf_params] + pdf = self._pdf_grid_func(pdf_params) + return self._pred_func(pdf, fk_tables), pdf + + +def forward_map(_pred_data, pdf_model, FIT_XGRID, extra_param_names=[]): + """Reportengine provider that builds the default FK-table forward map. + + Parameters + ---------- + _pred_data : callable + Prediction function of the form ``pred_func(pdf, fk_tables) -> predictions``. + pdf_model : object + The PDF model object; must expose ``param_names`` and ``grid_values_func``. + FIT_XGRID : array-like + The x-grid on which the PDF is evaluated. + extra_param_names : list[str], optional + Names of any additional fit parameters beyond the PDF parameters. + + """ + pdf_grid_func = pdf_model.grid_values_func(FIT_XGRID) + return FKTableForwardMap( + pred_func=_pred_data, + pdf_model=pdf_model, + pdf_grid_func=pdf_grid_func, + extra_param_names=extra_param_names, + ) diff --git a/colibri/likelihood.py b/colibri/likelihood.py index 7f3d2563..ff45eec9 100644 --- a/colibri/likelihood.py +++ b/colibri/likelihood.py @@ -23,7 +23,6 @@ def __init__( self, central_covmat_index, pdf_model, - fit_xgrid, forward_map, fast_kernel_arrays, positivity_fast_kernel_arrays, @@ -38,8 +37,6 @@ def __init__( pdf_model: pdf_model.PDFModel - fit_xgrid: np.ndarray - forward_map: Callable fast_kernel_arrays: tuple @@ -62,9 +59,7 @@ def __init__( self.positivity_penalty_settings = positivity_penalty_settings self.integrability_penalty = integrability_penalty - self.pred_and_pdf = pdf_model.pred_and_pdf_func( - fit_xgrid, forward_map=forward_map - ) + self.forward_map = forward_map self.fast_kernel_arrays = fast_kernel_arrays self.positivity_fast_kernel_arrays = positivity_fast_kernel_arrays @@ -126,7 +121,7 @@ def log_likelihood( jnp.ndarray jax array with the value of the log-likelihood. """ - predictions, pdf = self.pred_and_pdf(params, fast_kernel_arrays) + predictions, pdf = self.forward_map(fast_kernel_arrays, params) # Select only the data relevant for this likelihood # Especially important when using a training/validation split predictions = predictions[self.central_values_idx] @@ -168,8 +163,7 @@ def log_likelihood( def log_likelihood( central_covmat_index, pdf_model, - FIT_XGRID, - _pred_data, + forward_map, fast_kernel_arrays, positivity_fast_kernel_arrays, _penalty_posdata, @@ -185,8 +179,7 @@ def log_likelihood( return LogLikelihood( central_covmat_index, pdf_model, - FIT_XGRID, - _pred_data, + forward_map, fast_kernel_arrays, positivity_fast_kernel_arrays, _penalty_posdata, @@ -199,8 +192,7 @@ def mc_log_likelihood( mc_pseudodata, general_covariance_matrix, pdf_model, - FIT_XGRID, - _pred_data, + forward_map, fast_kernel_arrays, positivity_fast_kernel_arrays, _penalty_posdata, @@ -227,8 +219,7 @@ def mc_log_likelihood( train_loglike = LogLikelihood( central_covmat_index_train, pdf_model, - FIT_XGRID, - _pred_data, + forward_map, fast_kernel_arrays, positivity_fast_kernel_arrays, _penalty_posdata, @@ -253,8 +244,7 @@ def mc_log_likelihood( val_loglike = LogLikelihood( central_covmat_index_val, pdf_model, - FIT_XGRID, - _pred_data, + forward_map, fast_kernel_arrays, positivity_fast_kernel_arrays, _penalty_posdata, diff --git a/colibri/mc_utils.py b/colibri/mc_utils.py index 37f5f0fc..b8aa935c 100644 --- a/colibri/mc_utils.py +++ b/colibri/mc_utils.py @@ -126,9 +126,10 @@ def write_exportgrid_mc( # Create the exportgrid lhapdf_interpolator = pdf_model.grid_values_func(LHAPDF_XGRID) + n_pdf_params = len(pdf_model.param_names) # Rotate the grid from the evolution basis into the export grid basis - grid_for_writing = np.array(lhapdf_interpolator(parameters)) + grid_for_writing = np.array(lhapdf_interpolator(parameters[:n_pdf_params])) write_exportgrid( grid_for_writing=grid_for_writing, diff --git a/colibri/pdf_model.py b/colibri/pdf_model.py index f6183e46..961b39c9 100644 --- a/colibri/pdf_model.py +++ b/colibri/pdf_model.py @@ -6,7 +6,7 @@ """ from abc import ABC, abstractmethod -from typing import Callable, Tuple +from typing import Callable import jax.numpy as jnp from jax.typing import ArrayLike @@ -60,40 +60,3 @@ def func(params): return func """ pass - - def pred_and_pdf_func( - self, - xgrid: ArrayLike, - forward_map: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray], - ) -> Callable[[jnp.ndarray, jnp.ndarray], Tuple[jnp.ndarray, jnp.ndarray]]: - """Creates a function that returns a tuple of two arrays, given the model parameters and the fast kernel arrays as input. - - The returned function produces: - - The first array: 1D vector of theory predictions for the data. - - The second array: PDF values evaluated on the x-grid, using `self.grid_values_func`, with shape (Nfl, Nx). - - The `forward_map` is used to map the PDF values defined on the x-grid and the fast kernel arrays into the corresponding theory prediction vector. - """ - pdf_func = self.grid_values_func(xgrid) - - def pred_and_pdf(params, fast_kernel_arrays): - """ - Parameters - ---------- - params: jnp.array - The model parameters. - - fast_kernel_arrays: tuple - tuple of tuples of jnp.arrays - The FK tables to use. - - Returns - ------- - tuple - The predictions and the PDF values. - """ - pdf = pdf_func(params) - predictions = forward_map(pdf, fast_kernel_arrays) - return predictions, pdf - - return pred_and_pdf diff --git a/colibri/tests/conftest.py b/colibri/tests/conftest.py index 7baddcf1..e5bf20f0 100644 --- a/colibri/tests/conftest.py +++ b/colibri/tests/conftest.py @@ -10,6 +10,7 @@ from colibri.pdf_model import PDFModel from colibri.core import PriorSettings +from colibri.forward_map import FKTableForwardMap CONFIG_YML_PATH = "test_runcards/test_config.yaml" @@ -299,16 +300,7 @@ def wmin_param(params): MOCK_PDF_MODEL.grid_values_func = lambda xgrid: lambda params: jnp.sum( jnp.array([param * TEST_PDF_GRID for param in params]), axis=0 ) -""" -Mock PDF model with 2 parameters and grid_values_func simple mult add operation on np.ones grid. -""" -MOCK_PDF_MODEL.pred_and_pdf_func = ( - lambda xgrid, forward_map: lambda params, fast_kernel_arrays: ( - forward_map(MOCK_PDF_MODEL.grid_values_func(xgrid)(params), fast_kernel_arrays), - MOCK_PDF_MODEL.grid_values_func(xgrid)(params), - ) -) """ Mock prediction function of PDF model. """ @@ -340,7 +332,11 @@ def wmin_param(params): """ -TEST_FORWARD_MAP_DIS = lambda pdf, fk_arrays: jnp.einsum("ijk,jk->i", fk_arrays[0], pdf) +TEST_FORWARD_MAP_DIS = FKTableForwardMap( + lambda pdf, fk_arrays: jnp.einsum("ijk,jk->i", fk_arrays[0], pdf), + pdf_model=MOCK_PDF_MODEL, + pdf_grid_func=MOCK_PDF_MODEL.grid_values_func(TEST_XGRID), +) """ Mock DIS forward map function for testing purposes. Function expects a tuple of DIS-like fast kernel array of shape (N_data, TEST_N_FL, TEST_N_XGRID) and a PDF of shape (TEST_N_FL, TEST_N_XGRID). diff --git a/colibri/tests/test_analytic_fit.py b/colibri/tests/test_analytic_fit.py index 4067cdad..9120e831 100644 --- a/colibri/tests/test_analytic_fit.py +++ b/colibri/tests/test_analytic_fit.py @@ -11,16 +11,17 @@ 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 from colibri.tests.conftest import ( MOCK_CENTRAL_COVMAT_INDEX, MOCK_PDF_MODEL, TEST_FK_ARRAYS, - TEST_FORWARD_MAP_DIS, - TEST_PDF_GRID, TEST_PRIOR_SETTINGS_UNIFORM, TEST_XGRID, + TEST_DATASETS, ) analytic_settings = { @@ -29,53 +30,62 @@ "n_posterior_samples": 10, } +TEST_DATA = colibriAPI.data(**TEST_DATASETS) + def test_analytic_fit_flat_direction(): """ Tests that the analytic fit raises a ValueError when the - pred_and_pdf_func returns a flat direction in the parameter space. + forward_map returns a flat direction in the parameter space. """ - # override the pred_and_pdf_func to return a flat direction - # in the parameter space - MOCK_PDF_MODEL.pred_and_pdf_func = lambda xgrid, forward_map: ( - lambda params, fkarrs: (jnp.ones_like(params), TEST_PDF_GRID) - ) + n_params = len(MOCK_PDF_MODEL.param_names) - _pred_data = TEST_FORWARD_MAP_DIS + forward_map = FKTableForwardMap( + lambda pdf, fkarrs: jnp.ones(n_params), + pdf_model=MOCK_PDF_MODEL, + pdf_grid_func=MOCK_PDF_MODEL.grid_values_func(TEST_XGRID), + ) with pytest.raises(ValueError): # Run the analytic fit and make sure that the Value Error is raised analytic_fit( MOCK_CENTRAL_COVMAT_INDEX, - _pred_data, + forward_map, MOCK_PDF_MODEL, analytic_settings, TEST_PRIOR_SETTINGS_UNIFORM, TEST_XGRID, TEST_FK_ARRAYS, + TEST_DATA, ) -def test_analytic_fit(caplog): +def test_analytic_fit(caplog, monkeypatch): """ Tests basic functionality of the analytic fit function. """ - MOCK_PDF_MODEL.pred_and_pdf_func = lambda xgrid, forward_map: ( - lambda params, fkarrs: (params, TEST_PDF_GRID) + # Mock the grid_values_func of the PDF model within the test to return the input parameters as the PDF grid values + monkeypatch.setattr( + MOCK_PDF_MODEL, "grid_values_func", lambda xgrid: lambda params: params ) - _pred_data = TEST_FORWARD_MAP_DIS + forward_map = FKTableForwardMap( + lambda pdf, fkarrs: pdf, + pdf_model=MOCK_PDF_MODEL, + pdf_grid_func=MOCK_PDF_MODEL.grid_values_func(TEST_XGRID), + ) # Run the analytic fit result = analytic_fit( MOCK_CENTRAL_COVMAT_INDEX, - _pred_data, + forward_map, MOCK_PDF_MODEL, analytic_settings, TEST_PRIOR_SETTINGS_UNIFORM, TEST_XGRID, TEST_FK_ARRAYS, + TEST_DATA, ) assert isinstance(result, AnalyticFit) @@ -92,12 +102,13 @@ def test_analytic_fit(caplog): with caplog.at_level(logging.ERROR): # Set the log level to ERROR result_2 = analytic_fit( MOCK_CENTRAL_COVMAT_INDEX, - _pred_data, + 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 @@ -112,7 +123,7 @@ def test_analytic_fit(caplog): assert len(result_2.param_names) == len(MOCK_PDF_MODEL.param_names) -def test_analytic_fit_different_priors(caplog): +def test_analytic_fit_different_priors(caplog, monkeypatch): PRIOR_SETTINGS1 = PriorSettings( **{ @@ -121,21 +132,27 @@ def test_analytic_fit_different_priors(caplog): } ) - MOCK_PDF_MODEL.pred_and_pdf_func = lambda xgrid, forward_map: ( - lambda params, fkarrs: (params, TEST_PDF_GRID) + # Mock the grid_values_func of the PDF model within the test to return the input parameters as the PDF grid values + monkeypatch.setattr( + MOCK_PDF_MODEL, "grid_values_func", lambda xgrid: lambda params: params ) - _pred_data = None + forward_map = FKTableForwardMap( + lambda pdf, fkarrs: pdf, + pdf_model=MOCK_PDF_MODEL, + pdf_grid_func=MOCK_PDF_MODEL.grid_values_func(TEST_XGRID), + ) # Run the analytic fit result = analytic_fit( MOCK_CENTRAL_COVMAT_INDEX, - _pred_data, + forward_map, MOCK_PDF_MODEL, analytic_settings, PRIOR_SETTINGS1, TEST_XGRID, TEST_FK_ARRAYS, + TEST_DATA, ) assert isinstance(result, AnalyticFit) @@ -156,12 +173,13 @@ def test_analytic_fit_different_priors(caplog): # Run the analytic fit with custom uniform prior result = analytic_fit( MOCK_CENTRAL_COVMAT_INDEX, - _pred_data, + forward_map, MOCK_PDF_MODEL, analytic_settings, PRIOR_SETTINGS2, TEST_XGRID, TEST_FK_ARRAYS, + TEST_DATA, ) diff --git a/colibri/tests/test_bayes_prior.py b/colibri/tests/test_bayes_prior.py index bd20a4e2..331a43f3 100644 --- a/colibri/tests/test_bayes_prior.py +++ b/colibri/tests/test_bayes_prior.py @@ -17,14 +17,20 @@ from colibri.bayes_prior import bayesian_prior from colibri.core import PriorSettings from colibri.tests.conftest import MOCK_PDF_MODEL, TEST_PRIOR_SETTINGS_UNIFORM +from unittest.mock import Mock + +# Create a mock forward_map that exposes param_names matching MOCK_PDF_MODEL +MOCK_FORWARD_MAP = Mock() +MOCK_FORWARD_MAP.param_names = MOCK_PDF_MODEL.param_names +MOCK_FORWARD_MAP.pdf_param_names = MOCK_PDF_MODEL.param_names def test_uniform_prior(): """ Test the transformation of a uniform prior distribution. """ + prior_transform = bayesian_prior(TEST_PRIOR_SETTINGS_UNIFORM, MOCK_FORWARD_MAP) MOCK_PDF_MODEL.n_parameters = 2 - prior_transform = bayesian_prior(TEST_PRIOR_SETTINGS_UNIFORM, MOCK_PDF_MODEL) key = random.PRNGKey(0) n_params = MOCK_PDF_MODEL.n_parameters @@ -68,7 +74,7 @@ def test_uniform_prior(): } ) - prior_transform_bounds = bayesian_prior(prior_settings_bounds, MOCK_PDF_MODEL) + prior_transform_bounds = bayesian_prior(prior_settings_bounds, MOCK_FORWARD_MAP) cube_bounds = random.uniform(key, shape=(2,)) expected_bounds = jnp.array( @@ -98,7 +104,7 @@ def test_uniform_prior(): ) with pytest.raises(ValueError, match="Missing bounds for parameters"): - bayesian_prior(prior_settings_missing_bounds, MOCK_PDF_MODEL) + bayesian_prior(prior_settings_missing_bounds, MOCK_FORWARD_MAP) # ---- Test missing min_val/max_val and bounds ---- prior_settings_invalid = PriorSettings( @@ -109,7 +115,7 @@ def test_uniform_prior(): ) with pytest.raises(ValueError, match="prior_distribution_specs must define either"): - bayesian_prior(prior_settings_invalid, MOCK_PDF_MODEL) + bayesian_prior(prior_settings_invalid, MOCK_FORWARD_MAP) @patch("colibri.bayes_prior.get_full_posterior") @@ -135,7 +141,7 @@ def cov(self): } ) - prior_transform = bayesian_prior(prior_settings, MOCK_PDF_MODEL) + prior_transform = bayesian_prior(prior_settings, MOCK_FORWARD_MAP) key = random.PRNGKey(0) cube = random.uniform(key, shape=(10, 2)) @@ -162,4 +168,4 @@ def test_invalid_prior_type(): ) with pytest.raises(ValueError) as e: - bayesian_prior(prior_settings, MOCK_PDF_MODEL) + bayesian_prior(prior_settings, MOCK_FORWARD_MAP) diff --git a/colibri/tests/test_blackjax_fit.py b/colibri/tests/test_blackjax_fit.py index c8845666..61831471 100644 --- a/colibri/tests/test_blackjax_fit.py +++ b/colibri/tests/test_blackjax_fit.py @@ -21,6 +21,7 @@ from colibri.core import BlackJAXFit, BayesianPrior from colibri.blackjax_fit import blackjax_fit, run_blackjax_fit +from colibri.forward_map import FKTableForwardMap from colibri.likelihood import LogLikelihood jax.config.update("jax_enable_x64", True) @@ -42,7 +43,7 @@ def mock_sample(rng_key, n_samples): bayesian_prior = BayesianPrior( prior_transform=lambda x: x, log_prob=lambda x: -jnp.sum(x**2, axis=-1), - sample=lambda rng, n: jnp.zeros((n, MOCK_PDF_MODEL.n_parameters)), + sample=lambda rng, n: jnp.zeros((n, len(MOCK_PDF_MODEL.param_names))), ) integrability_penalty = lambda pdf: jnp.array([0.0]) @@ -61,12 +62,15 @@ def mock_sample(rng_key, n_samples): @pytest.mark.parametrize("pos_penalty", [True, False]) def test_blackjax_fit(pos_penalty): - _pred_data = lambda *args: jnp.array([0.0]) + forward_map = FKTableForwardMap( + lambda pdf, fk: jnp.zeros(len(MOCK_PDF_MODEL.param_names)), + pdf_model=MOCK_PDF_MODEL, + pdf_grid_func=MOCK_PDF_MODEL.grid_values_func(TEST_XGRID), + ) mock_log_likelihood = LogLikelihood( MOCK_CENTRAL_COVMAT_INDEX, MOCK_PDF_MODEL, - TEST_XGRID, - _pred_data, + forward_map, TEST_FK_ARRAYS, TEST_POS_FK_ARRAYS, MOCK_PENALTY_POSDATA, @@ -78,12 +82,10 @@ def test_blackjax_fit(pos_penalty): integrability_penalty=integrability_penalty, ) - MOCK_PDF_MODEL.n_parameters = len(MOCK_PDF_MODEL.param_names) - with patch("colibri.blackjax_fit.anesthetic.NestedSamples"): fit_result = blackjax_fit( - MOCK_PDF_MODEL, + forward_map, bayesian_prior, blackjax_settings, mock_log_likelihood, @@ -93,13 +95,14 @@ def test_blackjax_fit(pos_penalty): def test_blackjax_fit_truncates_posterior_and_warns(caplog): - # --- ensure pdf_model is consistent --- - MOCK_PDF_MODEL.n_parameters = len(MOCK_PDF_MODEL.param_names) + # --- build a forward_map with the right param_names --- + mock_forward_map = Mock() + mock_forward_map.param_names = ["param1", "param2"] bayesian_prior = BayesianPrior( prior_transform=lambda x: x, log_prob=lambda x: -jnp.sum(x**2, axis=-1), - sample=lambda rng, n: jnp.zeros((n, MOCK_PDF_MODEL.n_parameters)), + sample=lambda rng, n: jnp.zeros((n, len(mock_forward_map.param_names))), ) blackjax_settings = { @@ -143,7 +146,7 @@ def test_blackjax_fit_truncates_posterior_and_warns(caplog): caplog.set_level("WARNING") fit_result = blackjax_fit( - MOCK_PDF_MODEL, + mock_forward_map, bayesian_prior, blackjax_settings, log_likelihood, diff --git a/colibri/tests/test_checks.py b/colibri/tests/test_checks.py index 12495d1d..a0145512 100644 --- a/colibri/tests/test_checks.py +++ b/colibri/tests/test_checks.py @@ -9,6 +9,8 @@ import jax.numpy as jnp import pytest +from colibri.forward_map import FKTableForwardMap + from colibri.checks import check_pdf_model_is_linear, check_pdf_models_equal from colibri.core import PriorSettings @@ -19,11 +21,8 @@ read_data="theoryid: 123\nt0pdfset: t0pdfset1", ) @patch("os.path.exists", return_value=True) -@patch("colibri.checks.get_pdf_model", return_value="model1") -@patch("colibri.checks.pdf_models_equal") -def test_check_pdf_models_equal_true( - mock_pdf_models_equal, mock_get_pdf_model, mock_exists, mock_open -): +@patch("colibri.checks.get_pdf_model") +def test_check_pdf_models_equal_true(mock_get_pdf_model, mock_exists, mock_open): # Setup prior_settings = PriorSettings( **{ @@ -31,16 +30,20 @@ def test_check_pdf_models_equal_true( "prior_distribution_specs": {"prior_fit": "fit1"}, } ) - pdf_model = "model1" + + # The prior model returned by get_pdf_model must have matching param_names + mock_prior_model = MagicMock() + mock_prior_model.param_names = ["param1", "param2"] + mock_get_pdf_model.return_value = mock_prior_model + + forward_map = MagicMock() + forward_map.pdf_param_names = ["param1", "param2"] theoryid = MagicMock() theoryid.id = 123 - # Configure mock behavior - mock_pdf_models_equal.side_effect = lambda x, y: x == y - - # Act - check_pdf_models_equal.__wrapped__(prior_settings, pdf_model, theoryid) + # Act — should not raise + check_pdf_models_equal.__wrapped__(prior_settings, forward_map, theoryid) @patch( @@ -49,10 +52,9 @@ def test_check_pdf_models_equal_true( read_data="theoryid: 456\nt0pdfset: t0pdfset1", ) @patch("os.path.exists", return_value=True) -@patch("colibri.checks.get_pdf_model", return_value="model1") -@patch("colibri.checks.pdf_models_equal") +@patch("colibri.checks.get_pdf_model") def test_check_pdf_models_equal_false_theoryid( - mock_pdf_models_equal, mock_get_pdf_model, mock_exists, mock_open + mock_get_pdf_model, mock_exists, mock_open ): # Setup prior_settings = PriorSettings( @@ -61,21 +63,20 @@ def test_check_pdf_models_equal_false_theoryid( "prior_distribution_specs": {"prior_fit": "fit1"}, } ) - pdf_model = "model1" - theoryid = MagicMock() - theoryid.id = 123 + mock_prior_model = MagicMock() + mock_prior_model.param_names = ["param1", "param2"] + mock_get_pdf_model.return_value = mock_prior_model - t0pdfset = MagicMock() - t0pdfset.name = "t0pdfset1" + forward_map = MagicMock() + forward_map.pdf_param_names = ["param1", "param2"] - # Configure mock behavior - mock_pdf_models_equal.side_effect = lambda x, y: x == y + theoryid = MagicMock() + theoryid.id = 123 + # Theory ID mismatch (file says 456, fit says 123) with pytest.raises(Exception): - check_pdf_models_equal.__wrapped__( - prior_settings, pdf_model, theoryid, t0pdfset - ) + check_pdf_models_equal.__wrapped__(prior_settings, forward_map, theoryid) @patch( @@ -84,31 +85,30 @@ def test_check_pdf_models_equal_false_theoryid( read_data="theoryid: 123\nt0pdfset: t0pdfset2", ) @patch("os.path.exists", return_value=True) -@patch("colibri.checks.get_pdf_model", return_value="model1") -@patch("colibri.checks.pdf_models_equal") -def test_check_pdf_models_equal_false_t0pdf( - mock_pdf_models_equal, mock_get_pdf_model, mock_exists, mock_open +@patch("colibri.checks.get_pdf_model") +def test_check_pdf_models_equal_false_param_names( + mock_get_pdf_model, mock_exists, mock_open ): - # Setup - prior_settings = { - "prior_distribution": "prior_from_gauss_posterior", - "prior_distribution_specs": {"prior_fit": "fit1"}, - } - pdf_model = "model1" + # Setup — param names mismatch between prior model and forward_map + prior_settings = PriorSettings( + **{ + "prior_distribution": "prior_from_gauss_posterior", + "prior_distribution_specs": {"prior_fit": "fit1"}, + } + ) - theoryid = MagicMock() - theoryid.id = 123 + mock_prior_model = MagicMock() + mock_prior_model.param_names = ["param1", "param2", "param3"] # different + mock_get_pdf_model.return_value = mock_prior_model - t0pdfset = MagicMock() - t0pdfset.name = "t0pdfset1" + forward_map = MagicMock() + forward_map.pdf_param_names = ["param1", "param2"] - # Configure mock behavior - mock_pdf_models_equal.side_effect = lambda x, y: x == y + theoryid = MagicMock() + theoryid.id = 123 - with pytest.raises(Exception): - check_pdf_models_equal.__wrapped__( - prior_settings, pdf_model, theoryid, t0pdfset - ) + with pytest.raises(ValueError): + check_pdf_models_equal.__wrapped__(prior_settings, forward_map, theoryid) @patch("colibri.checks.make_pred_data") @@ -125,24 +125,31 @@ def test_check_pdf_model_is_linear(mock_fast_kernel_arrays, mock_make_pred_data) mock_pdf_model = MagicMock() mock_pdf_model.param_names = ["a", "b", "c"] - # Mock the behavior of pred_and_pdf_func to return a linear model - def linear_model(params, fk): - # Simulating a simple linear model: f(x) = a*x + b*y + c*z + 3.0, where params = [a, b, c] - return (jnp.dot(params, fk) + 3.0, params) + # Mock the behavior of pdf_grid to return a linear model + def pdf_linear_model(params): + return params + + # Set the mock's grid_values_func to return the linear_model function + mock_pdf_model.grid_values_func.return_value = pdf_linear_model - # Set the mock's pred_and_pdf_func to return the linear_model function - mock_pdf_model.pred_and_pdf_func.return_value = linear_model + forward_map_lin = FKTableForwardMap( + # Simulating a simple linear model: f(x) = a*x + b*y + c*z + 3.0, where pdf = [a, b, c] + lambda pdf, fk: jnp.dot(pdf, fk) + 3.0, + pdf_model=mock_pdf_model, + pdf_grid_func=mock_pdf_model.grid_values_func(FIT_XGRID), + ) # Test for linear model (should not raise an exception) - check_pdf_model_is_linear.__wrapped__(mock_pdf_model, FIT_XGRID, data) + check_pdf_model_is_linear(mock_pdf_model, forward_map_lin, FIT_XGRID, data) # Now mock a non-linear model to ensure the ValueError is raised - def non_linear_model(params, fk): + non_linear_model = FKTableForwardMap( # Introduce some non-linearity - return (jnp.dot(params**2, FIT_XGRID) + fk, params) - - mock_pdf_model.pred_and_pdf_func.return_value = non_linear_model + lambda pdf, fk: jnp.dot(pdf**2, FIT_XGRID) + fk, + pdf_model=mock_pdf_model, + pdf_grid_func=mock_pdf_model.grid_values_func(FIT_XGRID), + ) # Ensure ValueError is raised for non-linear model with pytest.raises(ValueError): - check_pdf_model_is_linear.__wrapped__(mock_pdf_model, FIT_XGRID, data) + check_pdf_model_is_linear(mock_pdf_model, non_linear_model, FIT_XGRID, data) diff --git a/colibri/tests/test_forward_map.py b/colibri/tests/test_forward_map.py new file mode 100644 index 00000000..75601f5b --- /dev/null +++ b/colibri/tests/test_forward_map.py @@ -0,0 +1,396 @@ +""" +colibri.tests.test_forward_map + +Tests for the ForwardMap abstract base class, FKTableForwardMap, and the +forward_map provider function. +""" + +import pytest +import numpy as np +import jax.numpy as jnp +from unittest.mock import Mock +from numpy.testing import assert_array_almost_equal + +from colibri.forward_map import ForwardMap, FKTableForwardMap, forward_map +from colibri.tests.conftest import ( + TEST_FK_ARRAYS, + TEST_PDF_GRID, + TEST_N_DATA, + TEST_N_FL, + TEST_N_XGRID, + TEST_XGRID, + MOCK_PDF_MODEL, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _simple_pred_func(pdf, fk_tables): + """DIS-like prediction: einsum over the first FK table.""" + return jnp.einsum("ijk,jk->i", fk_tables[0], pdf) + + +def _make_pdf_grid_func(pdf_grid): + """Return a callable that ignores params and always returns pdf_grid.""" + return lambda params: pdf_grid + + +# --------------------------------------------------------------------------- +# ForwardMap (abstract base class) +# --------------------------------------------------------------------------- + + +def _mock_pdf_model(param_names): + """Create a mock pdf_model with the given param_names.""" + model = Mock() + model.param_names = param_names + return model + + +def test_forward_map_cannot_be_instantiated(): + """ForwardMap is abstract; direct instantiation must raise TypeError.""" + with pytest.raises(TypeError): + ForwardMap(pdf_model=_mock_pdf_model(["a", "b"])) + + +def test_forward_map_subclass_without_call_cannot_be_instantiated(): + """A subclass that does not implement __call__ must also raise TypeError.""" + + class NoCallSubclass(ForwardMap): + pass + + with pytest.raises(TypeError): + NoCallSubclass(pdf_model=_mock_pdf_model(["a", "b"])) + + +def test_forward_map_abstract_call_raises_not_implemented(): + """Calling super().__call__() must hit the raise NotImplementedError body.""" + + class SuperCallingForwardMap(ForwardMap): + def __call__(self, fk_tables, params): + return super().__call__(fk_tables, params) + + fm = SuperCallingForwardMap(pdf_model=_mock_pdf_model(["a", "b"])) + with pytest.raises(NotImplementedError): + fm(TEST_FK_ARRAYS, jnp.array([1.0, 2.0])) + + +def test_forward_map_none_pdf_model_sets_empty_param_names(): + """When pdf_model is None, pdf_param_names must be set to an empty list.""" + + class MinimalForwardMap(ForwardMap): + def __call__(self, fk_tables, params): + return None + + fm = MinimalForwardMap(pdf_model=None) + assert fm.pdf_param_names == [] + assert fm.n_pdf_params == 0 + assert fm.pdf_model is None + + +def test_forward_map_subclass_stores_pdf_param_names(): + """pdf_param_names passed to super().__init__ must be stored on the instance.""" + + class MinimalForwardMap(ForwardMap): + def __call__(self, fk_tables, params): + return None + + mock_model = _mock_pdf_model(["p0", "p1", "p2", "p3", "p4"]) + fm = MinimalForwardMap(pdf_model=mock_model) + assert fm.pdf_param_names == ["p0", "p1", "p2", "p3", "p4"] + assert fm.pdf_model is mock_model + assert fm.n_pdf_params == 5 + + +def test_forward_map_extra_param_names_default(): + """extra_param_names defaults to an empty tuple.""" + + class MinimalForwardMap(ForwardMap): + def __call__(self, fk_tables, params): + return None + + fm = MinimalForwardMap(pdf_model=_mock_pdf_model(["a", "b"])) + assert list(fm.extra_param_names) == [] + assert fm.param_names == ["a", "b"] + + +def test_forward_map_extra_param_names(): + """extra_param_names are stored and appear in param_names after pdf_param_names.""" + + class MinimalForwardMap(ForwardMap): + def __call__(self, fk_tables, params): + return None + + fm = MinimalForwardMap( + pdf_model=_mock_pdf_model(["a", "b"]), extra_param_names=["norm", "scale"] + ) + assert list(fm.extra_param_names) == ["norm", "scale"] + assert fm.param_names == ["a", "b", "norm", "scale"] + assert fm.n_pdf_params == 2 + + +# --------------------------------------------------------------------------- +# FKTableForwardMap.__init__ +# --------------------------------------------------------------------------- + + +def test_fktable_forward_map_stores_pdf_param_names(): + """FKTableForwardMap.__init__ must store pdf_param_names via the base class.""" + pdf_grid_func = _make_pdf_grid_func(TEST_PDF_GRID) + fm = FKTableForwardMap( + pred_func=_simple_pred_func, + pdf_model=_mock_pdf_model(["a", "b", "c"]), + pdf_grid_func=pdf_grid_func, + ) + assert fm.pdf_param_names == ["a", "b", "c"] + assert fm.n_pdf_params == 3 + + +def test_fktable_forward_map_extra_param_names(): + """FKTableForwardMap must accept and store extra_param_names.""" + pdf_grid_func = _make_pdf_grid_func(TEST_PDF_GRID) + fm = FKTableForwardMap( + pred_func=_simple_pred_func, + pdf_model=_mock_pdf_model(["a", "b"]), + pdf_grid_func=pdf_grid_func, + extra_param_names=["norm"], + ) + assert fm.param_names == ["a", "b", "norm"] + assert fm.n_pdf_params == 2 + assert list(fm.extra_param_names) == ["norm"] + + +def test_fktable_forward_map_stores_pred_func(): + """FKTableForwardMap.__init__ must store the pred_func.""" + pdf_grid_func = _make_pdf_grid_func(TEST_PDF_GRID) + fm = FKTableForwardMap( + pred_func=_simple_pred_func, + pdf_model=_mock_pdf_model(["a", "b", "c"]), + pdf_grid_func=pdf_grid_func, + ) + assert fm._pred_func is _simple_pred_func + + +def test_fktable_forward_map_stores_pdf_grid_func(): + """FKTableForwardMap.__init__ must store the pdf_grid_func.""" + pdf_grid_func = _make_pdf_grid_func(TEST_PDF_GRID) + fm = FKTableForwardMap( + pred_func=_simple_pred_func, + pdf_model=_mock_pdf_model(["a", "b"]), + pdf_grid_func=pdf_grid_func, + ) + assert fm._pdf_grid_func is pdf_grid_func + + +# --------------------------------------------------------------------------- +# FKTableForwardMap.__call__ +# --------------------------------------------------------------------------- + + +def test_fktable_forward_map_returns_tuple(): + """__call__ must return a 2-tuple (predictions, pdf).""" + pdf_grid_func = _make_pdf_grid_func(TEST_PDF_GRID) + fm = FKTableForwardMap( + pred_func=_simple_pred_func, + pdf_model=_mock_pdf_model(["a", "b"]), + pdf_grid_func=pdf_grid_func, + ) + params = jnp.array([1.0, 2.0]) + + result = fm(TEST_FK_ARRAYS, params) + + assert isinstance(result, tuple) + assert len(result) == 2 + + +def test_fktable_forward_map_predictions_shape(): + """Predictions returned by __call__ must have shape (N_data,).""" + pdf_grid_func = _make_pdf_grid_func(TEST_PDF_GRID) + fm = FKTableForwardMap( + pred_func=_simple_pred_func, + pdf_model=_mock_pdf_model(["a", "b"]), + pdf_grid_func=pdf_grid_func, + ) + params = jnp.array([1.0, 2.0]) + + predictions, _ = fm(TEST_FK_ARRAYS, params) + + assert predictions.shape == (TEST_N_DATA,) + + +def test_fktable_forward_map_pdf_shape(): + """PDF returned by __call__ must have shape (N_fl, N_x).""" + pdf_grid_func = _make_pdf_grid_func(TEST_PDF_GRID) + fm = FKTableForwardMap( + pred_func=_simple_pred_func, + pdf_model=_mock_pdf_model(["a", "b"]), + pdf_grid_func=pdf_grid_func, + ) + params = jnp.array([1.0, 2.0]) + + _, pdf = fm(TEST_FK_ARRAYS, params) + + assert pdf.shape == (TEST_N_FL, TEST_N_XGRID) + + +def test_fktable_forward_map_slices_pdf_params(): + """ + __call__ must pass only params[:n_pdf_params] to pdf_grid_func; extra + parameters appended to params must not affect the PDF or predictions. + """ + pdf_grid_func = _make_pdf_grid_func(TEST_PDF_GRID) + fm = FKTableForwardMap( + pred_func=_simple_pred_func, + pdf_model=_mock_pdf_model(["a", "b"]), + pdf_grid_func=pdf_grid_func, + ) + + pdf_params = jnp.array([1.0, 2.0]) + extra_params = jnp.array([99.0, -99.0]) # should be ignored + + params_no_extra = pdf_params + params_with_extra = jnp.concatenate([pdf_params, extra_params]) + + preds_no_extra, pdf_no_extra = fm(TEST_FK_ARRAYS, params_no_extra) + preds_with_extra, pdf_with_extra = fm(TEST_FK_ARRAYS, params_with_extra) + + assert_array_almost_equal(preds_no_extra, preds_with_extra) + assert_array_almost_equal(pdf_no_extra, pdf_with_extra) + + +def test_fktable_forward_map_uses_pdf_grid_func(): + """ + __call__ must feed the pdf returned by the bound pdf_grid_func into pred_func. + We verify this by constructing two forward maps with differently scaled pdf_grid_funcs. + """ + scale = 3.0 + base_pdf_grid_func = _make_pdf_grid_func(TEST_PDF_GRID) + scaled_pdf_grid_func = lambda p: scale * base_pdf_grid_func(p) # noqa: E731 + + fm_base = FKTableForwardMap( + pred_func=_simple_pred_func, + pdf_model=_mock_pdf_model(["a", "b"]), + pdf_grid_func=base_pdf_grid_func, + ) + fm_scaled = FKTableForwardMap( + pred_func=_simple_pred_func, + pdf_model=_mock_pdf_model(["a", "b"]), + pdf_grid_func=scaled_pdf_grid_func, + ) + + params = jnp.array([1.0, 2.0]) + preds_base, _ = fm_base(TEST_FK_ARRAYS, params) + preds_scaled, _ = fm_scaled(TEST_FK_ARRAYS, params) + + np.testing.assert_allclose(preds_scaled, scale * preds_base, rtol=1e-5) + + +def test_fktable_forward_map_correct_values(): + """ + __call__ must produce predictions equal to pred_func(pdf_grid_func(params), fk). + """ + pdf_grid_func = _make_pdf_grid_func(TEST_PDF_GRID) + fm = FKTableForwardMap( + pred_func=_simple_pred_func, + pdf_model=_mock_pdf_model(["a", "b"]), + pdf_grid_func=pdf_grid_func, + ) + + params = jnp.array([1.0, 2.0]) + predictions, pdf = fm(TEST_FK_ARRAYS, params) + + expected_pdf = pdf_grid_func(params) + expected_preds = _simple_pred_func(expected_pdf, TEST_FK_ARRAYS) + + assert_array_almost_equal(predictions, expected_preds) + assert_array_almost_equal(pdf, expected_pdf) + + +# --------------------------------------------------------------------------- +# forward_map provider function +# --------------------------------------------------------------------------- + + +def test_forward_map_provider_returns_fktable_forward_map(): + """forward_map() must return an FKTableForwardMap instance.""" + result = forward_map( + _pred_data=_simple_pred_func, pdf_model=MOCK_PDF_MODEL, FIT_XGRID=TEST_XGRID + ) + assert isinstance(result, FKTableForwardMap) + + +def test_forward_map_provider_infers_pdf_param_names(): + """ + forward_map() must set pdf_param_names equal to pdf_model.param_names, + and n_pdf_params must equal len(pdf_model.param_names). + """ + result = forward_map( + _pred_data=_simple_pred_func, pdf_model=MOCK_PDF_MODEL, FIT_XGRID=TEST_XGRID + ) + assert result.pdf_param_names == MOCK_PDF_MODEL.param_names + assert result.n_pdf_params == len(MOCK_PDF_MODEL.param_names) + + +def test_forward_map_provider_stores_pred_func(): + """forward_map() must wire _pred_data into the FKTableForwardMap.""" + result = forward_map( + _pred_data=_simple_pred_func, pdf_model=MOCK_PDF_MODEL, FIT_XGRID=TEST_XGRID + ) + assert result._pred_func is _simple_pred_func + + +def test_forward_map_provider_binds_pdf_grid_func(): + """forward_map() must call pdf_model.grid_values_func(FIT_XGRID) and store the result.""" + result = forward_map( + _pred_data=_simple_pred_func, pdf_model=MOCK_PDF_MODEL, FIT_XGRID=TEST_XGRID + ) + assert callable(result._pdf_grid_func) + + +def test_forward_map_provider_functional(): + """ + The FKTableForwardMap built by forward_map() must produce correct results + when called. + """ + fm = forward_map( + _pred_data=_simple_pred_func, pdf_model=MOCK_PDF_MODEL, FIT_XGRID=TEST_XGRID + ) + params = jnp.array([1.0, 2.0]) + + predictions, pdf = fm(TEST_FK_ARRAYS, params) + + assert predictions.shape == (TEST_N_DATA,) + assert pdf.shape == (TEST_N_FL, TEST_N_XGRID) + + +def test_forward_map_provider_with_different_param_counts(): + """ + forward_map() must correctly handle pdf_models with different numbers of + parameters. + """ + for n in [1, 3, 7]: + mock_model = Mock() + mock_model.param_names = [f"p_{i}" for i in range(n)] + fm = forward_map( + _pred_data=_simple_pred_func, pdf_model=mock_model, FIT_XGRID=TEST_XGRID + ) + assert fm.pdf_param_names == mock_model.param_names + assert fm.n_pdf_params == n + + +def test_forward_map_provider_with_extra_param_names(): + """ + forward_map() must forward extra_param_names to FKTableForwardMap + and expose them via param_names. + """ + extra = ["norm", "scale"] + fm = forward_map( + _pred_data=_simple_pred_func, + pdf_model=MOCK_PDF_MODEL, + FIT_XGRID=TEST_XGRID, + extra_param_names=extra, + ) + assert fm.param_names == MOCK_PDF_MODEL.param_names + extra + assert list(fm.extra_param_names) == extra diff --git a/colibri/tests/test_likelihood.py b/colibri/tests/test_likelihood.py index 4d00e77c..72c79f94 100644 --- a/colibri/tests/test_likelihood.py +++ b/colibri/tests/test_likelihood.py @@ -18,7 +18,6 @@ TEST_FK_ARRAYS, TEST_FORWARD_MAP_DIS, TEST_POS_FK_ARRAYS, - TEST_XGRID, ) from colibri.data_batch import BatchSpec @@ -38,7 +37,6 @@ def test_LogLikelihood_class(pos_penalty): log_likelihood_class = LogLikelihood( central_covmat_index=MOCK_CENTRAL_COVMAT_INDEX, pdf_model=MOCK_PDF_MODEL, - fit_xgrid=TEST_XGRID, forward_map=TEST_FORWARD_MAP_DIS, fast_kernel_arrays=TEST_FK_ARRAYS, positivity_fast_kernel_arrays=TEST_POS_FK_ARRAYS, @@ -66,8 +64,8 @@ def test_LogLikelihood_class(pos_penalty): ] ) # Compute expected value using actual prediction and covariance - predictions, pdf = log_likelihood_class.pred_and_pdf( - params, log_likelihood_class.fast_kernel_arrays + predictions, pdf = log_likelihood_class.forward_map( + log_likelihood_class.fast_kernel_arrays, params ) predictions = predictions[log_likelihood_class.central_values_idx] diff = predictions - log_likelihood_class.central_values @@ -104,7 +102,6 @@ def test_log_likelihood(pos_penalty): log_likelihood_class = LogLikelihood( central_covmat_index=MOCK_CENTRAL_COVMAT_INDEX, pdf_model=MOCK_PDF_MODEL, - fit_xgrid=TEST_XGRID, forward_map=TEST_FORWARD_MAP_DIS, fast_kernel_arrays=TEST_FK_ARRAYS, positivity_fast_kernel_arrays=TEST_POS_FK_ARRAYS, @@ -115,7 +112,6 @@ def test_log_likelihood(pos_penalty): log_like = log_likelihood( MOCK_CENTRAL_COVMAT_INDEX, MOCK_PDF_MODEL, - TEST_XGRID, TEST_FORWARD_MAP_DIS, TEST_FK_ARRAYS, TEST_POS_FK_ARRAYS, @@ -143,7 +139,6 @@ def test_log_likelihood_with_and_without_pos_penalty(): log_likelihood_class = LogLikelihood( MOCK_CENTRAL_COVMAT_INDEX, MOCK_PDF_MODEL, - TEST_XGRID, TEST_FORWARD_MAP_DIS, TEST_FK_ARRAYS, TEST_POS_FK_ARRAYS, @@ -165,8 +160,8 @@ def test_log_likelihood_with_and_without_pos_penalty(): ) # Compute expectation directly: -0.5 * (chi2 + pos_pen + integ_pen) - predictions, pdf = log_likelihood_class.pred_and_pdf( - params, log_likelihood_class.fast_kernel_arrays + predictions, pdf = log_likelihood_class.forward_map( + log_likelihood_class.fast_kernel_arrays, params ) predictions = predictions[log_likelihood_class.central_values_idx] diff = predictions - log_likelihood_class.central_values @@ -195,7 +190,6 @@ def test_log_likelihood_with_and_without_pos_penalty(): log_likelihood_class = LogLikelihood( MOCK_CENTRAL_COVMAT_INDEX, MOCK_PDF_MODEL, - TEST_XGRID, TEST_FORWARD_MAP_DIS, TEST_FK_ARRAYS, TEST_POS_FK_ARRAYS, @@ -213,8 +207,8 @@ def test_log_likelihood_with_and_without_pos_penalty(): ) # Expectation: Only chi2 value (penalties zeroed) - predictions, pdf = log_likelihood_class.pred_and_pdf( - params, log_likelihood_class.fast_kernel_arrays + predictions, pdf = log_likelihood_class.forward_map( + log_likelihood_class.fast_kernel_arrays, params ) predictions = predictions[log_likelihood_class.central_values_idx] diff = predictions - log_likelihood_class.central_values @@ -255,7 +249,6 @@ def test_mc_log_likelihood_with_split(pos_penalty): mc_pd, general_covariance_matrix, MOCK_PDF_MODEL, - TEST_XGRID, TEST_FORWARD_MAP_DIS, TEST_FK_ARRAYS, TEST_POS_FK_ARRAYS, @@ -275,7 +268,7 @@ def test_mc_log_likelihood_with_split(pos_penalty): # Compute expected for train and validation independently def compute_expected(ll_obj): - preds, pdf = ll_obj.pred_and_pdf(params, ll_obj.fast_kernel_arrays) + preds, pdf = ll_obj.forward_map(ll_obj.fast_kernel_arrays, params) preds = preds[ll_obj.central_values_idx] diff = preds - ll_obj.central_values inv = ll_obj.inv_covmat @@ -333,7 +326,6 @@ def test_mc_log_likelihood_without_split_returns_nan_for_validation(pos_penalty) mc_pd, general_covariance_matrix, MOCK_PDF_MODEL, - TEST_XGRID, TEST_FORWARD_MAP_DIS, TEST_FK_ARRAYS, TEST_POS_FK_ARRAYS, @@ -348,8 +340,8 @@ def test_mc_log_likelihood_without_split_returns_nan_for_validation(pos_penalty) params = jnp.array([0.3, 0.4]) train_val = train_loglike(params) # Compute expected train value - predictions, pdf = train_loglike.pred_and_pdf( - params, train_loglike.fast_kernel_arrays + predictions, pdf = train_loglike.forward_map( + train_loglike.fast_kernel_arrays, params ) predictions = predictions[train_loglike.central_values_idx] diff = predictions - train_loglike.central_values @@ -392,7 +384,6 @@ def test_LogLikelihood_call_with_batch_idx(pos_penalty): log_likelihood_class = LogLikelihood( central_covmat_index=MOCK_CENTRAL_COVMAT_INDEX, pdf_model=MOCK_PDF_MODEL, - fit_xgrid=TEST_XGRID, forward_map=TEST_FORWARD_MAP_DIS, fast_kernel_arrays=TEST_FK_ARRAYS, positivity_fast_kernel_arrays=TEST_POS_FK_ARRAYS, @@ -409,8 +400,8 @@ def test_LogLikelihood_call_with_batch_idx(pos_penalty): ll_value_batched = log_likelihood_class(params, batch=batch) # Compute expected on the batch index: recompute inv_covmat on the sub-covmat - predictions, pdf = log_likelihood_class.pred_and_pdf( - params, log_likelihood_class.fast_kernel_arrays + predictions, pdf = log_likelihood_class.forward_map( + log_likelihood_class.fast_kernel_arrays, params ) predictions = predictions[log_likelihood_class.central_values_idx] predictions_b = predictions[batch.idx] @@ -454,7 +445,6 @@ def test_LogLikelihood_call_with_batch_with_inv_cov(pos_penalty): log_likelihood_class = LogLikelihood( central_covmat_index=MOCK_CENTRAL_COVMAT_INDEX, pdf_model=MOCK_PDF_MODEL, - fit_xgrid=TEST_XGRID, forward_map=TEST_FORWARD_MAP_DIS, fast_kernel_arrays=TEST_FK_ARRAYS, positivity_fast_kernel_arrays=TEST_POS_FK_ARRAYS, @@ -476,8 +466,8 @@ def test_LogLikelihood_call_with_batch_with_inv_cov(pos_penalty): ll_value_batched = log_likelihood_class(params, batch=batch) # Compute expected value using the provided inv_b (should be identical) - predictions, pdf = log_likelihood_class.pred_and_pdf( - params, log_likelihood_class.fast_kernel_arrays + predictions, pdf = log_likelihood_class.forward_map( + log_likelihood_class.fast_kernel_arrays, params ) predictions = predictions[log_likelihood_class.central_values_idx] predictions_b = predictions[batch.idx] diff --git a/colibri/tests/test_pdf_model.py b/colibri/tests/test_pdf_model.py index e21127c1..e170b9ab 100644 --- a/colibri/tests/test_pdf_model.py +++ b/colibri/tests/test_pdf_model.py @@ -42,17 +42,3 @@ def test_grid_values_func(): expected_output = sum([param * TEST_PDF_GRID for param in params]) assert_array_equal(func(params), expected_output) - - -def test_pred_and_pdf_func(): - """ - Tests that the pred_and_pdf_func returns the correct values. - """ - pred_and_pdf = model.pred_and_pdf_func(TEST_XGRID, TEST_FORWARD_MAP_DIS) - - params = jnp.array([2, 3]) - predictions, pdf = pred_and_pdf(params, TEST_FK_ARRAYS) - - expected_predictions = jnp.einsum("ijk,jk->i", TEST_FK_ARRAYS[0], pdf) - - assert jnp.allclose(predictions, expected_predictions) diff --git a/colibri/tests/test_ultranest_fit.py b/colibri/tests/test_ultranest_fit.py index 4b6df815..912a1f16 100644 --- a/colibri/tests/test_ultranest_fit.py +++ b/colibri/tests/test_ultranest_fit.py @@ -22,6 +22,7 @@ from colibri.ultranest_fit import UltranestFit, run_ultranest_fit, ultranest_fit from colibri.likelihood import LogLikelihood from colibri.core import BayesianPrior +from colibri.forward_map import FKTableForwardMap jax.config.update("jax_enable_x64", True) @@ -66,11 +67,15 @@ def mock_sample(rng_key, n_samples): def test_ultranest_fit(pos_penalty): _pred_data = lambda *args: jnp.array([0.0]) + forward_map = FKTableForwardMap( + _pred_data, + pdf_model=MOCK_PDF_MODEL, + pdf_grid_func=MOCK_PDF_MODEL.grid_values_func(TEST_XGRID), + ) mock_log_likelihood = LogLikelihood( MOCK_CENTRAL_COVMAT_INDEX, MOCK_PDF_MODEL, - TEST_XGRID, - _pred_data, + forward_map, TEST_FK_ARRAYS, TEST_POS_FK_ARRAYS, MOCK_PENALTY_POSDATA, @@ -83,7 +88,7 @@ def test_ultranest_fit(pos_penalty): ) fit_result = ultranest_fit( - MOCK_PDF_MODEL, + forward_map, bayesian_prior, ultranest_settings, mock_log_likelihood, @@ -103,13 +108,17 @@ def test_ultranest_fit(pos_penalty): def test_ultranest_fit_vectorized(pos_penalty): _pred_data = lambda *args: jnp.array([0.0]) + forward_map = FKTableForwardMap( + _pred_data, + pdf_model=MOCK_PDF_MODEL, + pdf_grid_func=MOCK_PDF_MODEL.grid_values_func(TEST_XGRID), + ) ultranest_settings["ReactiveNS_settings"]["vectorized"] = True mock_log_likelihood = LogLikelihood( MOCK_CENTRAL_COVMAT_INDEX, MOCK_PDF_MODEL, - TEST_XGRID, - _pred_data, + forward_map, TEST_FK_ARRAYS, TEST_POS_FK_ARRAYS, MOCK_PENALTY_POSDATA, @@ -122,7 +131,7 @@ def test_ultranest_fit_vectorized(pos_penalty): ) fit_result = ultranest_fit( - MOCK_PDF_MODEL, + forward_map, bayesian_prior, ultranest_settings, mock_log_likelihood, @@ -152,12 +161,16 @@ def test_ultranest_fit_with_SliceSampler(pos_penalty): } _pred_data = lambda *args: jnp.array([0.0]) + forward_map = FKTableForwardMap( + _pred_data, + pdf_model=MOCK_PDF_MODEL, + pdf_grid_func=MOCK_PDF_MODEL.grid_values_func(TEST_XGRID), + ) mock_log_likelihood = LogLikelihood( MOCK_CENTRAL_COVMAT_INDEX, MOCK_PDF_MODEL, - TEST_XGRID, - _pred_data, + forward_map, TEST_FK_ARRAYS, TEST_POS_FK_ARRAYS, MOCK_PENALTY_POSDATA, @@ -170,7 +183,7 @@ def test_ultranest_fit_with_SliceSampler(pos_penalty): ) fit_result = ultranest_fit( - MOCK_PDF_MODEL, + forward_map, bayesian_prior, ultranest_settings, mock_log_likelihood, @@ -200,12 +213,16 @@ def test_ultranest_fit_with_popSliceSampler(pos_penalty): } _pred_data = lambda *args: jnp.array([0.0]) + forward_map = FKTableForwardMap( + _pred_data, + pdf_model=MOCK_PDF_MODEL, + pdf_grid_func=MOCK_PDF_MODEL.grid_values_func(TEST_XGRID), + ) mock_log_likelihood = LogLikelihood( MOCK_CENTRAL_COVMAT_INDEX, MOCK_PDF_MODEL, - TEST_XGRID, - _pred_data, + forward_map, TEST_FK_ARRAYS, TEST_POS_FK_ARRAYS, MOCK_PENALTY_POSDATA, @@ -218,7 +235,7 @@ def test_ultranest_fit_with_popSliceSampler(pos_penalty): ) fit_result = ultranest_fit( - MOCK_PDF_MODEL, + forward_map, bayesian_prior, ultranest_settings, mock_log_likelihood, @@ -252,12 +269,16 @@ def test_ultranest_fit_with_sampler_plot(mock_sampler_class, pos_penalty): } _pred_data = lambda *args: jnp.array([0.0]) + forward_map = FKTableForwardMap( + _pred_data, + pdf_model=MOCK_PDF_MODEL, + pdf_grid_func=MOCK_PDF_MODEL.grid_values_func(TEST_XGRID), + ) mock_log_likelihood = LogLikelihood( MOCK_CENTRAL_COVMAT_INDEX, MOCK_PDF_MODEL, - TEST_XGRID, - _pred_data, + forward_map, TEST_FK_ARRAYS, TEST_POS_FK_ARRAYS, MOCK_PENALTY_POSDATA, @@ -285,7 +306,7 @@ def test_ultranest_fit_with_sampler_plot(mock_sampler_class, pos_penalty): mock_sampler_instance.plot = Mock() fit_result = ultranest_fit( - MOCK_PDF_MODEL, + forward_map, bayesian_prior, ultranest_settings_with_plot, mock_log_likelihood, diff --git a/colibri/tests/test_utils.py b/colibri/tests/test_utils.py index 9a4981ce..d7a7459c 100644 --- a/colibri/tests/test_utils.py +++ b/colibri/tests/test_utils.py @@ -31,6 +31,7 @@ TEST_DATASET, TEST_DATASET_HAD, ) +from colibri.forward_map import FKTableForwardMap from colibri.utils import ( cast_to_numpy, closest_indices, @@ -342,6 +343,11 @@ def test_likelihood_float_type( len(MOCK_CENTRAL_COVMAT_INDEX.central_values) ) # Mock _pred_data FIT_XGRID = jnp.linspace(0, 1, 10) # Mock FIT_XGRID + forward_map = FKTableForwardMap( + _pred_data, + pdf_model=MOCK_PDF_MODEL, + pdf_grid_func=MOCK_PDF_MODEL.grid_values_func(FIT_XGRID), + ) # Mock forward_map output_path = tmp_path fast_kernel_arrays = jax.random.uniform( @@ -350,7 +356,7 @@ def test_likelihood_float_type( # Call the function under test likelihood_float_type( - _pred_data=_pred_data, + forward_map=forward_map, pdf_model=MOCK_PDF_MODEL, FIT_XGRID=FIT_XGRID, bayesian_prior=mock_bayesian_prior, diff --git a/colibri/ultranest_fit.py b/colibri/ultranest_fit.py index 9896a218..ac6e396c 100644 --- a/colibri/ultranest_fit.py +++ b/colibri/ultranest_fit.py @@ -41,7 +41,7 @@ def ultranest_fit( - pdf_model, + forward_map, bayesian_prior, ultranest_settings, log_likelihood, @@ -51,8 +51,8 @@ def ultranest_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 The prior object containing prior_transform, log_prob, and sample functions. @@ -74,7 +74,7 @@ def ultranest_fit( # set the ultranest seed np.random.seed(ultranest_settings["ultranest_seed"]) - parameters = pdf_model.param_names + parameters = forward_map.param_names if ultranest_settings["ReactiveNS_settings"]["vectorized"]: log.info("Vectorized likelihood for ultranest fit.") diff --git a/colibri/utils.py b/colibri/utils.py index 19157426..aee86d0d 100644 --- a/colibri/utils.py +++ b/colibri/utils.py @@ -291,7 +291,7 @@ def wrapper(*args, **kwargs): def likelihood_float_type( - _pred_data, + forward_map, pdf_model, FIT_XGRID, bayesian_prior, @@ -309,10 +309,8 @@ def likelihood_float_type( central_values = central_covmat_index.central_values covmat = central_covmat_index.covmat - pred_and_pdf = pdf_model.pred_and_pdf_func(FIT_XGRID, forward_map=_pred_data) - def log_likelihood(params, central_values, inv_covmat, fast_kernel_arrays): - predictions, _ = pred_and_pdf(params, fast_kernel_arrays) + predictions, pdf = forward_map(fast_kernel_arrays, params) return -0.5 * loss_function(central_values, predictions, inv_covmat) params = bayesian_prior.prior_transform( @@ -454,6 +452,7 @@ def write_resampled_bayesian_fit( # overwrite old ns_result.csv with resampled posterior parameters = pdf_model.param_names + n_pdf_params = len(pdf_model.param_names) df = pd.DataFrame(resampled_posterior, columns=parameters) df.to_csv(str(resampled_fit_path) + f"/{csv_results_name}.csv", float_format="%.5e") @@ -466,7 +465,7 @@ def write_resampled_bayesian_fit( for i, parameters in enumerate(resampled_posterior): # Get the PDF grid in the evolution basis lhapdf_interpolator = pdf_model.grid_values_func(LHAPDF_XGRID) - grid_for_writing = np.array(lhapdf_interpolator(parameters)) + grid_for_writing = np.array(lhapdf_interpolator(parameters[:n_pdf_params])) replica_index = i + 1