kulprit 0.2.0__tar.gz → 0.3.0__tar.gz

{kulprit-0.2.0 → kulprit-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: kulprit
-Version: 0.2.0
+Version: 0.3.0
 Summary: Kullback-Leibler projections for Bayesian model selection.
 Maintainer-email: Osvaldo Martin <aloctavodia@gmail.com>, Tomás Capretto <tomicapretto@gmail.com>
 Description-Content-Type: text/markdown
@@ -18,7 +18,6 @@ Requires-Dist: arviz>=0.17.1
 Requires-Dist: bambi>=0.12.0
 Requires-Dist: scikit-learn>=1.0.2
 Requires-Dist: numba>=0.56.0
-Requires-Dist: preliz>=0.4.1
 Project-URL: source, https://github.com/bambinos/kulprit
 Project-URL: tracker, https://github.com/bambinos/kulprit/issues
 

{kulprit-0.2.0 → kulprit-0.3.0}/kulprit/__init__.py

@@ -4,11 +4,11 @@ Kulprit.
 Kullback-Leibler projections for Bayesian model selection.
 """
 
-from kulprit.reference import ProjectionPredictive
+from kulprit.projector import ProjectionPredictive
 from kulprit.datasets import *
 
 
-__version__ = "0.2.0"
+__version__ = "0.3.0"
 
 
 __all__ = ["ProjectionPredictive"]

{kulprit-0.2.0 → kulprit-0.3.0}/kulprit/plots/plots.py

@@ -4,7 +4,9 @@ import matplotlib.pyplot as plt
 import numpy as np
 
 
-def plot_compare(cmp_df, legend=True, title=True, figsize=None, plot_kwargs=None):
+def plot_compare(
+    elpd_info, label_terms=None, legend=True, title=True, figsize=None, plot_kwargs=None
+):
     """
     Plot model comparison.
 
@@ -14,6 +16,8 @@ def plot_compare(cmp_df, legend=True, title=True, figsize=None, plot_kwargs=None
         Dataframe containing the comparison data. Should have columns
         `elpd_loo` and `elpd_diff` containing the ELPD values and the
         differences to the reference model.
+    label_terms : list
+        List of the labels for the submodels.
     legend : bool
         Flag for plotting the legend, default True.
     title : bool
@@ -22,29 +26,26 @@ def plot_compare(cmp_df, legend=True, title=True, figsize=None, plot_kwargs=None
         Figure size. If None it will be defined automatically.
     plot_kwargs : dict
     """
-
     if plot_kwargs is None:
         plot_kwargs = {}
 
     if figsize is None:
-        figsize = (len(cmp_df) - 1, 10)
+        figsize = (10, 4)
 
     figsize, ax_labelsize, _, xt_labelsize, linewidth, _ = _scale_fig_size(figsize, None, 1, 1)
 
-    xticks_pos, step = np.linspace(0, -1, ((cmp_df.shape[0]) * 2) - 2, retstep=True)
-    xticks_pos[1::2] = xticks_pos[1::2] - step * 1.5
-
-    labels = cmp_df.index.values[1:]
-    xticks_labels = [""] * len(xticks_pos)
-    xticks_labels[0] = labels[0]
-    xticks_labels[2::2] = labels[1:]
+    xticks_pos = np.linspace(0, 1, len(elpd_info) - 1)
+    xticks_num_labels = [value[0] for value in elpd_info[1:]]
+    xticks_name_labels = [f"\n\n{term}" for term in label_terms]
+    elpd_loo = [value[1] for value in elpd_info]
+    elpd_se = [value[2] for value in elpd_info]
 
     fig, axes = plt.subplots(1, figsize=figsize)
 
     axes.errorbar(
-        y=cmp_df["elpd_loo"][1:],
-        x=xticks_pos[::2],
-        yerr=cmp_df.se[1:],
+        y=elpd_loo[1:],
+        x=xticks_pos,
+        yerr=elpd_se[1:],
         label="Submodels",
         color=plot_kwargs.get("color_eldp", "k"),
         fmt=plot_kwargs.get("marker_eldp", "o"),
@@ -55,7 +56,7 @@ def plot_compare(cmp_df, legend=True, title=True, figsize=None, plot_kwargs=None
     )
 
     axes.axhline(
-        cmp_df["elpd_loo"].iloc[0],
+        elpd_loo[0],
         ls=plot_kwargs.get("ls_reference", "--"),
         color=plot_kwargs.get("color_ls_reference", "grey"),
         lw=linewidth,
@@ -63,16 +64,16 @@ def plot_compare(cmp_df, legend=True, title=True, figsize=None, plot_kwargs=None
     )
 
     axes.fill_between(
-        [-2, 1],
-        cmp_df["elpd_loo"].iloc[0] + cmp_df["se"].iloc[0],
-        cmp_df["elpd_loo"].iloc[0] - cmp_df["se"].iloc[0],
+        [-0.15, 1.15],
+        elpd_loo[0] + elpd_se[0],
+        elpd_loo[0] - elpd_se[0],
         alpha=0.1,
         color=plot_kwargs.get("color_ls_reference", "grey"),
     )
 
     if legend:
         fig.legend(
-            bbox_to_anchor=(0.9, 0.1),
+            bbox_to_anchor=(0.9, 0.3),
             loc="lower right",
             ncol=1,
             fontsize=ax_labelsize * 0.6,
@@ -84,15 +85,18 @@ def plot_compare(cmp_df, legend=True, title=True, figsize=None, plot_kwargs=None
             fontsize=ax_labelsize * 0.6,
         )
 
-    # remove double ticks
-    xticks_pos, xticks_labels = xticks_pos[::2], xticks_labels[::2]
+    sec0 = axes.secondary_xaxis(location=0)
+    sec0.set_xticks(xticks_pos, xticks_num_labels)
+    sec0.tick_params("x", length=0, labelsize=xt_labelsize * 0.6)
 
-    # set axes
-    axes.set_xticks(xticks_pos)
+    sec1 = axes.secondary_xaxis(location=0)
+    sec1.set_xticks(xticks_pos, xticks_name_labels, rotation=plot_kwargs.get("xlabel_rotation", 0))
+    sec1.tick_params("x", length=0, labelsize=xt_labelsize * 0.6)
+    sec1.set_xlabel("Submodels", fontsize=ax_labelsize * 0.6)
+
+    axes.set_xticks([])
     axes.set_ylabel("ELPD", fontsize=ax_labelsize * 0.6)
-    axes.set_xlabel("Submodel size", fontsize=ax_labelsize * 0.6)
-    axes.set_xticklabels(xticks_labels)
-    axes.set_xlim(-1 + step, 0 - step)
+    axes.set_xlim(-0.1, 1.1)
     axes.tick_params(labelsize=xt_labelsize * 0.6)
 
     return axes
@@ -100,47 +104,35 @@ def plot_compare(cmp_df, legend=True, title=True, figsize=None, plot_kwargs=None
 
 def plot_densities(
     model,
-    path,
     idata,
+    submodels,
     var_names=None,
-    submodels=None,
     include_reference=True,
-    labels="formula",
+    labels="size",
     kind="density",
     figsize=None,
     plot_kwargs=None,
 ):
     """Compare the projected posterior densities of the submodels"""
-
     if plot_kwargs is None:
         plot_kwargs = {}
-    plot_kwargs.setdefault("figsize", figsize)
 
     if kind not in ["density", "forest"]:
         raise ValueError("kind must be one of 'density' or 'forest'")
 
     # set default variable names to the reference model terms
     if not var_names:
-        var_names = list(
-            set(model.components[model.family.likelihood.parent].common_terms.keys())
-            - set([model.response_component.term.name])
-        )
+        var_names = [fvar.name for fvar in model.backend.model.free_RVs]
 
     if include_reference:
         data = [idata]
         l_labels = ["Reference"]
-        var_names.append(f"~{model.family.likelihood.parent}")
     else:
         data = []
         l_labels = []
 
-    if submodels is None:
-        submodels = path.values()
-    else:
-        submodels = [path[key] for key in submodels]
-
     if labels == "formula":
-        l_labels.extend([submodel.model.formula for submodel in submodels])
+        l_labels.extend([",".join(submodel.term_names) for submodel in submodels])
     else:
         l_labels.extend([submodel.size for submodel in submodels])
 
@@ -149,6 +141,7 @@ def plot_densities(
     if kind == "density":
         plot_kwargs.setdefault("outline", False)
         plot_kwargs.setdefault("shade", 0.4)
+        plot_kwargs.setdefault("figsize", figsize)
 
         axes = plot_density(
             data=data,
@@ -159,6 +152,8 @@ def plot_densities(
 
     elif kind == "forest":
         plot_kwargs.setdefault("combined", True)
+        plot_kwargs.setdefault("figsize", (10, 2 + len(var_names) ** 0.5))
+
         axes = plot_forest(
             data=data,
             model_names=l_labels,
@@ -167,13 +162,3 @@ def plot_densities(
         )
 
     return axes
-
-
-def align_yaxis(axes, v_1, ax2, v_2):
-    """adjust ax2 ylimit so that v2 in ax2 is aligned to v1 in axes"""
-    _, y_1 = axes.transData.transform((0, v_1))
-    _, y_2 = ax2.transData.transform((0, v_2))
-    inv = ax2.transData.inverted()
-    _, d_y = inv.transform((0, 0)) - inv.transform((0, y_1 - y_2))
-    miny, maxy = ax2.get_ylim()
-    ax2.set_ylim(miny + d_y, maxy + d_y)

kulprit-0.3.0/kulprit/projection/arviz_io.py

@@ -0,0 +1,44 @@
+"""Functions to interact with InferenceData objects"""
+import warnings
+from arviz import convert_to_dataset, extract, loo
+
+
+def get_observed_data(idata, response_name):
+    """Extract the observed data from the reference model."""
+    observed_data = {response_name: idata.observed_data.get(response_name).to_dict().get("data")}
+    return convert_to_dataset(observed_data), idata.observed_data.get(response_name).values
+
+
+def compute_pps(model, idata):
+    """Compute posterior predictive samples from the reference model."""
+    if "posterior_predictive" not in idata.groups():
+        model.predict(idata, kind="response", inplace=True)
+
+
+def get_pps(idata, response_name, num_samples):
+    """Extract samples posterior predictive samples from the reference model."""
+    pps = extract(
+        idata,
+        group="posterior_predictive",
+        var_names=[response_name],
+        num_samples=num_samples,
+        rng=1,
+    ).values.T
+    return pps
+
+
+def compute_loo(submodel=None, idata=None):
+    """Compute the PSIS-LOO-CV for a submodel or InferenceData object."""
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", message="Estimated shape parameter of Pareto distribution"
+        )
+        if submodel is not None:
+            elpd = loo(submodel.idata)
+            submodel.elpd_loo = elpd.elpd_loo
+            submodel.elpd_se = elpd.se
+
+        if idata is not None:
+            return loo(idata)
+
+    return None

kulprit-0.3.0/kulprit/projection/pymc_io.py

@@ -0,0 +1,87 @@
+"""Functions to interact with PyMC models"""
+import numpy as np
+
+from pymc import do, compute_log_likelihood
+from pymc.util import is_transformed_name, get_untransformed_name
+from pymc.pytensorf import join_nonshared_inputs
+
+try:
+    # Try to import compile_pymc as compile for older versions of pymc
+    from pymc.pytensorf import compile_pymc as compile_
+except ImportError:
+    # Fallback to importing compile for newer versions of pymc
+    from pymc.pytensorf import compile as compile_
+
+from pytensor import function
+from pytensor.tensor import matrix
+
+
+def compile_mllk(model):
+    """
+    Compile the log-likelihood function for the model to be able to condition on both
+    data and parameters.
+    """
+    obs_rvs = model.observed_RVs[0]
+    old_y_value = model.rvs_to_values[obs_rvs]
+    new_y_value = obs_rvs.type()
+    model.rvs_to_values[obs_rvs] = new_y_value
+
+    vars_ = model.value_vars
+    initial_point = model.initial_point()
+
+    [logp], raveled_inp = join_nonshared_inputs(
+        point=initial_point, outputs=[model.datalogp], inputs=vars_
+    )
+
+    rv_logp_fn = compile_([raveled_inp, new_y_value], logp)
+    rv_logp_fn.trust_input = True
+
+    def fmodel(params, obs):
+        return -rv_logp_fn(params, obs)
+
+    return fmodel, old_y_value, obs_rvs
+
+
+def compute_new_model(model, ref_var_info, all_terms, term_names):
+    """
+    Compute a new model by excluding the terms not in term_names.
+    """
+    # get all the terms not in term_names
+    exclude_terms = {term: 0 for term in set(all_terms) - set(term_names)}
+
+    for term in exclude_terms.keys():
+        shape = ref_var_info[term][0]
+        exclude_terms[term] = np.zeros(shape)
+
+    return do(model, exclude_terms)
+
+
+def compute_llk(idata, model):
+    """Compute log-likelihood for the submodel."""
+    return compute_log_likelihood(idata, model=model, progressbar=False, extend_inferencedata=False)
+
+
+def get_model_information(model):
+    """
+    Get the size and transformation of each variable in a PyMC model.
+    """
+
+    var_info = {}
+    initial_point = model.initial_point()
+    for v_var in model.value_vars:
+        name = v_var.name
+        if is_transformed_name(name):
+            name = get_untransformed_name(name)
+            x_var = matrix(f"{name}_transformed")
+            z_var = model.rvs_to_transforms[model.values_to_rvs[v_var]].backward(x_var)
+            transformation = function(inputs=[x_var], outputs=z_var)
+        else:
+            transformation = None
+
+        var_info[name] = (
+            initial_point[v_var.name].shape,
+            initial_point[v_var.name].size,
+            transformation,
+        )
+
+    return var_info

kulprit-0.3.0/kulprit/projection/search_strategies.py

@@ -0,0 +1,188 @@
+"""This module contains the search strategies"""
+
+import numpy as np
+from sklearn.linear_model import lasso_path
+from kulprit.projection.arviz_io import compute_loo
+
+
+def user_path(_project, path):
+    submodels = []
+    for term_names in path:
+        submodel = _project(term_names)
+        compute_loo(submodel=submodel)
+        submodels.append(submodel)
+    return submodels
+
+
+def forward_search(_project, ref_terms, max_terms, elpd_ref, early_stop):
+    """Method for performing forward search.
+
+    Parameters:
+    ----------
+    _project : function
+        A function that takes a list of term names and returns a submodel object.
+    ref_terms : list
+        A list of all terms in the model.
+    max_terms : int
+        The maximum number of terms to include in the submodel.
+    elpd_ref : float
+        The expected log pointwise predictive density of the reference model.
+    early_stop : str
+        The early stopping criterion. Either "mean" or "se".
+
+    Returns:
+    -------
+    List: A list of submodels, each containing the terms of the submodel and its ELPD.
+    """
+
+    # initial intercept-only subset
+    submodel_size = 0
+    term_names = []
+    submodel = _project(term_names)
+    compute_loo(submodel=submodel)
+    submodels = [submodel]
+
+    while submodel_size < max_terms:
+        # increment submodel size
+        submodel_size += 1
+
+        # get list of candidate submodels, project onto them, and compute
+        # their distances
+        candidates = _get_candidates(submodel.term_names, ref_terms)
+        projections = [_project(candidate) for candidate in candidates]
+
+        # identify the best candidate by loss (equivalent to KL min)
+        submodel = min(projections, key=lambda projection: projection.loss)
+
+        # compute loo for the best candidate and update inplace
+        compute_loo(submodel=submodel)
+
+        if _early_stopping(submodel, elpd_ref, early_stop):
+            submodels.append(submodel)
+            break
+
+        # add best candidate to the list of selected submodels
+        submodels.append(submodel)
+
+    return submodels
+
+
+def l1_search(_project, model, ref_terms, max_terms, elpd_ref, early_stop):
+    """Method for performing l1 search.
+
+    Parameters:
+    ----------
+    _project : function
+        A function that takes a list of term names and returns a submodel object.
+    model : object
+        A Bambi model.
+    ref_terms : list
+        A list of all terms in the model.
+    max_terms : int
+        The maximum number of terms to include in the submodel.
+    elpd_ref : float
+        The expected log pointwise predictive density of the reference model.
+    early_stop : str
+        The early stopping criterion. Either "mean" or "se".
+
+    Returns:
+    -------
+    List: A list of submodels, each containing the terms of the submodel and its ELPD.
+    """
+
+    d_component = model.distributional_components[model.family.likelihood.parent]
+    X = np.column_stack([d_component.design.common[term] for term in ref_terms])
+    # XXX we need to make this more general  # pylint: disable=fixme
+    mean_param_name = list(model.family.link.keys())[0]
+    eta = model.family.link[mean_param_name].link(
+        model.components[model.family.likelihood.parent].design.response.design_matrix
+    )
+    # compute L1 path in the latent space
+    _, coef_path, *_ = lasso_path(X, eta)
+    cov_order = _first_non_zero_idx(coef_path)
+
+    # sort the covariates according to their L1 ordering
+    cov_lasso = dict(sorted(cov_order.items(), key=lambda item: item[1]))
+    sorted_covs = [ref_terms[k] for k in cov_lasso]
+
+    submodel_size = 0
+    submodels = []
+
+    while submodel_size < max_terms + 1:
+        term_names = sorted_covs[:submodel_size]
+        submodel = _project(term_names)
+
+        # compute loo for the best candidate and update inplace
+        compute_loo(submodel=submodel)
+
+        if _early_stopping(submodel, elpd_ref, early_stop):
+            submodels.append(submodel)
+            break
+
+        # add best candidate to the list of selected submodels
+        submodels.append(submodel)
+
+        submodel_size += 1
+    return submodels
+
+
+def _get_candidates(prev_subset, ref_terms):
+    """Method for extracting a list of all candidate submodels.
+
+    Parameters:
+    ----------
+    prev_subset : list
+        The terms of the previous submodel.
+    ref_terms : list
+        The terms of the reference model.
+
+    Returns:
+    -------
+        List: A list of lists, each containing the terms of all candidate submodels
+    """
+
+    candidate_additions = list(set(ref_terms).difference(prev_subset))
+    candidates = [prev_subset + [addition] for addition in candidate_additions]
+    return candidates
+
+
+def _first_non_zero_idx(arr):
+    """Find the index of the first non-zero element in each row of a matrix.
+
+    Parameters:
+    ----------
+    arr : np.ndarray
+        A matrix.
+
+    Returns:
+    -------
+    dict: Dictionary keyed by the row number where each value is the index of the first
+    non-zero element in that row.
+    """
+
+    # initialise dictionary of indices
+    idx_dict = {}
+
+    # loop through each row and find first non-zero element
+    for i, j in zip(*np.where(arr != 0)):
+        if i in idx_dict:
+            continue
+        idx_dict[i] = j
+
+    # identify which keys are missing and set their values to infinity
+    if len(idx_dict) < arr.shape[0]:
+        missing_keys = set(range(arr.shape[0])) - set(idx_dict.keys())
+        for key in missing_keys:
+            idx_dict[key] = np.inf
+
+    return idx_dict
+
+
+def _early_stopping(submodel, elpd_ref, early_stop):
+    if early_stop == "mean":
+        if elpd_ref.elpd_loo - submodel.elpd_loo <= 4:
+            return True
+    elif early_stop == "se":
+        if submodel.elpd_loo + submodel.elpd_se >= elpd_ref.elpd_loo:
+            return True
+    return False

kulprit-0.3.0/kulprit/projection/solver.py

@@ -0,0 +1,63 @@
+"""optimization module."""
+from arviz import from_dict
+import numpy as np
+from scipy.optimize import minimize
+
+
+def solve(model_log_likelihood, pps, initial_guess, var_info, tolerance):
+    """The primary projection method in the procedure.
+
+    Parameters:
+    ----------
+    model_log_likelihood: Callable
+        The log-likelihood function of the model
+    pps: array
+        The predictions of the reference model
+    initial_guess: array
+        The initial guess for the optimization
+    var_info: dict
+        The dictionary containing information about the size and transformation of the variables
+
+    Returns:
+    -------
+        new_idata: arviz.InferenceData
+        loss: float
+    """
+    num_samples = len(pps)
+    posterior_array = np.zeros((num_samples, len(initial_guess)))
+    posterior_dict = {}
+    objectives = []
+
+    opt = minimize(
+        model_log_likelihood,
+        args=(pps[0]),
+        x0=initial_guess,
+    )
+    initial_guess = opt.x
+
+    for idx, obs in enumerate(pps):
+        opt = minimize(
+            model_log_likelihood,
+            args=(obs),
+            x0=initial_guess,
+            tol=tolerance,
+            method="SLSQP",
+        )
+
+        posterior_array[idx] = opt.x
+        objectives.append(opt.fun)
+        initial_guess = opt.x
+
+    size = 0
+    for key, values in var_info.items():
+        shape, new_size, transformation = values
+        posterior_dict[key] = posterior_array[:, size : size + new_size].reshape(
+            1, num_samples, *shape
+        )
+        if transformation is not None:
+            posterior_dict[key] = transformation(posterior_dict[key])
+        size += new_size
+
+    new_idata = from_dict(posterior=posterior_dict)
+    loss = np.mean(objectives)
+    return new_idata, loss