Fix for simulation-based shadow pricing segmentation

activitysim/abm/tables/shadow_pricing.py

@@ -6,11 +6,13 @@
 import logging
 import multiprocessing
 import time
+import warnings
 from collections import OrderedDict
 from typing import Any, Literal
 
 import numpy as np
 import pandas as pd
+from pydantic import model_validator
 
 from activitysim.abm.tables.size_terms import size_terms as get_size_terms
 from activitysim.abm.tables.size_terms import tour_destination_size_terms
@@ -62,12 +64,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,11 +130,20 @@ 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."""
+    SEGMENT_TO_NAME: dict[str, str] | None = None
+    """Deprecated compatibility setting; chooser segmentation comes from model settings."""
+
+    @model_validator(mode="before")
+    @classmethod
+    def warn_deprecated_segment_to_name(cls, data: Any) -> Any:
+        if isinstance(data, dict) and "SEGMENT_TO_NAME" in data:
+            warnings.warn(
+                "SEGMENT_TO_NAME is deprecated and ignored; use "
+                "CHOOSER_SEGMENT_COLUMN_NAME in location choice settings instead",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        return data
 
 
 class ShadowPriceCalculator:
@@ -176,6 +181,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,13 +875,34 @@ 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"
+                )
 
             if type(self.choices_synced) != pd.DataFrame:
                 self.choices_synced = self.choices_synced.to_frame()
 
+            # 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]
+                .reindex(self.choices_synced.index)
+                .dropna()
+                .unique()
+            )
+            if not segment_values.intersection(chooser_values):
+                raise SystemConfigurationError(
+                    f"No overlap between SEGMENT_IDS values ({sorted(segment_values)}) and "
+                    f"modeled chooser values in persons_merged['{segment_name}'] "
+                    f"for {self.model_selector} "
+                    "simulation shadow pricing"
+                )
+
             choices_synced = self.choices_synced.merge(
                 persons_merged[segment_name],
                 how="left",

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,163 @@ 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, network_los
+):
+    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 "
+            "modeled chooser values in persons_merged['school_segment_string_99'] for school "
+            "simulation shadow pricing"
+        )
+    finally:
+        persons_merged.pop(custom_segment_column)
+
+
+def test_shadow_pricing_simulate_segment_overlap_ignores_unmodeled_choosers(
+    state, model_settings, network_los
+):
+    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")
+    persons_merged[custom_segment_column] = "99"
+    unmodeled_chooser = persons_merged.index[0]
+    persons_merged.loc[unmodeled_chooser, custom_segment_column] = "1"
+    modeled_choosers = persons_merged.drop(index=unmodeled_chooser)
+    state.settings.use_shadow_pricing = True
+
+    try:
+        spc = shadow_pricing.load_shadow_price_calculator(state, custom_model_settings)
+        choices = pd.Series(
+            22660, index=modeled_choosers.index, name="choice", dtype=int
+        )
+        spc.set_choices(choices, modeled_choosers[custom_segment_column])
+
+        with pytest.raises(
+            SystemConfigurationError, match="No overlap between SEGMENT_IDS"
+        ):
+            spc.update_shadow_prices(state)
+    finally:
+        persons_merged.pop(custom_segment_column)
+
+
+def test_shadow_price_settings_accept_deprecated_segment_to_name():
+    legacy_setting = {"school": "school_segment"}
+
+    with pytest.warns(DeprecationWarning, match="SEGMENT_TO_NAME is deprecated"):
+        settings = shadow_pricing.ShadowPriceSettings.model_validate(
+            {"SEGMENT_TO_NAME": legacy_setting}
+        )
+
+    assert settings.SEGMENT_TO_NAME == legacy_setting
+
+
 def test_shadow_pricing_dedicated_rng_channel_eet_only(
     state, model_settings, network_los
 ):