From 8aee9996b52adfe16e171e06ac3bde1d4bcf4ec8 Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Thu, 18 Apr 2024 12:44:54 -0500 Subject: [PATCH] protect additional columns not in spec --- activitysim/abm/models/joint_tour_participation.py | 11 ++++++++++- activitysim/core/configuration/base.py | 3 +++ activitysim/core/simulate.py | 1 + activitysim/core/util.py | 9 ++++++++- 4 files changed, 22 insertions(+), 2 deletions(-) diff --git a/activitysim/abm/models/joint_tour_participation.py b/activitysim/abm/models/joint_tour_participation.py index 98a8c70c7..55d5367b3 100644 --- a/activitysim/abm/models/joint_tour_participation.py +++ b/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, diff --git a/activitysim/core/configuration/base.py b/activitysim/core/configuration/base.py index 94c1d0de0..5b1cbc22f 100644 --- a/activitysim/core/configuration/base.py +++ b/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): diff --git a/activitysim/core/simulate.py b/activitysim/core/simulate.py index 0647883ce..d97b87a0d 100644 --- a/activitysim/core/simulate.py +++ b/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: diff --git a/activitysim/core/util.py b/activitysim/core/util.py index 08d1a3475..eb874f284 100644 --- a/activitysim/core/util.py +++ b/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))