Global Sensitivity Analysis for Categorical Features (#2357)

Summary:

Sobol sensitivity analysis and the sensitivity plots currently do not support categorical features. The plots in particular error out when categorical features are present because the sensitivity analysis (both first order and gradient-based).

This commit adds support for categorical features by
- introducing a uniformly random integer-valued point distribution for categorical and ordinal features in `SobolSensitivity`, in line with the usual integer-valued encoding,
-  ignoring the sign of the derivative-based sensitivity analysis of the categorical features, since the `CategoricalKernel` is non-differentiable, and the "direction" of categorical features is not well defined,
- adding a separate case for categorical features to the sensitivity analysis plot, stating that the categorical features "affect" but don't "increase" or "decrease" the metric.

Note that the results for the categorical features are still first order or total Sobol indices, so all the results in the plot are on the same scale.

Differential Revision: D56070326

ax/plot/feature_importances.py

@@ -12,6 +12,7 @@
 import numpy as np
 import pandas as pd
 import plotly.graph_objs as go
+from ax.core.parameter import ChoiceParameter
 from ax.exceptions.core import NoDataError
 from ax.modelbridge import ModelBridge
 from ax.plot.base import AxPlotConfig, AxPlotTypes
@@ -143,6 +144,14 @@ def plot_feature_importance_by_feature_plotly(
         }
     traces = []
     dropdown = []
+    categorical_features = []
+    if model is not None:
+        categorical_features = [
+            name
+            for name, par in model.model_space.parameters.items()
+            if isinstance(par, ChoiceParameter)
+        ]
+
     for i, metric_name in enumerate(sorted(sensitivity_values.keys())):
         importances = sensitivity_values[metric_name]
         factor_col = "Factor"
@@ -157,7 +166,11 @@ def plot_feature_importance_by_feature_plotly(
                         factor_col: factor,
                         importance_col: np.asarray(importance)[0],
                         importance_col_se: np.asarray(importance)[2],
-                        sign_col: np.sign(np.asarray(importance)[0]).astype(int),
+                        sign_col: (
+                            0
+                            if factor in categorical_features
+                            else 2 * (np.asarray(importance)[0] >= 0).astype(int) - 1
+                        ),
                     }
                     for factor, importance in importances.items()
                 ]
@@ -172,7 +185,11 @@ def plot_feature_importance_by_feature_plotly(
                     {
                         factor_col: factor,
                         importance_col: importance,
-                        sign_col: np.sign(importance).astype(int),
+                        sign_col: (
+                            0
+                            if factor in categorical_features
+                            else 2 * (importance >= 0).astype(int) - 1
+                        ),
                     }
                     for factor, importance in importances.items()
                 ]
@@ -183,9 +200,13 @@ def plot_feature_importance_by_feature_plotly(
         if relative:
             df[importance_col] = df[importance_col].div(df[importance_col].sum())
 
-        colors = {-1: "darkorange", 1: "steelblue"}
-        names = {-1: "Decreases metric", 1: "Increases metric"}
-        legend_counter = {-1: 0, 1: 0}
+        colors = {-1: "darkorange", 0: "gray", 1: "steelblue"}
+        names = {
+            -1: "Decreases metric",
+            0: "Affects metric (categorical choice)",
+            1: "Increases metric",
+        }
+        legend_counter = {-1: 0, 0: 0, 1: 0}
         all_positive = all(df[sign_col] >= 0)
         for _, row in df.iterrows():
             traces.append(

ax/utils/sensitivity/sobol_measures.py

@@ -6,6 +6,7 @@
 # pyre-strict
 
 from copy import deepcopy
+from functools import partial
 
 from typing import Any, Callable, Dict, List, Optional, Union
 
@@ -36,6 +37,7 @@ def __init__(
         second_order: bool = False,
         num_bootstrap_samples: int = 1,
         bootstrap_array: bool = False,
+        discrete_features: Optional[List[int]] = None,
     ) -> None:
         r"""Computes three types of Sobol indices:
         first order indices, total indices and second order indices (if specified ).
@@ -51,6 +53,9 @@ def __init__(
                 to be specified.
             bootstrap_array: If true, all the num_bootstrap_samples extimated indices
                 are returned instead of their mean and Var.
+            discrete_features: If specified, the inputs associated with the indices in
+                this list are generated using an integer-valued uniform distribution,
+                rather than the default (pseudo-)random continuous uniform distribution.
         """
         self.input_function = input_function
         self.dim: int = bounds.shape[-1]
@@ -71,6 +76,16 @@ def __init__(
         else:
             self.A = unnormalize(torch.rand(num_mc_samples, self.dim), bounds=bounds)
             self.B = unnormalize(torch.rand(num_mc_samples, self.dim), bounds=bounds)
+
+        # uniform integral distribution for discrete features
+        if discrete_features is not None:
+            all_low = bounds[0, discrete_features].to(dtype=torch.int).tolist()
+            all_high = (bounds[1, discrete_features]).to(dtype=torch.int).tolist()
+            for i, low, high in zip(discrete_features, all_low, all_high):
+                randint = partial(torch.randint, low=low, high=high + 1)
+                self.A[:, i] = randint(size=self.A.shape[:-1])
+                self.B[:, i] = randint(size=self.B.shape[:-1])
+
         # pyre-fixme[4]: Attribute must be annotated.
         self.A_B_ABi = self.generate_all_input_matrix().to(torch.double)
 
@@ -395,6 +410,7 @@ def __init__(
             [torch.Tensor, torch.Tensor], torch.Tensor
         ] = GaussianLinkMean,
         mini_batch_size: int = 128,
+        discrete_features: Optional[List[int]] = None,
     ) -> None:
         r"""Computes three types of Sobol indices:
         first order indices, total indices and second order indices (if specified ).
@@ -411,6 +427,9 @@ def __init__(
                 to be specified.
             mini_batch_size: The size of the mini-batches used while evaluating the
                 model posterior. Increasing this will increase the memory usage.
+            discrete_features: If specified, the inputs associated with the indices in
+                this list are generated using an integer-valued uniform distribution,
+                rather than the default (pseudo-)random continuous uniform distribution.
         """
         self.model = model
         self.second_order = second_order
@@ -438,6 +457,7 @@ def input_function(x: Tensor) -> Tensor:
             second_order=self.second_order,
             input_qmc=self.input_qmc,
             num_bootstrap_samples=self.num_bootstrap_samples,
+            discrete_features=discrete_features,
         )
         self.sensitivity.evalute_function()
 
@@ -486,6 +506,7 @@ def __init__(
         input_qmc: bool = False,
         gp_sample_qmc: bool = False,
         num_bootstrap_samples: int = 1,
+        discrete_features: Optional[List[int]] = None,
     ) -> None:
         r"""Computes three types of Sobol indices:
         first order indices, total indices and second order indices (if specified ).
@@ -502,6 +523,9 @@ def __init__(
                 SobolQMCNormalSampler.
             num_bootstrap_samples: If bootstrap is true, the number of bootstraps has
                 to be specified.
+            discrete_features: If specified, the inputs associated with the indices in
+                this list are generated using an integer-valued uniform distribution,
+                rather than the default (pseudo-)random continuous uniform distribution.
         """
         self.model = model
         self.second_order = second_order
@@ -519,6 +543,7 @@ def __init__(
             input_qmc=self.input_qmc,
             num_bootstrap_samples=self.num_bootstrap_samples,
             bootstrap_array=True,
+            discrete_features=discrete_features,
         )
         # TODO: Ideally, we would reduce the memory consumption here as well
         # but this is a tricky since it uses joint posterior sampling.
@@ -717,6 +742,7 @@ def compute_sobol_indices_from_model_list(
     model_list: List[Model],
     bounds: Tensor,
     order: str = "first",
+    discrete_features: Optional[List[int]] = None,
     **sobol_kwargs: Any,
 ) -> Tensor:
     """
@@ -728,6 +754,9 @@ def compute_sobol_indices_from_model_list(
         bounds: A 2 x d Tensor of lower and upper bounds of the domain of the models.
         order: A string specifying the order of the Sobol indices to be computed.
             Supports "first" and "total" and defaults to "first".
+        discrete_features: If specified, the inputs associated with the indices in
+            this list are generated using an integer-valued uniform distribution,
+            rather than the default (pseudo-)random continuous uniform distribution.
         sobol_kwargs: keyword arguments passed on to SobolSensitivityGPMean.
 
     Returns:
@@ -739,6 +768,7 @@ def compute_sobol_indices_from_model_list(
         sens_class = SobolSensitivityGPMean(
             model=model,
             bounds=bounds,
+            discrete_features=discrete_features,
             **sobol_kwargs,
         )
         indices.append(method(sens_class))
@@ -789,6 +819,7 @@ def ax_parameter_sens(
         model_list=model_list,
         bounds=bounds,
         order=order,
+        discrete_features=digest.categorical_features + digest.ordinal_features,
         **sobol_kwargs,
     )
     if signed:
@@ -797,6 +828,14 @@ def ax_parameter_sens(
             bounds=bounds,
             **sobol_kwargs,
         )
+        # categorical features don't have a direction, so we set the derivative to 1.0
+        # in order not to zero our their sensitivity. We treat categorical features
+        # separately in the sensitivity analysis plot as well, to make clear that they
+        # are affecting the metric, but neither increasing nor decreasing. Note that the
+        # orginal variables have a well defined direction, so we do not need to treat
+        # them differently here.
+        for i in digest.categorical_features:
+            ind_deriv[:, i] = 1.0
         ind *= torch.sign(ind_deriv)
     return _array_with_string_indices_to_dict(
         rows=metrics, cols=digest.feature_names, A=ind.numpy()

ax/utils/sensitivity/tests/test_sensitivity.py

@@ -6,6 +6,7 @@
 # pyre-strict
 
 
+import copy
 import math
 from typing import cast
 from unittest.mock import patch, PropertyMock
@@ -270,35 +271,52 @@ def test_SobolGPMean(self) -> None:
                             )
         # Test with signed
         model_bridge = get_modelbridge(modular=True)
-        # Unsigned
+
+        # adding a categorical feature
+        cat_model_bridge = copy.deepcopy(model_bridge)
+        digest = cat_model_bridge.model.search_space_digest
+        digest.categorical_features = [0]
+
         sobol_kwargs = {"input_qmc": True, "num_mc_samples": 10}
-        ind_dict = ax_parameter_sens(
-            model_bridge,  # pyre-ignore
-            order="total",
-            signed=False,
-            **sobol_kwargs,  # pyre-ignore
-        )
-        ind_deriv = compute_derivatives_from_model_list(
-            model_list=[model_bridge.model.surrogate.model],
-            bounds=torch.tensor(model_bridge.model.search_space_digest.bounds).T,
-            **sobol_kwargs,
-        )
-        ind_dict_signed = ax_parameter_sens(
-            model_bridge,  # pyre-ignore
-            order="total",
-            # signed=True
-            **sobol_kwargs,  # pyre-ignore
-        )
-        for i, pname in enumerate(["x1", "x2"]):
-            self.assertEqual(
-                torch.sign(ind_deriv[0, i]).item(),
-                math.copysign(1, ind_dict_signed["branin"][pname]),
-            )
-            self.assertAlmostEqual(
-                (torch.sign(ind_deriv[0, i]) * ind_dict["branin"][pname]).item(),
-                ind_dict_signed["branin"][pname],
-            )  # signed
-            self.assertTrue(ind_dict["branin"][pname] >= 0)  # unsigned
+        seed = 1234
+        for bridge in [model_bridge, cat_model_bridge]:
+            with self.subTest(model_bridge=bridge):
+                torch.manual_seed(seed)
+                # Unsigned
+                ind_dict = ax_parameter_sens(
+                    bridge,  # pyre-ignore
+                    order="total",
+                    signed=False,
+                    **sobol_kwargs,  # pyre-ignore
+                )
+                ind_deriv = compute_derivatives_from_model_list(
+                    model_list=[bridge.model.surrogate.model],
+                    bounds=torch.tensor(bridge.model.search_space_digest.bounds).T,
+                    **sobol_kwargs,
+                )
+                torch.manual_seed(seed)  # reset seed to keep discrete features the same
+                cat_indices = bridge.model.search_space_digest.categorical_features
+                ind_dict_signed = ax_parameter_sens(
+                    bridge,  # pyre-ignore
+                    order="total",
+                    # signed=True
+                    **sobol_kwargs,  # pyre-ignore
+                )
+                for i, pname in enumerate(["x1", "x2"]):
+                    if i in cat_indices:  # special case for categorical features
+                        expected_sign = 1
+                    else:
+                        expected_sign = torch.sign(ind_deriv[0, i]).item()
+
+                    self.assertEqual(
+                        expected_sign,
+                        math.copysign(1, ind_dict_signed["branin"][pname]),
+                    )
+                    self.assertAlmostEqual(
+                        (expected_sign * ind_dict["branin"][pname]).item(),
+                        ind_dict_signed["branin"][pname],
+                    )  # signed
+                    self.assertTrue(ind_dict["branin"][pname] >= 0)  # unsigned
 
     def test_SobolGPSampling(self) -> None:
         bounds = torch.tensor([(0.0, 1.0) for _ in range(2)]).t()
@@ -349,6 +367,32 @@ def test_SobolGPSampling(self) -> None:
             total_order = sensitivity_sampling.total_order_indices()
             second_order = sensitivity_sampling.second_order_indices()
 
+        discrete_feature = 0
+        sensitivity_sampling_discrete = SobolSensitivityGPSampling(
+            self.model,
+            num_mc_samples=10,
+            num_gp_samples=10,
+            bounds=bounds,
+            second_order=True,
+            discrete_features=[discrete_feature],
+        )
+        sens = sensitivity_sampling_discrete.sensitivity
+        A = sens.A
+        B = sens.B
+        Arnd = A.round()
+        Brnd = B.round()
+        # testing that the discrete feature is integer valued
+        self.assertTrue(
+            torch.allclose(Arnd[:, discrete_feature], A[:, discrete_feature])
+        )
+        self.assertTrue(
+            torch.allclose(Brnd[:, discrete_feature], B[:, discrete_feature])
+        )
+
+        # testing that the other features are not integer valued
+        self.assertFalse(torch.allclose(Arnd, A))
+        self.assertFalse(torch.allclose(Brnd, B))
+
     def test_DerivativeGp(self) -> None:
         test_x = torch.rand(2, 2)
         posterior = posterior_derivative(self.model, test_x, kernel_type="matern")