feat: add validation for static covariates in tsx baseline

Author: rvasahu-amazon

src/sagemaker/clarify.py

@@ -408,7 +408,7 @@ class DatasetType(Enum):
 class TimeSeriesJSONDatasetFormat(Enum):
     """Possible dataset formats for JSON time series data files.
 
-    Below is an example ``COLUMNS`` dataset for time series explainability.::
+    Below is an example ``COLUMNS`` dataset for time series explainability::
 
         {
             "ids": [1, 2],
@@ -420,15 +420,15 @@ class TimeSeriesJSONDatasetFormat(Enum):
             "scv2": [30, 40]
         }
 
-    For this example, JMESPaths are specified when creating ``TimeSeriesDataConfig`` as follows.::
+    For this example, JMESPaths are specified when creating ``TimeSeriesDataConfig`` as follows::
 
         item_id="ids"
         timestamp="timestamps"
         target_time_series="target_ts"
         related_time_series=["rts1", "rts2"]
         static_covariates=["scv1", "scv2"]
 
-    Below is an example ``ITEM_RECORDS`` dataset for time series explainability.::
+    Below is an example ``ITEM_RECORDS`` dataset for time series explainability::
 
         [
             {
@@ -452,15 +452,15 @@ class TimeSeriesJSONDatasetFormat(Enum):
             }
         ]
 
-    For this example, JMESPaths are specified when creating ``TimeSeriesDataConfig`` as follows.::
+    For this example, JMESPaths are specified when creating ``TimeSeriesDataConfig`` as follows::
 
         item_id="[*].id"
         timestamp="[*].timeseries[].timestamp"
         target_time_series="[*].timeseries[].target_ts"
         related_time_series=["[*].timeseries[].rts1", "[*].timeseries[].rts2"]
         static_covariates=["[*].scv1", "[*].scv2"]
 
-    Below is an example ``TIMESTAMP_RECORDS`` dataset for time series explainability.::
+    Below is an example ``TIMESTAMP_RECORDS`` dataset for time series explainability::
 
         [
             {"id": 1, "timestamp": 1, "target_ts": 5, "scv1": 10, "rts1": 0.25},
@@ -469,7 +469,7 @@ class TimeSeriesJSONDatasetFormat(Enum):
             {"id": 2, "timestamp": 5, "target_ts": 10, "scv1": 20, "rts1": 1}
         ]
 
-    For this example, JMESPaths are specified when creating ``TimeSeriesDataConfig`` as follows.::
+    For this example, JMESPaths are specified when creating ``TimeSeriesDataConfig`` as follows::
 
         item_id="[*].id"
         timestamp="[*].timestamp"
@@ -1784,7 +1784,7 @@ def __init__(
                 values will be replaced with the average of a time series. For static data
                 (static covariates), a baseline value for each covariate should be provided for
                 each possible item_id. An example config follows, where ``item1`` and ``item2``
-                are item ids.::
+                are item ids::
 
                     {
                         "target_time_series": "zero",
@@ -2548,14 +2548,16 @@ def explainability(
         explainability_config: Union[ExplainabilityConfig, List[ExplainabilityConfig]],
     ):
         """Generates a config for Explainability"""
-        # determine if this is a timeseries explainability case by checking
+        # determine if this is a time series explainability case by checking
         # if *both* TimeSeriesDataConfig and TimeSeriesModelConfig were given
         ts_data_config_present = "time_series_data_config" in data_config.analysis_config
         ts_model_config_present = "time_series_predictor_config" in model_config.predictor_config
 
         if isinstance(explainability_config, AsymmetricShapleyValueConfig):
             assert ts_data_config_present, "Please provide a TimeSeriesDataConfig to DataConfig."
             assert ts_model_config_present, "Please provide a TimeSeriesModelConfig to ModelConfig."
+            # Check static covariates baseline matches number of provided static covariate columns
+
         else:
             if ts_data_config_present:
                 raise ValueError(
@@ -2759,6 +2761,44 @@ def _merge_explainability_configs(
             return explainability_methods
         return explainability_config.get_explainability_config()
 
+    @classmethod
+    def _validate_time_series_static_covariates_baseline(
+        cls,
+        explainability_config: AsymmetricShapleyValueConfig,
+        data_config: DataConfig,
+    ):
+        """Validates static covariates in baseline for asymmetric shapley value (for time series).
+
+        Checks that baseline values set for static covariate columns are
+        consistent between every item_id and the number of static covariate columns
+        provided in DataConfig.
+        """
+        baseline = explainability_config.get_explainability_config()[
+            "asymmetric_shapley_value"
+        ].get("baseline")
+        if baseline and "static_covariates" in baseline:
+            covariate_count = len(
+                data_config.get_config()["time_series_data_config"].get("static_covariates", [])
+            )
+            if covariate_count > 0:
+                for item_id in baseline.get("static_covariates", []):
+                    baseline_entry = baseline["static_covariates"][item_id]
+                    assert isinstance(baseline_entry, list), (
+                        f"Baseline entry for {item_id} must be a list, is "
+                        f"{type(baseline_entry)}."
+                    )
+                    assert len(baseline_entry) == covariate_count, (
+                        f"Length of baseline entry for {item_id} does not match number "
+                        f"of static covariate columns. Please ensure every covariate "
+                        f"has a baseline value for every item id."
+                    )
+            else:
+                raise ValueError(
+                    "Static covariate baselines are provided in AsymmetricShapleyValueConfig "
+                    "when no static covariate columns are provided in TimeSeriesDataConfig. "
+                    "Please check these configs."
+                )
+
 
 def _upload_analysis_config(analysis_config_file, s3_output_path, sagemaker_session, kms_key):
     """Uploads the local ``analysis_config_file`` to the ``s3_output_path``.

tests/unit/test_clarify.py

@@ -16,6 +16,7 @@
 import copy
 
 import pytest
+from dataclasses import dataclass
 from mock import ANY, MagicMock, Mock, patch
 from typing import Any, Dict, List, NamedTuple, Optional, Union
 
@@ -2574,7 +2575,8 @@ def _build_pdp_config_mock():
 
 def _build_asymmetric_shapley_value_config_mock():
     asym_shap_val_config_dict = {
-        "explanation_type": "fine_grained",
+        "direction": "chronological",
+        "granularity": "fine_grained",
         "num_samples": 20,
     }
     asym_shap_val_config = Mock(spec=AsymmetricShapleyValueConfig)
@@ -2613,6 +2615,14 @@ def _build_model_config_mock():
     return model_config
 
 
+@dataclass
+class ValidateTSXBaselineCase:
+    explainability_config: AsymmetricShapleyValueConfig
+    data_config: DataConfig
+    error: Optional[Exception] = None
+    error_msg: Optional[str] = None
+
+
 class TestAnalysisConfigGeneratorForTimeSeriesExplainability:
     @patch("sagemaker.clarify._AnalysisConfigGenerator._add_methods")
     @patch("sagemaker.clarify._AnalysisConfigGenerator._add_predictor")
@@ -2794,6 +2804,170 @@ def test_merge_explainability_configs_with_timeseries_invalid(
                 explainability_config=mock_config,
             )
 
+    @pytest.mark.parametrize(
+        "case",
+        [
+            ValidateTSXBaselineCase(
+                explainability_config=AsymmetricShapleyValueConfig(
+                    direction="chronological",
+                    granularity="timewise",
+                    baseline={
+                        "target_time_series": "zero",
+                        "related_time_series": "zero",
+                        "static_covariates": {
+                            "item1": [0.0, 0.5, 1.0],
+                            "item2": [0.3, 0.6, 0.9],
+                            "item3": [0.0, 1.0, 1.0],
+                            "item4": [0.9, 0.6, 0.3],
+                            "item5": [1.0, 0.5, 0.0],
+                        },
+                    },
+                ),
+                data_config=DataConfig(
+                    s3_data_input_path="s3://data/input",
+                    s3_output_path="s3://data/output",
+                    headers=["id", "time", "tts", "rts_1", "rts_2", "scv1", "scv2", "scv3"],
+                    dataset_type="application/json",
+                    time_series_data_config=TimeSeriesDataConfig(
+                        item_id="[].id",
+                        timestamp="[].temporal[].timestamp",
+                        target_time_series="[].temporal[].target",
+                        related_time_series=["[].temporal[].rts_1", "[].temporal[].rts_2"],
+                        static_covariates=["[].cov_1", "[].cov_2", "[].cov_3"],
+                        dataset_format=TimeSeriesJSONDatasetFormat.ITEM_RECORDS,
+                    ),
+                ),
+            ),
+        ],
+    )
+    def test_time_series_baseline_valid_static_covariates(self, case: ValidateTSXBaselineCase):
+        """
+        GIVEN AsymmetricShapleyValueConfig and TimeSeriesDataConfig are created and a baseline
+            is provided
+        WHEN AnalysisConfigGenerator._validate_time_series_static_covariates_baseline() is called
+        THEN no error is raised
+        """
+        _AnalysisConfigGenerator._validate_time_series_static_covariates_baseline(
+            explainability_config=case.explainability_config,
+            data_config=case.data_config,
+        )
+
+    @pytest.mark.parametrize(
+        "case",
+        [
+            ValidateTSXBaselineCase(  # some item ids are missing baseline values
+                explainability_config=AsymmetricShapleyValueConfig(
+                    direction="chronological",
+                    granularity="timewise",
+                    baseline={
+                        "target_time_series": "zero",
+                        "related_time_series": "zero",
+                        "static_covariates": {
+                            "item1": [0.0, 0.5, 1.0],
+                            "item2": [0.3, 0.6, 0.9],
+                            "item3": [0.0],
+                            "item4": [0.9, 0.6, 0.3],
+                            "item5": [1.0],
+                        },
+                    },
+                ),
+                data_config=DataConfig(
+                    s3_data_input_path="s3://data/input",
+                    s3_output_path="s3://data/output",
+                    headers=["id", "time", "tts", "rts_1", "rts_2", "scv1", "scv2", "scv3"],
+                    dataset_type="application/json",
+                    time_series_data_config=TimeSeriesDataConfig(
+                        item_id="[].id",
+                        timestamp="[].temporal[].timestamp",
+                        target_time_series="[].temporal[].target",
+                        related_time_series=["[].temporal[].rts_1", "[].temporal[].rts_2"],
+                        static_covariates=["[].cov_1", "[].cov_2", "[].cov_3"],
+                        dataset_format=TimeSeriesJSONDatasetFormat.ITEM_RECORDS,
+                    ),
+                ),
+                error=AssertionError,
+                error_msg="baseline entry for item3 does not match number",
+            ),
+            ValidateTSXBaselineCase(  # no static covariates are in data config
+                explainability_config=AsymmetricShapleyValueConfig(
+                    direction="chronological",
+                    granularity="timewise",
+                    baseline={
+                        "target_time_series": "zero",
+                        "related_time_series": "zero",
+                        "static_covariates": {
+                            "item1": [0.0, 0.5, 1.0],
+                            "item2": [0.3, 0.6, 0.9],
+                            "item3": [0.0, 1.0, 1.0],
+                            "item4": [0.9, 0.6, 0.3],
+                            "item5": [1.0, 0.5, 0.0],
+                        },
+                    },
+                ),
+                data_config=DataConfig(
+                    s3_data_input_path="s3://data/input",
+                    s3_output_path="s3://data/output",
+                    headers=["id", "time", "tts", "rts_1", "rts_2"],
+                    dataset_type="application/json",
+                    time_series_data_config=TimeSeriesDataConfig(
+                        item_id="[].id",
+                        timestamp="[].temporal[].timestamp",
+                        target_time_series="[].temporal[].target",
+                        related_time_series=["[].temporal[].rts_1", "[].temporal[].rts_2"],
+                        dataset_format=TimeSeriesJSONDatasetFormat.ITEM_RECORDS,
+                    ),
+                ),
+                error=ValueError,
+                error_msg="no static covariate columns are provided in TimeSeriesDataConfig",
+            ),
+            ValidateTSXBaselineCase(  # some item ids do not have a list as their baseline
+                explainability_config=AsymmetricShapleyValueConfig(
+                    direction="chronological",
+                    granularity="timewise",
+                    baseline={
+                        "target_time_series": "zero",
+                        "related_time_series": "zero",
+                        "static_covariates": {
+                            "item1": [0.0, 0.5, 1.0],
+                            "item2": [0.3, 0.6, 0.9],
+                            "item3": [0.0, 1.0, 1.0],
+                            "item4": [0.9, 0.6, 0.3],
+                            "item5": {"cov_1": 1.0, "cov_2": 0.5, "cov_3": 0.0},
+                        },
+                    },
+                ),
+                data_config=DataConfig(
+                    s3_data_input_path="s3://data/input",
+                    s3_output_path="s3://data/output",
+                    headers=["id", "time", "tts", "rts_1", "rts_2", "scv1", "scv2", "scv3"],
+                    dataset_type="application/json",
+                    time_series_data_config=TimeSeriesDataConfig(
+                        item_id="[].id",
+                        timestamp="[].temporal[].timestamp",
+                        target_time_series="[].temporal[].target",
+                        related_time_series=["[].temporal[].rts_1", "[].temporal[].rts_2"],
+                        static_covariates=["[].cov_1", "[].cov_2", "[].cov_3"],
+                        dataset_format=TimeSeriesJSONDatasetFormat.ITEM_RECORDS,
+                    ),
+                ),
+                error=AssertionError,
+                error_msg="Baseline entry for item5 must be a list",
+            ),
+        ],
+    )
+    def test_time_series_baseline_invalid_static_covariates(self, case: ValidateTSXBaselineCase):
+        """
+        GIVEN AsymmetricShapleyValueConfig and TimeSeriesDataConfig are created and a baseline
+            is provided where the static covariates baseline values are misconfigured
+        WHEN AnalysisConfigGenerator._validate_time_series_static_covariates_baseline() is called
+        THEN the appropriate error is raised
+        """
+        with pytest.raises(case.error, match=case.error_msg):
+            _AnalysisConfigGenerator._validate_time_series_static_covariates_baseline(
+                explainability_config=case.explainability_config,
+                data_config=case.data_config,
+            )
+
 
 class TestProcessingOutputHandler:
     def test_get_s3_upload_mode_image(self):