Source code for contextualized.analysis.effects

"""
Utilities for plotting effects learned by Contextualized models.
"""
from typing import *

import numpy as np
import matplotlib.pyplot as plt

from contextualized.easy.wrappers import SKLearnWrapper


def simple_plot(
    x_vals: List[Union[float, int]],
    y_vals: List[Union[float, int]],
    **kwargs,
) -> None:
    """
    Simple plotting of y vs x with kwargs passed to matplotlib helpers.

    Args:
        x_vals: x values to plot
        y_vals: y values to plot
        **kwargs: kwargs passed to matplotlib helpers (fill_alpha, fill_color, y_lowers, y_uppers, x_label, y_label, x_ticks, x_ticklabels, y_ticks, y_ticklabels)

    Returns:
        None
    """
    plt.figure(figsize=kwargs.get("figsize", (8, 8)))
    if "y_lowers" in kwargs and "y_uppers" in kwargs:
        plt.fill_between(
            x_vals,
            np.squeeze(kwargs["y_lowers"]),
            np.squeeze(kwargs["y_uppers"]),
            alpha=kwargs.get("fill_alpha", 0.2),
            color=kwargs.get("fill_color", "blue"),
        )
    plt.plot(x_vals, y_vals)
    plt.xlabel(kwargs.get("x_label", "X"))
    plt.ylabel(kwargs.get("y_label", "Y"))
    if (
        kwargs.get("x_ticks", None) is not None
        and kwargs.get("x_ticklabels", None) is not None
    ):
        plt.xticks(kwargs["x_ticks"], kwargs["x_ticklabels"])
    if (
        kwargs.get("y_ticks", None) is not None
        and kwargs.get("y_ticklabels", None) is not None
    ):
        plt.yticks(kwargs["y_ticks"], kwargs["y_ticklabels"])
    plt.show()


def plot_effect(x_vals, y_means, y_lowers=None, y_uppers=None, **kwargs):
    """Plots a single effect."""
    min_val = np.min(y_means)
    y_means -= min_val
    if y_lowers is not None and y_uppers is not None:
        y_lowers -= min_val
        y_uppers -= min_val
    if kwargs.get("should_exponentiate", False):
        y_means = np.exp(y_means)
        if y_lowers is not None and y_uppers is not None:
            y_lowers = np.exp(y_lowers)
            y_uppers = np.exp(y_uppers)
    try:
        if "x_encoder" in kwargs and kwargs["x_encoder"] is not None:
            x_classes = kwargs["x_encoder"].classes_
            # Line up class values with centered values.
            x_ticks = np.array(list(range(len(x_classes))))
            if (
                kwargs.get("x_means", None) is not None
                and kwargs.get("x_stds", None) is not None
            ):
                x_ticks = (x_ticks - kwargs["x_means"]) / kwargs["x_stds"]
        else:
            x_ticks = None
            x_classes = None
    except:
        x_classes = None
        x_ticks = None

    if np.max(y_means) > kwargs.get("min_effect_size", 0.0):
        simple_plot(
            x_vals,
            y_means,
            x_label=kwargs.get("xlabel", "X"),
            y_label=kwargs.get("ylabel", "Odds Ratio of Outcome"),
            y_lowers=y_lowers,
            y_uppers=y_uppers,
            x_ticks=x_ticks,
            x_ticklabels=x_classes,
        )


def get_homogeneous_context_effects(
    model: SKLearnWrapper, C: np.ndarray, **kwargs
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Get the homogeneous (context-invariant) effects of context.

    Args:
        model (SKLearnWrapper): a fitted ``contextualized.easy`` model
        C: the context values to use to estimate the effects
        verbose (bool, optional): print progess. Default True.
        individual_preds (bool, optional): whether to use plot each bootstrap. Default True.
        C_vis (np.ndarray, optional): Context bins used to visualize context (n_vis, n_contexts). Default None to construct anew.
        n_vis (int, optional): Number of bins to use to visualize context. Default 1000.

    Returns:
        Tuple[np.ndarray, np.ndarray]:
            c_vis: the context values that were used to estimate the effects
            effects: array of effects, one for each context. Each homogeneous effect is a matrix of shape:
                (n_bootstraps, n_context_vals, n_outcomes).
    """
    if kwargs.get("verbose", True):
        print("Estimating Homogeneous Contextual Effects.")
    c_vis = maybe_make_c_vis(C, **kwargs)

    effects = []
    for j in range(C.shape[1]):
        c_j = np.zeros_like(c_vis)
        c_j[:, j] = c_vis[:, j]
        try:
            (_, mus) = model.predict_params(
                c_j, individual_preds=kwargs.get("individual_preds", True)
            )
        except ValueError:
            (_, mus) = model.predict_params(c_j)
        effects.append(mus)
    return c_vis, np.array(effects)


def get_homogeneous_predictor_effects(model, C, **kwargs):
    """
    Get the homogeneous (context-invariant) effects of predictors.
    :param model:
    :param C:

    returns:
        c_vis: the context values that were used to estimate the effects
        effects: np array of effects, one for each predictor. Each homogeneous effect is a matrix of shape:
            (n_bootstraps, n_outcomes).

    """
    if kwargs.get("verbose", True):
        print("Estimating Homogeneous Predictor Effects.")
    c_vis = maybe_make_c_vis(C, **kwargs)
    c_idx = 0
    try:
        (betas, _) = model.predict_params(
            c_vis, individual_preds=kwargs.get("individual_preds", True)
        )
        # bootstraps x C_vis x outcomes x predictors
        if len(betas.shape) == 4:
            c_idx = 1
    except ValueError:
        (betas, _) = model.predict_params(c_vis)
    betas = np.mean(
        betas, axis=c_idx
    )  # homogeneous predictor effect is context-invariant
    return c_vis, np.transpose(betas, (2, 0, 1))


def get_heterogeneous_predictor_effects(model, C, **kwargs):
    """
    Get the heterogeneous (context-variant) effects of predictors.
    :param model:
    :param C:

    returns:
        c_vis: the context values that were used to estimate the effects
        effects: np array of effects, one for each context x predictor pair.
            Each heterogeneous effect is a matrix of shape:
                (n_predictors, n_bootstraps, n_context_vals, n_outcomes).
    """
    if kwargs.get("verbose", True):
        print("Estimating Heterogeneous Predictor Effects.")
    c_vis = maybe_make_c_vis(C, **kwargs)

    effects = []
    for j in range(C.shape[1]):
        c_j = np.zeros_like(c_vis)
        c_j[:, j] = c_vis[:, j]
        c_idx = 0
        try:
            (betas, _) = model.predict_params(
                c_j, individual_preds=kwargs.get("individual_preds", True)
            )
            # bootstraps x C_vis x outcomes x predictors
            if len(betas.shape) == 4:
                c_idx = 1
        except ValueError:
            (betas, _) = model.predict_params(c_j)
        # Heterogeneous Effects are mean-centered wrt C
        effect = np.transpose(
            np.transpose(betas, (0, 2, 3, 1))
            - np.tile(
                np.expand_dims(np.mean(betas, axis=c_idx), -1),
                (1, 1, 1, betas.shape[1]),
            ),
            (0, 3, 1, 2),
        )
        effects.append(effect)
    effects = np.array(effects)
    if len(effects.shape) == 5:
        effects = np.transpose(
            effects, (0, 4, 1, 2, 3)
        )  # (n_contexts, n_predictors, n_bootstraps, n_context_vals, n_outcomes)
    else:
        effects = np.transpose(
            effects, (0, 3, 1, 2)
        )  # (n_contexts, n_predictors, n_context_vals, n_outcomes)
    return c_vis, effects


def plot_boolean_vars(names, y_mean, y_err, **kwargs):
    """
    Plots Boolean variables.
    """
    plt.figure(figsize=kwargs.get("figsize", (12, 8)))
    sorted_i = np.argsort(y_mean)
    if kwargs.get("classification", True):
        y_mean = np.exp(y_mean)
        y_err = np.exp(y_err)
    for counter, i in enumerate(sorted_i):
        plt.bar(
            counter,
            y_mean[i],
            width=0.5,
            color=kwargs.get("fill_color", "blue"),
            edgecolor=kwargs.get("edge_color", "black"),
            yerr=y_err,
        )
    plt.xticks(
        range(len(names)),
        np.array(names)[sorted_i],
        rotation=60,
        fontsize=kwargs.get("boolean_x_ticksize", 18),
        ha="right",
    )
    plt.ylabel(
        kwargs.get("ylabel", "Odds Ratio of Outcome"),
        fontsize=kwargs.get("ylabel_fontsize", 32),
    )
    plt.yticks(fontsize=kwargs.get("ytick_fontsize", 18))
    if kwargs.get("bool_figname", None) is not None:
        plt.savefig(kwargs.get("bool_figname"), dpi=300, bbox_inches="tight")
    else:
        plt.show()


[docs]def plot_homogeneous_context_effects( model: SKLearnWrapper, C: np.ndarray, **kwargs, ) -> None: """ Plot the direct effect of context on outcomes, disregarding other features. This context effect is homogeneous in that it is a static function of context (context-invariant). Args: model (SKLearnWrapper): a fitted ``contextualized.easy`` model C: the context values to use to estimate the effects verbose (bool, optional): print progess. Default True. individual_preds (bool, optional): whether to use plot each bootstrap. Default True. C_vis (np.ndarray, optional): Context bins used to visualize context (n_vis, n_contexts). Default None to construct anew. n_vis (int, optional): Number of bins to use to visualize context. Default 1000. lower_pct (int, optional): Lower percentile for bootstraps. Default 2.5. upper_pct (int, optional): Upper percentile for bootstraps. Default 97.5. classification (bool, optional): Whether to exponentiate the effects. Default True. C_encoders (List[sklearn.preprocessing.LabelEncoder], optional): encoders for each context. Default None. C_means (np.ndarray, optional): means for each context. Default None. C_stds (np.ndarray, optional): standard deviations for each context. Default None. xlabel_prefix (str, optional): prefix for x label. Default "". figname (str, optional): name of figure to save. Default None. Returns: None """ c_vis, effects = get_homogeneous_context_effects(model, C, **kwargs) # effects.shape is (n_context, n_bootstraps, n_context_vals, n_outcomes) for outcome in range(effects.shape[-1]): for j in range(effects.shape[0]): try: mus = effects[j, :, :, outcome] means = np.mean(mus, axis=0) lowers = np.percentile(mus, kwargs.get("lower_pct", 2.5), axis=0) uppers = np.percentile(mus, kwargs.get("upper_pct", 97.5), axis=0) except ValueError: mus = effects[j, :, outcome] means = mus # no bootstraps were provided. lowers, uppers = None, None if "C_encoders" in kwargs: encoder = kwargs["C_encoders"][j] else: encoder = None if "C_means" in kwargs: c_means = kwargs["C_means"][j] else: c_means = None if "C_stds" in kwargs: c_stds = kwargs["C_stds"][j] else: c_stds = None plot_effect( c_vis[:, j], means, lowers, uppers, should_exponentiate=kwargs.get("classification", True), x_encoder=encoder, x_means=c_means, x_stds=c_stds, xlabel=C.columns.tolist()[j], **kwargs, )
[docs]def plot_homogeneous_predictor_effects( model: SKLearnWrapper, C: np.ndarray, X: np.ndarray, **kwargs, ) -> None: """ Plot the effect of predictors on outcomes that do not change with context (homogeneous). Args: model (SKLearnWrapper): a fitted ``contextualized.easy`` model C: the context values to use to estimate the effects X: the predictor values to use to estimate the effects max_classes_for_discrete (int, optional): maximum number of classes to treat as discrete. Default 10. min_effect_size (float, optional): minimum effect size to plot. Default 0.003. ylabel (str, optional): y label for plot. Default "Influence of ". xlabel_prefix (str, optional): prefix for x label. Default "". X_names (List[str], optional): names of predictors. Default None. X_encoders (List[sklearn.preprocessing.LabelEncoder], optional): encoders for each predictor. Default None. X_means (np.ndarray, optional): means for each predictor. Default None. X_stds (np.ndarray, optional): standard deviations for each predictor. Default None. verbose (bool, optional): print progess. Default True. lower_pct (int, optional): Lower percentile for bootstraps. Default 2.5. upper_pct (int, optional): Upper percentile for bootstraps. Default 97.5. classification (bool, optional): Whether to exponentiate the effects. Default True. figname (str, optional): name of figure to save. Default None. Returns: None """ c_vis = np.zeros_like(C.values) x_vis = make_grid_mat(X.values, 1000) (betas, _) = model.predict_params( c_vis, individual_preds=True ) # bootstraps x C_vis x outcomes x predictors homogeneous_betas = np.mean(betas, axis=1) # bootstraps x outcomes x predictors for outcome in range(homogeneous_betas.shape[1]): betas = homogeneous_betas[:, outcome, :] # bootstraps x predictors my_avg_betas = np.mean(betas, axis=0) lowers = np.percentile(betas, kwargs.get("lower_pct", 2.5), axis=0) uppers = np.percentile(betas, kwargs.get("upper_pct", 97.5), axis=0) max_impacts = [] # Calculate the max impact of each effect. for k in range(my_avg_betas.shape[0]): effect_range = my_avg_betas[k] * np.ptp(x_vis[:, k]) max_impacts.append(effect_range) effects_by_desc_impact = np.argsort(max_impacts)[::-1] boolean_vars = [j for j in range(X.shape[-1]) if len(set(X.values[:, j])) == 2] if len(boolean_vars) > 0: plot_boolean_vars( [X.columns[j] for j in boolean_vars], [max_impacts[j] for j in boolean_vars], [np.max(uppers[j]) - max_impacts[j] for j in boolean_vars], **kwargs, ) for j in effects_by_desc_impact: if j in boolean_vars: continue means = my_avg_betas[j] * x_vis[:, j] my_lowers = lowers[j] * x_vis[:, j] my_uppers = uppers[j] * x_vis[:, j] if "X_encoders" in kwargs: encoder = kwargs["X_encoders"][j] else: encoder = None if "X_means" in kwargs: x_means = kwargs["X_means"][j] else: x_means = None if "X_stds" in kwargs: x_stds = kwargs["X_stds"][j] else: x_stds = None plot_effect( x_vis[:, j], means, my_lowers, my_uppers, should_exponentiate=kwargs.get("classification", True), x_encoder=encoder, x_means=x_means, x_stds=x_stds, xlabel=f"{kwargs.get('xlabel_prefix', '')} {X.columns[j]}", **kwargs, )
[docs]def plot_heterogeneous_predictor_effects(model, C, X, **kwargs): """ Plot how the effect of predictors on outcomes changes with context (heterogeneous). Args: model (SKLearnWrapper): a fitted ``contextualized.easy`` model C: the context values to use to estimate the effects X: the predictor values to use to estimate the effects max_classes_for_discrete (int, optional): maximum number of classes to treat as discrete. Default 10. min_effect_size (float, optional): minimum effect size to plot. Default 0.003. y_prefix (str, optional): y prefix for plot. Default "Influence of ". X_names (List[str], optional): names of predictors. Default None. verbose (bool, optional): print progess. Default True. individual_preds (bool, optional): whether to use plot each bootstrap. Default True. C_vis (np.ndarray, optional): Context bins used to visualize context (n_vis, n_contexts). Default None to construct anew. n_vis (int, optional): Number of bins to use to visualize context. Default 1000. lower_pct (int, optional): Lower percentile for bootstraps. Default 2.5. upper_pct (int, optional): Upper percentile for bootstraps. Default 97.5. classification (bool, optional): Whether to exponentiate the effects. Default True. C_encoders (List[sklearn.preprocessing.LabelEncoder], optional): encoders for each context. Default None. C_means (np.ndarray, optional): means for each context. Default None. C_stds (np.ndarray, optional): standard deviations for each context. Default None. xlabel_prefix (str, optional): prefix for x label. Default "". figname (str, optional): name of figure to save. Default None. Returns: None """ c_vis = maybe_make_c_vis(C, **kwargs) n_vis = c_vis.shape[0] # c_names = C.columns.tolist() for j in range(C.shape[1]): c_j = c_vis.copy() c_j[:, :j] = 0.0 c_j[:, j + 1 :] = 0.0 (models, _) = model.predict_params( c_j, individual_preds=True ) # n_bootstraps x n_vis x outcomes x predictors homogeneous_effects = np.mean( models, axis=1 ) # n_bootstraps x outcomes x predictors heterogeneous_effects = models.copy() for i in range(n_vis): heterogeneous_effects[:, i] -= homogeneous_effects # n_bootstraps x n_vis x outcomes x predictors for outcome in range(heterogeneous_effects.shape[2]): my_effects = heterogeneous_effects[ :, :, outcome, : ] # n_bootstraps x n_vis x predictors means = np.mean(my_effects, axis=0) # n_vis x predictors my_lowers = np.percentile(my_effects, kwargs.get("lower_pct", 2.5), axis=0) my_uppers = np.percentile(my_effects, kwargs.get("upper_pct", 97.5), axis=0) x_ticks = None x_ticklabels = None try: x_classes = kwargs["encoders"][j].classes_ if len(x_classes) <= kwargs.get("max_classes_for_discrete", 10): x_ticks = np.array(list(range(len(x_classes)))) if "c_means" in kwargs: x_ticks -= kwargs["c_means"][j] if "c_stds" in kwargs: x_ticks /= kwargs["c_stds"][j] x_ticklabels = x_classes except KeyError: pass for k in range(my_effects.shape[-1]): if np.max(heterogeneous_effects[:, k]) > kwargs.get( "min_effect_size", 0.0 ): simple_plot( c_vis[:, j], means[:, k], x_label=C.columns[j], y_label=f"{kwargs.get('y_prefix', 'Influence of')} {X.columns[k]}", y_lowers=my_lowers[:, k], y_uppers=my_uppers[:, k], x_ticks=x_ticks, x_ticklabels=x_ticklabels, **kwargs, )
def make_grid_mat(observation_mat, n_vis): """ :param observation_mat: defines the domain for each feature. :param n_vis: returns a matrix of n_vis x n_features that can be used to visualize the effects of the features. """ ar_vis = np.zeros((n_vis, observation_mat.shape[1])) for j in range(observation_mat.shape[1]): ar_vis[:, j] = np.linspace( np.min(observation_mat[:, j]), np.max(observation_mat[:, j]), n_vis ) return ar_vis def make_c_vis(C, n_vis): """ :param C: :param n_vis: returns a matrix of n_vis x n_contexts that can be used to visualize the effects of the context variables. """ if isinstance(C, np.ndarray): return make_grid_mat(C, n_vis) return make_grid_mat(C.values, n_vis) def maybe_make_c_vis(C, **kwargs): """ :param C: :param n_vis: returns a matrix of n_vis x n_contexts that can be used to visualize the effects of the context variables. if C_vis is supplied, then we use that instead. """ if kwargs.get("C_vis", None) is None: if kwargs.get("verbose", True): print( """Generating datapoints for visualization by assuming the encoder is an additive model and thus doesn't require sampling on a manifold. If the encoder has interactions, please supply C_vis so that we can visualize these effects on the correct data manifold.""" ) return make_c_vis(C, kwargs.get("n_vis", 1000)) return kwargs["C_vis"]