Merge branch 'main' into explicit_chunking

activitysim/abm/models/parking_location_choice.py

@@ -187,9 +187,19 @@ def choose_parking_location(
     locals_dict["PARKING"] = skims["op_skims"].dest_key
 
     spec = get_spec_for_segment(state, model_settings, segment_name)
-    trips = drop_unused_columns(trips, spec, locals_dict, custom_chooser=None)
+    trips = drop_unused_columns(
+        trips,
+        spec,
+        locals_dict,
+        custom_chooser=None,
+        additional_columns=model_settings.compute_settings.protect_columns,
+    )
     alternatives = drop_unused_columns(
-        alternatives, spec, locals_dict, custom_chooser=None
+        alternatives,
+        spec,
+        locals_dict,
+        custom_chooser=None,
+        additional_columns=model_settings.compute_settings.protect_columns,
     )
 
     destination_sample = logit.interaction_dataset(

activitysim/abm/models/trip_destination.py

@@ -60,7 +60,8 @@ class TripDestinationSettings(LocationComponentSettings, extra="forbid"):
     PRIMARY_DEST: str = "tour_leg_dest"  # must be created in preprocessor
     REDUNDANT_TOURS_MERGED_CHOOSER_COLUMNS: list[str] | None = None
     preprocessor: PreprocessorSettings | None = None
-    alts_preprocessor: PreprocessorSettings | None = None
+    alts_preprocessor_sample: PreprocessorSettings | None = None
+    alts_preprocessor_simulate: PreprocessorSettings | None = None
     CLEANUP: bool
     fail_some_trips_for_testing: bool = False
     """This setting is used by testing code to force failed trip_destination."""
@@ -202,6 +203,15 @@ def _destination_sample(
 
     log_alt_losers = state.settings.log_alt_losers
 
+    if model_settings.alts_preprocessor_sample:
+        expressions.assign_columns(
+            state,
+            df=alternatives,
+            model_settings=model_settings.alts_preprocessor_sample,
+            locals_dict=locals_dict,
+            trace_label=tracing.extend_trace_label(trace_label, "alts"),
+        )
+
     choices = interaction_sample(
         state,
         choosers=trips,
@@ -941,6 +951,15 @@ def trip_destination_simulate(
     )
     locals_dict.update(skims)
 
+    if model_settings.alts_preprocessor_simulate:
+        expressions.assign_columns(
+            state,
+            df=destination_sample,
+            model_settings=model_settings.alts_preprocessor_simulate,
+            locals_dict=locals_dict,
+            trace_label=tracing.extend_trace_label(trace_label, "alts"),
+        )
+
     log_alt_losers = state.settings.log_alt_losers
     destinations = interaction_sample_simulate(
         state,
@@ -1252,7 +1271,6 @@ def run_trip_destination(
             state.filesystem, model_settings_file_name
         )
     preprocessor_settings = model_settings.preprocessor
-    alts_preprocessor_settings = model_settings.alts_preprocessor
     logsum_settings = state.filesystem.read_model_settings(
         model_settings.LOGSUM_SETTINGS
     )
@@ -1375,15 +1393,6 @@ def run_trip_destination(
                     trace_label=nth_trace_label,
                 )
 
-            if alts_preprocessor_settings:
-                expressions.assign_columns(
-                    state,
-                    df=alternatives,
-                    model_settings=alts_preprocessor_settings,
-                    locals_dict=locals_dict,
-                    trace_label=tracing.extend_trace_label(nth_trace_label, "alts"),
-                )
-
             if isinstance(
                 nth_trips["trip_period"].dtype, pd.api.types.CategoricalDtype
             ):

activitysim/abm/models/trip_scheduling_choice.py

@@ -13,7 +13,11 @@
     get_time_windows,
 )
 from activitysim.core import chunk, expressions, simulate, tracing, workflow
-from activitysim.core.configuration.base import PreprocessorSettings, PydanticReadable
+from activitysim.core.configuration.base import (
+    ComputeSettings,
+    PreprocessorSettings,
+    PydanticReadable,
+)
 from activitysim.core.interaction_sample_simulate import _interaction_sample_simulate
 from activitysim.core.skim_dataset import SkimDataset
 from activitysim.core.skim_dictionary import SkimDict
@@ -223,6 +227,7 @@ def run_trip_scheduling_choice(
     skims,
     locals_dict: Mapping,
     trace_label: str,
+    model_settings: TripSchedulingChoiceSettings,
 ):
     NUM_TOUR_LEGS = 3
     trace_label = tracing.extend_trace_label(trace_label, "interaction_sample_simulate")
@@ -296,6 +301,7 @@ def run_trip_scheduling_choice(
                 trace_choice_name="trip_schedule_stage_1",
                 estimator=None,
                 chunk_sizer=chunk_sizer,
+                compute_settings=model_settings.compute_settings,
             )
 
             assert len(choices.index) == len(choosers.index)
@@ -338,6 +344,9 @@ class TripSchedulingChoiceSettings(PydanticReadable, extra="forbid"):
     SPECIFICATION: str
     """file name of specification file"""
 
+    compute_settings: ComputeSettings = ComputeSettings()
+    """Compute settings for this component."""
+
 
 @workflow.step
 def trip_scheduling_choice(
@@ -419,7 +428,13 @@ def trip_scheduling_choice(
         )
 
     tours_df = run_trip_scheduling_choice(
-        state, spec, tours_df, skims, locals_dict, trace_label
+        state,
+        spec,
+        tours_df,
+        skims,
+        locals_dict,
+        trace_label,
+        model_settings,
     )
 
     state.add_table("tours", tours_df)

activitysim/core/configuration/base.py

@@ -243,6 +243,7 @@ def subcomponent_settings(self, subcomponent: str) -> ComputeSettings:
             use_numexpr=self.use_numexpr,
             use_numba=self.use_numba,
             drop_unused_columns=self.drop_unused_columns,
+            protect_columns=self.protect_columns,
         )
 
 

activitysim/core/interaction_sample.py

@@ -257,6 +257,7 @@ def _interaction_sample(
             locals_d,
             custom_chooser=None,
             sharrow_enabled=sharrow_enabled,
+            additional_columns=compute_settings.protect_columns,
         )
 
         alternatives = util.drop_unused_columns(
@@ -265,7 +266,7 @@ def _interaction_sample(
             locals_d,
             custom_chooser=None,
             sharrow_enabled=sharrow_enabled,
-            additional_columns=["tdd", "origin_destination"],
+            additional_columns=["tdd"] + compute_settings.protect_columns,
         )
 
     if sharrow_enabled:

activitysim/core/interaction_sample_simulate.py

@@ -158,6 +158,7 @@ def _interaction_sample_simulate(
             locals_d,
             custom_chooser=None,
             sharrow_enabled=sharrow_enabled,
+            additional_columns=compute_settings.protect_columns,
         )
 
         alternatives = util.drop_unused_columns(
@@ -166,7 +167,7 @@ def _interaction_sample_simulate(
             locals_d,
             custom_chooser=None,
             sharrow_enabled=sharrow_enabled,
-            additional_columns=["tdd", "origin_destination"],
+            additional_columns=["tdd"] + compute_settings.protect_columns,
         )
 
     interaction_df = alternatives.join(choosers, how="left", rsuffix="_chooser")

activitysim/core/interaction_simulate.py

@@ -729,6 +729,7 @@ def _interaction_simulate(
             locals_d,
             custom_chooser=None,
             sharrow_enabled=sharrow_enabled,
+            additional_columns=compute_settings.protect_columns,
         )
 
     if (

activitysim/core/simulate.py

@@ -1989,6 +1989,7 @@ def _simple_simulate_logsums(
             locals_d,
             custom_chooser=None,
             sharrow_enabled=state.settings.sharrow,
+            additional_columns=compute_settings.protect_columns,
         )
 
     if nest_spec is None:

activitysim/examples/prototype_sandag_xborder/configs/tour_od_choice.yaml

@@ -55,3 +55,7 @@ SEGMENTS:
 
 LOGSUM_SETTINGS: tour_mode_choice
 LOGSUM_PREPROCESSOR: preprocessor
+
+compute_settings:
+  protect_columns:
+  - origin_destination