reduce memory usage in parking_location_choice

activitysim/abm/models/parking_location_choice.py

@@ -21,7 +21,7 @@
 from activitysim.core.configuration.logit import LogitComponentSettings
 from activitysim.core.interaction_sample_simulate import interaction_sample_simulate
 from activitysim.core.tracing import print_elapsed_time
-from activitysim.core.util import assign_in_place
+from activitysim.core.util import assign_in_place, drop_unused_chooser_columns
 
 logger = logging.getLogger(__name__)
 
@@ -106,6 +106,7 @@ def parking_destination_simulate(
     destination_sample,
     model_settings: ParkingLocationSettings,
     skims,
+    locals_dict,
     chunk_size,
     trace_hh_id,
     trace_label,
@@ -132,11 +133,6 @@ def parking_destination_simulate(
 
     logger.info("Running parking_destination_simulate with %d trips", len(trips))
 
-    locals_dict = config.get_model_constants(model_settings).copy()
-    locals_dict.update(skims)
-    locals_dict["timeframe"] = "trip"
-    locals_dict["PARKING"] = skims["op_skims"].dest_key
-
     parking_locations = interaction_sample_simulate(
         state,
         choosers=trips,
@@ -181,6 +177,19 @@ def choose_parking_location(
     t0 = print_elapsed_time()
 
     alt_dest_col_name = model_settings.ALT_DEST_COL_NAME
+
+    # remove trips and alts columns that are not used in spec
+    locals_dict = config.get_model_constants(model_settings).copy()
+    locals_dict.update(skims)
+    locals_dict["timeframe"] = "trip"
+    locals_dict["PARKING"] = skims["op_skims"].dest_key
+
+    spec = get_spec_for_segment(state, model_settings, segment_name)
+    trips = drop_unused_chooser_columns(trips, spec, locals_dict, custom_chooser=None)
+    alternatives = drop_unused_chooser_columns(
+        alternatives, spec, locals_dict, custom_chooser=None
+    )
+
     destination_sample = logit.interaction_dataset(
         state, trips, alternatives, alt_index_id=alt_dest_col_name
     )
@@ -194,6 +203,7 @@ def choose_parking_location(
         destination_sample=destination_sample,
         model_settings=model_settings,
         skims=skims,
+        locals_dict=locals_dict,
         chunk_size=chunk_size,
         trace_hh_id=trace_hh_id,
         trace_label=trace_label,

activitysim/core/util.py

@@ -638,3 +638,78 @@ def zarr_file_modification_time(zarr_dir: Path):
     if t == 0:
         raise FileNotFoundError(zarr_dir)
     return t
+
+
+def drop_unused_chooser_columns(
+    choosers, spec, locals_d, custom_chooser, sharrow_enabled=False
+):
+    """
+    Drop unused columns from the chooser table, based on the spec and custom_chooser function.
+    """
+    # keep only variables needed for spec
+    import re
+
+    # define a regular expression to find variables in spec
+    pattern = r"[a-zA-Z_][a-zA-Z0-9_]*"
+
+    unique_variables_in_spec = set(
+        spec.reset_index()["Expression"].apply(lambda x: re.findall(pattern, x)).sum()
+    )
+
+    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))
+        if locals_d.get("timeframe") == "trip":
+            orig_col_name = locals_d.get("ORIGIN", None)
+            dest_col_name = locals_d.get("DESTINATION", None)
+            stop_col_name = None
+            parking_col_name = locals_d.get("PARKING", None)
+            primary_origin_col_name = None
+            if orig_col_name is None and "od_skims" in locals_d:
+                orig_col_name = locals_d["od_skims"].orig_key
+            if dest_col_name is None and "od_skims" in locals_d:
+                dest_col_name = locals_d["od_skims"].dest_key
+            if stop_col_name is None and "dp_skims" in locals_d:
+                stop_col_name = locals_d["dp_skims"].dest_key
+            if primary_origin_col_name is None and "dnt_skims" in locals_d:
+                primary_origin_col_name = locals_d["dnt_skims"].dest_key
+            unique_variables_in_spec.add(orig_col_name)
+            unique_variables_in_spec.add(dest_col_name)
+            unique_variables_in_spec.add(parking_col_name)
+            unique_variables_in_spec.add(primary_origin_col_name)
+            unique_variables_in_spec.add(stop_col_name)
+            unique_variables_in_spec.add("trip_period")
+        # when using trip_scheduling_choice for trup scheduling
+        unique_variables_in_spec.add("last_outbound_stop")
+        unique_variables_in_spec.add("last_inbound_stop")
+
+    # when sharrow mode, need to keep the following columns in the choosers table
+    if sharrow_enabled:
+        unique_variables_in_spec.add("out_period")
+        unique_variables_in_spec.add("in_period")
+        unique_variables_in_spec.add("purpose_index_num")
+
+    if custom_chooser:
+        import inspect
+
+        custom_chooser_lines = inspect.getsource(custom_chooser)
+        unique_variables_in_spec.update(re.findall(pattern, custom_chooser_lines))
+
+    logger.info("Dropping unused variables in chooser table")
+
+    logger.info(
+        "before dropping, the choosers table has {} columns: {}".format(
+            len(choosers.columns), choosers.columns
+        )
+    )
+
+    # keep only variables needed for spec
+    choosers = choosers[[c for c in choosers.columns if c in unique_variables_in_spec]]
+
+    logger.info(
+        "after dropping, the choosers table has {} columns: {}".format(
+            len(choosers.columns), choosers.columns
+        )
+    )
+
+    return choosers