Models

The contextualized.easy module contains Contextualized models, which are implemented with a simple SKLearn-style import-fit-predict pattern. All models can be loaded directly from the module, e.g. from contextualized.easy import ContextualizedRegressor.

ContextualizedRegressor.ContextualizedRegressor	Contextualized Linear Regression quantifies context-varying linear relationships.
ContextualizedClassifier.ContextualizedClassifier	Contextualized Logistic Regression reveals context-dependent decisions and decision boundaries.
ContextualGAM.ContextualGAMRegressor	The Contextual GAM Regressor separates and interprets the effect of context in context-varying relationships, such as heterogeneous treatment effects.
ContextualGAM.ContextualGAMClassifier	The Contextual GAM Classifier separates and interprets the effect of context in context-varying decisions and classifiers, such as heterogeneous disease diagnoses.
ContextualizedNetworks.ContextualizedCorrelationNetworks	Contextualized Correlation Networks reveal context-varying feature correlations, interaction strengths, dependencies in feature groups.
ContextualizedNetworks.ContextualizedMarkovNetworks	Contextualized Markov Networks reveal context-varying feature dependencies, cliques, and modules.
ContextualizedNetworks.ContextualizedBayesianNetworks	Contextualized Bayesian Networks and Directed Acyclic Graphs (DAGs) reveal context-dependent causal relationships, effect sizes, and variable ordering.

class ContextualizedRegressor(**kwargs)

Contextualized Linear Regression quantifies context-varying linear relationships.

Parameters:

• n_bootstraps (int, optional) – Number of bootstraps to use. Defaults to 1.

• num_archetypes (int, optional) – Number of archetypes to use. Defaults to 0, which used the NaiveMetaModel. If > 0, uses archetypes in the ContextualizedMetaModel.

• encoder_type (str, optional) – Type of encoder to use (“mlp”, “ngam”, “linear”). Defaults to “mlp”.

• loss_fn (torch.nn.Module, optional) – Loss function. Defaults to LOSSES[“mse”].

• link_fn (torch.nn.Module, optional) – Link function. Defaults to LINK_FUNCTIONS[“identity”].

• alpha (float, optional) – Regularization strength. Defaults to 0.0.

• mu_ratio (float, optional) – Float in range (0.0, 1.0), governs how much the regularization applies to context-specific parameters or context-specific offsets. Defaults to 0.0.

• l1_ratio (float, optional) – Float in range (0.0, 1.0), governs how much the regularization penalizes l1 vs l2 parameter norms. Defaults to 0.0.

fit(*args, **kwargs) → None

Fit contextualized model to data.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• Y (np.ndarray, optional) – Target array of shape (N, n_targets). Defaults to None, where X will be used as targets such as in Contextualized Networks.

• max_epochs (int, optional) – Maximum number of epochs to train for. Defaults to 1.

• learning_rate (float, optional) – Learning rate for optimizer. Defaults to 1e-3.

• val_split (float, optional) – Proportion of data to use for validation and early stopping. Defaults to 0.2.

• n_bootstraps (int, optional) – Number of bootstraps to use. Defaults to 1.

• train_batch_size (int, optional) – Batch size for training. Defaults to 1.

• val_batch_size (int, optional) – Batch size for validation. Defaults to 16.

• test_batch_size (int, optional) – Batch size for testing. Defaults to 16.

• es_patience (int, optional) – Number of epochs to wait before early stopping. Defaults to 1.

• es_monitor (str, optional) – Metric to monitor for early stopping. Defaults to “val_loss”.

• es_mode (str, optional) – Mode for early stopping. Defaults to “min”.

• es_verbose (bool, optional) – Whether to print early stopping updates. Defaults to False.

predict(C: ndarray, X: ndarray, individual_preds: bool = False, **kwargs) → Union[ndarray, List[ndarray]]

Predict outcomes from context C and predictors X.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• individual_preds (bool, optional) – Whether to return individual predictions for each model. Defaults to False.

Returns:

The outcomes predicted by the context-specific models (n_samples, y_dim). Returned as lists of individual bootstraps if individual_preds is True.

Return type:

Union[np.ndarray, List[np.ndarray]]

predict_params(C: ndarray, individual_preds: bool = False, model_includes_mus: bool = True, **kwargs) → Union[ndarray, List[ndarray], Tuple[ndarray, ndarray], Tuple[List[ndarray], List[ndarray]]]

Predict context-specific model parameters from context C.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• individual_preds (bool, optional) – Whether to return individual model predictions for each bootstrap. Defaults to False, averaging across bootstraps.

• model_includes_mus (bool, optional) – Whether the model includes context-specific offsets (mu). Defaults to True.

Returns:

The parameters of the predicted context-specific models. Returned as lists of individual bootstraps if individual_preds is True, otherwise averages the bootstraps for a better estimate. If model_includes_mus is True, returns both coefficients and offsets as a tuple of (betas, mus). Otherwise, returns coefficients (betas) only. For model_includes_mus=True, ([betas], [mus]) if individual_preds is True, otherwise (betas, mus). For model_includes_mus=False, [betas] if individual_preds is True, otherwise betas. betas is shape (n_samples, x_dim, y_dim) or (n_samples, x_dim) if y_dim = 1. mus is shape (n_samples, y_dim) or (n_samples,) if y_dim = 1.

Return type:

Union[np.ndarray, List[np.ndarray], Tuple[np.ndarray, np.ndarray], Tuple[List[np.ndarray], List[np.ndarray]]

class ContextualizedClassifier(**kwargs)

Contextualized Logistic Regression reveals context-dependent decisions and decision boundaries. Implemented as a ContextualizedRegressor with logistic link function and binary cross-entropy loss.

Parameters:

• n_bootstraps (int, optional) – Number of bootstraps to use. Defaults to 1.

• num_archetypes (int, optional) – Number of archetypes to use. Defaults to 0, which used the NaiveMetaModel. If > 0, uses archetypes in the ContextualizedMetaModel.

• encoder_type (str, optional) – Type of encoder to use (“mlp”, “ngam”, “linear”). Defaults to “mlp”.

• alpha (float, optional) – Regularization strength. Defaults to 0.0.

• mu_ratio (float, optional) – Float in range (0.0, 1.0), governs how much the regularization applies to context-specific parameters or context-specific offsets. Defaults to 0.0.

• l1_ratio (float, optional) – Float in range (0.0, 1.0), governs how much the regularization penalizes l1 vs l2 parameter norms. Defaults to 0.0.

fit(*args, **kwargs) → None

Fit contextualized model to data.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• Y (np.ndarray, optional) – Target array of shape (N, n_targets). Defaults to None, where X will be used as targets such as in Contextualized Networks.

• max_epochs (int, optional) – Maximum number of epochs to train for. Defaults to 1.

• learning_rate (float, optional) – Learning rate for optimizer. Defaults to 1e-3.

• val_split (float, optional) – Proportion of data to use for validation and early stopping. Defaults to 0.2.

• n_bootstraps (int, optional) – Number of bootstraps to use. Defaults to 1.

• train_batch_size (int, optional) – Batch size for training. Defaults to 1.

• val_batch_size (int, optional) – Batch size for validation. Defaults to 16.

• test_batch_size (int, optional) – Batch size for testing. Defaults to 16.

• es_patience (int, optional) – Number of epochs to wait before early stopping. Defaults to 1.

• es_monitor (str, optional) – Metric to monitor for early stopping. Defaults to “val_loss”.

• es_mode (str, optional) – Mode for early stopping. Defaults to “min”.

• es_verbose (bool, optional) – Whether to print early stopping updates. Defaults to False.

predict(C, X, individual_preds=False, **kwargs)

Predict binary outcomes from context C and predictors X.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• individual_preds (bool, optional) – Whether to return individual predictions for each model. Defaults to False.

Returns:

The binary outcomes predicted by the context-specific models (n_samples, y_dim). Returned as lists of individual bootstraps if individual_preds is True.

Return type:

Union[np.ndarray, List[np.ndarray]]

predict_params(C: ndarray, individual_preds: bool = False, model_includes_mus: bool = True, **kwargs) → Union[ndarray, List[ndarray], Tuple[ndarray, ndarray], Tuple[List[ndarray], List[ndarray]]]

Predict context-specific model parameters from context C.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• individual_preds (bool, optional) – Whether to return individual model predictions for each bootstrap. Defaults to False, averaging across bootstraps.

• model_includes_mus (bool, optional) – Whether the model includes context-specific offsets (mu). Defaults to True.

Returns:

The parameters of the predicted context-specific models. Returned as lists of individual bootstraps if individual_preds is True, otherwise averages the bootstraps for a better estimate. If model_includes_mus is True, returns both coefficients and offsets as a tuple of (betas, mus). Otherwise, returns coefficients (betas) only. For model_includes_mus=True, ([betas], [mus]) if individual_preds is True, otherwise (betas, mus). For model_includes_mus=False, [betas] if individual_preds is True, otherwise betas. betas is shape (n_samples, x_dim, y_dim) or (n_samples, x_dim) if y_dim = 1. mus is shape (n_samples, y_dim) or (n_samples,) if y_dim = 1.

Return type:

Union[np.ndarray, List[np.ndarray], Tuple[np.ndarray, np.ndarray], Tuple[List[np.ndarray], List[np.ndarray]]

predict_proba(C, X, **kwargs)

Predict probabilities of outcomes from context C and predictors X.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• individual_preds (bool, optional) – Whether to return individual predictions for each model. Defaults to False.

Returns:

The outcome probabilities predicted by the context-specific models (n_samples, y_dim, 2). Returned as lists of individual bootstraps if individual_preds is True.

Return type:

Union[np.ndarray, List[np.ndarray]]

class ContextualGAMRegressor(**kwargs)

The Contextual GAM Regressor separates and interprets the effect of context in context-varying relationships, such as heterogeneous treatment effects. Implemented as a Contextual Generalized Additive Model with a linear regressor on top. Always uses a Neural Additive Model (“ngam”) encoder for interpretability. See this paper for more details.

Parameters:

• n_bootstraps (int, optional) – Number of bootstraps to use. Defaults to 1.

• num_archetypes (int, optional) – Number of archetypes to use. Defaults to 0, which used the NaiveMetaModel. If > 0, uses archetypes in the ContextualizedMetaModel.

• alpha (float, optional) – Regularization strength. Defaults to 0.0.

• mu_ratio (float, optional) – Float in range (0.0, 1.0), governs how much the regularization applies to context-specific parameters or context-specific offsets. Defaults to 0.0.

• l1_ratio (float, optional) – Float in range (0.0, 1.0), governs how much the regularization penalizes l1 vs l2 parameter norms. Defaults to 0.0.

fit(*args, **kwargs) → None

Fit contextualized model to data.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• Y (np.ndarray, optional) – Target array of shape (N, n_targets). Defaults to None, where X will be used as targets such as in Contextualized Networks.

• max_epochs (int, optional) – Maximum number of epochs to train for. Defaults to 1.

• learning_rate (float, optional) – Learning rate for optimizer. Defaults to 1e-3.

• val_split (float, optional) – Proportion of data to use for validation and early stopping. Defaults to 0.2.

• n_bootstraps (int, optional) – Number of bootstraps to use. Defaults to 1.

• train_batch_size (int, optional) – Batch size for training. Defaults to 1.

• val_batch_size (int, optional) – Batch size for validation. Defaults to 16.

• test_batch_size (int, optional) – Batch size for testing. Defaults to 16.

• es_patience (int, optional) – Number of epochs to wait before early stopping. Defaults to 1.

• es_monitor (str, optional) – Metric to monitor for early stopping. Defaults to “val_loss”.

• es_mode (str, optional) – Mode for early stopping. Defaults to “min”.

• es_verbose (bool, optional) – Whether to print early stopping updates. Defaults to False.

predict(C: ndarray, X: ndarray, individual_preds: bool = False, **kwargs) → Union[ndarray, List[ndarray]]

Predict outcomes from context C and predictors X.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• individual_preds (bool, optional) – Whether to return individual predictions for each model. Defaults to False.

Returns:

The outcomes predicted by the context-specific models (n_samples, y_dim). Returned as lists of individual bootstraps if individual_preds is True.

Return type:

Union[np.ndarray, List[np.ndarray]]

predict_params(C: ndarray, individual_preds: bool = False, model_includes_mus: bool = True, **kwargs) → Union[ndarray, List[ndarray], Tuple[ndarray, ndarray], Tuple[List[ndarray], List[ndarray]]]

Predict context-specific model parameters from context C.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• individual_preds (bool, optional) – Whether to return individual model predictions for each bootstrap. Defaults to False, averaging across bootstraps.

• model_includes_mus (bool, optional) – Whether the model includes context-specific offsets (mu). Defaults to True.

Returns:

The parameters of the predicted context-specific models. Returned as lists of individual bootstraps if individual_preds is True, otherwise averages the bootstraps for a better estimate. If model_includes_mus is True, returns both coefficients and offsets as a tuple of (betas, mus). Otherwise, returns coefficients (betas) only. For model_includes_mus=True, ([betas], [mus]) if individual_preds is True, otherwise (betas, mus). For model_includes_mus=False, [betas] if individual_preds is True, otherwise betas. betas is shape (n_samples, x_dim, y_dim) or (n_samples, x_dim) if y_dim = 1. mus is shape (n_samples, y_dim) or (n_samples,) if y_dim = 1.

Return type:

Union[np.ndarray, List[np.ndarray], Tuple[np.ndarray, np.ndarray], Tuple[List[np.ndarray], List[np.ndarray]]

class ContextualGAMClassifier(**kwargs)

The Contextual GAM Classifier separates and interprets the effect of context in context-varying decisions and classifiers, such as heterogeneous disease diagnoses. Implemented as a Contextual Generalized Additive Model with a classifier on top. Always uses a Neural Additive Model (“ngam”) encoder for interpretability. See this paper for more details.

Parameters:

• n_bootstraps (int, optional) – Number of bootstraps to use. Defaults to 1.

• num_archetypes (int, optional) – Number of archetypes to use. Defaults to 0, which used the NaiveMetaModel. If > 0, uses archetypes in the ContextualizedMetaModel.

• alpha (float, optional) – Regularization strength. Defaults to 0.0.

• mu_ratio (float, optional) – Float in range (0.0, 1.0), governs how much the regularization applies to context-specific parameters or context-specific offsets. Defaults to 0.0.

• l1_ratio (float, optional) – Float in range (0.0, 1.0), governs how much the regularization penalizes l1 vs l2 parameter norms. Defaults to 0.0.

fit(*args, **kwargs) → None

Fit contextualized model to data.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• Y (np.ndarray, optional) – Target array of shape (N, n_targets). Defaults to None, where X will be used as targets such as in Contextualized Networks.

• max_epochs (int, optional) – Maximum number of epochs to train for. Defaults to 1.

• learning_rate (float, optional) – Learning rate for optimizer. Defaults to 1e-3.

• val_split (float, optional) – Proportion of data to use for validation and early stopping. Defaults to 0.2.

• n_bootstraps (int, optional) – Number of bootstraps to use. Defaults to 1.

• train_batch_size (int, optional) – Batch size for training. Defaults to 1.

• val_batch_size (int, optional) – Batch size for validation. Defaults to 16.

• test_batch_size (int, optional) – Batch size for testing. Defaults to 16.

• es_patience (int, optional) – Number of epochs to wait before early stopping. Defaults to 1.

• es_monitor (str, optional) – Metric to monitor for early stopping. Defaults to “val_loss”.

• es_mode (str, optional) – Mode for early stopping. Defaults to “min”.

• es_verbose (bool, optional) – Whether to print early stopping updates. Defaults to False.

predict(C, X, individual_preds=False, **kwargs)

Predict binary outcomes from context C and predictors X.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• individual_preds (bool, optional) – Whether to return individual predictions for each model. Defaults to False.

Returns:

The binary outcomes predicted by the context-specific models (n_samples, y_dim). Returned as lists of individual bootstraps if individual_preds is True.

Return type:

Union[np.ndarray, List[np.ndarray]]

predict_params(C: ndarray, individual_preds: bool = False, model_includes_mus: bool = True, **kwargs) → Union[ndarray, List[ndarray], Tuple[ndarray, ndarray], Tuple[List[ndarray], List[ndarray]]]

Predict context-specific model parameters from context C.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• individual_preds (bool, optional) – Whether to return individual model predictions for each bootstrap. Defaults to False, averaging across bootstraps.

• model_includes_mus (bool, optional) – Whether the model includes context-specific offsets (mu). Defaults to True.

Returns:

The parameters of the predicted context-specific models. Returned as lists of individual bootstraps if individual_preds is True, otherwise averages the bootstraps for a better estimate. If model_includes_mus is True, returns both coefficients and offsets as a tuple of (betas, mus). Otherwise, returns coefficients (betas) only. For model_includes_mus=True, ([betas], [mus]) if individual_preds is True, otherwise (betas, mus). For model_includes_mus=False, [betas] if individual_preds is True, otherwise betas. betas is shape (n_samples, x_dim, y_dim) or (n_samples, x_dim) if y_dim = 1. mus is shape (n_samples, y_dim) or (n_samples,) if y_dim = 1.

Return type:

Union[np.ndarray, List[np.ndarray], Tuple[np.ndarray, np.ndarray], Tuple[List[np.ndarray], List[np.ndarray]]

predict_proba(C, X, **kwargs)

Predict probabilities of outcomes from context C and predictors X.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• individual_preds (bool, optional) – Whether to return individual predictions for each model. Defaults to False.

Returns:

The outcome probabilities predicted by the context-specific models (n_samples, y_dim, 2). Returned as lists of individual bootstraps if individual_preds is True.

Return type:

Union[np.ndarray, List[np.ndarray]]

class ContextualizedCorrelationNetworks(**kwargs)

Contextualized Correlation Networks reveal context-varying feature correlations, interaction strengths, dependencies in feature groups. Uses the Contextualized Networks model, see the paper for detailed estimation procedures.

Parameters:

• n_bootstraps (int, optional) – Number of bootstraps to use. Defaults to 1.

• num_archetypes (int, optional) – Number of archetypes to use. Defaults to 10. Always uses archetypes in the ContextualizedMetaModel.

• encoder_type (str, optional) – Type of encoder to use (“mlp”, “ngam”, “linear”). Defaults to “mlp”.

• alpha (float, optional) – Regularization strength. Defaults to 0.0.

• mu_ratio (float, optional) – Float in range (0.0, 1.0), governs how much the regularization applies to context-specific parameters or context-specific offsets. Defaults to 0.0.

• l1_ratio (float, optional) – Float in range (0.0, 1.0), governs how much the regularization penalizes l1 vs l2 parameter norms. Defaults to 0.0.

fit(*args, **kwargs) → None

Fit contextualized model to data.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• Y (np.ndarray, optional) – Target array of shape (N, n_targets). Defaults to None, where X will be used as targets such as in Contextualized Networks.

• max_epochs (int, optional) – Maximum number of epochs to train for. Defaults to 1.

• learning_rate (float, optional) – Learning rate for optimizer. Defaults to 1e-3.

• val_split (float, optional) – Proportion of data to use for validation and early stopping. Defaults to 0.2.

• n_bootstraps (int, optional) – Number of bootstraps to use. Defaults to 1.

• train_batch_size (int, optional) – Batch size for training. Defaults to 1.

• val_batch_size (int, optional) – Batch size for validation. Defaults to 16.

• test_batch_size (int, optional) – Batch size for testing. Defaults to 16.

• es_patience (int, optional) – Number of epochs to wait before early stopping. Defaults to 1.

• es_monitor (str, optional) – Metric to monitor for early stopping. Defaults to “val_loss”.

• es_mode (str, optional) – Mode for early stopping. Defaults to “min”.

• es_verbose (bool, optional) – Whether to print early stopping updates. Defaults to False.

measure_mses(C: ndarray, X: ndarray, individual_preds: bool = False) → Union[ndarray, List[ndarray]]

Measures mean-squared errors.

Parameters:

• C (np.ndarray) – Contextual features for each sample (n_samples, n_context_features)

• X (np.ndarray) – The data matrix (n_samples, n_features)

• individual_preds (bool, optional) – If True, returns the predictions for each bootstrap. Defaults to False.

Returns:

The mean-squared errors for each sample, or for each bootstrap if individual_preds is True (n_samples).

Return type:

Union[np.ndarray, List[np.ndarray]]

predict(C: ndarray, X: ndarray, individual_preds: bool = False, **kwargs) → Union[ndarray, List[ndarray]]

Predict outcomes from context C and predictors X.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• individual_preds (bool, optional) – Whether to return individual predictions for each model. Defaults to False.

Returns:

The outcomes predicted by the context-specific models (n_samples, y_dim). Returned as lists of individual bootstraps if individual_preds is True.

Return type:

Union[np.ndarray, List[np.ndarray]]

predict_X(C: ndarray, X: ndarray, individual_preds: bool = False, **kwargs) → Union[ndarray, List[ndarray]]

Reconstructs the data matrix based on predicted contextualized networks and the true data matrix. Useful for measuring reconstruction error or for imputation.

Parameters:

• C (np.ndarray) – Contextual features for each sample (n_samples, n_context_features)

• X (np.ndarray) – The data matrix (n_samples, n_features)

• individual_preds (bool, optional) – If True, returns the predictions for each bootstrap. Defaults to False.

• **kwargs – Keyword arguments for the Lightning trainer’s predict_y method.

Returns:

The predicted data matrix, or matrices for each bootstrap if individual_preds is True (n_samples, n_features).

Return type:

Union[np.ndarray, List[np.ndarray]]

predict_correlation(C: ndarray, individual_preds: bool = True, squared: bool = True) → Union[ndarray, List[ndarray]]

Predicts context-specific correlations between features.

Parameters:

• C (Numpy ndarray) – Contextual features for each sample (n_samples, n_context_features)

• individual_preds (bool, optional) – If True, returns the predictions for each bootstrap. Defaults to True.

• squared (bool, optional) – If True, returns the squared correlations. Defaults to True.

Returns:

The predicted context-specific correlation matrices, or matrices for each bootstrap if individual_preds is True (n_samples, n_features, n_features).

Return type:

Union[np.ndarray, List[np.ndarray]]

predict_networks(C: ndarray, with_offsets: bool = False, individual_preds: bool = False, **kwargs) → Union[ndarray, List[ndarray], Tuple[ndarray, ndarray], Tuple[List[ndarray], List[ndarray]]]

Predicts context-specific networks given contextual features.

Parameters:

• C (np.ndarray) – Contextual features for each sample (n_samples, n_context_features)

• with_offsets (bool, optional) – If True, returns both the network parameters and offsets. Defaults to False.

• individual_preds (bool, optional) – If True, returns the predictions for each bootstrap. Defaults to False.

Returns:

The predicted network parameters (and offsets if with_offsets is True). Returned as lists of individual bootstraps if individual_preds is True.

Return type:

Union[np.ndarray, List[np.ndarray], Tuple[np.ndarray, np.ndarray], Tuple[List[np.ndarray], List[np.ndarray]]]

predict_params(C: ndarray, individual_preds: bool = False, model_includes_mus: bool = True, **kwargs) → Union[ndarray, List[ndarray], Tuple[ndarray, ndarray], Tuple[List[ndarray], List[ndarray]]]

Predict context-specific model parameters from context C.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• individual_preds (bool, optional) – Whether to return individual model predictions for each bootstrap. Defaults to False, averaging across bootstraps.

• model_includes_mus (bool, optional) – Whether the model includes context-specific offsets (mu). Defaults to True.

Returns:

The parameters of the predicted context-specific models. Returned as lists of individual bootstraps if individual_preds is True, otherwise averages the bootstraps for a better estimate. If model_includes_mus is True, returns both coefficients and offsets as a tuple of (betas, mus). Otherwise, returns coefficients (betas) only. For model_includes_mus=True, ([betas], [mus]) if individual_preds is True, otherwise (betas, mus). For model_includes_mus=False, [betas] if individual_preds is True, otherwise betas. betas is shape (n_samples, x_dim, y_dim) or (n_samples, x_dim) if y_dim = 1. mus is shape (n_samples, y_dim) or (n_samples,) if y_dim = 1.

Return type:

Union[np.ndarray, List[np.ndarray], Tuple[np.ndarray, np.ndarray], Tuple[List[np.ndarray], List[np.ndarray]]

class ContextualizedMarkovNetworks(**kwargs)

Contextualized Markov Networks reveal context-varying feature dependencies, cliques, and modules. Implemented as Contextualized Gaussian Precision Matrices, directly interpretable as Markov Networks. Uses the Contextualized Networks model, see the paper for detailed estimation procedures.

Parameters:

• n_bootstraps (int, optional) – Number of bootstraps to use. Defaults to 1.

• num_archetypes (int, optional) – Number of archetypes to use. Defaults to 10. Always uses archetypes in the ContextualizedMetaModel.

• encoder_type (str, optional) – Type of encoder to use (“mlp”, “ngam”, “linear”). Defaults to “mlp”.

• alpha (float, optional) – Regularization strength. Defaults to 0.0.

• mu_ratio (float, optional) – Float in range (0.0, 1.0), governs how much the regularization applies to context-specific parameters or context-specific offsets. Defaults to 0.0.

• l1_ratio (float, optional) – Float in range (0.0, 1.0), governs how much the regularization penalizes l1 vs l2 parameter norms. Defaults to 0.0.

fit(*args, **kwargs) → None

Fit contextualized model to data.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• Y (np.ndarray, optional) – Target array of shape (N, n_targets). Defaults to None, where X will be used as targets such as in Contextualized Networks.

• max_epochs (int, optional) – Maximum number of epochs to train for. Defaults to 1.

• learning_rate (float, optional) – Learning rate for optimizer. Defaults to 1e-3.

• val_split (float, optional) – Proportion of data to use for validation and early stopping. Defaults to 0.2.

• n_bootstraps (int, optional) – Number of bootstraps to use. Defaults to 1.

• train_batch_size (int, optional) – Batch size for training. Defaults to 1.

• val_batch_size (int, optional) – Batch size for validation. Defaults to 16.

• test_batch_size (int, optional) – Batch size for testing. Defaults to 16.

• es_patience (int, optional) – Number of epochs to wait before early stopping. Defaults to 1.

• es_monitor (str, optional) – Metric to monitor for early stopping. Defaults to “val_loss”.

• es_mode (str, optional) – Mode for early stopping. Defaults to “min”.

• es_verbose (bool, optional) – Whether to print early stopping updates. Defaults to False.

measure_mses(C: ndarray, X: ndarray, individual_preds: bool = False) → Union[ndarray, List[ndarray]]

Measures mean-squared errors.

Parameters:

• C (np.ndarray) – Contextual features for each sample (n_samples, n_context_features)

• X (np.ndarray) – The data matrix (n_samples, n_features)

• individual_preds (bool, optional) – If True, returns the predictions for each bootstrap. Defaults to False.

Returns:

The mean-squared errors for each sample, or for each bootstrap if individual_preds is True (n_samples).

Return type:

Union[np.ndarray, List[np.ndarray]]

predict(C: ndarray, X: ndarray, individual_preds: bool = False, **kwargs) → Union[ndarray, List[ndarray]]

Predict outcomes from context C and predictors X.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• individual_preds (bool, optional) – Whether to return individual predictions for each model. Defaults to False.

Returns:

The outcomes predicted by the context-specific models (n_samples, y_dim). Returned as lists of individual bootstraps if individual_preds is True.

Return type:

Union[np.ndarray, List[np.ndarray]]

predict_X(C: ndarray, X: ndarray, individual_preds: bool = False, **kwargs) → Union[ndarray, List[ndarray]]

Reconstructs the data matrix based on predicted contextualized networks and the true data matrix. Useful for measuring reconstruction error or for imputation.

Parameters:

• C (np.ndarray) – Contextual features for each sample (n_samples, n_context_features)

• X (np.ndarray) – The data matrix (n_samples, n_features)

• individual_preds (bool, optional) – If True, returns the predictions for each bootstrap. Defaults to False.

• **kwargs – Keyword arguments for the Lightning trainer’s predict_y method.

Returns:

The predicted data matrix, or matrices for each bootstrap if individual_preds is True (n_samples, n_features).

Return type:

Union[np.ndarray, List[np.ndarray]]

predict_networks(C: ndarray, with_offsets: bool = False, individual_preds: bool = False, **kwargs) → Union[ndarray, List[ndarray], Tuple[ndarray, ndarray], Tuple[List[ndarray], List[ndarray]]]

Predicts context-specific networks given contextual features.

Parameters:

• C (np.ndarray) – Contextual features for each sample (n_samples, n_context_features)

• with_offsets (bool, optional) – If True, returns both the network parameters and offsets. Defaults to False.

• individual_preds (bool, optional) – If True, returns the predictions for each bootstrap. Defaults to False.

Returns:

The predicted network parameters (and offsets if with_offsets is True). Returned as lists of individual bootstraps if individual_preds is True.

Return type:

Union[np.ndarray, List[np.ndarray], Tuple[np.ndarray, np.ndarray], Tuple[List[np.ndarray], List[np.ndarray]]]

predict_params(C: ndarray, individual_preds: bool = False, model_includes_mus: bool = True, **kwargs) → Union[ndarray, List[ndarray], Tuple[ndarray, ndarray], Tuple[List[ndarray], List[ndarray]]]

Predict context-specific model parameters from context C.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• individual_preds (bool, optional) – Whether to return individual model predictions for each bootstrap. Defaults to False, averaging across bootstraps.

• model_includes_mus (bool, optional) – Whether the model includes context-specific offsets (mu). Defaults to True.

Returns:

The parameters of the predicted context-specific models. Returned as lists of individual bootstraps if individual_preds is True, otherwise averages the bootstraps for a better estimate. If model_includes_mus is True, returns both coefficients and offsets as a tuple of (betas, mus). Otherwise, returns coefficients (betas) only. For model_includes_mus=True, ([betas], [mus]) if individual_preds is True, otherwise (betas, mus). For model_includes_mus=False, [betas] if individual_preds is True, otherwise betas. betas is shape (n_samples, x_dim, y_dim) or (n_samples, x_dim) if y_dim = 1. mus is shape (n_samples, y_dim) or (n_samples,) if y_dim = 1.

Return type:

Union[np.ndarray, List[np.ndarray], Tuple[np.ndarray, np.ndarray], Tuple[List[np.ndarray], List[np.ndarray]]

predict_precisions(C: ndarray, individual_preds: bool = True) → Union[ndarray, List[ndarray]]

Predicts context-specific precision matrices. Can be converted to context-specific Markov networks by binarizing the networks and setting all non-zero entries to 1. Can be converted to context-specific covariance matrices by taking the inverse.

Parameters:

• C (np.ndarray) – Contextual features for each sample (n_samples, n_context_features)

• individual_preds (bool, optional) – If True, returns the predictions for each bootstrap. Defaults to True.

Returns:

The predicted context-specific Markov networks as precision matrices, or matrices for each bootstrap if individual_preds is True (n_samples, n_features, n_features).

Return type:

Union[np.ndarray, List[np.ndarray]]

class ContextualizedBayesianNetworks(**kwargs)

Contextualized Bayesian Networks and Directed Acyclic Graphs (DAGs) reveal context-dependent causal relationships, effect sizes, and variable ordering. Uses the NOTMAD model, see the paper for detailed estimation procedures.

Parameters:

• n_bootstraps (int, optional) – Number of bootstraps to use. Defaults to 1.

• num_archetypes (int, optional) – Number of archetypes to use. Defaults to 16. Always uses archetypes in the ContextualizedMetaModel.

• encoder_type (str, optional) – Type of encoder to use (“mlp”, “ngam”, “linear”). Defaults to “mlp”.

• archetype_dag_loss_type (str, optional) – The type of loss to use for the archetype loss. Defaults to “l1”.

• archetype_l1 (float, optional) – The strength of the l1 regularization for the archetype loss. Defaults to 0.0.

• archetype_dag_params (dict, optional) – Parameters for the archetype loss. Defaults to {“loss_type”: “l1”, “params”: {“alpha”: 0.0, “rho”: 0.0, “s”: 0.0, “tol”: 1e-4}}.

• archetype_dag_loss_params (dict, optional) – Parameters for the archetype loss. Defaults to {“alpha”: 0.0, “rho”: 0.0, “s”: 0.0, “tol”: 1e-4}.

• archetype_alpha (float, optional) – The strength of the alpha regularization for the archetype loss. Defaults to 0.0.

• archetype_rho (float, optional) – The strength of the rho regularization for the archetype loss. Defaults to 0.0.

• archetype_s (float, optional) – The strength of the s regularization for the archetype loss. Defaults to 0.0.

• archetype_tol (float, optional) – The tolerance for the archetype loss. Defaults to 1e-4.

• archetype_use_dynamic_alpha_rho (bool, optional) – Whether to use dynamic alpha and rho for the archetype loss. Defaults to False.

• init_mat (np.ndarray, optional) – The initial adjacency matrix for the archetype loss. Defaults to None.

• num_factors (int, optional) – The number of factors for the archetype loss. Defaults to 0.

• factor_mat_l1 (float, optional) – The strength of the l1 regularization for the factor matrix for the archetype loss. Defaults to 0.

• sample_specific_dag_loss_type (str, optional) – The type of loss to use for the sample-specific loss. Defaults to “l1”.

• sample_specific_alpha (float, optional) – The strength of the alpha regularization for the sample-specific loss. Defaults to 0.0.

• sample_specific_rho (float, optional) – The strength of the rho regularization for the sample-specific loss. Defaults to 0.0.

• sample_specific_s (float, optional) – The strength of the s regularization for the sample-specific loss. Defaults to 0.0.

• sample_specific_tol (float, optional) – The tolerance for the sample-specific loss. Defaults to 1e-4.

• sample_specific_use_dynamic_alpha_rho (bool, optional) – Whether to use dynamic alpha and rho for the sample-specific loss. Defaults to False.

fit(*args, **kwargs) → None

Fit contextualized model to data.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• Y (np.ndarray, optional) – Target array of shape (N, n_targets). Defaults to None, where X will be used as targets such as in Contextualized Networks.

• max_epochs (int, optional) – Maximum number of epochs to train for. Defaults to 1.

• learning_rate (float, optional) – Learning rate for optimizer. Defaults to 1e-3.

• val_split (float, optional) – Proportion of data to use for validation and early stopping. Defaults to 0.2.

• n_bootstraps (int, optional) – Number of bootstraps to use. Defaults to 1.

• train_batch_size (int, optional) – Batch size for training. Defaults to 1.

• val_batch_size (int, optional) – Batch size for validation. Defaults to 16.

• test_batch_size (int, optional) – Batch size for testing. Defaults to 16.

• es_patience (int, optional) – Number of epochs to wait before early stopping. Defaults to 1.

• es_monitor (str, optional) – Metric to monitor for early stopping. Defaults to “val_loss”.

• es_mode (str, optional) – Mode for early stopping. Defaults to “min”.

• es_verbose (bool, optional) – Whether to print early stopping updates. Defaults to False.

measure_mses(C: ndarray, X: ndarray, individual_preds: bool = False, **kwargs) → Union[ndarray, List[ndarray]]

Measures mean-squared errors.

Parameters:

• C (np.ndarray) – Contextual features for each sample (n_samples, n_context_features)

• X (np.ndarray) – The data matrix (n_samples, n_features)

• individual_preds (bool, optional) – If True, returns the predictions for each bootstrap. Defaults to False.

• **kwargs – Keyword arguments for the contextualized.dags.GraphTrainer’s predict_params method.

Returns:

The mean-squared errors for each sample, or for each bootstrap if individual_preds is True (n_samples).

Return type:

Union[np.ndarray, List[np.ndarray]]

predict(C: ndarray, X: ndarray, individual_preds: bool = False, **kwargs) → Union[ndarray, List[ndarray]]

Predict outcomes from context C and predictors X.

Parameters:

• C (np.ndarray) – Context array of shape (n_samples, n_context_features)

• X (np.ndarray) – Predictor array of shape (N, n_features)

• individual_preds (bool, optional) – Whether to return individual predictions for each model. Defaults to False.

Returns:

The outcomes predicted by the context-specific models (n_samples, y_dim). Returned as lists of individual bootstraps if individual_preds is True.

Return type:

Union[np.ndarray, List[np.ndarray]]

predict_X(C: ndarray, X: ndarray, individual_preds: bool = False, **kwargs) → Union[ndarray, List[ndarray]]

Reconstructs the data matrix based on predicted contextualized networks and the true data matrix. Useful for measuring reconstruction error or for imputation.

Parameters:

• C (np.ndarray) – Contextual features for each sample (n_samples, n_context_features)

• X (np.ndarray) – The data matrix (n_samples, n_features)

• individual_preds (bool, optional) – If True, returns the predictions for each bootstrap. Defaults to False.

• **kwargs – Keyword arguments for the Lightning trainer’s predict_y method.

Returns:

The predicted data matrix, or matrices for each bootstrap if individual_preds is True (n_samples, n_features).

Return type:

Union[np.ndarray, List[np.ndarray]]

predict_networks(C: ndarray, project_to_dag: bool = True, **kwargs) → Union[ndarray, List[ndarray]]

Predicts context-specific Bayesian networks.

Parameters:

• C (np.ndarray) – Contextual features for each sample (n_samples, n_context_features)

• project_to_dag (bool, optional) – If True, guarantees returned graphs are DAGs by trimming edges until acyclicity is satisified. Defaults to True.

• **kwargs – Keyword arguments for the contextualized.dags.GraphTrainer’s predict_params method.

Returns:

The linear coefficients of the predicted context-specific Bayesian network parameters (n_samples, n_features, n_features). Returned as lists of individual bootstraps if individual_preds is True.

Return type:

Union[np.ndarray, List[np.ndarray]]

predict_params(C: ndarray, **kwargs) → Union[ndarray, List[ndarray]]

Predicts context-specific Bayesian network parameters as linear coefficients in a linear structural equation model (SEM).

Parameters:

• C (np.ndarray) – Contextual features for each sample (n_samples, n_context_features)

• **kwargs – Keyword arguments for the contextualized.dags.GraphTrainer’s predict_params method.

Returns:

The linear coefficients of the predicted context-specific Bayesian network parameters (n_samples, n_features, n_features). Returned as lists of individual bootstraps if individual_preds is True.

Return type:

Union[np.ndarray, List[np.ndarray]]