Preprocessing & Annotation functionality (#957)

* initial batch of models with pre and post processing added

* nonmand sched and transit pass models

* first pass at preprocessing and annotate functionality in all models

* fixing bugs in jtf and trip purpose

* adding persons back in to locals_d in jtc

* model name missing in tour scheduling

* missing expressions import in tour sched prob

* ci unit tests & fixing estimation test error

* addressing review comments

---------

Co-authored-by: Ali Etezady <[email protected]>

Author: dhensle

activitysim/abm/models/atwork_subtour_destination.py

@@ -7,7 +7,7 @@
 import pandas as pd
 
 from activitysim.abm.models.util import tour_destination
-from activitysim.core import config, estimation, los, tracing, workflow
+from activitysim.core import config, estimation, los, tracing, workflow, expressions
 from activitysim.core.configuration.logit import TourLocationComponentSettings
 from activitysim.core.util import assign_in_place
 
@@ -120,3 +120,11 @@ def atwork_subtour_destination(
         state.tracing.trace_df(
             tours, label="atwork_subtour_destination", columns=["destination"]
         )
+
+    expressions.annotate_tables(
+        state,
+        locals_dict={},
+        skims=None,
+        model_settings=model_settings,
+        trace_label=trace_label,
+    )

activitysim/abm/models/atwork_subtour_frequency.py

@@ -37,8 +37,8 @@ class AtworkSubtourFrequencySettings(LogitComponentSettings, extra="forbid"):
     Settings for the `atwork_subtour_frequency` component.
     """
 
-    preprocessor: PreprocessorSettings | None = None
-    """Setting for the preprocessor."""
+    # no additional fields are required for this component
+    pass
 
 
 @workflow.step
@@ -92,15 +92,15 @@ def atwork_subtour_frequency(
     nest_spec = config.get_logit_model_settings(model_settings)
     constants = config.get_model_constants(model_settings)
 
-    # - preprocessor
-    preprocessor_settings = model_settings.preprocessor
-    if preprocessor_settings:
-        expressions.assign_columns(
-            state,
-            df=work_tours,
-            model_settings=preprocessor_settings,
-            trace_label=trace_label,
-        )
+    # preprocess choosers
+    expressions.annotate_preprocessors(
+        state,
+        df=work_tours,
+        locals_dict=constants,
+        skims=None,
+        model_settings=model_settings,
+        trace_label=trace_label,
+    )
 
     if estimator:
         estimator.write_spec(model_settings)
@@ -164,3 +164,11 @@ def atwork_subtour_frequency(
 
     if trace_hh_id:
         state.tracing.trace_df(tours, label="atwork_subtour_frequency.tours")
+
+    expressions.annotate_tables(
+        state,
+        locals_dict=constants,
+        skims=None,
+        model_settings=model_settings,
+        trace_label=trace_label,
+    )

activitysim/abm/models/atwork_subtour_mode_choice.py

@@ -195,21 +195,18 @@ def atwork_subtour_mode_choice(
     )
     state.add_table("tours", tours)
 
-    # - annotate tours table
-    if model_settings.annotate_tours:
-        tours = state.get_dataframe("tours")
-        expressions.assign_columns(
-            state,
-            df=tours,
-            model_settings=model_settings.annotate_tours,
-            trace_label=tracing.extend_trace_label(trace_label, "annotate_tours"),
-        )
-        state.add_table("tours", tours)
-
     if trace_hh_id:
         state.tracing.trace_df(
             tours[tours.tour_category == "atwork"],
             label=tracing.extend_trace_label(trace_label, mode_column_name),
             slicer="tour_id",
             index_label="tour_id",
         )
+
+    expressions.annotate_tables(
+        state,
+        locals_dict=constants,
+        skims=skims,
+        model_settings=model_settings,
+        trace_label=trace_label,
+    )

activitysim/abm/models/atwork_subtour_scheduling.py

@@ -141,3 +141,11 @@ def atwork_subtour_scheduling(
             trace_label,
             "tour_map",
         )
+
+    expressions.annotate_tables(
+        state,
+        locals_dict=constants,
+        skims=skims,
+        model_settings=model_settings,
+        trace_label=trace_label,
+    )

activitysim/abm/models/auto_ownership.py

@@ -18,8 +18,6 @@
 from activitysim.core.configuration.base import PreprocessorSettings, PydanticReadable
 from activitysim.core.configuration.logit import LogitComponentSettings
 
-from .util import annotate
-
 logger = logging.getLogger(__name__)
 
 
@@ -28,8 +26,8 @@ class AutoOwnershipSettings(LogitComponentSettings, extra="forbid"):
     Settings for the `auto_ownership` component.
     """
 
-    preprocessor: PreprocessorSettings | None = None
-    annotate_households: PreprocessorSettings | None = None
+    # no additional fields are required for this component
+    pass
 
 
 @workflow.step
@@ -69,20 +67,14 @@ def auto_ownership_simulate(
 
     logger.info("Running %s with %d households", trace_label, len(choosers))
 
-    # - preprocessor
-    preprocessor_settings = model_settings.preprocessor
-    if preprocessor_settings:
-        locals_d = {}
-        if constants is not None:
-            locals_d.update(constants)
-
-        expressions.assign_columns(
-            state,
-            df=choosers,
-            model_settings=preprocessor_settings,
-            locals_dict=locals_d,
-            trace_label=trace_label,
-        )
+    expressions.annotate_preprocessors(
+        state,
+        df=choosers,
+        locals_dict=constants,
+        skims=None,
+        model_settings=model_settings,
+        trace_label=trace_label,
+    )
 
     if estimator:
         estimator.write_model_settings(model_settings, model_settings_file_name)
@@ -120,8 +112,13 @@ def auto_ownership_simulate(
         "auto_ownership", households.auto_ownership, value_counts=True
     )
 
-    if model_settings.annotate_households:
-        annotate.annotate_households(state, model_settings, trace_label)
-
     if trace_hh_id:
         state.tracing.trace_df(households, label="auto_ownership", warn_if_empty=True)
+
+    expressions.annotate_tables(
+        state,
+        locals_dict=constants,
+        skims=None,
+        model_settings=model_settings,
+        trace_label=trace_label,
+    )

activitysim/abm/models/cdap.py

@@ -36,12 +36,17 @@ class CdapSettings(PydanticReadable, extra="forbid"):
     JOINT_TOUR_COEFFICIENTS: str = "cdap_joint_tour_coefficients.csv"
     JOINT_TOUR_USEFUL_COLUMNS: list[str] | None = None
     """Columns to include from the persons table that will be need to calculate household joint tour utility."""
-    annotate_persons: PreprocessorSettings | None = None
-    annotate_households: PreprocessorSettings | None = None
     COEFFICIENTS: Path
     CONSTANTS: dict[str, Any] = {}
     compute_settings: ComputeSettings | None = None
 
+    preprocessor: PreprocessorSettings | None = None
+    """Preprocess choosers tables before running the model."""
+    annotate_persons: PreprocessorSettings | None = None
+    """Postprocess persons table after model completion."""
+    annotate_households: PreprocessorSettings | None = None
+    """Postprocess households table after model completion."""
+
 
 @workflow.step
 def cdap_simulate(
@@ -171,6 +176,16 @@ def cdap_simulate(
                     index=True,
                 )
 
+    # preprocess choosers
+    expressions.annotate_preprocessors(
+        state,
+        df=persons_merged,
+        locals_dict=constants,
+        skims=None,
+        model_settings=model_settings,
+        trace_label=trace_label,
+    )
+
     if estimator:
         estimator.write_model_settings(model_settings, "cdap.yaml")
         estimator.write_spec(model_settings, tag="INDIV_AND_HHSIZE1_SPEC")
@@ -241,31 +256,25 @@ def cdap_simulate(
     cap_cat_type = pd.api.types.CategoricalDtype(["", "M", "N", "H"], ordered=False)
     choices = choices.astype(cap_cat_type)
     persons["cdap_activity"] = choices
-
-    expressions.assign_columns(
-        state,
-        df=persons,
-        model_settings=model_settings.annotate_persons,
-        trace_label=tracing.extend_trace_label(trace_label, "annotate_persons"),
-    )
-
     state.add_table("persons", persons)
 
     # - annotate households table
     if add_joint_tour_utility:
         hh_joint = hh_joint.reindex(households.index)
         households["has_joint_tour"] = hh_joint
 
-    expressions.assign_columns(
-        state,
-        df=households,
-        model_settings=model_settings.annotate_households,
-        trace_label=tracing.extend_trace_label(trace_label, "annotate_households"),
-    )
     state.add_table("households", households)
 
     tracing.print_summary("cdap_activity", persons.cdap_activity, value_counts=True)
     logger.info(
         "cdap crosstabs:\n%s"
         % pd.crosstab(persons.ptype, persons.cdap_activity, margins=True)
     )
+
+    expressions.annotate_tables(
+        state,
+        locals_dict=constants,
+        skims=None,
+        model_settings=model_settings,
+        trace_label=trace_label,
+    )

activitysim/abm/models/free_parking.py

@@ -26,9 +26,6 @@ class FreeParkingSettings(LogitComponentSettings, extra="forbid"):
     Settings for the `free_parking` component.
     """
 
-    preprocessor: PreprocessorSettings | None = None
-    """Setting for the preprocessor."""
-
     FREE_PARKING_ALT: int
     """The code for free parking."""
 
@@ -78,21 +75,6 @@ def free_parking(
 
     constants = model_settings.CONSTANTS or {}
 
-    # - preprocessor
-    preprocessor_settings = model_settings.preprocessor
-    if preprocessor_settings:
-        locals_d = {}
-        if constants is not None:
-            locals_d.update(constants)
-
-        expressions.assign_columns(
-            state,
-            df=choosers,
-            model_settings=preprocessor_settings,
-            locals_dict=locals_d,
-            trace_label=trace_label,
-        )
-
     model_spec = state.filesystem.read_model_spec(file_name=model_settings.SPEC)
     coefficients_df = state.filesystem.read_model_coefficients(model_settings)
     model_spec = simulate.eval_coefficients(
@@ -101,6 +83,15 @@ def free_parking(
 
     nest_spec = config.get_logit_model_settings(model_settings)
 
+    expressions.annotate_preprocessors(
+        state,
+        df=choosers,
+        locals_dict=constants,
+        skims=None,
+        model_settings=model_settings,
+        trace_label=trace_label,
+    )
+
     if estimator:
         estimator.write_model_settings(model_settings, model_settings_file_name)
         estimator.write_spec(file_name=model_settings.SPEC)
@@ -144,3 +135,11 @@ def free_parking(
 
     if state.settings.trace_hh_id:
         state.tracing.trace_df(persons, label=trace_label, warn_if_empty=True)
+
+    expressions.annotate_tables(
+        state,
+        locals_dict=constants,
+        skims=None,
+        model_settings=model_settings,
+        trace_label=trace_label,
+    )

activitysim/abm/models/joint_tour_composition.py

@@ -18,6 +18,7 @@
 from activitysim.core.configuration.base import PreprocessorSettings
 from activitysim.core.configuration.logit import LogitComponentSettings
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -36,8 +37,7 @@ class JointTourCompositionSettings(LogitComponentSettings, extra="forbid"):
     Settings for the `joint_tour_composition` component.
     """
 
-    preprocessor: PreprocessorSettings | None = None
-    """Setting for the preprocessor."""
+    pass
 
 
 @workflow.step
@@ -78,26 +78,6 @@ def joint_tour_composition(
         "Running joint_tour_composition with %d joint tours" % joint_tours.shape[0]
     )
 
-    # - run preprocessor
-    preprocessor_settings = model_settings.preprocessor
-    if preprocessor_settings:
-        locals_dict = {
-            "persons": persons,
-            "hh_time_window_overlap": lambda *x: hh_time_window_overlap(state, *x),
-        }
-
-        expressions.assign_columns(
-            state,
-            df=households,
-            model_settings=preprocessor_settings,
-            locals_dict=locals_dict,
-            trace_label=trace_label,
-        )
-
-    joint_tours_merged = pd.merge(
-        joint_tours, households, left_on="household_id", right_index=True, how="left"
-    )
-
     # - simple_simulate
     model_spec = state.filesystem.read_model_spec(file_name=model_settings.SPEC)
     coefficients_df = state.filesystem.read_model_coefficients(model_settings)
@@ -108,6 +88,25 @@ def joint_tour_composition(
     nest_spec = config.get_logit_model_settings(model_settings)
     constants = config.get_model_constants(model_settings)
 
+    locals_dict = {
+        "persons": persons,
+        "hh_time_window_overlap": lambda *x: hh_time_window_overlap(state, *x),
+    }
+    locals_dict.update(constants)
+
+    expressions.annotate_preprocessors(
+        state,
+        df=households,
+        locals_dict=locals_dict,
+        skims=None,
+        model_settings=model_settings,
+        trace_label=trace_label,
+    )
+
+    joint_tours_merged = pd.merge(
+        joint_tours, households, left_on="household_id", right_index=True, how="left"
+    )
+
     if estimator:
         estimator.write_spec(model_settings)
         estimator.write_model_settings(model_settings, model_settings_file_name)
@@ -156,3 +155,11 @@ def joint_tour_composition(
             label="joint_tour_composition.joint_tours",
             slicer="household_id",
         )
+
+    expressions.annotate_tables(
+        state,
+        locals_dict=locals_dict,
+        skims=None,
+        model_settings=model_settings,
+        trace_label=trace_label,
+    )