Fix for simulation-based shadow pricing segmentation

activitysim/abm/tables/shadow_pricing.py

@@ -62,12 +62,6 @@
 TALLY_CHECKOUT = (1, -1)
 TALLY_PENDING_PERSONS = (2, -1)
 
-default_segment_to_name_dict = {
-    # model_selector : persons_segment_name
-    "school": "school_segment",
-    "workplace": "income_segment",
-}
-
 
 def size_table_name(model_selector):
     """
@@ -134,12 +128,6 @@ class ShadowPriceSettings(PydanticReadable, extra="forbid"):
 
     WRITE_ITERATION_CHOICES: bool = False
 
-    SEGMENT_TO_NAME: dict[str, str] = {
-        "school": "school_segment",
-        "workplace": "income_segment",
-    }  # pydantic uses deep copy, so mutable default value is ok here
-    """Mapping from model_selector to persons_segment_name."""
-
 
 class ShadowPriceCalculator:
     def __init__(
@@ -176,6 +164,7 @@ def __init__(
         )
 
         self.model_selector = model_settings.MODEL_SELECTOR
+        self.chooser_segment_column = model_settings.CHOOSER_SEGMENT_COLUMN_NAME
 
         if (self.num_processes > 1) and not state.settings.fail_fast:
             # if we are multiprocessing, then fail_fast should be true or we will wait forever for failed processes
@@ -869,9 +858,24 @@ def update_shadow_prices(self, state):
             sampled_persons = pd.DataFrame()
             persons_merged = state.get_dataframe("persons_merged")
 
-            # need to join the segment to the choices to sample correct persons
-            segment_to_name_dict = self.shadow_settings.SEGMENT_TO_NAME
-            segment_name = segment_to_name_dict[self.model_selector]
+            # Use the model chooser segmentation to keep shadow-pricing resampling
+            # consistent with segment_ids in location choice settings.
+            segment_name = self.chooser_segment_column
+            if segment_name not in persons_merged.columns:
+                raise SystemConfigurationError(
+                    f"Missing chooser segment column '{segment_name}' in persons_merged "
+                    f"for {self.model_selector} simulation shadow pricing"
+                )
+
+            # Fail fast on obvious misconfiguration instead of silently sampling no one.
+            segment_values = set(self.segment_ids.values())
+            chooser_values = set(persons_merged[segment_name].dropna().unique())
+            if not segment_values.intersection(chooser_values):
+                raise SystemConfigurationError(
+                    f"No overlap between SEGMENT_IDS values ({sorted(segment_values)}) and "
+                    f"persons_merged['{segment_name}'] values for {self.model_selector} "
+                    "simulation shadow pricing"
+                )
 
             if type(self.choices_synced) != pd.DataFrame:
                 self.choices_synced = self.choices_synced.to_frame()

activitysim/abm/test/test_misc/test_shadow_pricing_simulate.py

@@ -12,6 +12,7 @@
 from activitysim.abm.tables import shadow_pricing
 from activitysim.core import los, workflow
 from activitysim.core.configuration.logit import TourLocationComponentSettings
+from activitysim.core.exceptions import SystemConfigurationError
 
 LAND_USE_FIELDS = [
     "e01_nrm",
@@ -333,6 +334,9 @@ def persons() -> pd.DataFrame:
         }
     )
 
+    persons["school_segment_string"] = persons["school_segment"].astype(str)
+    persons["school_segment_string_99"] = "99"
+
     return persons
 
 
@@ -572,6 +576,114 @@ def test_shadow_pricing_simulate(state, model_settings, network_los):
         )
 
 
+def test_shadow_pricing_simulate_custom_segment(state, model_settings, network_los):
+    """Run simulation shadow pricing with string-valued chooser segments."""
+    segment_ids = {
+        "university": "3",
+        "highschool": "2",
+        "gradeschool": "1",
+    }
+    custom_model_settings = model_settings.model_copy(
+        update={
+            "CHOOSER_SEGMENT_COLUMN_NAME": "school_segment_string",
+            "SEGMENT_IDS": segment_ids,
+        }
+    )
+    custom_model_settings.LOGSUM_SETTINGS = None
+
+    spc = shadow_pricing.load_shadow_price_calculator(state, custom_model_settings)
+
+    max_iterations = 5
+    chooser_segment_column = "school_segment_string"
+    save_sample_df = choices_df = None
+    persons_merged = state.get_dataframe("persons_merged")
+
+    for iteration in range(1, max_iterations + 1):
+        old_shadow_prices = spc.shadow_prices["highschool"].values
+        persons_merged_df_ = persons_merged.copy()
+
+        if spc.use_shadow_pricing and iteration > 1:
+            spc.update_shadow_prices(state)
+
+            if spc.shadow_settings.SHADOW_PRICE_METHOD == "simulation":
+                persons_merged_df_ = persons_merged_df_[
+                    persons_merged_df_.index.isin(spc.sampled_persons.index)
+                ].sort_index()
+
+        choices_df_, save_sample_df = run_location_choice(
+            state,
+            persons_merged_df_,
+            network_los,
+            shadow_price_calculator=spc,
+            want_logsums=False,
+            want_sample_table=False,
+            estimator=None,
+            model_settings=custom_model_settings,
+            chunk_size=0,
+            chunk_tag="school_location_string_segment",
+            trace_label=f"school_location_string_segment_{iteration}",
+        )
+
+        if spc.use_shadow_pricing:
+            if (
+                spc.shadow_settings.SHADOW_PRICE_METHOD == "simulation"
+                and iteration > 1
+            ):
+                if len(choices_df_) != 0:
+                    choices_df = pd.concat([choices_df, choices_df_], axis=0)
+                    choices_df_index = choices_df_.index.name
+                    choices_df = choices_df.reset_index()
+                    choices_df = choices_df.drop_duplicates(
+                        subset=[choices_df_index], keep="last"
+                    )
+                    choices_df = choices_df.set_index(choices_df_index).sort_index()
+            else:
+                choices_df = choices_df_.copy()
+
+        new_shadow_prices = spc.shadow_prices["highschool"].values
+        assert not any((old_shadow_prices == -999) & (new_shadow_prices != -999))
+        check_shadow_prices(spc, iteration)
+
+        spc.set_choices(
+            choices=choices_df["choice"],
+            segment_ids=persons_merged[chooser_segment_column].reindex(
+                choices_df.index
+            ),
+        )
+
+
+def test_shadow_pricing_simulate_segment_values_do_not_overlap(state, model_settings):
+    custom_segment_column = "school_segment_string_99"
+    custom_model_settings = model_settings.model_copy(
+        update={
+            "CHOOSER_SEGMENT_COLUMN_NAME": custom_segment_column,
+            "SEGMENT_IDS": {
+                "university": "3",
+                "highschool": "2",
+                "gradeschool": "1",
+            },
+        }
+    )
+    persons_merged = state.get_dataframe("persons_merged")
+    state.settings.use_shadow_pricing = True
+
+    try:
+        spc = shadow_pricing.load_shadow_price_calculator(state, custom_model_settings)
+        choices = pd.Series(22660, index=persons_merged.index, name="choice", dtype=int)
+        spc.set_choices(choices, persons_merged[custom_segment_column])
+
+        with pytest.raises(SystemConfigurationError) as error:
+            spc.update_shadow_prices(state)
+
+        assert str(error.value) == (
+            "No overlap between SEGMENT_IDS values (['1', '2', '3']) and "
+            "persons_merged['school_segment_string_99'] values for school "
+            "simulation shadow pricing"
+        )
+    finally:
+        persons_merged.pop(custom_segment_column)
+
+
 def test_shadow_pricing_dedicated_rng_channel_eet_only(
     state, model_settings, network_los
 ):