From c4c6ab4724948e019afd1bad369d8c723896a5fa Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Mon, 7 Sep 2026 09:31:15 -0500 Subject: [PATCH 1/7] Reduce memory pressure in destination choice models Chunk logsum joins and release allocator pages between model segments. Preserve tracing run IDs for programmatic multiprocess runs and include example resources in built packages. --- activitysim/abm/models/location_choice.py | 67 ++- activitysim/abm/models/trip_destination.py | 453 +++++++++++++----- activitysim/abm/models/trip_mode_choice.py | 59 +-- activitysim/abm/models/trip_scheduling.py | 39 +- activitysim/abm/models/util/logsums.py | 8 +- .../abm/models/util/tour_destination.py | 75 ++- activitysim/abm/test/test_location_choice.py | 58 +++ .../test_tour_destination_sampling.py | 135 ++++++ .../test/trip_dest/test_trip_destination.py | 42 +- activitysim/core/mem.py | 76 ++- activitysim/core/test/test_mem.py | 7 + activitysim/core/test/test_workflow_runner.py | 37 ++ activitysim/core/workflow/runner.py | 18 +- pyproject.toml | 6 + 14 files changed, 871 insertions(+), 209 deletions(-) create mode 100644 activitysim/core/test/test_mem.py create mode 100644 activitysim/core/test/test_workflow_runner.py diff --git a/activitysim/abm/models/location_choice.py b/activitysim/abm/models/location_choice.py index c955f4a606..a3587e4976 100644 --- a/activitysim/abm/models/location_choice.py +++ b/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,52 @@ 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. + 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, + chunk_size, + chunk_tag, + chunk_trace_label, + explicit_chunk_size=0, + ) + ) + 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 diff --git a/activitysim/abm/models/trip_destination.py b/activitysim/abm/models/trip_destination.py index 7dc6154cc5..ddae622764 100644 --- a/activitysim/abm/models/trip_destination.py +++ b/activitysim/abm/models/trip_destination.py @@ -26,11 +26,11 @@ estimation, expressions, los, + mem, simulate, tracing, workflow, ) -from activitysim.core.configuration.base import PreprocessorSettings from activitysim.core.configuration.logit import LocationComponentSettings from activitysim.core.exceptions import DuplicateWorkflowTableError, InvalidTravelError from activitysim.core.interaction_sample import ( @@ -82,8 +82,7 @@ def deprecated_destination_prefix(cls, values): if values[badkey] != values[goodkey]: # both keys are given, with different values -> error raise ValueError( - f"Deprecated `{badkey}` field must have the " - f"same value as `{goodkey}` if both are provided." + f"Deprecated `{badkey}` field must have the same value as `{goodkey}` if both are provided." ) else: # both keys are given, with same values -> warning @@ -98,9 +97,7 @@ def deprecated_destination_prefix(cls, values): else: # only the wrong key is given -> warning warnings.warn( - f"Use of the field `{badkey}` in the " - "trip_destination configuration file is deprecated, use " - f"`{goodkey}` instead.", + f"Use of the field `{badkey}` in the trip_destination configuration file is deprecated, use `{goodkey}` instead.", FutureWarning, stacklevel=2, ) @@ -149,6 +146,7 @@ def _destination_sample( chunk_tag: str, trace_label: str, zone_layer=None, + preprocess_alternatives: bool = True, ): """ @@ -209,16 +207,17 @@ def _destination_sample( log_alt_losers = state.settings.log_alt_losers - # preprocessing alternatives - expressions.annotate_preprocessors( - state, - df=alternatives, - locals_dict=locals_dict, - skims=skims, - model_settings=model_settings, - trace_label=trace_label, - preprocessor_setting_name="alts_preprocessor_sample", - ) + if preprocess_alternatives: + # preprocessing alternatives + expressions.annotate_preprocessors( + state, + df=alternatives, + locals_dict=locals_dict, + skims=skims, + model_settings=model_settings, + trace_label=trace_label, + preprocessor_setting_name="alts_preprocessor_sample", + ) # Trip destination keeps the alternative universe here so stable_alt_positions is not needed. choices = interaction_sample( @@ -694,33 +693,62 @@ def destination_presample( skims = skim_hotel.sample_skims(presample=True) - taz_sample = _destination_sample( + # Bound the entire two-stage sampling pipeline, not just utility + # evaluation inside interaction_sample. choose_MAZ_for_TAZ temporarily + # expands every sampled TAZ by its constituent MAZs; doing that for a full + # trip-purpose segment can consume many GiB even when interaction_sample is + # itself chunked. Keeping both stages inside this outer chooser loop lets + # each expanded MAZ frame be released before the next chunk. + maz_sample_chunks = [] + preprocess_alternatives = True + for ( + _i, + trips_taz_chunk, + chunk_trace_label, + _chunk_sizer, + ) in chunk.adaptive_chunked_choosers( state, - primary_purpose, trips_taz, - alternatives, - model_settings, - TAZ_size_term_matrix, - skims, - alt_dest_col_name, - estimator, - chunk_tag=chunk_tag, - trace_label=trace_label, - zone_layer="taz", - ) - - # choose a MAZ for each DEST_TAZ choice, choice probability based on MAZ size_term fraction of TAZ total - maz_sample = choose_MAZ_for_TAZ( - state, - taz_sample, - size_term_matrix, - trips, - network_los, - alt_dest_col_name, trace_label, - model_settings, - full_taz_index=full_taz_index, - ) + f"{chunk_tag}.pipeline", + chunk_size=state.settings.chunk_size, + explicit_chunk_size=model_settings.explicit_chunk, + ): + taz_sample = _destination_sample( + state, + primary_purpose, + trips_taz_chunk, + alternatives, + model_settings, + TAZ_size_term_matrix, + skims, + alt_dest_col_name, + estimator, + chunk_tag=chunk_tag, + trace_label=chunk_trace_label, + zone_layer="taz", + preprocess_alternatives=preprocess_alternatives, + ) + preprocess_alternatives = False + + # Choose a MAZ for each DEST_TAZ choice, with probability based on its + # share of the TAZ's purpose-specific size term. + maz_sample_chunk = choose_MAZ_for_TAZ( + state, + taz_sample, + size_term_matrix, + trips, + network_los, + alt_dest_col_name, + chunk_trace_label, + model_settings, + full_taz_index=full_taz_index, + ) + maz_sample_chunks.append(maz_sample_chunk) + del taz_sample, maz_sample_chunk + mem.release_memory() + + maz_sample = pd.concat(maz_sample_chunks) assert alt_dest_col_name in maz_sample @@ -766,8 +794,7 @@ def trip_destination_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: @@ -878,27 +905,12 @@ def compute_logsums( # chunk usage is uniform so better to combine chunk_tag = "trip_destination.compute_logsums" - # FIXME should pass this in? - network_los = state.get_injectable("network_los") - # - trips_merged - merge trips and tours_merged trips_merged = pd.merge( trips, tours_merged, left_on="tour_id", right_index=True, how="left" ) assert trips_merged.index.equals(trips.index) - # - choosers - merge destination_sample and trips_merged - # re/set index because pandas merge does not preserve left index if it has duplicate values! - choosers = pd.merge( - destination_sample, - trips_merged.reset_index(), - left_index=True, - right_on="trip_id", - how="left", - suffixes=("", "_r"), - ).set_index("trip_id") - assert choosers.index.equals(destination_sample.index) - logsum_settings = state.filesystem.read_model_settings( model_settings.LOGSUM_SETTINGS ) @@ -932,20 +944,6 @@ def compute_logsums( "timeframe": "trip", } - destination_sample["od_logsum"] = compute_ood_logsums( - state, - choosers, - logsum_settings, - nest_spec, - logsum_spec, - od_skims, - locals_dict, - state.settings.chunk_size, - trace_label=tracing.extend_trace_label(trace_label, "od"), - chunk_tag=chunk_tag, - explicit_chunk_size=model_settings.explicit_chunk, - ) - # - dp_logsums dp_skims = { "ORIGIN": model_settings.ALT_DEST_COL_NAME, @@ -955,19 +953,83 @@ def compute_logsums( "od_skims": skims["dp_skims"], } - destination_sample["dp_logsum"] = compute_ood_logsums( + # Merge sampled alternatives with chooser attributes inside the chunk loop. + # Previously this merge and both preprocessors ran on the full sample table + # before ``simple_simulate_logsums`` chunked utility evaluation. For large + # trip models, that briefly retained millions of rows and dozens of derived + # columns per process, defeating the purpose of explicit chunking. + od_logsum_chunks = [] + dp_logsum_chunks = [] + for ( + _i, + trips_chunk, + destination_sample_chunk, + chunk_trace_label, + chunk_sizer, + ) in chunk.adaptive_chunked_choosers_and_alts( state, - choosers, - logsum_settings, - nest_spec, - logsum_spec, - dp_skims, - locals_dict, - state.settings.chunk_size, - trace_label=tracing.extend_trace_label(trace_label, "dp"), - chunk_tag=chunk_tag, + trips_merged, + destination_sample, + trace_label, + chunk_tag, + chunk_size=state.settings.chunk_size, explicit_chunk_size=model_settings.explicit_chunk, - ) + ): + # Re/set the index because pandas merge does not preserve the left index + # when it contains the repeated trip ids of sampled alternatives. + choosers = pd.merge( + destination_sample_chunk, + trips_chunk.reset_index(), + left_index=True, + right_on="trip_id", + how="left", + suffixes=("", "_r"), + ).set_index("trip_id") + assert choosers.index.equals(destination_sample_chunk.index) + chunk_sizer.log_df(chunk_trace_label, "logsum_choosers", choosers) + + # The outer loop now owns chunking. Disable nested chunking so each + # merged/preprocessed sample chunk can be released before constructing + # the next one. + od_logsum_chunk = compute_ood_logsums( + state, + choosers, + logsum_settings, + nest_spec, + logsum_spec, + od_skims, + locals_dict, + state.settings.chunk_size, + trace_label=tracing.extend_trace_label(chunk_trace_label, "od"), + chunk_tag=chunk_tag, + explicit_chunk_size=0, + ) + dp_logsum_chunk = compute_ood_logsums( + state, + choosers, + logsum_settings, + nest_spec, + logsum_spec, + dp_skims, + locals_dict, + state.settings.chunk_size, + trace_label=tracing.extend_trace_label(chunk_trace_label, "dp"), + chunk_tag=chunk_tag, + explicit_chunk_size=0, + ) + od_logsum_chunks.append(od_logsum_chunk) + dp_logsum_chunks.append(dp_logsum_chunk) + chunk_sizer.log_df(chunk_trace_label, "logsum_choosers", None) + del choosers, od_logsum_chunk, dp_logsum_chunk + mem.release_memory() + + od_logsums = pd.concat(od_logsum_chunks) + dp_logsums = pd.concat(dp_logsum_chunks) + assert od_logsums.index.equals(destination_sample.index) + assert dp_logsums.index.equals(destination_sample.index) + + destination_sample["od_logsum"] = od_logsums + destination_sample["dp_logsum"] = dp_logsums return destination_sample @@ -1026,9 +1088,6 @@ def trip_destination_simulate( trip_period_idx = skims["odt_skims"].map_time_periods(trips) if trip_period_idx is not None: trips["trip_period"] = trip_period_idx - else: - None - locals_dict = model_settings.CONSTANTS.copy() locals_dict.update( { @@ -1040,38 +1099,94 @@ def trip_destination_simulate( ) locals_dict.update(skims) - # preprocessing alternatives - expressions.annotate_preprocessors( - state, - df=destination_sample, - locals_dict=locals_dict, - skims=skims, - model_settings=model_settings, - trace_label=trace_label, - preprocessor_setting_name="alts_preprocessor_simulate", - ) - log_alt_losers = state.settings.log_alt_losers - destinations = interaction_sample_simulate( - state, - choosers=trips, - alternatives=destination_sample, - spec=spec, - choice_column=alt_dest_col_name, - log_alt_losers=log_alt_losers, - want_logsums=want_logsums, - allow_zero_probs=True, - zero_prob_choice_val=NO_DESTINATION, - skims=skims, - locals_d=locals_dict, - chunk_size=state.settings.chunk_size, - chunk_tag=chunk_tag, - trace_label=trace_label, - trace_choice_name="trip_dest", - estimator=estimator, - explicit_chunk_size=model_settings.explicit_chunk, - alts_context=alts_context, - ) + if estimator: + # Preserve the estimator's single-call lifecycle and output bundle. + expressions.annotate_preprocessors( + state, + df=destination_sample, + locals_dict=locals_dict, + skims=skims, + model_settings=model_settings, + trace_label=trace_label, + preprocessor_setting_name="alts_preprocessor_simulate", + ) + destinations = interaction_sample_simulate( + state, + choosers=trips, + alternatives=destination_sample, + spec=spec, + choice_column=alt_dest_col_name, + log_alt_losers=log_alt_losers, + want_logsums=want_logsums, + allow_zero_probs=True, + zero_prob_choice_val=NO_DESTINATION, + skims=skims, + locals_d=locals_dict, + chunk_size=state.settings.chunk_size, + chunk_tag=chunk_tag, + trace_label=trace_label, + trace_choice_name="trip_dest", + estimator=estimator, + explicit_chunk_size=model_settings.explicit_chunk, + alts_context=alts_context, + ) + else: + # The alternative preprocessor can add several derived columns. Running + # it on the complete sampled table before interaction_sample_simulate + # briefly materializes millions of rows per worker and defeats the + # evaluator's internal chunking. Preprocess and simulate each sampled + # chooser chunk end-to-end instead. + destination_chunks = [] + for ( + _i, + trips_chunk, + destination_sample_chunk, + chunk_trace_label, + _chunk_sizer, + ) in chunk.adaptive_chunked_choosers_and_alts( + state, + trips, + destination_sample, + trace_label, + chunk_tag, + chunk_size=state.settings.chunk_size, + explicit_chunk_size=model_settings.explicit_chunk, + ): + expressions.annotate_preprocessors( + state, + df=destination_sample_chunk, + locals_dict=locals_dict, + skims=skims, + model_settings=model_settings, + trace_label=chunk_trace_label, + preprocessor_setting_name="alts_preprocessor_simulate", + ) + destination_chunk = interaction_sample_simulate( + state, + choosers=trips_chunk, + alternatives=destination_sample_chunk, + spec=spec, + choice_column=alt_dest_col_name, + log_alt_losers=log_alt_losers, + want_logsums=want_logsums, + allow_zero_probs=True, + zero_prob_choice_val=NO_DESTINATION, + skims=skims, + locals_d=locals_dict, + chunk_size=state.settings.chunk_size, + chunk_tag=chunk_tag, + trace_label=chunk_trace_label, + trace_choice_name="trip_dest", + estimator=None, + explicit_chunk_size=0, + alts_context=alts_context, + ) + destination_chunks.append(destination_chunk) + del destination_chunk + mem.release_memory() + + destinations = pd.concat(destination_chunks) if not want_logsums: # for consistency, always return a dataframe with canonical column name @@ -1095,8 +1210,7 @@ def trip_destination_simulate( return destinations -@workflow.func -def choose_trip_destination( +def _choose_trip_destination_unchunked( state: workflow.State, primary_purpose, trips, @@ -1136,8 +1250,7 @@ def choose_trip_destination( dropped_trips = ~trips.index.isin(destination_sample.index.unique()) if dropped_trips.any(): logger.warning( - "%s trip_destination_sample %s trips " - "without viable destination alternatives" + "%s trip_destination_sample %s trips without viable destination alternatives" % (trace_label, dropped_trips.sum()) ) trips = trips[~dropped_trips] @@ -1186,8 +1299,7 @@ def choose_trip_destination( dropped_trips = ~trips.index.isin(destinations.index) if dropped_trips.any(): logger.warning( - "%s trip_destination_simulate %s trips " - "without viable destination alternatives" + "%s trip_destination_simulate %s trips without viable destination alternatives" % (trace_label, dropped_trips.sum()) ) @@ -1206,6 +1318,90 @@ def choose_trip_destination( return destinations, destination_sample +@workflow.func +def choose_trip_destination( + state: workflow.State, + primary_purpose, + trips, + alternatives, + tours_merged, + model_settings: TripDestinationSettings, + want_logsums, + want_sample_table, + size_term_matrix, + skim_hotel, + estimator, + chunk_size, + trace_label, +): + """Run the complete destination-choice pipeline in bounded chooser chunks. + + Sampling, logsum calculation, and final simulation used to be chunked + independently. That bounded each temporary calculation, but the complete + sampled-alternative table for a purpose remained live between stages. For + large trip-purpose segments this retained millions of rows per worker and + made memory grow again as the next trip number began. + + When explicit chunking is configured, stream each chooser chunk through + all three stages and retain only its final one-row-per-trip choice. The + estimator and sample-table paths keep their original whole-segment + lifecycle because they intentionally consume the complete sample table. + """ + + if estimator or want_sample_table or not model_settings.explicit_chunk: + return _choose_trip_destination_unchunked( + state, + primary_purpose, + trips, + alternatives, + tours_merged, + model_settings, + want_logsums, + want_sample_table, + size_term_matrix, + skim_hotel, + estimator, + chunk_size, + trace_label, + ) + + destination_chunks = [] + for ( + _i, + trips_chunk, + chunk_trace_label, + _chunk_sizer, + ) in chunk.adaptive_chunked_choosers( + state, + trips, + trace_label, + "trip_destination.pipeline", + chunk_size=state.settings.chunk_size, + explicit_chunk_size=model_settings.explicit_chunk, + ): + destinations_chunk, destination_sample = _choose_trip_destination_unchunked( + state, + primary_purpose, + trips_chunk, + alternatives, + tours_merged, + model_settings, + want_logsums, + want_sample_table=False, + size_term_matrix=size_term_matrix, + skim_hotel=skim_hotel, + estimator=None, + chunk_size=chunk_size, + trace_label=chunk_trace_label, + ) + assert destination_sample is None + destination_chunks.append(destinations_chunk) + del destinations_chunk, destination_sample + mem.release_memory() + + return pd.concat(destination_chunks), None + + class SkimHotel: def __init__( self, @@ -1331,8 +1527,6 @@ def run_trip_destination( model_settings = TripDestinationSettings.read_settings_file( state.filesystem, model_settings_file_name ) - preprocessor_settings = model_settings.preprocessor - # read in logsum settings if they exist, otherwise logsum calculations are skipped if model_settings.LOGSUM_SETTINGS: logsum_settings = state.filesystem.read_model_settings( @@ -1481,9 +1675,6 @@ def run_trip_destination( ) if trip_period_idx is not None: nth_trips["trip_period"] = trip_period_idx - else: - None - logger.debug( "Running %s with %d trips", nth_trace_label, nth_trips.shape[0] ) @@ -1535,6 +1726,14 @@ def run_trip_destination( assert destination_sample is not None sample_list.append(destination_sample) + # Each purpose can create very large temporary TAZ-to-MAZ and + # logsum arrays. Explicit chunking bounds one allocation, while + # allocator pressure relief prevents freed chunks from + # accumulating in a long-lived multiprocess worker before the + # next purpose begins. + if model_settings.explicit_chunk: + mem.release_memory() + destinations_df = pd.concat(choices_list) if fail_some_trips_for_testing: diff --git a/activitysim/abm/models/trip_mode_choice.py b/activitysim/abm/models/trip_mode_choice.py index 4aad9dda62..2e69af97df 100644 --- a/activitysim/abm/models/trip_mode_choice.py +++ b/activitysim/abm/models/trip_mode_choice.py @@ -3,24 +3,21 @@ from __future__ import annotations import logging -from typing import Any -import numpy as np import pandas as pd from activitysim.abm.models.util import school_escort_tours_trips from activitysim.abm.models.util.mode import mode_choice_simulate from activitysim.core import ( - chunk, config, estimation, expressions, los, + mem, simulate, tracing, workflow, ) -from activitysim.core.configuration.base import PreprocessorSettings, PydanticReadable from activitysim.core.configuration.logit import TemplatedLogitComponentSettings from activitysim.core.util import assign_in_place @@ -44,7 +41,7 @@ class TripModeChoiceSettings(TemplatedLogitComponentSettings, extra="forbid"): FORCE_ESCORTEE_CHAUFFEUR_MODE_MATCH: bool = True """ If True, overwrite the trip mode of escortee trips to match the mode selected - by the chauffeur. This is useful for school escort tours where the escortee trip + by the chauffeur. This is useful for school escort tours where the escortee trip mode (e.g., "transit") should match the chauffeur trip mode. """ @@ -96,21 +93,12 @@ def trip_mode_choice( else: tours_merged = pd.DataFrame() - # - trips_merged - merge trips and tours_merged - trips_merged = pd.merge( - trips_df, tours_merged, left_on="tour_id", right_index=True, how="left" - ) - assert trips_merged.index.equals(trips.index) - tracing.print_summary( "primary_purpose", trips_df.primary_purpose, value_counts=True ) # setup skim keys - assert "trip_period" not in trips_merged - trips_merged["trip_period"] = network_los.skim_time_period_label( - trips_merged.depart, as_cat=True - ) + assert "trip_period" not in trips_df orig_col = "origin" dest_col = "destination" @@ -138,13 +126,6 @@ def trip_mode_choice( ) od_skim_wrapper = skim_dict.wrap("origin", "destination") - if hasattr(skim_dict, "map_time_periods_from_series"): - trip_period_idx = skim_dict.map_time_periods_from_series( - trips_merged["trip_period"] - ) - if trip_period_idx is not None: - trips_merged["trip_period"] = trip_period_idx - skims = { "odt_skims": odt_skim_stack_wrapper, "dot_skims": dot_skim_stack_wrapper, @@ -169,11 +150,35 @@ def trip_mode_choice( choices_list = [] cols_to_keep_list = [] - for primary_purpose, trips_segment in trips_merged.groupby( + for primary_purpose, base_trips_segment in trips_df.groupby( "primary_purpose", observed=True ): segment_trace_label = tracing.extend_trace_label(trace_label, primary_purpose) + # A full trips/tours merge duplicates several gigabytes at production + # scale. Materialize only the chooser rows for the active purpose. + if len(tours_cols) > 0: + trips_segment = pd.merge( + base_trips_segment, + tours_merged, + left_on="tour_id", + right_index=True, + how="left", + ) + else: + trips_segment = base_trips_segment.copy() + assert trips_segment.index.equals(base_trips_segment.index) + + trips_segment["trip_period"] = network_los.skim_time_period_label( + trips_segment.depart, as_cat=True + ) + if hasattr(skim_dict, "map_time_periods_from_series"): + trip_period_idx = skim_dict.map_time_periods_from_series( + trips_segment["trip_period"] + ) + if trip_period_idx is not None: + trips_segment["trip_period"] = trip_period_idx + logger.info( "trip_mode_choice tour_type '%s' (%s trips)" % ( @@ -266,6 +271,9 @@ def trip_mode_choice( ), "{cols_not_in_choosers} from CHOOSER_COLS_TO_KEEP is not in the choosers dataframe" cols_to_keep_list.append(trips_segment[cols_to_keep]) + del trips_segment, base_trips_segment + mem.release_memory() + choices_df = pd.concat(choices_list) if estimator: @@ -295,8 +303,6 @@ def trip_mode_choice( ) ) - tracing.print_summary("trip_modes", trips_merged.tour_mode, value_counts=True) - tracing.print_summary( "trip_mode_choice choices", trips_df[mode_column_name], value_counts=True ) @@ -330,9 +336,6 @@ def trip_mode_choice( # need to update locals_dict to access skims that are the same .shape as trips table locals_dict = {} locals_dict.update(constants) - if state.settings.skip_failed_choices: - trips_merged = trips_merged.loc[~mask_skipped] - simulate.set_skim_wrapper_targets(trips_merged, skims) locals_dict.update(skims) locals_dict["timeframe"] = "trip" expressions.annotate_tables( diff --git a/activitysim/abm/models/trip_scheduling.py b/activitysim/abm/models/trip_scheduling.py index 0ce4f2ffff..da4a595770 100644 --- a/activitysim/abm/models/trip_scheduling.py +++ b/activitysim/abm/models/trip_scheduling.py @@ -15,10 +15,18 @@ split_out_school_escorting_trips, ) from activitysim.abm.models.util.trip import cleanup_failed_trips, failed_trip_cohorts -from activitysim.core import chunk, config, estimation, expressions, tracing, workflow +from activitysim.core import ( + chunk, + config, + estimation, + expressions, + mem, + tracing, + workflow, +) from activitysim.core.configuration.base import PreprocessorSettings, PydanticReadable -from activitysim.core.util import reindex from activitysim.core.exceptions import InvalidTravelError, PipelineError +from activitysim.core.util import reindex logger = logging.getLogger(__name__) @@ -521,6 +529,7 @@ def trip_scheduling( ) trips_df = trips.copy() + original_trip_columns = trips.columns if state.is_table("school_escort_trips"): school_escort_trips = state.get_dataframe("school_escort_trips") @@ -630,14 +639,36 @@ def trip_scheduling( choices_list.append(choices) - trips_df = trips.copy() - if state.is_table("school_escort_trips"): + # The working frame contains only non-school-escort trips, so rebuild + # the complete output frame as before. Release it first to avoid + # briefly retaining two full trip-table copies. + del trips_chunk, trips_df + mem.release_memory() + trips_df = trips.copy() + # separate out school escorting trips to exclude them from the model and estimation data bundle trips_df, se_trips_df, full_trips_index = split_out_school_escorting_trips( trips_df, school_escort_trips ) non_se_trips_df = trips_df + else: + # ``trips_df`` began as an exact copy of ``trips`` and scheduling only + # added working columns to this outer frame. Reuse it instead of + # allocating a second complete copy of a potentially 30-million-row + # trip table merely to discard those columns. + temporary_columns = trips_df.columns.difference(original_trip_columns) + if len(temporary_columns): + trips_df.drop(columns=temporary_columns, inplace=True) + + # Preserve the old reset-to-``trips`` semantics for the unlikely case + # where a caller supplied columns that trip scheduling also uses as + # scratch space. + overwritten_columns = original_trip_columns.intersection( + ["earliest", "latest", "tour_hour", "stop_num", "chunk_id"] + ) + for column in overwritten_columns: + trips_df[column] = trips[column] choices = pd.concat(choices_list) choices = choices.reindex(trips_df.index) diff --git a/activitysim/abm/models/util/logsums.py b/activitysim/abm/models/util/logsums.py index a0ec5b8457..ccba509fc2 100644 --- a/activitysim/abm/models/util/logsums.py +++ b/activitysim/abm/models/util/logsums.py @@ -7,12 +7,12 @@ import pandas as pd from pydantic import BaseModel as PydanticBase +from activitysim.abm.models.park_and_ride_lot_choice import run_park_and_ride_lot_choice from activitysim.core import config, expressions, los, simulate, tracing, workflow from activitysim.core.configuration.logit import ( TourLocationComponentSettings, TourModeComponentSettings, ) -from activitysim.abm.models.park_and_ride_lot_choice import run_park_and_ride_lot_choice logger = logging.getLogger(__name__) @@ -192,6 +192,7 @@ def compute_location_choice_logsums( in_period_col: str | None = None, out_period_col: str | None = None, duration_col: str | None = None, + explicit_chunk_size: float | None = None, ): """ @@ -351,6 +352,9 @@ def compute_location_choice_logsums( trace_label=trace_label, ) + if explicit_chunk_size is None: + explicit_chunk_size = model_settings.explicit_chunk + logsums = simulate.simple_simulate_logsums( state, choosers, @@ -361,7 +365,7 @@ def compute_location_choice_logsums( chunk_size=chunk_size, chunk_tag=chunk_tag, trace_label=trace_label, - explicit_chunk_size=model_settings.explicit_chunk, + explicit_chunk_size=explicit_chunk_size, compute_settings=logsum_settings.compute_settings, ) diff --git a/activitysim/abm/models/util/tour_destination.py b/activitysim/abm/models/util/tour_destination.py index 11bdb8ecdb..d129ae4a58 100644 --- a/activitysim/abm/models/util/tour_destination.py +++ b/activitysim/abm/models/util/tour_destination.py @@ -12,6 +12,7 @@ from activitysim.abm.models.util.maz_sampling import draw_maz_rands from activitysim.abm.tables.size_terms import tour_destination_size_terms from activitysim.core import ( + chunk, config, estimation, expressions, @@ -712,8 +713,7 @@ def run_destination_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: @@ -793,31 +793,59 @@ def run_destination_logsums( chunk_tag = "tour_destination.logsums" - # merge persons into tours - choosers = pd.merge( - destination_sample, - persons_merged, - left_on=chooser_id_column, - right_index=True, - how="left", - ) - - logger.debug("Running %s with %s rows", trace_label, len(choosers)) - state.tracing.dump_df(DUMP, persons_merged, trace_label, "persons_merged") - state.tracing.dump_df(DUMP, choosers, trace_label, "choosers") - logsums = logsum.compute_location_choice_logsums( + # One chooser row per tour, aligned with the repeated tour index of the + # sampled alternatives. Merge person attributes and run the logsum + # preprocessor inside this outer chunk so the full sampled table is never + # materialized with all person and derived columns at once. + tour_choosers = destination_sample.loc[ + ~destination_sample.index.duplicated(keep="first"), [chooser_id_column] + ] + logsum_chunks = [] + for ( + _i, + tour_choosers_chunk, + destination_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, + tour_choosers, + destination_sample, trace_label, - ) + chunk_tag, + chunk_size=chunk_size, + explicit_chunk_size=model_settings.explicit_chunk, + ): + choosers = pd.merge( + destination_sample_chunk, + persons_merged, + left_on=chooser_id_column, + right_index=True, + how="left", + ) + assert choosers.index.equals(destination_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, + chunk_size, + chunk_tag, + chunk_trace_label, + explicit_chunk_size=0, + ) + ) + chunk_sizer.log_df(chunk_trace_label, "logsum_choosers", None) + + logsums = pd.concat(logsum_chunks) + assert logsums.index.equals(destination_sample.index) destination_sample["mode_choice_logsum"] = logsums @@ -951,6 +979,7 @@ def run_destination_simulate( trace_choice_name="destination", estimator=estimator, skip_choice=skip_choice, + explicit_chunk_size=model_settings.explicit_chunk, compute_settings=model_settings.compute_settings, alts_context=alts_context, ) diff --git a/activitysim/abm/test/test_location_choice.py b/activitysim/abm/test/test_location_choice.py index 9bcdd7711f..4146321122 100644 --- a/activitysim/abm/test/test_location_choice.py +++ b/activitysim/abm/test/test_location_choice.py @@ -88,3 +88,61 @@ def test_estimation_override_preserves_destination_choice_logsum(monkeypatch): ) pdt.assert_frame_equal(choices, expected) assert sample is None + + +def test_location_logsums_join_person_attributes_inside_chunks(monkeypatch): + person_index = pd.Index([1, 2], name="person_id") + persons = pd.DataFrame({"income": [10, 20]}, index=person_index) + sample = pd.DataFrame( + {"alt_dest": [101, 102, 201, 202, 203]}, + index=pd.Index([1, 1, 2, 2, 2], name="person_id"), + ) + model_settings = SimpleNamespace( + LOGSUM_SETTINGS="tour_mode_choice.yaml", + LOGSUM_TOUR_PURPOSE="work", + explicit_chunk=0.5, + ) + state = SimpleNamespace(filesystem=Mock()) + chunk_sizer = Mock() + + monkeypatch.setattr( + location_choice.TourModeComponentSettings, + "read_settings_file", + lambda *args, **kwargs: SimpleNamespace(), + ) + + def chunked(*args, **kwargs): + assert args[1] is persons + assert args[2] is sample + assert kwargs["explicit_chunk_size"] == 0.5 + yield 1, persons.iloc[:1], sample.iloc[:2], "logsums.i1", chunk_sizer + yield 2, persons.iloc[1:], sample.iloc[2:], "logsums.i2", chunk_sizer + + monkeypatch.setattr( + location_choice.chunk, "adaptive_chunked_choosers_and_alts", chunked + ) + chooser_lengths = [] + + def compute_logsums(_state, choosers, *args, **kwargs): + chooser_lengths.append(len(choosers)) + assert kwargs["explicit_chunk_size"] == 0 + return choosers["alt_dest"] + choosers["income"] + + monkeypatch.setattr( + location_choice.logsum, "compute_location_choice_logsums", compute_logsums + ) + + result = location_choice.run_location_logsums( + state, + "work", + persons, + Mock(), + sample, + model_settings, + chunk_size=0, + chunk_tag="school_location.logsums", + trace_label="school_location.logsums.work", + ) + + assert chooser_lengths == [2, 3] + assert result[location_choice.ALT_LOGSUM].tolist() == [111, 112, 221, 222, 223] diff --git a/activitysim/abm/test/test_misc/test_tour_destination_sampling.py b/activitysim/abm/test/test_misc/test_tour_destination_sampling.py index 2157fdcadb..1ab2eab4f3 100644 --- a/activitysim/abm/test/test_misc/test_tour_destination_sampling.py +++ b/activitysim/abm/test/test_misc/test_tour_destination_sampling.py @@ -1,6 +1,7 @@ from __future__ import annotations from types import SimpleNamespace +from unittest.mock import Mock import numpy as np import pandas as pd @@ -9,6 +10,140 @@ from activitysim.core import workflow +def test_destination_logsums_join_persons_inside_chunks(monkeypatch): + persons = pd.DataFrame( + {"income": [10, 20]}, index=pd.Index([1, 2], name="person_id") + ) + sample = pd.DataFrame( + { + "person_id": [1, 1, 2, 2, 2], + "alt_dest": [101, 102, 201, 202, 203], + }, + index=pd.Index([11, 11, 22, 22, 22], name="tour_id"), + ) + model_settings = SimpleNamespace( + LOGSUM_SETTINGS="tour_mode_choice.yaml", + CHOOSER_ID_COLUMN="person_id", + explicit_chunk=0.5, + ) + state = SimpleNamespace( + filesystem=SimpleNamespace( + read_model_settings=lambda *args, **kwargs: SimpleNamespace() + ), + tracing=SimpleNamespace(dump_df=lambda *args, **kwargs: None), + ) + chunk_sizer = Mock() + + def chunked(*args, **kwargs): + assert args[1].index.tolist() == [11, 22] + assert args[2] is sample + assert kwargs["explicit_chunk_size"] == 0.5 + yield 1, args[1].iloc[:1], sample.iloc[:2], "logsums.i1", chunk_sizer + yield 2, args[1].iloc[1:], sample.iloc[2:], "logsums.i2", chunk_sizer + + monkeypatch.setattr( + tour_destination.chunk, "adaptive_chunked_choosers_and_alts", chunked + ) + chooser_lengths = [] + + def compute_logsums(_state, choosers, *args, **kwargs): + chooser_lengths.append(len(choosers)) + assert kwargs["explicit_chunk_size"] == 0 + return choosers["alt_dest"] + choosers["income"] + + monkeypatch.setattr( + tour_destination.logsum, + "compute_location_choice_logsums", + compute_logsums, + ) + + result = tour_destination.run_destination_logsums( + state, + "shopping", + persons, + sample, + model_settings, + Mock(), + chunk_size=0, + trace_label="non_mandatory.shopping.logsums", + ) + + assert chooser_lengths == [2, 3] + assert result["mode_choice_logsum"].tolist() == [111, 112, 221, 222, 223] + + +def test_destination_simulate_forwards_explicit_chunk_size(monkeypatch): + tours = pd.DataFrame( + {"person_id": [1, 2], "home_zone_id": [10, 20]}, + index=pd.Index([11, 22], name="tour_id"), + ) + persons = pd.DataFrame( + {"income": [100, 200]}, index=pd.Index([1, 2], name="person_id") + ) + sample = pd.DataFrame( + {"alt_dest": [101, 102, 201]}, + index=pd.Index([11, 11, 22], name="tour_id"), + ) + destination_size_terms = pd.DataFrame( + {"size_term": [1.0, 2.0, 3.0]}, + index=pd.Index([101, 102, 201], name="alt_dest"), + ) + model_settings = SimpleNamespace( + SPEC="destination.csv", + COEFFICIENTS="coefficients.csv", + CHOOSER_ID_COLUMN="person_id", + ALT_DEST_COL_NAME="alt_dest", + CHOOSER_ORIG_COL_NAME="home_zone_id", + CONSTANTS=None, + explicit_chunk=0.25, + compute_settings=SimpleNamespace(), + ) + state = SimpleNamespace( + settings=SimpleNamespace(log_alt_losers=False, use_explicit_error_terms=False), + tracing=SimpleNamespace(dump_df=lambda *args, **kwargs: None), + ) + network_los = SimpleNamespace(get_default_skim_dict=lambda: _DummySkimDict()) + monkeypatch.setattr( + tour_destination.simulate, + "spec_for_segment", + lambda *args, **kwargs: pd.DataFrame({"coefficient": [1.0]}), + ) + monkeypatch.setattr( + tour_destination.expressions, + "annotate_preprocessors", + lambda *args, **kwargs: None, + ) + captured = {} + + def simulate_sampled(*args, **kwargs): + captured.update(kwargs) + return pd.DataFrame( + {"choice": [101, 201], "logsum": [1.0, 2.0]}, index=tours.index + ) + + monkeypatch.setattr( + tour_destination, "interaction_sample_simulate", simulate_sampled + ) + + result = tour_destination.run_destination_simulate( + state, + "shopping", + tours, + persons, + sample, + want_logsums=True, + model_settings=model_settings, + network_los=network_los, + destination_size_terms=destination_size_terms, + estimator=None, + chunk_size=0, + trace_label="non_mandatory.shopping.simulate", + ) + + assert captured["explicit_chunk_size"] == 0.25 + assert result["choice"].tolist() == [101, 201] + + class _DummySkimDict: def wrap(self, orig_key, dest_key): return type("WrappedSkims", (), {"orig_key": orig_key, "dest_key": dest_key})() diff --git a/activitysim/abm/test/trip_dest/test_trip_destination.py b/activitysim/abm/test/trip_dest/test_trip_destination.py index fdfd9b6de9..7e5686ef6c 100644 --- a/activitysim/abm/test/trip_dest/test_trip_destination.py +++ b/activitysim/abm/test/trip_dest/test_trip_destination.py @@ -9,18 +9,31 @@ from activitysim.core import workflow as wf -def test_trip_destination(tmp_path: Path): +def run_trip_destination( + tmp_path: Path, + explicit_chunk: float | None = None, + repeat_work_tours: int = 1, +): shutil.copytree( Path(__file__).parent.joinpath("configs"), tmp_path.joinpath("configs") ) shutil.copytree(Path(__file__).parent.joinpath("data"), tmp_path.joinpath("data")) + if explicit_chunk is not None: + with (tmp_path / "configs" / "trip_destination.yaml").open("a") as stream: + stream.write(f"\nexplicit_chunk: {explicit_chunk}\n") + state = wf.State.make_default(working_dir=tmp_path) # init tours tours = pd.read_csv( tmp_path / state.filesystem.data_dir[0] / "tours.csv" ).set_index("tour_id") + base_tour = tours.loc[[500]] + tours = pd.concat( + [tours] + + [base_tour.rename(index={500: 510 + i}) for i in range(1, repeat_work_tours)] + ).sort_index() state.add_table("tours", tours) state.tracing.register_traceable_table("tours", tours) state.get_rn_generator().add_channel("tours", tours) @@ -29,13 +42,25 @@ def test_trip_destination(tmp_path: Path): trips = pd.read_csv( tmp_path / state.filesystem.data_dir[0] / "trips.csv" ).set_index("trip_id") + base_trips = trips[trips.tour_id == 500] + repeated_trips = [trips] + for i in range(1, repeat_work_tours): + trip_copy = base_trips.copy() + trip_copy.index += 100_000 * i + trip_copy["tour_id"] = 510 + i + repeated_trips.append(trip_copy) + trips = pd.concat(repeated_trips).sort_index() state.add_table("trips", trips) state.tracing.register_traceable_table("trips", trips) state.get_rn_generator().add_channel("trips", trips) state.run.all() - out_trips = state.get_dataframe("trips") + return state.get_dataframe("trips") + + +def test_trip_destination(tmp_path: Path): + out_trips = run_trip_destination(tmp_path) # logsums are generated for intermediate trips only assert out_trips["destination_logsum"].isna().tolist() == [ @@ -46,3 +71,16 @@ def test_trip_destination(tmp_path: Path): False, True, ] + + +def test_trip_destination_chunked_logsums_match_unchunked(tmp_path: Path): + # Repeat one tour so each work-purpose presample has enough chooser rows to + # exercise the outer TAZ-to-MAZ pipeline chunker. + unchunked = run_trip_destination(tmp_path / "unchunked", repeat_work_tours=4) + chunked = run_trip_destination( + tmp_path / "chunked", + explicit_chunk=0.5, + repeat_work_tours=4, + ) + + pd.testing.assert_frame_equal(chunked, unchunked) diff --git a/activitysim/core/mem.py b/activitysim/core/mem.py index fbfa11fbe5..19440f24cd 100644 --- a/activitysim/core/mem.py +++ b/activitysim/core/mem.py @@ -2,12 +2,14 @@ # See full license in LICENSE.txt. from __future__ import annotations +import ctypes import datetime import gc import glob import logging import multiprocessing import os +import sys import threading import time @@ -15,7 +17,7 @@ import pandas as pd import psutil -from activitysim.core import config, util, workflow +from activitysim.core import util, workflow logger = logging.getLogger(__name__) @@ -29,13 +31,79 @@ MEM_TICK = 0 MEM_LOG_FILE_NAME = "mem.csv" -OMNIBUS_LOG_FILE_NAME = f"omnibus_mem.csv" +OMNIBUS_LOG_FILE_NAME = "omnibus_mem.csv" SUMMARY_BIN_SIZE_IN_SECONDS = 15 mem_log_lock = threading.Lock() +def release_memory(): + """Return unused allocator pages to the operating system when possible. + + Garbage collection destroys unreachable Python objects, but native arrays + created by pandas, NumPy, and compiled model evaluators can leave free pages + in the process allocator. Long, segmented model steps may therefore retain + a high resident set even after each segment's temporary frames are gone. + + The allocator-pressure calls below are advisory and platform specific. A + failed or unavailable call is harmless: collection has still occurred and + the model continues normally. + """ + + was_disabled = not gc.isenabled() + if was_disabled: + gc.enable() + gc.collect() + if was_disabled: + gc.disable() + + try: + if sys.platform.startswith("linux"): + libc = ctypes.CDLL(None) + malloc_trim = getattr(libc, "malloc_trim", None) + if malloc_trim is not None: + malloc_trim.argtypes = [ctypes.c_size_t] + malloc_trim.restype = ctypes.c_int + return bool(malloc_trim(0)) + + elif sys.platform == "darwin": + libc = ctypes.CDLL(None) + malloc_default_zone = getattr(libc, "malloc_default_zone", None) + pressure_relief = getattr(libc, "malloc_zone_pressure_relief", None) + if malloc_default_zone is not None and pressure_relief is not None: + malloc_default_zone.argtypes = [] + malloc_default_zone.restype = ctypes.c_void_p + pressure_relief.argtypes = [ctypes.c_void_p, ctypes.c_size_t] + pressure_relief.restype = ctypes.c_size_t + return bool(pressure_relief(malloc_default_zone(), 0)) + + elif sys.platform == "win32": + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + get_current_process = kernel32.GetCurrentProcess + get_current_process.argtypes = [] + get_current_process.restype = ctypes.c_void_p + set_working_set_size = kernel32.SetProcessWorkingSetSize + set_working_set_size.argtypes = [ + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_size_t, + ] + set_working_set_size.restype = ctypes.c_int + maximum_size = ctypes.c_size_t(-1).value + return bool( + set_working_set_size( + get_current_process(), + maximum_size, + maximum_size, + ) + ) + except (AttributeError, OSError, TypeError, ValueError): + logger.debug("Platform allocator did not release memory", exc_info=True) + + return False + + def time_bin(timestamps): bins_size_in_seconds = SUMMARY_BIN_SIZE_IN_SECONDS epoch = pd.Timestamp("1970-01-01") @@ -217,7 +285,7 @@ def trace_memory_info(event, trace_ticks=0, force_garbage_collect=False, *, stat child_info = child.memory_info() full_rss += child_info.rss num_children += 1 - except (psutil.NoSuchProcess, psutil.AccessDenied) as e: + except (psutil.NoSuchProcess, psutil.AccessDenied): pass noteworthy = ( @@ -296,7 +364,7 @@ def shared_memory_size(data_buffers): if data_buffers is None: data_buffers = {} - for k, data_buffer in data_buffers.items(): + for _k, data_buffer in data_buffers.items(): if isinstance(data_buffer, str) and data_buffer.startswith("sh.Dataset:"): from sharrow import Dataset diff --git a/activitysim/core/test/test_mem.py b/activitysim/core/test/test_mem.py new file mode 100644 index 0000000000..9f3ffb7e9a --- /dev/null +++ b/activitysim/core/test/test_mem.py @@ -0,0 +1,7 @@ +from __future__ import annotations + +from activitysim.core import mem + + +def test_release_memory_is_advisory(): + assert isinstance(mem.release_memory(), bool) diff --git a/activitysim/core/test/test_workflow_runner.py b/activitysim/core/test/test_workflow_runner.py new file mode 100644 index 0000000000..89c01ec3a0 --- /dev/null +++ b/activitysim/core/test/test_workflow_runner.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from activitysim.core import mp_tasks +from activitysim.core.run_id import RunId +from activitysim.core.workflow.runner import Runner + + +class ProgrammaticState: + def __init__(self): + self.settings = SimpleNamespace(memory_profile=False, multiprocess=True) + self.tracing = SimpleNamespace(run_id=RunId("abc123")) + + def __contains__(self, key): + return key == "preload_injectables" + + def get_injectable(self, key): + assert key != "run_id" + return f"value-for-{key}" + + +def test_programmatic_multiprocess_run_gets_tracing_run_id(monkeypatch): + state = ProgrammaticState() + runner = Runner(state) + received = {} + + def capture_run_multiprocess(passed_state, injectables): + assert passed_state is state + received.update(injectables) + + monkeypatch.setattr(mp_tasks, "run_multiprocess", capture_run_multiprocess) + + runner.all(config_logger=False, filter_warnings=False) + + assert received["run_id"] == "abc123" + assert received["settings"] is state.settings diff --git a/activitysim/core/workflow/runner.py b/activitysim/core/workflow/runner.py index 79ecd0ed4f..df56f4cbc4 100644 --- a/activitysim/core/workflow/runner.py +++ b/activitysim/core/workflow/runner.py @@ -5,6 +5,7 @@ import time from collections.abc import Callable, Iterable from datetime import timedelta +from typing import Any from activitysim.core import tracing from activitysim.core.exceptions import DuplicateWorkflowNameError @@ -266,8 +267,7 @@ def _pre_run_step(self, model_name: str) -> bool | None: if self._obj.settings.duplicate_step_execution == "error": checkpointed_model_bullets = "\n - ".join(checkpointed_models) raise DuplicateWorkflowNameError( - f"Checkpointed Models:\n - {checkpointed_model_bullets}\n" - f"Cannot run model '{model_name}' more than once" + f"Checkpointed Models:\n - {checkpointed_model_bullets}\nCannot run model '{model_name}' more than once" ) self._obj.rng().begin_step(model_name) @@ -412,7 +412,19 @@ def all( from activitysim.cli.run import INJECTABLES from activitysim.core import mp_tasks - injectables = {k: self._obj.get_injectable(k) for k in INJECTABLES} + # ``run_id`` is initialized lazily by the tracing accessor and + # is not necessarily present as a top-level injectable when a + # caller constructs State programmatically. The CLI mirrors it + # explicitly in ``handle_standard_args``; do the equivalent + # here so ``State.run.all()`` works for multiprocessing too. + injectables = { + key: ( + self._obj.tracing.run_id + if key == "run_id" + else self._obj.get_injectable(key) + ) + for key in INJECTABLES + } injectables["settings"] = self._obj.settings # injectables["settings_package"] = state.settings.dict() mp_tasks.run_multiprocess(self._obj, injectables) diff --git a/pyproject.toml b/pyproject.toml index df86451c8f..4ac291fb62 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,12 @@ exclude = [ "activitysim/abm/models/util/test*", ] +# Explicitly include the example resources used by ``activitysim create``. +# Relying on setuptools-scm's implicit file list works in a Git checkout, but +# produces an incomplete wheel in Docker and other source-only build contexts. +[tool.setuptools.package-data] +activitysim = ["examples/**/*"] + [tool.setuptools_scm] fallback_version = "999" write_to = "activitysim/_generated_version.py" From 6d3714fc2bb910ef8392127cb88ba18b985b6335 Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Mon, 7 Sep 2026 15:32:22 -0500 Subject: [PATCH 2/7] Preserve legacy unchunked destination sampling Apply outer chooser chunking and memory release only when explicit chunking is configured. --- activitysim/abm/models/trip_destination.py | 45 +++++++++++++--------- 1 file changed, 27 insertions(+), 18 deletions(-) diff --git a/activitysim/abm/models/trip_destination.py b/activitysim/abm/models/trip_destination.py index ddae622764..eedf07b462 100644 --- a/activitysim/abm/models/trip_destination.py +++ b/activitysim/abm/models/trip_destination.py @@ -693,27 +693,34 @@ def destination_presample( skims = skim_hotel.sample_skims(presample=True) + explicit_chunk_size = getattr(model_settings, "explicit_chunk", 0) + if explicit_chunk_size: + chooser_chunks = chunk.adaptive_chunked_choosers( + state, + trips_taz, + trace_label, + f"{chunk_tag}.pipeline", + chunk_size=state.settings.chunk_size, + explicit_chunk_size=explicit_chunk_size, + ) + else: + # Preserve the legacy unchunked call contract, including compatibility + # with callers that supply a lightweight settings object. + chooser_chunks = ((0, trips_taz, trace_label, None),) + # Bound the entire two-stage sampling pipeline, not just utility - # evaluation inside interaction_sample. choose_MAZ_for_TAZ temporarily + # evaluation inside interaction_sample. choose_MAZ_for_TAZ temporarily # expands every sampled TAZ by its constituent MAZs; doing that for a full # trip-purpose segment can consume many GiB even when interaction_sample is - # itself chunked. Keeping both stages inside this outer chooser loop lets + # itself chunked. Keeping both stages inside this outer chooser loop lets # each expanded MAZ frame be released before the next chunk. maz_sample_chunks = [] preprocess_alternatives = True - for ( - _i, - trips_taz_chunk, - chunk_trace_label, - _chunk_sizer, - ) in chunk.adaptive_chunked_choosers( - state, - trips_taz, - trace_label, - f"{chunk_tag}.pipeline", - chunk_size=state.settings.chunk_size, - explicit_chunk_size=model_settings.explicit_chunk, - ): + for _i, trips_taz_chunk, chunk_trace_label, _chunk_sizer in chooser_chunks: + sample_kwargs = {} + if explicit_chunk_size: + sample_kwargs["preprocess_alternatives"] = preprocess_alternatives + taz_sample = _destination_sample( state, primary_purpose, @@ -727,9 +734,10 @@ def destination_presample( chunk_tag=chunk_tag, trace_label=chunk_trace_label, zone_layer="taz", - preprocess_alternatives=preprocess_alternatives, + **sample_kwargs, ) - preprocess_alternatives = False + if explicit_chunk_size: + preprocess_alternatives = False # Choose a MAZ for each DEST_TAZ choice, with probability based on its # share of the TAZ's purpose-specific size term. @@ -746,7 +754,8 @@ def destination_presample( ) maz_sample_chunks.append(maz_sample_chunk) del taz_sample, maz_sample_chunk - mem.release_memory() + if explicit_chunk_size: + mem.release_memory() maz_sample = pd.concat(maz_sample_chunks) From 67c3d385a131a7641e32059cffb459fb60e2cc94 Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Mon, 7 Sep 2026 20:11:54 -0500 Subject: [PATCH 3/7] Fix trip-period annotations and logsum chunking Pass zero chunk sizes to logsum calculations, preserve configured outer chunking, and add trip-period data during post-choice annotations with regression coverage. --- activitysim/abm/models/location_choice.py | 2 +- activitysim/abm/models/trip_destination.py | 6 +- activitysim/abm/models/trip_mode_choice.py | 46 ++++-- .../abm/models/util/tour_destination.py | 2 +- activitysim/abm/test/test_location_choice.py | 4 +- .../test_tour_destination_sampling.py | 4 +- activitysim/abm/test/test_trip_mode_choice.py | 138 ++++++++++++++++++ 7 files changed, 179 insertions(+), 23 deletions(-) create mode 100644 activitysim/abm/test/test_trip_mode_choice.py diff --git a/activitysim/abm/models/location_choice.py b/activitysim/abm/models/location_choice.py index a3587e4976..04766014de 100644 --- a/activitysim/abm/models/location_choice.py +++ b/activitysim/abm/models/location_choice.py @@ -658,7 +658,7 @@ def run_location_logsums( logsum_settings, model_settings, network_los, - chunk_size, + 0, chunk_tag, chunk_trace_label, explicit_chunk_size=0, diff --git a/activitysim/abm/models/trip_destination.py b/activitysim/abm/models/trip_destination.py index eedf07b462..386d5c68de 100644 --- a/activitysim/abm/models/trip_destination.py +++ b/activitysim/abm/models/trip_destination.py @@ -1008,7 +1008,7 @@ def compute_logsums( logsum_spec, od_skims, locals_dict, - state.settings.chunk_size, + 0, trace_label=tracing.extend_trace_label(chunk_trace_label, "od"), chunk_tag=chunk_tag, explicit_chunk_size=0, @@ -1021,7 +1021,7 @@ def compute_logsums( logsum_spec, dp_skims, locals_dict, - state.settings.chunk_size, + 0, trace_label=tracing.extend_trace_label(chunk_trace_label, "dp"), chunk_tag=chunk_tag, explicit_chunk_size=0, @@ -1183,7 +1183,7 @@ def trip_destination_simulate( zero_prob_choice_val=NO_DESTINATION, skims=skims, locals_d=locals_dict, - chunk_size=state.settings.chunk_size, + chunk_size=0, chunk_tag=chunk_tag, trace_label=chunk_trace_label, trace_choice_name="trip_dest", diff --git a/activitysim/abm/models/trip_mode_choice.py b/activitysim/abm/models/trip_mode_choice.py index 2e69af97df..6169baaa7f 100644 --- a/activitysim/abm/models/trip_mode_choice.py +++ b/activitysim/abm/models/trip_mode_choice.py @@ -118,6 +118,17 @@ def trip_mode_choice( skim_dict = network_los.get_default_skim_dict() + def add_trip_period(choosers): + choosers["trip_period"] = network_los.skim_time_period_label( + choosers.depart, as_cat=True + ) + if hasattr(skim_dict, "map_time_periods_from_series"): + trip_period_idx = skim_dict.map_time_periods_from_series( + choosers["trip_period"] + ) + if trip_period_idx is not None: + choosers["trip_period"] = trip_period_idx + odt_skim_stack_wrapper = skim_dict.wrap_3d( orig_key=orig_col, dest_key=dest_col, dim3_key="trip_period" ) @@ -169,15 +180,7 @@ def trip_mode_choice( trips_segment = base_trips_segment.copy() assert trips_segment.index.equals(base_trips_segment.index) - trips_segment["trip_period"] = network_los.skim_time_period_label( - trips_segment.depart, as_cat=True - ) - if hasattr(skim_dict, "map_time_periods_from_series"): - trip_period_idx = skim_dict.map_time_periods_from_series( - trips_segment["trip_period"] - ) - if trip_period_idx is not None: - trips_segment["trip_period"] = trip_period_idx + add_trip_period(trips_segment) logger.info( "trip_mode_choice tour_type '%s' (%s trips)" @@ -338,10 +341,21 @@ def trip_mode_choice( locals_dict.update(constants) locals_dict.update(skims) locals_dict["timeframe"] = "trip" - expressions.annotate_tables( - state, - locals_dict=locals_dict, - skims=skims, - model_settings=model_settings, - trace_label=trace_label, - ) + # Three-dimensional skim wrappers require trip_period. It is normally only + # needed in the purpose-sized chooser frames above, but post-choice table + # annotators may also use those skims. Add it to the full trips table only + # for annotation, then restore the original table schema. + add_trip_period(trips_df) + try: + expressions.annotate_tables( + state, + locals_dict=locals_dict, + skims=skims, + model_settings=model_settings, + trace_label=trace_label, + ) + finally: + trips_df.drop(columns="trip_period", inplace=True) + state_trips = state.get_dataframe("trips", as_copy=False) + if state_trips is not trips_df: + state_trips.drop(columns="trip_period", inplace=True) diff --git a/activitysim/abm/models/util/tour_destination.py b/activitysim/abm/models/util/tour_destination.py index d129ae4a58..5401ad4a22 100644 --- a/activitysim/abm/models/util/tour_destination.py +++ b/activitysim/abm/models/util/tour_destination.py @@ -836,7 +836,7 @@ def run_destination_logsums( logsum_settings, model_settings, network_los, - chunk_size, + 0, chunk_tag, chunk_trace_label, explicit_chunk_size=0, diff --git a/activitysim/abm/test/test_location_choice.py b/activitysim/abm/test/test_location_choice.py index 4146321122..c8cf1b5023 100644 --- a/activitysim/abm/test/test_location_choice.py +++ b/activitysim/abm/test/test_location_choice.py @@ -114,6 +114,7 @@ def test_location_logsums_join_person_attributes_inside_chunks(monkeypatch): def chunked(*args, **kwargs): assert args[1] is persons assert args[2] is sample + assert kwargs["chunk_size"] == 123 assert kwargs["explicit_chunk_size"] == 0.5 yield 1, persons.iloc[:1], sample.iloc[:2], "logsums.i1", chunk_sizer yield 2, persons.iloc[1:], sample.iloc[2:], "logsums.i2", chunk_sizer @@ -125,6 +126,7 @@ def chunked(*args, **kwargs): def compute_logsums(_state, choosers, *args, **kwargs): chooser_lengths.append(len(choosers)) + assert args[4] == 0 assert kwargs["explicit_chunk_size"] == 0 return choosers["alt_dest"] + choosers["income"] @@ -139,7 +141,7 @@ def compute_logsums(_state, choosers, *args, **kwargs): Mock(), sample, model_settings, - chunk_size=0, + chunk_size=123, chunk_tag="school_location.logsums", trace_label="school_location.logsums.work", ) diff --git a/activitysim/abm/test/test_misc/test_tour_destination_sampling.py b/activitysim/abm/test/test_misc/test_tour_destination_sampling.py index 1ab2eab4f3..0dfcc0eeac 100644 --- a/activitysim/abm/test/test_misc/test_tour_destination_sampling.py +++ b/activitysim/abm/test/test_misc/test_tour_destination_sampling.py @@ -37,6 +37,7 @@ def test_destination_logsums_join_persons_inside_chunks(monkeypatch): def chunked(*args, **kwargs): assert args[1].index.tolist() == [11, 22] assert args[2] is sample + assert kwargs["chunk_size"] == 123 assert kwargs["explicit_chunk_size"] == 0.5 yield 1, args[1].iloc[:1], sample.iloc[:2], "logsums.i1", chunk_sizer yield 2, args[1].iloc[1:], sample.iloc[2:], "logsums.i2", chunk_sizer @@ -48,6 +49,7 @@ def chunked(*args, **kwargs): def compute_logsums(_state, choosers, *args, **kwargs): chooser_lengths.append(len(choosers)) + assert args[4] == 0 assert kwargs["explicit_chunk_size"] == 0 return choosers["alt_dest"] + choosers["income"] @@ -64,7 +66,7 @@ def compute_logsums(_state, choosers, *args, **kwargs): sample, model_settings, Mock(), - chunk_size=0, + chunk_size=123, trace_label="non_mandatory.shopping.logsums", ) diff --git a/activitysim/abm/test/test_trip_mode_choice.py b/activitysim/abm/test/test_trip_mode_choice.py new file mode 100644 index 0000000000..56009eea06 --- /dev/null +++ b/activitysim/abm/test/test_trip_mode_choice.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pandas as pd + +from activitysim.abm.models import trip_mode_choice as trip_mode_choice_module + + +class _DummySkimDict: + def wrap_3d(self, **_kwargs): + return object() + + def wrap(self, *_args): + return object() + + def map_time_periods_from_series(self, periods): + return periods.map({"AM": 0, "PM": 1}) + + +class _DummyState: + current_model_name = "test_trip_mode_choice" + + def __init__(self, trips): + self.settings = SimpleNamespace( + downcast_int=False, + downcast_float=False, + skip_failed_choices=False, + trace_hh_id=None, + ) + self.filesystem = SimpleNamespace( + read_model_spec=lambda **_kwargs: pd.DataFrame(), + get_segment_coefficients=lambda *_args: {}, + ) + self.tables = {"trips": trips} + + def add_table(self, name, df): + self.tables[name] = df + + def get_dataframe(self, name, columns=None, as_copy=True): + df = self.tables[name] + if columns is not None: + df = df[columns] + return df.copy() if as_copy else df + + def is_table(self, _name): + return False + + +def test_post_choice_annotations_receive_full_trip_period_then_remove_it(monkeypatch): + trips = pd.DataFrame( + { + "tour_id": [11, 12, 13], + "household_id": [1, 2, 3], + "primary_purpose": ["work", "shopping", "work"], + "depart": [8, 17, 9], + "origin": [1, 2, 3], + "destination": [2, 3, 1], + }, + index=pd.Index([101, 102, 103], name="trip_id"), + ) + state = _DummyState(trips) + skim_dict = _DummySkimDict() + network_los = SimpleNamespace( + skim_time_periods=SimpleNamespace(period_minutes=60), + skim_time_period_label=lambda depart, as_cat=True: depart.map( + lambda value: "AM" if value < 12 else "PM" + ), + get_default_skim_dict=lambda: skim_dict, + ) + model_settings = SimpleNamespace( + MODE_CHOICE_LOGSUM_COLUMN_NAME="mode_choice_logsum", + TOURS_MERGED_CHOOSER_COLUMNS=[], + CHOOSER_COLS_TO_KEEP=[], + FORCE_ESCORTEE_CHAUFFEUR_MODE_MATCH=False, + SPEC="trip_mode_choice.csv", + explicit_chunk=None, + compute_settings=None, + ) + + monkeypatch.setattr( + trip_mode_choice_module.tracing, "print_summary", lambda *_args, **_kwargs: None + ) + monkeypatch.setattr( + trip_mode_choice_module.config, "get_model_constants", lambda *_args: {} + ) + monkeypatch.setattr( + trip_mode_choice_module.config, "get_logit_model_settings", lambda *_args: None + ) + monkeypatch.setattr( + trip_mode_choice_module.simulate, + "eval_coefficients", + lambda _state, spec, *_args: spec, + ) + monkeypatch.setattr( + trip_mode_choice_module.simulate, + "eval_nest_coefficients", + lambda *_args: None, + ) + monkeypatch.setattr( + trip_mode_choice_module.expressions, + "annotate_preprocessors", + lambda *_args, **_kwargs: None, + ) + + def choose_mode(_state, choosers, **_kwargs): + return pd.DataFrame( + { + "trip_mode": "DRIVE", + "mode_choice_logsum": 1.0, + }, + index=choosers.index, + ) + + monkeypatch.setattr(trip_mode_choice_module, "mode_choice_simulate", choose_mode) + + def annotate_tables(_state, **_kwargs): + annotated = _state.get_dataframe("trips") + assert annotated.index.equals(trips.index) + assert annotated["trip_period"].tolist() == [0, 1, 0] + annotated["post_choice_skim_value"] = [10.0, 20.0, 30.0] + _state.add_table("trips", annotated) + + monkeypatch.setattr( + trip_mode_choice_module.expressions, "annotate_tables", annotate_tables + ) + + trip_mode_choice_module.trip_mode_choice( + state, + trips, + network_los, + model_settings=model_settings, + ) + + result = state.get_dataframe("trips", as_copy=False) + assert "trip_period" not in trips + assert "trip_period" not in result + assert result["post_choice_skim_value"].tolist() == [10.0, 20.0, 30.0] From 2c0f9895c17369756dd900f7c80d5f9c344147e6 Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Tue, 8 Sep 2026 11:04:54 -0500 Subject: [PATCH 4/7] address copilot review comments --- activitysim/abm/models/trip_destination.py | 19 ++++- activitysim/abm/models/trip_mode_choice.py | 5 ++ activitysim/abm/test/test_trip_mode_choice.py | 23 ++++- .../test/trip_dest/test_trip_destination.py | 83 +++++++++++++++++++ activitysim/core/test/test_mem.py | 80 ++++++++++++++++++ 5 files changed, 206 insertions(+), 4 deletions(-) diff --git a/activitysim/abm/models/trip_destination.py b/activitysim/abm/models/trip_destination.py index 386d5c68de..5f8157a0a1 100644 --- a/activitysim/abm/models/trip_destination.py +++ b/activitysim/abm/models/trip_destination.py @@ -694,6 +694,13 @@ def destination_presample( skims = skim_hotel.sample_skims(presample=True) explicit_chunk_size = getattr(model_settings, "explicit_chunk", 0) + if ( + explicit_chunk_size + and state.settings.chunk_training_mode != chunk.MODE_EXPLICIT + ): + # Adaptive sampling already owns a ledger; do not nest an additional + # pipeline ledger merely because the model has an explicit setting. + explicit_chunk_size = 0 if explicit_chunk_size: chooser_chunks = chunk.adaptive_chunked_choosers( state, @@ -1357,7 +1364,12 @@ def choose_trip_destination( lifecycle because they intentionally consume the complete sample table. """ - if estimator or want_sample_table or not model_settings.explicit_chunk: + if ( + state.settings.chunk_training_mode != chunk.MODE_EXPLICIT + or estimator + or want_sample_table + or not model_settings.explicit_chunk + ): return _choose_trip_destination_unchunked( state, primary_purpose, @@ -1374,6 +1386,9 @@ def choose_trip_destination( trace_label, ) + # The outer pipeline owns the chunk boundary. In particular, a fractional + # size must not be applied again to each already bounded chooser chunk. + inner_settings = model_settings.model_copy(update={"explicit_chunk": 0}) destination_chunks = [] for ( _i, @@ -1394,7 +1409,7 @@ def choose_trip_destination( trips_chunk, alternatives, tours_merged, - model_settings, + inner_settings, want_logsums, want_sample_table=False, size_term_matrix=size_term_matrix, diff --git a/activitysim/abm/models/trip_mode_choice.py b/activitysim/abm/models/trip_mode_choice.py index 6169baaa7f..acdc7e85eb 100644 --- a/activitysim/abm/models/trip_mode_choice.py +++ b/activitysim/abm/models/trip_mode_choice.py @@ -274,6 +274,11 @@ def add_trip_period(choosers): ), "{cols_not_in_choosers} from CHOOSER_COLS_TO_KEEP is not in the choosers dataframe" cols_to_keep_list.append(trips_segment[cols_to_keep]) + # Wrappers retain their last chooser frame (and possibly array views). + # Retarget to an independent empty frame before releasing this purpose. + simulate.set_skim_wrapper_targets( + trips_segment.iloc[:0].copy(), skims, allow_partial_success=False + ) del trips_segment, base_trips_segment mem.release_memory() diff --git a/activitysim/abm/test/test_trip_mode_choice.py b/activitysim/abm/test/test_trip_mode_choice.py index 56009eea06..4aa901afa2 100644 --- a/activitysim/abm/test/test_trip_mode_choice.py +++ b/activitysim/abm/test/test_trip_mode_choice.py @@ -1,5 +1,6 @@ from __future__ import annotations +import weakref from types import SimpleNamespace import pandas as pd @@ -7,12 +8,19 @@ from activitysim.abm.models import trip_mode_choice as trip_mode_choice_module +class _DummySkimWrapper: + df = None + + def set_df(self, df): + self.df = df + + class _DummySkimDict: def wrap_3d(self, **_kwargs): - return object() + return _DummySkimWrapper() def wrap(self, *_args): - return object() + return _DummySkimWrapper() def map_time_periods_from_series(self, periods): return periods.map({"AM": 0, "PM": 1}) @@ -103,7 +111,16 @@ def test_post_choice_annotations_receive_full_trip_period_then_remove_it(monkeyp lambda *_args, **_kwargs: None, ) + chooser_refs = [] + wrappers = [] + def choose_mode(_state, choosers, **_kwargs): + assert all(ref() is None for ref in chooser_refs) + chooser_refs.append(weakref.ref(choosers)) + wrappers[:] = _kwargs["skims"].values() + trip_mode_choice_module.simulate.set_skim_wrapper_targets( + choosers, _kwargs["skims"] + ) return pd.DataFrame( { "trip_mode": "DRIVE", @@ -132,6 +149,8 @@ def annotate_tables(_state, **_kwargs): model_settings=model_settings, ) + assert all(ref() is None for ref in chooser_refs) + assert all(wrapper.df.empty for wrapper in wrappers) result = state.get_dataframe("trips", as_copy=False) assert "trip_period" not in trips assert "trip_period" not in result diff --git a/activitysim/abm/test/trip_dest/test_trip_destination.py b/activitysim/abm/test/trip_dest/test_trip_destination.py index 7e5686ef6c..f3ae7ca64f 100644 --- a/activitysim/abm/test/trip_dest/test_trip_destination.py +++ b/activitysim/abm/test/trip_dest/test_trip_destination.py @@ -1,9 +1,14 @@ from __future__ import annotations import shutil +from unittest.mock import Mock from pathlib import Path import pandas as pd +import pytest + +from activitysim.abm.models import trip_destination as td +from activitysim.core import chunk from activitysim import abm # noqa: F401 from activitysim.core import workflow as wf @@ -13,6 +18,7 @@ def run_trip_destination( tmp_path: Path, explicit_chunk: float | None = None, repeat_work_tours: int = 1, + chunk_training_mode: str = chunk.MODE_EXPLICIT, ): shutil.copytree( Path(__file__).parent.joinpath("configs"), tmp_path.joinpath("configs") @@ -25,6 +31,10 @@ def run_trip_destination( state = wf.State.make_default(working_dir=tmp_path) + state.settings.chunk_training_mode = chunk_training_mode + if chunk_training_mode != chunk.MODE_EXPLICIT: + state.settings.chunk_size = 1_000_000 + # init tours tours = pd.read_csv( tmp_path / state.filesystem.data_dir[0] / "tours.csv" @@ -84,3 +94,76 @@ def test_trip_destination_chunked_logsums_match_unchunked(tmp_path: Path): ) pd.testing.assert_frame_equal(chunked, unchunked) + + +@pytest.mark.parametrize("mode", chunk.TRAINING_MODES) +@pytest.mark.parametrize("explicit_chunk", [0, 0.5, 2]) +@pytest.mark.parametrize("legacy_reason", [None, "estimator", "sample_table"]) +def test_destination_pipeline_chunk_boundary( + tmp_path, monkeypatch, mode, explicit_chunk, legacy_reason +): + state = wf.State.make_default( + working_dir=tmp_path, + configs_dir=Path(__file__).parent / "configs", + data_dir=Path(__file__).parent / "data", + ) + state.settings.chunk_training_mode = mode + state.settings.chunk_size = 1_000_000 + settings = td.TripDestinationSettings.model_construct(explicit_chunk=explicit_chunk) + trips = pd.DataFrame({"value": range(4)}, index=pd.Index(range(4), name="trip_id")) + calls = [] + + def choose( + _state, purpose, chooser, alternatives, tours, model_settings, *args, **kwargs + ): + calls.append((chooser.index.tolist(), model_settings)) + return chooser.copy(), None + + monkeypatch.setattr(td, "_choose_trip_destination_unchunked", choose) + monkeypatch.setattr(td.mem, "release_memory", lambda: False) + use_outer = mode == chunk.MODE_EXPLICIT and explicit_chunk and legacy_reason is None + if not use_outer: + outer = Mock( + side_effect=AssertionError("legacy path must not create an outer chunker") + ) + monkeypatch.setattr(td.chunk, "adaptive_chunked_choosers", outer) + result, sample = td.choose_trip_destination( + state, + "work", + trips, + None, + None, + settings, + False, + legacy_reason == "sample_table", + None, + None, + object() if legacy_reason == "estimator" else None, + 1_000_000, + "test", + ) + pd.testing.assert_frame_equal(result, trips) + assert sample is None + assert settings.explicit_chunk == explicit_chunk + if use_outer: + assert [rows for rows, _ in calls] == [[0, 1], [2, 3]] + assert all( + inner.explicit_chunk == 0 and inner is not settings for _, inner in calls + ) + else: + assert calls == [([0, 1, 2, 3], settings)] + + +def test_trip_destination_training_ignores_explicit_chunk(tmp_path): + unchunked = run_trip_destination( + tmp_path / "default", + repeat_work_tours=4, + chunk_training_mode=chunk.MODE_RETRAIN, + ) + explicit = run_trip_destination( + tmp_path / "explicit", + explicit_chunk=0.5, + repeat_work_tours=4, + chunk_training_mode=chunk.MODE_RETRAIN, + ) + pd.testing.assert_frame_equal(explicit, unchunked) diff --git a/activitysim/core/test/test_mem.py b/activitysim/core/test/test_mem.py index 9f3ffb7e9a..f1e87f6390 100644 --- a/activitysim/core/test/test_mem.py +++ b/activitysim/core/test/test_mem.py @@ -1,7 +1,87 @@ from __future__ import annotations +import ctypes +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + from activitysim.core import mem def test_release_memory_is_advisory(): assert isinstance(mem.release_memory(), bool) + + +@pytest.mark.parametrize("platform", ["linux", "darwin", "win32", "unsupported"]) +@pytest.mark.parametrize("gc_enabled", [True, False]) +@pytest.mark.parametrize( + "outcome", ["success", "no_release", "missing", "load_error", "call_error"] +) +def test_release_memory_platforms(monkeypatch, platform, gc_enabled, outcome): + # Mock the platform module itself so pytest and dependencies retain the host + # platform, and never invoke a foreign allocator on the test machine. + monkeypatch.setattr(mem, "sys", SimpleNamespace(platform=platform)) + gc = Mock() + gc.isenabled.return_value = gc_enabled + monkeypatch.setattr(mem, "gc", gc) + release = Mock(return_value=1 if outcome == "success" else 0) + handle = Mock(return_value=1234) + library = SimpleNamespace() + if outcome != "missing": + library = SimpleNamespace( + malloc_trim=release, + malloc_default_zone=handle, + malloc_zone_pressure_relief=release, + GetCurrentProcess=handle, + SetProcessWorkingSetSize=release, + ) + loader = Mock(return_value=library) + if outcome == "load_error": + loader.side_effect = OSError("allocator unavailable") + elif outcome == "call_error": + release.side_effect = OSError("allocator call failed") + monkeypatch.setattr(mem.ctypes, "CDLL", loader) + monkeypatch.setattr(mem.ctypes, "WinDLL", loader, raising=False) + + assert mem.release_memory() is (outcome == "success" and platform != "unsupported") + gc.collect.assert_called_once_with() + if gc_enabled: + gc.enable.assert_not_called() + gc.disable.assert_not_called() + else: + assert [call[0] for call in gc.mock_calls] == [ + "isenabled", + "enable", + "collect", + "disable", + ] + if platform == "unsupported": + loader.assert_not_called() + else: + if platform == "win32": + loader.assert_called_once_with("kernel32", use_last_error=True) + else: + loader.assert_called_once_with(None) + if outcome in ("success", "no_release", "call_error"): + if platform == "linux": + release.assert_called_once_with(0) + assert release.argtypes == [ctypes.c_size_t] + assert release.restype == ctypes.c_int + elif platform == "darwin": + handle.assert_called_once_with() + release.assert_called_once_with(1234, 0) + assert handle.restype == ctypes.c_void_p + assert release.argtypes == [ctypes.c_void_p, ctypes.c_size_t] + assert release.restype == ctypes.c_size_t + else: + handle.assert_called_once_with() + maximum_size = ctypes.c_size_t(-1).value + release.assert_called_once_with(1234, maximum_size, maximum_size) + assert handle.restype == ctypes.c_void_p + assert release.argtypes == [ + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_size_t, + ] + assert release.restype == ctypes.c_int From 060a8f54d966a83693587d4c1656bb13b0355d70 Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Tue, 8 Sep 2026 11:30:21 -0500 Subject: [PATCH 5/7] Fix nested PNR chunker failures by forwarding logsum chunk overrides --- .../abm/models/park_and_ride_lot_choice.py | 9 +- activitysim/abm/models/util/logsums.py | 4 + .../abm/test/test_park_and_ride_logsums.py | 145 ++++++++++++++++++ activitysim/core/interaction_simulate.py | 10 +- 4 files changed, 166 insertions(+), 2 deletions(-) create mode 100644 activitysim/abm/test/test_park_and_ride_logsums.py diff --git a/activitysim/abm/models/park_and_ride_lot_choice.py b/activitysim/abm/models/park_and_ride_lot_choice.py index 4cebc4020a..eff96595c8 100644 --- a/activitysim/abm/models/park_and_ride_lot_choice.py +++ b/activitysim/abm/models/park_and_ride_lot_choice.py @@ -182,6 +182,8 @@ 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, ) -> pd.Series: """ Run the park-and-ride lot choice model. @@ -345,7 +347,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, ) diff --git a/activitysim/abm/models/util/logsums.py b/activitysim/abm/models/util/logsums.py index ccba509fc2..8908ed4622 100644 --- a/activitysim/abm/models/util/logsums.py +++ b/activitysim/abm/models/util/logsums.py @@ -300,6 +300,10 @@ def compute_location_choice_logsums( estimator=None, pnr_capacity_cls=None, trace_label=tracing.extend_trace_label(trace_label, "pnr_lot_choice"), + # The caller may already own the logsum chunk ledger. Propagate + # both overrides through lot choice, not just mode simulation. + chunk_size=chunk_size, + explicit_chunk_size=explicit_chunk_size, ) logsum_spec = state.filesystem.read_model_spec(file_name=logsum_settings.SPEC) diff --git a/activitysim/abm/test/test_park_and_ride_logsums.py b/activitysim/abm/test/test_park_and_ride_logsums.py new file mode 100644 index 0000000000..72428027e0 --- /dev/null +++ b/activitysim/abm/test/test_park_and_ride_logsums.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pandas as pd +import pytest + +from activitysim.abm.models import location_choice, park_and_ride_lot_choice as pnr +from activitysim.abm.models.util import logsums, tour_destination +from activitysim.core import interaction_simulate, workflow + + +@pytest.fixture +def pnr_context(monkeypatch): + state = workflow.State.make_temp() + state.settings.chunk_size = 1_000_000_000 + state.settings.sharrow = False + state.rng().begin_step("destination") + land_use = pd.DataFrame({"pnr_spaces": [10, 10]}, index=[1, 2]) + state.add_table("land_use", land_use) + settings = SimpleNamespace( + SPEC="unused.csv", + CONSTANTS={}, + LANDUSE_PNR_SPACES_COLUMN="pnr_spaces", + preprocessor=None, + CHOOSER_FILTER_EXPR=None, + explicit_chunk=0.5, + compute_settings=None, + ) + spec = pd.DataFrame({"coefficient": [1.0]}, index=["1"]) + fs = type(state.filesystem) + monkeypatch.setattr(fs, "read_model_spec", lambda *a, **k: spec) + monkeypatch.setattr(fs, "read_model_coefficients", lambda *a, **k: {}) + monkeypatch.setattr(fs, "get_segment_coefficients", lambda *a, **k: {}) + monkeypatch.setattr( + pnr.ParkAndRideLotChoiceSettings, "read_settings_file", lambda *a, **k: settings + ) + monkeypatch.setattr( + pnr.simulate, "eval_coefficients", lambda state, spec, *a, **k: spec + ) + monkeypatch.setattr( + pnr, + "filter_chooser_to_transit_accessible_destinations", + lambda state, choosers, *a: choosers, + ) + monkeypatch.setattr(logsums, "setup_skims", lambda *a, **k: {}) + monkeypatch.setattr(pnr.expressions, "annotate_preprocessors", lambda *a, **k: None) + monkeypatch.setattr(pnr.util, "drop_unused_columns", lambda df, *a, **k: df) + network = SimpleNamespace( + skim_time_period_label=lambda period, as_cat, broadcast_to: pd.Series( + "AM", index=broadcast_to + ) + ) + return state, land_use, network, settings + + +@pytest.mark.parametrize("mode", ["training", "adaptive", "explicit"]) +@pytest.mark.parametrize("component", ["tour", "location"]) +def test_pnr_logsums_share_outer_chunk_boundary( + pnr_context, monkeypatch, mode, component +): + state, land_use, network, pnr_settings = pnr_context + state.settings.chunk_training_mode = mode + persons = pd.DataFrame( + {"home_zone_id": [1, 2]}, index=pd.Index([1, 2], name="person_id") + ) + sample = pd.DataFrame( + {"person_id": [1, 1, 2, 2], "alt_dest": [1, 2, 1, 2]}, + index=pd.Index( + [1, 1, 2, 2], name="person_id" if component == "location" else "tour_id" + ), + ) + if component == "location": + sample = sample.drop(columns="person_id") + model = SimpleNamespace( + LOGSUM_SETTINGS="unused.yaml", + LOGSUM_TOUR_PURPOSE="work", + CHOOSER_ID_COLUMN="person_id", + explicit_chunk=0.5, + CHOOSER_ORIG_COL_NAME="home_zone_id", + ALT_DEST_COL_NAME="alt_dest", + IN_PERIOD=17, + OUT_PERIOD=8, + LOGSUM_PREPROCESSOR="preprocessor", + ) + logsum_settings = SimpleNamespace( + include_pnr_for_logsums=True, + SPEC="unused.csv", + preprocessor=None, + compute_settings=None, + ) + monkeypatch.setattr( + type(state.filesystem), "read_model_settings", lambda *a, **k: logsum_settings + ) + monkeypatch.setattr( + location_choice.TourModeComponentSettings, + "read_settings_file", + lambda *a, **k: logsum_settings, + ) + monkeypatch.setattr(logsums.config, "get_logit_model_settings", lambda *a: None) + monkeypatch.setattr(logsums.config, "get_model_constants", lambda *a: {}) + monkeypatch.setattr( + logsums.simulate, + "simple_simulate_logsums", + lambda state, choosers, *a, **k: choosers.pnr_zone_id.astype(float), + ) + sizes = [] + + def evaluate(state, choosers, *args, **kwargs): + sizes.append(len(choosers)) + return pd.Series(1, index=choosers.index) + + # Keep the actual interaction simulator and its ChunkSizer, replacing only + # utility evaluation so the test exercises the full ledger call chain. + monkeypatch.setattr(interaction_simulate, "_interaction_simulate", evaluate) + if component == "tour": + result = tour_destination.run_destination_logsums( + state, + "work", + persons, + sample, + model, + network, + state.settings.chunk_size, + "test", + ) + else: + result = location_choice.run_location_logsums( + state, + "work", + persons, + network, + sample, + model, + state.settings.chunk_size, + "logsums", + "test", + ) + assert result.mode_choice_logsum.tolist() == [1.0] * 4 + assert sum(sizes) == 4 + if mode == "explicit": + assert sizes == [2, 2] # PNR's fractional setting must not split these again. + assert pnr_settings.explicit_chunk == 0.5 + assert not state.chunk.CHUNK_SIZERS + assert not state.chunk.CHUNK_LEDGERS diff --git a/activitysim/core/interaction_simulate.py b/activitysim/core/interaction_simulate.py index 496abf79e0..c36791179c 100644 --- a/activitysim/core/interaction_simulate.py +++ b/activitysim/core/interaction_simulate.py @@ -993,6 +993,7 @@ def interaction_simulate( estimator=None, explicit_chunk_size=0, compute_settings: ComputeSettings | None = None, + chunk_size: int | None = None, ): """ Run a simulation in the situation in which alternatives must @@ -1032,6 +1033,9 @@ def interaction_simulate( when household tracing enabled. No tracing occurs if label is empty or None. trace_choice_name: str This is the column label to be used in trace file csv dump of choices + chunk_size : int, optional + Adaptive memory budget; zero disables inner chunking. Defaults to the + global setting when omitted. explicit_chunk_size : float, optional If > 0, specifies the chunk size to use when chunking the interaction simulation. If < 1, specifies the fraction of the total number of choosers. @@ -1054,7 +1058,11 @@ def interaction_simulate( chunk_trace_label, chunk_sizer, ) in chunk.adaptive_chunked_choosers( - state, choosers, trace_label, explicit_chunk_size=explicit_chunk_size + state, + choosers, + trace_label, + explicit_chunk_size=explicit_chunk_size, + chunk_size=chunk_size, ): choices = _interaction_simulate( state, From 2776bfd85ad1feb10e01e0335a1e602c0252bc28 Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Tue, 8 Sep 2026 11:32:04 -0500 Subject: [PATCH 6/7] Fix chunk-dependent PNR draws by preserving segment-wide chooser IDs --- activitysim/abm/models/location_choice.py | 4 + .../abm/models/park_and_ride_lot_choice.py | 13 ++- activitysim/abm/models/util/logsums.py | 20 ++++ .../abm/models/util/tour_destination.py | 4 + .../abm/test/test_park_and_ride_logsums.py | 92 +++++++++++++++++++ 5 files changed, 129 insertions(+), 4 deletions(-) diff --git a/activitysim/abm/models/location_choice.py b/activitysim/abm/models/location_choice.py index 04766014de..662f7bf264 100644 --- a/activitysim/abm/models/location_choice.py +++ b/activitysim/abm/models/location_choice.py @@ -630,6 +630,9 @@ def run_location_logsums( # 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, @@ -662,6 +665,7 @@ def run_location_logsums( 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) diff --git a/activitysim/abm/models/park_and_ride_lot_choice.py b/activitysim/abm/models/park_and_ride_lot_choice.py index eff96595c8..d99895a388 100644 --- a/activitysim/abm/models/park_and_ride_lot_choice.py +++ b/activitysim/abm/models/park_and_ride_lot_choice.py @@ -184,6 +184,7 @@ def run_park_and_ride_lot_choice( 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,7 +220,7 @@ 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 @@ -227,9 +228,13 @@ def run_park_and_ride_lot_choice( 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() ) diff --git a/activitysim/abm/models/util/logsums.py b/activitysim/abm/models/util/logsums.py index 8908ed4622..3fe2e67d41 100644 --- a/activitysim/abm/models/util/logsums.py +++ b/activitysim/abm/models/util/logsums.py @@ -179,6 +179,24 @@ def filter_chooser_columns( return choosers +def get_pnr_index_multiplier(choosers, logsum_settings): + """Keep PNR synthetic chooser IDs stable across complete chooser chunks. + + Lot choice disambiguates repeated destination-sample indices using a + multiplier rounded up to a multiple of ten. Compute it over the full + segment so changing chunk size cannot change the random-number keys. + """ + include_pnr = ( + logsum_settings.get("include_pnr_for_logsums", False) + if isinstance(logsum_settings, dict) + else getattr(logsum_settings, "include_pnr_for_logsums", False) + ) + if not include_pnr or choosers.index.is_unique: + return None + max_count = int(choosers.groupby(level=0).size().max()) + return ((max_count + 9) // 10) * 10 + + def compute_location_choice_logsums( state: workflow.State, choosers: pd.DataFrame, @@ -193,6 +211,7 @@ def compute_location_choice_logsums( out_period_col: str | None = None, duration_col: str | None = None, explicit_chunk_size: float | None = None, + pnr_index_multiplier: int | None = None, ): """ @@ -304,6 +323,7 @@ def compute_location_choice_logsums( # both overrides through lot choice, not just mode simulation. chunk_size=chunk_size, explicit_chunk_size=explicit_chunk_size, + chooser_index_multiplier=pnr_index_multiplier, ) logsum_spec = state.filesystem.read_model_spec(file_name=logsum_settings.SPEC) diff --git a/activitysim/abm/models/util/tour_destination.py b/activitysim/abm/models/util/tour_destination.py index 5401ad4a22..068732d5f9 100644 --- a/activitysim/abm/models/util/tour_destination.py +++ b/activitysim/abm/models/util/tour_destination.py @@ -802,6 +802,9 @@ def run_destination_logsums( tour_choosers = destination_sample.loc[ ~destination_sample.index.duplicated(keep="first"), [chooser_id_column] ] + pnr_index_multiplier = logsum.get_pnr_index_multiplier( + destination_sample, logsum_settings + ) logsum_chunks = [] for ( _i, @@ -840,6 +843,7 @@ def run_destination_logsums( 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) diff --git a/activitysim/abm/test/test_park_and_ride_logsums.py b/activitysim/abm/test/test_park_and_ride_logsums.py index 72428027e0..8f9a3ac0ab 100644 --- a/activitysim/abm/test/test_park_and_ride_logsums.py +++ b/activitysim/abm/test/test_park_and_ride_logsums.py @@ -143,3 +143,95 @@ def evaluate(state, choosers, *args, **kwargs): assert pnr_settings.explicit_chunk == 0.5 assert not state.chunk.CHUNK_SIZERS assert not state.chunk.CHUNK_LEDGERS + + +@pytest.mark.parametrize("component", ["tour", "location"]) +def test_pnr_random_draws_and_logsums_do_not_depend_on_chunk_size( + pnr_context, monkeypatch, component +): + state, land_use, network, pnr_settings = pnr_context + state.settings.chunk_training_mode = "explicit" + persons = pd.DataFrame( + {"home_zone_id": [1, 2, 1]}, index=pd.Index([11, 22, 33], name="person_id") + ) + # Different multiplicities cross a rounding boundary; the final chooser + # also exercises a chunk whose index is unique despite the segment's not + # being unique. Every segment uses the same canonical PNR random channel. + ids = [11] * 11 + [22] * 2 + [33] + sample = pd.DataFrame( + {"person_id": ids, "alt_dest": list(range(11)) + [1, 2, 1]}, + index=pd.Index(ids, name="person_id" if component == "location" else "tour_id"), + ) + if component == "location": + sample = sample.drop(columns="person_id") + model = SimpleNamespace( + LOGSUM_SETTINGS="unused.yaml", + LOGSUM_TOUR_PURPOSE="work", + CHOOSER_ID_COLUMN="person_id", + explicit_chunk=0, + CHOOSER_ORIG_COL_NAME="home_zone_id", + ALT_DEST_COL_NAME="alt_dest", + IN_PERIOD=17, + OUT_PERIOD=8, + LOGSUM_PREPROCESSOR="preprocessor", + ) + logsum_settings = SimpleNamespace( + include_pnr_for_logsums=True, + SPEC="unused.csv", + preprocessor=None, + compute_settings=None, + ) + monkeypatch.setattr( + type(state.filesystem), "read_model_settings", lambda *a, **k: logsum_settings + ) + monkeypatch.setattr( + location_choice.TourModeComponentSettings, + "read_settings_file", + lambda *a, **k: logsum_settings, + ) + monkeypatch.setattr(logsums.config, "get_logit_model_settings", lambda *a: None) + monkeypatch.setattr(logsums.config, "get_model_constants", lambda *a: {}) + monkeypatch.setattr( + logsums.simulate, + "simple_simulate_logsums", + lambda state, choosers, *a, **k: choosers.pnr_zone_id.astype(float), + ) + draws = [] + + def evaluate(state, choosers, *args, **kwargs): + rands = state.rng().random_for_df(choosers).ravel() + draws.extend(zip(choosers.index, rands)) + return pd.Series(1 + (rands > 0.5).astype(int), index=choosers.index) + + monkeypatch.setattr(interaction_simulate, "_interaction_simulate", evaluate) + baseline = baseline_draws = None + for size in [0, 0.5, 1, 2]: + model.explicit_chunk = size + draws.clear() + if component == "tour": + result = tour_destination.run_destination_logsums( + state, "work", persons, sample.copy(), model, network, 0, "test" + ) + else: + result = location_choice.run_location_logsums( + state, + "work", + persons, + network, + sample.copy(), + model, + 0, + "logsums", + "test", + ) + if baseline is None: + baseline, baseline_draws = result, list(draws) + assert [idx for idx, rand in draws] == list(range(220, 231)) + [ + 440, + 441, + 660, + ] + else: + pd.testing.assert_frame_equal(result, baseline) + assert draws == baseline_draws + assert "pnr_lot_choice" not in state.rng().channels From 51e298a84276813946e1d623c9a5785e078e022f Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Tue, 8 Sep 2026 11:33:04 -0500 Subject: [PATCH 7/7] Fix discarded trip_period output by limiting cleanup to temporary columns --- activitysim/abm/models/trip_mode_choice.py | 15 +++++--- activitysim/abm/test/test_trip_mode_choice.py | 36 +++++++++++++------ 2 files changed, 35 insertions(+), 16 deletions(-) diff --git a/activitysim/abm/models/trip_mode_choice.py b/activitysim/abm/models/trip_mode_choice.py index acdc7e85eb..9ef9d7cfd7 100644 --- a/activitysim/abm/models/trip_mode_choice.py +++ b/activitysim/abm/models/trip_mode_choice.py @@ -350,7 +350,9 @@ def add_trip_period(choosers): # needed in the purpose-sized chooser frames above, but post-choice table # annotators may also use those skims. Add it to the full trips table only # for annotation, then restore the original table schema. - add_trip_period(trips_df) + temporary_trip_period = "trip_period" not in trips_df.columns + if temporary_trip_period: + add_trip_period(trips_df) try: expressions.annotate_tables( state, @@ -360,7 +362,10 @@ def add_trip_period(choosers): trace_label=trace_label, ) finally: - trips_df.drop(columns="trip_period", inplace=True) - state_trips = state.get_dataframe("trips", as_copy=False) - if state_trips is not trips_df: - state_trips.drop(columns="trip_period", inplace=True) + # CHOOSER_COLS_TO_KEEP may have made trip_period an output column. + # Remove it only when this annotation block created it temporarily. + if temporary_trip_period: + trips_df.drop(columns="trip_period", inplace=True) + state_trips = state.get_dataframe("trips", as_copy=False) + if state_trips is not trips_df: + state_trips.drop(columns="trip_period", inplace=True) diff --git a/activitysim/abm/test/test_trip_mode_choice.py b/activitysim/abm/test/test_trip_mode_choice.py index 4aa901afa2..f40a7dc7f4 100644 --- a/activitysim/abm/test/test_trip_mode_choice.py +++ b/activitysim/abm/test/test_trip_mode_choice.py @@ -1,9 +1,11 @@ from __future__ import annotations import weakref +from contextlib import nullcontext from types import SimpleNamespace import pandas as pd +import pytest from activitysim.abm.models import trip_mode_choice as trip_mode_choice_module @@ -55,7 +57,12 @@ def is_table(self, _name): return False -def test_post_choice_annotations_receive_full_trip_period_then_remove_it(monkeypatch): +@pytest.mark.parametrize("keep_trip_period", [False, True]) +@pytest.mark.parametrize("copy_annotation_table", [False, True]) +@pytest.mark.parametrize("annotation_error", [False, True]) +def test_post_choice_annotations_preserve_requested_trip_period( + monkeypatch, keep_trip_period, copy_annotation_table, annotation_error +): trips = pd.DataFrame( { "tour_id": [11, 12, 13], @@ -79,7 +86,7 @@ def test_post_choice_annotations_receive_full_trip_period_then_remove_it(monkeyp model_settings = SimpleNamespace( MODE_CHOICE_LOGSUM_COLUMN_NAME="mode_choice_logsum", TOURS_MERGED_CHOOSER_COLUMNS=[], - CHOOSER_COLS_TO_KEEP=[], + CHOOSER_COLS_TO_KEEP=["trip_period"] if keep_trip_period else [], FORCE_ESCORTEE_CHAUFFEUR_MODE_MATCH=False, SPEC="trip_mode_choice.csv", explicit_chunk=None, @@ -132,26 +139,33 @@ def choose_mode(_state, choosers, **_kwargs): monkeypatch.setattr(trip_mode_choice_module, "mode_choice_simulate", choose_mode) def annotate_tables(_state, **_kwargs): - annotated = _state.get_dataframe("trips") + annotated = _state.get_dataframe("trips", as_copy=copy_annotation_table) assert annotated.index.equals(trips.index) assert annotated["trip_period"].tolist() == [0, 1, 0] annotated["post_choice_skim_value"] = [10.0, 20.0, 30.0] _state.add_table("trips", annotated) + if annotation_error: + raise RuntimeError("annotation failed") monkeypatch.setattr( trip_mode_choice_module.expressions, "annotate_tables", annotate_tables ) - trip_mode_choice_module.trip_mode_choice( - state, - trips, - network_los, - model_settings=model_settings, - ) + with pytest.raises( + RuntimeError, match="annotation failed" + ) if annotation_error else nullcontext(): + trip_mode_choice_module.trip_mode_choice( + state, + trips, + network_los, + model_settings=model_settings, + ) assert all(ref() is None for ref in chooser_refs) assert all(wrapper.df.empty for wrapper in wrappers) result = state.get_dataframe("trips", as_copy=False) - assert "trip_period" not in trips - assert "trip_period" not in result + assert ("trip_period" in trips) == keep_trip_period + assert ("trip_period" in result) == keep_trip_period + if keep_trip_period: + assert result["trip_period"].tolist() == [0, 1, 0] assert result["post_choice_skim_value"].tolist() == [10.0, 20.0, 30.0]