Fix discarded trip_period output by limiting cleanup to temporary columns

Author: jpn--

activitysim/abm/models/trip_mode_choice.py

@@ -350,7 +350,9 @@ def add_trip_period(choosers):
     # needed in the purpose-sized chooser frames above, but post-choice table
     # annotators may also use those skims. Add it to the full trips table only
     # for annotation, then restore the original table schema.
-    add_trip_period(trips_df)
+    temporary_trip_period = "trip_period" not in trips_df.columns
+    if temporary_trip_period:
+        add_trip_period(trips_df)
     try:
         expressions.annotate_tables(
             state,
@@ -360,7 +362,10 @@ def add_trip_period(choosers):
             trace_label=trace_label,
         )
     finally:
-        trips_df.drop(columns="trip_period", inplace=True)
-        state_trips = state.get_dataframe("trips", as_copy=False)
-        if state_trips is not trips_df:
-            state_trips.drop(columns="trip_period", inplace=True)
+        # CHOOSER_COLS_TO_KEEP may have made trip_period an output column.
+        # Remove it only when this annotation block created it temporarily.
+        if temporary_trip_period:
+            trips_df.drop(columns="trip_period", inplace=True)
+            state_trips = state.get_dataframe("trips", as_copy=False)
+            if state_trips is not trips_df:
+                state_trips.drop(columns="trip_period", inplace=True)

activitysim/abm/test/test_trip_mode_choice.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 import weakref
+from contextlib import nullcontext
 from types import SimpleNamespace
 
 import pandas as pd
+import pytest
 
 from activitysim.abm.models import trip_mode_choice as trip_mode_choice_module
 
@@ -55,7 +57,12 @@ def is_table(self, _name):
         return False
 
 
-def test_post_choice_annotations_receive_full_trip_period_then_remove_it(monkeypatch):
+@pytest.mark.parametrize("keep_trip_period", [False, True])
+@pytest.mark.parametrize("copy_annotation_table", [False, True])
+@pytest.mark.parametrize("annotation_error", [False, True])
+def test_post_choice_annotations_preserve_requested_trip_period(
+    monkeypatch, keep_trip_period, copy_annotation_table, annotation_error
+):
     trips = pd.DataFrame(
         {
             "tour_id": [11, 12, 13],
@@ -79,7 +86,7 @@ def test_post_choice_annotations_receive_full_trip_period_then_remove_it(monkeyp
     model_settings = SimpleNamespace(
         MODE_CHOICE_LOGSUM_COLUMN_NAME="mode_choice_logsum",
         TOURS_MERGED_CHOOSER_COLUMNS=[],
-        CHOOSER_COLS_TO_KEEP=[],
+        CHOOSER_COLS_TO_KEEP=["trip_period"] if keep_trip_period else [],
         FORCE_ESCORTEE_CHAUFFEUR_MODE_MATCH=False,
         SPEC="trip_mode_choice.csv",
         explicit_chunk=None,
@@ -132,26 +139,33 @@ def choose_mode(_state, choosers, **_kwargs):
     monkeypatch.setattr(trip_mode_choice_module, "mode_choice_simulate", choose_mode)
 
     def annotate_tables(_state, **_kwargs):
-        annotated = _state.get_dataframe("trips")
+        annotated = _state.get_dataframe("trips", as_copy=copy_annotation_table)
         assert annotated.index.equals(trips.index)
         assert annotated["trip_period"].tolist() == [0, 1, 0]
         annotated["post_choice_skim_value"] = [10.0, 20.0, 30.0]
         _state.add_table("trips", annotated)
+        if annotation_error:
+            raise RuntimeError("annotation failed")
 
     monkeypatch.setattr(
         trip_mode_choice_module.expressions, "annotate_tables", annotate_tables
     )
 
-    trip_mode_choice_module.trip_mode_choice(
-        state,
-        trips,
-        network_los,
-        model_settings=model_settings,
-    )
+    with pytest.raises(
+        RuntimeError, match="annotation failed"
+    ) if annotation_error else nullcontext():
+        trip_mode_choice_module.trip_mode_choice(
+            state,
+            trips,
+            network_los,
+            model_settings=model_settings,
+        )
 
     assert all(ref() is None for ref in chooser_refs)
     assert all(wrapper.df.empty for wrapper in wrappers)
     result = state.get_dataframe("trips", as_copy=False)
-    assert "trip_period" not in trips
-    assert "trip_period" not in result
+    assert ("trip_period" in trips) == keep_trip_period
+    assert ("trip_period" in result) == keep_trip_period
+    if keep_trip_period:
+        assert result["trip_period"].tolist() == [0, 1, 0]
     assert result["post_choice_skim_value"].tolist() == [10.0, 20.0, 30.0]