protect additional columns not in spec

activitysim/abm/models/joint_tour_participation.py

@@ -18,7 +18,7 @@
     tracing,
     workflow,
 )
-from activitysim.core.configuration.base import PreprocessorSettings
+from activitysim.core.configuration.base import ComputeSettings, PreprocessorSettings
 from activitysim.core.configuration.logit import LogitComponentSettings
 from activitysim.core.util import assign_in_place, reindex
 
@@ -408,6 +408,15 @@ def joint_tour_participation(
     )
     candidates["chunk_id"] = reindex(household_chunk_ids, candidates.household_id)
 
+    # these hardcoded columns need to be protected from being dropped
+    assert model_settings is not None
+    if model_settings.compute_settings is None:
+        model_settings.compute_settings = ComputeSettings()
+    assert model_settings.compute_settings is not None
+    for i in ["person_is_preschool", "composition", "adult"]:
+        if i not in model_settings.compute_settings.protect_columns:
+            model_settings.compute_settings.protect_columns.append(i)
+
     choices = simulate.simple_simulate_by_chunk_id(
         state,
         choosers=candidates,

activitysim/core/configuration/base.py

@@ -208,6 +208,9 @@ class ComputeSettings(PydanticBase):
     Default to True. If set to False, all columns in the data table will be kept.
     """
 
+    protect_columns: list[str] = []
+    """Protect these columns from being dropped from the chooser table."""
+
     def should_skip(self, subcomponent: str) -> bool:
         """Check if sharrow should be skipped for a particular subcomponent."""
         if isinstance(self.sharrow_skip, dict):

activitysim/core/simulate.py

@@ -1550,6 +1550,7 @@ def _simple_simulate(
             locals_d,
             custom_chooser,
             sharrow_enabled=sharrow_enabled,
+            additional_columns=compute_settings.protect_columns,
         )
 
     if nest_spec is None:

activitysim/core/util.py

@@ -641,7 +641,12 @@ def zarr_file_modification_time(zarr_dir: Path):
 
 
 def drop_unused_columns(
-    choosers, spec, locals_d, custom_chooser, sharrow_enabled=False
+    choosers,
+    spec,
+    locals_d,
+    custom_chooser,
+    sharrow_enabled=False,
+    additional_columns=None,
 ):
     """
     Drop unused columns from the chooser table, based on the spec and custom_chooser function.
@@ -656,6 +661,8 @@ def drop_unused_columns(
         spec.reset_index()["Expression"].apply(lambda x: re.findall(pattern, x)).sum()
     )
 
+    unique_variables_in_spec |= set(additional_columns or [])
+
     if locals_d:
         unique_variables_in_spec.add(locals_d.get("orig_col_name", None))
         unique_variables_in_spec.add(locals_d.get("dest_col_name", None))