Add verbose parameter to GCM auto assign function

When using the auto assignment function, it now provides more details into the fitting process when verbose is set to True (default).

Signed-off-by: Patrick Bloebaum <bloebp@amazon.com>

Author: bloebp

dowhy/gcm/auto.py

@@ -15,6 +15,7 @@
 from dowhy.gcm import config
 from dowhy.gcm.causal_mechanisms import AdditiveNoiseModel, ClassifierFCM
 from dowhy.gcm.causal_models import CAUSAL_MECHANISM, ProbabilisticCausalModel, validate_causal_model_assignment
+from dowhy.gcm.config import add_info_log_msg
 from dowhy.gcm.ml import (
     ClassificationModel,
     PredictionModel,
@@ -92,6 +93,7 @@ def assign_causal_mechanisms(
     based_on: pd.DataFrame,
     quality: AssignmentQuality = AssignmentQuality.GOOD,
     override_models: bool = False,
+    verbose: bool = True,
 ) -> None:
     """Automatically assigns appropriate causal models. If causal models are already assigned to nodes and
     override_models is set to False, this function only validates the assignments with respect to the graph structure.
@@ -129,41 +131,70 @@ def assign_causal_mechanisms(
 
     :return: None
     """
+    add_info_log_msg("----- Starting automatic model assignment -----", verbose)
     for node in causal_model.graph.nodes:
+        add_info_log_msg("--Node: %s" % node, verbose)
         if not override_models and CAUSAL_MECHANISM in causal_model.graph.nodes[node]:
+            add_info_log_msg(
+                "Node %s already has a model assigned and the override parameter is false. "
+                "Will skip this node." % node,
+                verbose,
+            )
             validate_causal_model_assignment(causal_model.graph, node)
             continue
-        assign_causal_mechanism_node(causal_model, node, based_on, quality)
+
+        assign_causal_mechanism_node(causal_model, node, based_on, quality, verbose)
+
+    add_info_log_msg("----- Finished automatic model assignment -----", verbose)
 
 
 def assign_causal_mechanism_node(
     causal_model: ProbabilisticCausalModel,
     node: str,
     based_on: pd.DataFrame,
     quality: AssignmentQuality = AssignmentQuality.GOOD,
+    verbose: bool = True,
 ) -> None:
     if is_root_node(causal_model.graph, node):
         causal_model.set_causal_mechanism(node, EmpiricalDistribution())
+
+        add_info_log_msg(
+            "Identified %s as a root node. Assigning %s to the node." % (node, causal_model.causal_mechanism(node)),
+            verbose,
+        )
     else:
+        add_info_log_msg("Identified %s as a non-root node." % node, verbose)
+
         prediction_model = select_model(
             based_on[get_ordered_predecessors(causal_model.graph, node)].to_numpy(),
             based_on[node].to_numpy(),
             quality,
+            verbose,
         )
 
         if isinstance(prediction_model, ClassificationModel):
             causal_model.set_causal_mechanism(node, ClassifierFCM(prediction_model))
+
+            add_info_log_msg("Assigning %s to the node %s." % (causal_model.causal_mechanism(node), node), verbose)
         else:
             causal_model.set_causal_mechanism(node, AdditiveNoiseModel(prediction_model))
 
+            add_info_log_msg(
+                "Assigning a %s to the node %s." % (causal_model.causal_mechanism(node), node),
+                verbose,
+            )
+
 
 def select_model(
-    X: np.ndarray, Y: np.ndarray, model_selection_quality: AssignmentQuality
+    X: np.ndarray, Y: np.ndarray, model_selection_quality: AssignmentQuality, verbose: bool
 ) -> Union[PredictionModel, ClassificationModel]:
+    add_info_log_msg("Looking for the best prediction model based on the %s." % model_selection_quality, verbose)
     if model_selection_quality == AssignmentQuality.BEST:
         try:
             from dowhy.gcm.ml.autogluon import AutoGluonClassifier, AutoGluonRegressor
 
+            add_info_log_msg("Using an autogluon model", verbose)
+
             if is_categorical(Y):
                 return AutoGluonClassifier()
             else:
@@ -190,8 +221,12 @@ def select_model(
         list_of_classifier += [partial(create_polynom_logistic_regression_classifier, max_iter=1000)]
 
     if is_categorical(Y):
+        add_info_log_msg("The node seems to be categorical. Checking classification models...", verbose)
+
         return find_best_model(list_of_classifier, X, Y, model_selection_splits=model_selection_splits)()
     else:
+        add_info_log_msg("The node seems to be continuous. Checking regression models....", verbose)
+
         return find_best_model(list_of_regressor, X, Y, model_selection_splits=model_selection_splits)()
 
 
@@ -253,19 +288,24 @@ def find_best_model(
     max_samples_per_split: int = 10000,
     model_selection_splits: int = 5,
     n_jobs: Optional[int] = None,
+    verbose: bool = True,
 ) -> Callable[[], PredictionModel]:
     n_jobs = config.default_n_jobs if n_jobs is None else n_jobs
 
     X, Y = shape_into_2d(X, Y)
 
     is_classification_problem = isinstance(prediction_model_factories[0](), ClassificationModel)
 
+    metric_name = "given"
+
     if metric is None:
+        metric_name = "(negative) F1"
         if is_classification_problem:
             metric = lambda y_true, y_preds: -metrics.f1_score(
                 y_true, y_preds, average="macro", zero_division=0
             )  # Higher score is better
         else:
+            metric_name = "mean squared error (MSE)"
             metric = metrics.mean_squared_error
 
     labelBinarizer = None
@@ -301,5 +341,22 @@ def estimate_average_score(prediction_model_factory: Callable[[], PredictionMode
         delayed(estimate_average_score)(prediction_model_factory, int(random_seed))
         for prediction_model_factory, random_seed in zip(prediction_model_factories, random_seeds)
     )
+    sorted_results = sorted(zip(prediction_model_factories, average_metric_scores), key=lambda x: x[1])
+    best_model = sorted_results[0]
+
+    add_info_log_msg(
+        "Using %d splits and the %s metric. The results are:\n-%s"
+        % (
+            model_selection_splits,
+            metric_name,
+            "\n-".join(["%s: %s" % (str(result[0]()).replace("()", ""), str(result[1])) for result in sorted_results]),
+        ),
+        verbose,
+    )
+    add_info_log_msg(
+        "Based on this, selecting %s as the best model to minimize the %s metric."
+        % (str(best_model[0]()).replace("()", ""), metric_name),
+        verbose,
+    )
 
-    return sorted(zip(prediction_model_factories, average_metric_scores), key=lambda x: x[1])[0][0]
+    return best_model[0]

dowhy/gcm/causal_mechanisms.py

@@ -10,7 +10,7 @@
 import numpy as np
 
 from dowhy.gcm.ml import ClassificationModel, PredictionModel
-from dowhy.gcm.ml.regression import InvertibleFunction
+from dowhy.gcm.ml.regression import InvertibleFunction, SklearnRegressionModel
 from dowhy.gcm.util.general import is_categorical, shape_into_2d
 
 
@@ -155,9 +155,14 @@ def evaluate(self, parent_samples: np.ndarray, noise_samples: np.ndarray) -> np.
         return self._invertible_function.evaluate(predictions + noise_samples)
 
     def __str__(self) -> str:
+        if isinstance(self._prediction_model, SklearnRegressionModel):
+            prediction_model_string = self._prediction_model.sklearn_model.__class__.__name__
+        else:
+            prediction_model_string = self._prediction_model.__class__.__name__
+
         return "%s with %s and an %s" % (
             self.__class__.__name__,
-            self._prediction_model.__class__.__name__,
+            prediction_model_string,
             self._invertible_function.__class__.__name__,
         )
 
@@ -207,6 +212,14 @@ def __init__(self, prediction_model: PredictionModel, noise_model: Optional[Stoc
     def clone(self):
         return AdditiveNoiseModel(prediction_model=self.prediction_model.clone(), noise_model=self.noise_model.clone())
 
+    def __str__(self) -> str:
+        if isinstance(self._prediction_model, SklearnRegressionModel):
+            prediction_model_string = self._prediction_model.sklearn_model.__class__.__name__
+        else:
+            prediction_model_string = self._prediction_model.__class__.__name__
+
+        return "AdditiveNoiseModel using %s" % prediction_model_string
+
 
 class ProbabilityEstimatorModel(ABC):
     @abstractmethod
@@ -291,3 +304,6 @@ def get_class_names(self, class_indices: np.ndarray) -> List[str]:
     @property
     def classifier_model(self) -> ClassificationModel:
         return self._classifier_model
+
+    def __repr__(self):
+        return "Classifier FCM based on %s" % self.classifier_model

dowhy/gcm/confidence_intervals_cms.py

@@ -53,6 +53,7 @@ def fit_and_compute(
     bootstrap_training_data: pd.DataFrame,
     bootstrap_data_subset_size_fraction: float = 0.75,
     auto_assign_quality: Optional[auto.AssignmentQuality] = None,
+    auto_assign_verbose: bool = False,
     *args,
     **kwargs,
 ):
@@ -78,6 +79,8 @@ def fit_and_compute(
     :param auto_assign_quality: If a quality is provided, then the existing causal mechanisms in the given causal_model
                                 are overridden by new automatically inferred mechanisms based on the provided
                                 AssignmentQuality. If None is given, the existing assigned mechanisms are used.
+    :param auto_assign_verbose: If True, the auto assignment logs additional information about the model selection
+                                process.
     :param args: Args passed through verbatim to the causal queries.
     :param kwargs: Keyword args passed through verbatim to the causal queries.
     :return: A tuple containing (1) the median of causal query results and (2) the confidence intervals.
@@ -94,7 +97,9 @@ def snapshot():
         ]
 
         if auto_assign_quality is not None:
-            auto.assign_causal_mechanisms(causal_model_copy, sampled_data, auto_assign_quality, override_models=True)
+            auto.assign_causal_mechanisms(
+                causal_model_copy, sampled_data, auto_assign_quality, override_models=True, verbose=auto_assign_verbose
+            )
 
         fit(causal_model_copy, sampled_data)
         return f(causal_model_copy, *args, **kwargs)

dowhy/gcm/config.py

@@ -1,6 +1,13 @@
+import logging
+import sys
+
 show_progress_bars = True
 default_n_jobs = -1
 
+_logger = logging.getLogger(__name__)
+_logger.setLevel(logging.INFO)
+_logger.addHandler(logging.StreamHandler(sys.stdout))
+
 
 def enable_progress_bars():
     global show_progress_bars
@@ -15,3 +22,8 @@ def disable_progress_bars():
 def set_default_n_jobs(n_jobs: int) -> None:
     global default_n_jobs
     default_n_jobs = n_jobs
+
+
+def add_info_log_msg(msg: str, verbose: bool) -> None:
+    if verbose:
+        _logger.info(msg)

dowhy/gcm/distribution_change.py

@@ -101,6 +101,7 @@ def distribution_change(
     mechanism_change_test_significance_level: float = 0.05,
     mechanism_change_test_fdr_control_method: Optional[str] = "fdr_bh",
     auto_assignment_quality: Optional[AssignmentQuality] = None,
+    auto_assignment_verbose: bool = False,
     return_additional_info: bool = False,
     shapley_config: Optional[ShapleyConfig] = None,
     graph_factory: Callable[[Any], DirectedGraph] = nx.DiGraph,
@@ -139,6 +140,8 @@ def distribution_change(
                                     old and new graph. However, they are re-fitted on the given data.
                                     If set to a valid assignment quality, new models are automatically assigned to the
                                     old and new graph based on the respective data.
+    :param auto_assignment_verbose: If True, the auto assignment logs additional information about the model selection
+                                    process.
     :param return_additional_info: If set to True, three additional items are returned: a dictionary indicating
                                    whether each node's mechanism changed, the causal DAG whose causal models are
                                    learned from old data, and the causal DAG whose causal models are learned from new
@@ -160,7 +163,13 @@ def distribution_change(
     if auto_assignment_quality is None:
         clone_causal_models(causal_model.graph, causal_model_old.graph)
     else:
-        assign_causal_mechanisms(causal_model_old, old_data, override_models=True, quality=auto_assignment_quality)
+        assign_causal_mechanisms(
+            causal_model_old,
+            old_data,
+            override_models=True,
+            quality=auto_assignment_quality,
+            verbose=auto_assignment_verbose,
+        )
     invariant_nodes = list(set(invariant_nodes).intersection(set(causal_graph_old.nodes)))
     _remove_invariant_nodes(invariant_nodes, causal_model_old, old_data, auto_assignment_quality)
 
@@ -169,7 +178,13 @@ def distribution_change(
     if auto_assignment_quality is None:
         clone_causal_models(causal_graph_old, causal_model_new.graph)
     else:
-        assign_causal_mechanisms(causal_model_new, new_data, override_models=True, quality=auto_assignment_quality)
+        assign_causal_mechanisms(
+            causal_model_new,
+            new_data,
+            override_models=True,
+            quality=auto_assignment_quality,
+            verbose=auto_assignment_verbose,
+        )
 
     mechanism_changes = _fit_accounting_for_mechanism_change(
         causal_model_old,

dowhy/gcm/fitting_sampling.py

@@ -2,7 +2,6 @@
 
 Functions in this module should be considered experimental, meaning there might be breaking API changes in the future.
 """
-
 from typing import Any
 
 import networkx as nx

dowhy/gcm/influence.py

@@ -226,6 +226,7 @@ def intrinsic_causal_influence(
     num_samples_baseline: int = 2000,
     max_batch_size: int = 250,
     auto_assign_quality: auto.AssignmentQuality = auto.AssignmentQuality.GOOD,
+    auto_assign_verbose: bool = False,
     shapley_config: Optional[ShapleyConfig] = None,
 ) -> Dict[Any, float]:
     """Computes the causal contribution of each upstream noise term of the target node (including the noise of the
@@ -268,6 +269,8 @@ def intrinsic_causal_influence(
                            significant impact on the overall memory usage. If set to -1, all samples are used in one
                            batch.
     :param auto_assign_quality: Auto assign quality for the 'approx' prediction_model option.
+    :param auto_assign_verbose: If True, the auto assignment logs additional information about the model selection
+                                process if selected 'approx' for the prediction_model option.
     :param shapley_config: :class:`~dowhy.gcm.shapley.ShapleyConfig` for the Shapley estimator.
     :return: Intrinsic causal contribution of each ancestor node to the statistical property defined by the
              attribution_func of the target node.
@@ -292,6 +295,7 @@ def intrinsic_causal_influence(
         target_samples,
         auto_assign_quality,
         target_is_categorical,
+        auto_assign_verbose,
     )
 
     if attribution_func is None:
@@ -332,6 +336,7 @@ def intrinsic_causal_influence_sample(
     num_noise_feature_samples: int = 5000,
     max_batch_size: int = 100,
     auto_assign_quality: auto.AssignmentQuality = auto.AssignmentQuality.GOOD,
+    auto_assign_verbose: bool = False,
     shapley_config: Optional[ShapleyConfig] = None,
 ) -> List[Dict[Any, Any]]:
     """Estimates the intrinsic causal impact of upstream nodes on a specified target_node, using the provided
@@ -367,6 +372,8 @@ def intrinsic_causal_influence_sample(
     :param max_batch_size: Maximum batch size for estimating multiple predictions at once. This has a significant influence on the
                           overall memory usage. If set to -1, all samples are used in one batch.
     :param auto_assign_quality: Auto assign quality for the 'approx' prediction_model option.
+    :param auto_assign_verbose: If True, the auto assignment logs additional information about the model selection
+                                process if selected 'approx' for the prediction_model option.
     :param shapley_config: :class:`~dowhy.gcm.shapley.ShapleyConfig` for the Shapley estimator.
     :return: A list of dictionaries indicating the intrinsic causal influence of a node on the target for a particular
              sample. This is, each dictionary belongs to one baseline sample.
@@ -403,6 +410,7 @@ def intrinsic_causal_influence_sample(
         target_samples,
         auto_assign_quality,
         False,  # Currently only supports continues target since we need to reconstruct its noise term.
+        auto_assign_verbose,
     )
 
     shapley_vales = feature_relevance_sample(
@@ -465,6 +473,7 @@ def _get_icc_noise_function(
     target_samples: np.ndarray,
     auto_assign_quality: auto.AssignmentQuality,
     target_is_categorical: bool,
+    auto_assign_verbose: bool,
 ) -> Callable[[np.ndarray], np.ndarray]:
     if isinstance(prediction_model, str) and prediction_model not in ("approx", "exact"):
         raise ValueError(
@@ -478,7 +487,9 @@ def _get_icc_noise_function(
         return prediction_model.predict
 
     if prediction_model == "approx":
-        prediction_model = auto.select_model(noise_samples, target_samples, auto_assign_quality)
+        prediction_model = auto.select_model(
+            noise_samples, target_samples, auto_assign_quality, verbose=auto_assign_verbose
+        )
         prediction_model.fit(noise_samples, target_samples)
 
         if target_is_categorical:

dowhy/gcm/ml/regression.py

@@ -57,6 +57,9 @@ def clone(self):
         """
         return SklearnRegressionModel(sklearn_mdl=sklearn.clone(self._sklearn_mdl))
 
+    def __str__(self):
+        return str(self._sklearn_mdl)
+
 
 def create_linear_regressor_with_given_parameters(
     coefficients: np.ndarray, intercept: float = 0, **kwargs

dowhy/gcm/stochastic_models.py

@@ -202,6 +202,9 @@ def draw_samples(self, num_samples: int) -> np.ndarray:
     def clone(self):
         return EmpiricalDistribution()
 
+    def __str__(self):
+        return "Empirical Distribution"
+
 
 class BayesianGaussianMixtureDistribution(StochasticModel):
     def __init__(self) -> None: