Reduce peak memory use in destination choice and large trip models

activitysim/abm/models/location_choice.py

@@ -11,7 +11,15 @@
 from activitysim.abm.models.util import tour_destination
 from activitysim.abm.models.util.bias_logsums import maybe_bias_logsums
 from activitysim.abm.tables import shadow_pricing
-from activitysim.core import estimation, expressions, los, simulate, tracing, workflow
+from activitysim.core import (
+    chunk,
+    estimation,
+    expressions,
+    los,
+    simulate,
+    tracing,
+    workflow,
+)
 from activitysim.core.configuration.logit import (
     TourLocationComponentSettings,
     TourModeComponentSettings,
@@ -520,8 +528,7 @@ def run_location_sample(
     full_dest_size_terms = dest_size_terms
 
     logger.debug(
-        f"dropping {(~(dest_size_terms.size_term > 0)).sum()} "
-        f"of {len(dest_size_terms)} rows where size_term is zero"
+        f"dropping {(~(dest_size_terms.size_term > 0)).sum()} of {len(dest_size_terms)} rows where size_term is zero"
     )
     dest_size_terms = dest_size_terms[dest_size_terms.size_term > 0]
 
@@ -532,8 +539,7 @@ def run_location_sample(
     if pre_sample_taz and not state.settings.want_dest_choice_presampling:
         pre_sample_taz = False
         logger.info(
-            f"Disabled destination zone presampling for {trace_label} "
-            f"because 'want_dest_choice_presampling' setting is False"
+            f"Disabled destination zone presampling for {trace_label} because 'want_dest_choice_presampling' setting is False"
         )
 
     if pre_sample_taz:
@@ -616,23 +622,56 @@ def run_location_logsums(
 
     logger.info(f"Running {trace_label} with {len(location_sample_df.index)} rows")
 
-    choosers = location_sample_df.join(persons_merged_df, how="left")
-
     tour_purpose = model_settings.LOGSUM_TOUR_PURPOSE
     if isinstance(tour_purpose, dict):
         tour_purpose = tour_purpose[segment_name]
 
-    logsums = logsum.compute_location_choice_logsums(
+    # Join sampled alternatives to person attributes inside the chooser chunk.
+    # The old full-table join could retain millions of rows and all derived
+    # logsum preprocessor columns before the utility evaluator began chunking.
+    # At production scale that defeated explicit chunking and exhausted memory.
+    pnr_index_multiplier = logsum.get_pnr_index_multiplier(
+        location_sample_df, logsum_settings
+    )
+    logsum_chunks = []
+    for (
+        _i,
+        persons_chunk,
+        location_sample_chunk,
+        chunk_trace_label,
+        chunk_sizer,
+    ) in chunk.adaptive_chunked_choosers_and_alts(
         state,
-        choosers,
-        tour_purpose,
-        logsum_settings,
-        model_settings,
-        network_los,
-        chunk_size,
-        chunk_tag,
+        persons_merged_df,
+        location_sample_df,
         trace_label,
-    )
+        chunk_tag,
+        chunk_size=chunk_size,
+        explicit_chunk_size=model_settings.explicit_chunk,
+    ):
+        choosers = location_sample_chunk.join(persons_chunk, how="left")
+        assert choosers.index.equals(location_sample_chunk.index)
+        chunk_sizer.log_df(chunk_trace_label, "logsum_choosers", choosers)
+
+        logsum_chunks.append(
+            logsum.compute_location_choice_logsums(
+                state,
+                choosers,
+                tour_purpose,
+                logsum_settings,
+                model_settings,
+                network_los,
+                0,
+                chunk_tag,
+                chunk_trace_label,
+                explicit_chunk_size=0,
+                pnr_index_multiplier=pnr_index_multiplier,
+            )
+        )
+        chunk_sizer.log_df(chunk_trace_label, "logsum_choosers", None)
+
+    logsums = pd.concat(logsum_chunks)
+    assert logsums.index.equals(location_sample_df.index)
 
     # "add_column series should have an index matching the table to which it is being added"
     # when the index has duplicates, however, in the special case that the series index exactly

activitysim/abm/models/park_and_ride_lot_choice.py

@@ -184,6 +184,9 @@ def run_park_and_ride_lot_choice(
     model_settings_file_name: str = "park_and_ride_lot_choice.yaml",
     pnr_capacity_cls: ParkAndRideCapacity | None = None,
     trace_label: str = "park_and_ride_lot_choice",
+    chunk_size: int | None = None,
+    explicit_chunk_size: float | None = None,
+    chooser_index_multiplier: int | None = None,
 ) -> pd.Series:
     """
     Run the park-and-ride lot choice model.
@@ -219,17 +222,21 @@ def run_park_and_ride_lot_choice(
         pnr_alts["pnr_lot_full"] = 0
 
     original_index = None
-    if not choosers.index.is_unique:
+    if chooser_index_multiplier is not None or not choosers.index.is_unique:
         # non-unique index will crash interaction_simulate
         # so we need to reset the index and add it to ActivitySim's rng
         # this happens while the disaggregate accessibility model is running pnr lot choice
         original_index = choosers.index
         oi_name = original_index.name
         oi_name = oi_name if oi_name else "index"
         choosers = choosers.reset_index(drop=False)
-        idx_multiplier = choosers.groupby(oi_name).size().max()
-        # round to the nearest 10's place
-        idx_multiplier = int(np.ceil(idx_multiplier / 10.0) * 10)
+        # A logsum caller supplies the full segment's multiplier. Use it even
+        # when this chunk happens to have unique indices, so all chunks use the
+        # same synthetic RNG channel and keys as the unchunked segment.
+        idx_multiplier = chooser_index_multiplier
+        if idx_multiplier is None:
+            idx_multiplier = choosers.groupby(oi_name).size().max()
+            idx_multiplier = int(np.ceil(idx_multiplier / 10.0) * 10)
         choosers.index = (
             original_index * idx_multiplier + choosers.groupby(oi_name).cumcount()
         )
@@ -347,7 +354,12 @@ def run_park_and_ride_lot_choice(
         trace_label=trace_label,
         trace_choice_name=trace_label,
         estimator=estimator,
-        explicit_chunk_size=model_settings.explicit_chunk,
+        chunk_size=chunk_size,
+        explicit_chunk_size=(
+            model_settings.explicit_chunk
+            if explicit_chunk_size is None
+            else explicit_chunk_size
+        ),
         compute_settings=model_settings.compute_settings,
     )