Easy-start Models#

The contextualized.easy module contains the most popular and useful model types, which have an SKLearn-like API for minimal overhead. 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)[source]#

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)[source]#

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)[source]#

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)[source]#

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). Returned as lists of individual bootstraps if individual_preds is True.

Return type:

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

class ContextualGAMRegressor(**kwargs)[source]#

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)[source]#

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). Returned as lists of individual bootstraps if individual_preds is True.

Return type:

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

class ContextualizedCorrelationNetworks(**kwargs)[source]#

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]][source]#

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]][source]#

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, 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)[source]#

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]][source]#

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, 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]][source]#

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)[source]#

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]][source]#

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]][source]#

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]][source]#

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]]