Reduce memory of summarize step by dropping unused variables before wide merge

activitysim/abm/models/summarize.py

@@ -4,6 +4,7 @@
 
 import logging
 import os
+import re
 from pathlib import Path
 
 import numpy as np
@@ -12,6 +13,7 @@
 from activitysim.core import expressions, timing, workflow
 from activitysim.core.configuration.base import PreprocessorSettings, PydanticReadable
 from activitysim.core.los import Network_LOS
+from activitysim.core.util import drop_unused_columns
 
 logger = logging.getLogger(__name__)
 
@@ -202,6 +204,62 @@ def manual_breaks(
         return bins
 
 
+def trim_unused_columns(
+    state: workflow.State,
+    model_settings: SummarizeSettings,
+    spec: pd.DataFrame,
+    *tables: pd.DataFrame,
+) -> tuple[pd.DataFrame, ...]:
+    """
+    Drop columns that are not referenced by the summarize specification.
+
+    Column references are found in the summarize spec, the preprocessor spec (if any),
+    and the BIN / AGGREGATE columns in the settings.  Because `trips` and `tours_merged`
+    are merged with ``_trip`` / ``_tour`` suffixes, the unsuffixed names of any suffixed
+    references are retained in both tables so that the merge produces the same columns.
+    """
+    expressions_df = spec[["Expression"]].dropna().astype(str)
+    if model_settings.preprocessor is not None:
+        preprocessor_spec = pd.read_csv(
+            state.filesystem.get_config_file_path(model_settings.preprocessor.SPEC),
+            comment="#",
+        )
+        expressions_df = pd.concat(
+            [expressions_df, preprocessor_spec[["Expression"]].dropna().astype(str)]
+        )
+
+    identifiers = set(
+        re.findall(r"[a-zA-Z_][a-zA-Z0-9_]*", " ".join(expressions_df["Expression"]))
+    )
+
+    # columns needed to merge trips and tours, and by wrap_skims
+    keep = {"tour_id", "person_id", "household_id", "start", "end", "depart"}
+    keep |= {"origin", "destination"}
+    keep |= {i.removesuffix(s) for i in identifiers for s in ("_trip", "_tour")}
+
+    for table_name in ("persons_merged", "trips_merged", "tours_merged", "land_use"):
+        meta = getattr(model_settings, table_name, None) or {}
+        for slicer in (meta.get("AGGREGATE") or []) + (meta.get("BIN") or []):
+            keep.add(slicer["column"])
+
+    return tuple(
+        drop_unused_columns(
+            table,
+            expressions_df,
+            None,
+            custom_chooser=None,
+            additional_columns=keep,
+        )
+        for table in tables
+    )
+
+
+def delete_temporary_variables(locals_d: dict[str, object], variables: str) -> None:
+    logger.debug(f"Deleting temporary variable(s): {variables}")
+    for var in variables.split(","):
+        locals_d.pop(var.strip(), None)
+
+
 class SummarizeSettings(PydanticReadable, extra="allow"):
     """
     Settings for the `summarize` component.
@@ -216,6 +274,12 @@ class SummarizeSettings(PydanticReadable, extra="allow"):
     EXPORT_PIPELINE_TABLES: bool = True
     """To export pipeline tables for expression development."""
 
+    DROP_UNUSED_COLUMNS: bool = True
+    """Drop trips and tours columns not referenced by the summarize specs before merging.
+
+    Ignored (no columns are dropped) when `EXPORT_PIPELINE_TABLES` is True.  Set to False
+    if expressions access columns indirectly, e.g. by names assembled at runtime."""
+
     preprocessor: PreprocessorSettings | None = None
 
 
@@ -261,6 +325,13 @@ def summarize(
         comment="#",
     )
 
+    # Trim the tables that are merged below to only the columns the specs use,
+    # unless the full tables are going to be exported for expression development
+    if model_settings.DROP_UNUSED_COLUMNS and not model_settings.EXPORT_PIPELINE_TABLES:
+        trips, tours_merged = trim_unused_columns(
+            state, model_settings, spec, trips, tours_merged
+        )
+
     # Load dataframes from pipeline
     tours = tours_merged
 
@@ -366,10 +437,8 @@ def summarize(
 
         # delete temporary variables listed in Expression when Output == "_del"
         if out_file == "_del":
-            logger.debug(f"Deleting temporary variable(s): {expr}")
             with performance_timer.time_expression(expr):
-                for var in str(expr).split(","):
-                    locals_d.pop(var.strip(), None)
+                delete_temporary_variables(locals_d, str(expr))
             continue
 
         # Save temporary variables starting with underscores in locals_d

test/summarize/test_summarize.py

@@ -150,3 +150,60 @@ def test_summarize(initialize_pipeline: workflow.State, caplog):
         )
     )
     assert temporary_dataframe_deleted["deleted"].tolist() == [True]
+
+
+def test_summarize_drop_unused_columns(
+    initialize_pipeline: workflow.State, monkeypatch
+):
+    from activitysim.abm.models import summarize as summarize_module
+    from activitysim.abm.models.summarize import SummarizeSettings
+
+    state = initialize_pipeline
+    state.settings.duplicate_step_execution = "allow"
+    read_settings_file = SummarizeSettings.read_settings_file.__func__
+
+    column_counts = []
+    trim_unused_columns = summarize_module.trim_unused_columns
+
+    def recording_trim(state, model_settings, spec, *tables):
+        trimmed = trim_unused_columns(state, model_settings, spec, *tables)
+        column_counts.append(
+            [(len(a.columns), len(b.columns)) for a, b in zip(tables, trimmed)]
+        )
+        return trimmed
+
+    monkeypatch.setattr(summarize_module, "trim_unused_columns", recording_trim)
+
+    output_dirs = {}
+    for drop_unused in (False, True):
+
+        def patched_read_settings_file(cls, *args, drop_unused=drop_unused, **kwargs):
+            model_settings = read_settings_file(cls, *args, **kwargs)
+            model_settings.EXPORT_PIPELINE_TABLES = False
+            model_settings.DROP_UNUSED_COLUMNS = drop_unused
+            model_settings.OUTPUT = f"summarize_drop_{drop_unused}"
+            return model_settings
+
+        monkeypatch.setattr(
+            SummarizeSettings,
+            "read_settings_file",
+            classmethod(patched_read_settings_file),
+        )
+        state.run.summarize()
+        output_dirs[drop_unused] = Path(
+            state.get_output_file_path(f"summarize_drop_{drop_unused}")
+        )
+
+    # trimming is only invoked when enabled, and must actually remove columns
+    assert len(column_counts) == 1
+    assert all(after < before for before, after in column_counts[0])
+
+    full_files = sorted(p.name for p in output_dirs[False].glob("*.csv"))
+    assert full_files
+    assert full_files == sorted(p.name for p in output_dirs[True].glob("*.csv"))
+    for name in full_files:
+        pd.testing.assert_frame_equal(
+            pd.read_csv(output_dirs[False] / name),
+            pd.read_csv(output_dirs[True] / name),
+            obj=name,
+        )