From 1c9850556dfe29bcaacc9a75a330feb3e4312870 Mon Sep 17 00:00:00 2001
From: David Hensle <hensle93@gmail.com>
Date: Fri, 15 Dec 2023 16:22:09 -0800
Subject: [PATCH] preprocessor and annotation in auto ownership

---
 activitysim/abm/models/auto_ownership.py | 22 ++++++++++++--
 activitysim/abm/models/util/annotate.py  | 38 ++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/activitysim/abm/models/auto_ownership.py b/activitysim/abm/models/auto_ownership.py
index 564d6f94b6..476bd18adf 100644
--- a/activitysim/abm/models/auto_ownership.py
+++ b/activitysim/abm/models/auto_ownership.py
@@ -2,9 +2,9 @@
 # See full license in LICENSE.txt.
 import logging
 
-from activitysim.core import config, inject, pipeline, simulate, tracing
+from activitysim.core import config, expressions, inject, pipeline, simulate, tracing
 
-from .util import estimation
+from .util import estimation, annotate
 
 logger = logging.getLogger(__name__)
 
@@ -32,6 +32,21 @@ def auto_ownership_simulate(households, households_merged, chunk_size, trace_hh_
 
     logger.info("Running %s with %d households", trace_label, len(choosers))
 
+    # - preprocessor
+    preprocessor_settings = model_settings.get("preprocessor", None)
+    if preprocessor_settings:
+
+        locals_d = {}
+        if constants is not None:
+            locals_d.update(constants)
+
+        expressions.assign_columns(
+            df=choosers,
+            model_settings=preprocessor_settings,
+            locals_dict=locals_d,
+            trace_label=trace_label,
+        )
+
     if estimator:
         estimator.write_model_settings(model_settings, model_settings_file_name)
         estimator.write_spec(model_settings)
@@ -69,5 +84,8 @@ def auto_ownership_simulate(households, households_merged, chunk_size, trace_hh_
         "auto_ownership", households.auto_ownership, value_counts=True
     )
 
+    if model_settings.get("annotate_households"):
+        annotate.annotate_households(model_settings, trace_label)
+
     if trace_hh_id:
         tracing.trace_df(households, label="auto_ownership", warn_if_empty=True)
diff --git a/activitysim/abm/models/util/annotate.py b/activitysim/abm/models/util/annotate.py
index e50519b38e..727ef4292f 100644
--- a/activitysim/abm/models/util/annotate.py
+++ b/activitysim/abm/models/util/annotate.py
@@ -15,6 +15,44 @@
 logger = logging.getLogger(__name__)
 
 
+def annotate_households(model_settings, trace_label, locals_dict={}):
+    """
+    Add columns to the households table in the pipeline according to spec.
+
+    Parameters
+    ----------
+    model_settings : dict
+    trace_label : str
+    """
+    households = inject.get_table("households").to_frame()
+    expressions.assign_columns(
+        df=households,
+        model_settings=model_settings.get("annotate_households"),
+        locals_dict=locals_dict,
+        trace_label=tracing.extend_trace_label(trace_label, "annotate_households"),
+    )
+    pipeline.replace_table("households", households)
+
+
+def annotate_persons(model_settings, trace_label, locals_dict={}):
+    """
+    Add columns to the persons table in the pipeline according to spec.
+
+    Parameters
+    ----------
+    model_settings : dict
+    trace_label : str
+    """
+    persons = inject.get_table("persons").to_frame()
+    expressions.assign_columns(
+        df=persons,
+        model_settings=model_settings.get("annotate_persons"),
+        locals_dict=locals_dict,
+        trace_label=tracing.extend_trace_label(trace_label, "annotate_persons"),
+    )
+    pipeline.replace_table("persons", persons)
+
+
 def annotate_tours(model_settings, trace_label, locals_dict={}):
     """
     Add columns to the tours table in the pipeline according to spec.