Common settings for fastmath, sharrow_skip, and other compute controls across components

activitysim/abm/models/atwork_subtour_frequency.py

@@ -117,6 +117,7 @@ def atwork_subtour_frequency(
         trace_label=trace_label,
         trace_choice_name="atwork_subtour_frequency",
         estimator=estimator,
+        fastmath=model_settings.sharrow_fastmath,
     )
 
     # convert indexes to alternative names

activitysim/abm/models/auto_ownership.py

@@ -9,14 +9,15 @@
 
 from activitysim.core import (
     config,
-    expressions,
     estimation,
+    expressions,
     simulate,
     tracing,
     workflow,
 )
 from activitysim.core.configuration.base import PreprocessorSettings, PydanticReadable
 from activitysim.core.configuration.logit import LogitComponentSettings
+
 from .util import annotate
 
 logger = logging.getLogger(__name__)
@@ -99,6 +100,7 @@ def auto_ownership_simulate(
         trace_choice_name="auto_ownership",
         log_alt_losers=log_alt_losers,
         estimator=estimator,
+        fastmath=model_settings.sharrow_fastmath,
     )
 
     if estimator:

activitysim/abm/models/free_parking.py

@@ -118,6 +118,7 @@ def free_parking(
         trace_label=trace_label,
         trace_choice_name="free_parking_at_work",
         estimator=estimator,
+        fastmath=model_settings.sharrow_fastmath,
     )
 
     free_parking_alt = model_settings.FREE_PARKING_ALT

activitysim/abm/models/joint_tour_composition.py

@@ -123,6 +123,7 @@ def joint_tour_composition(
         trace_label=trace_label,
         trace_choice_name="composition",
         estimator=estimator,
+        fastmath=model_settings.sharrow_fastmath,
     )
 
     # convert indexes to alternative names

activitysim/abm/models/joint_tour_frequency.py

@@ -112,6 +112,7 @@ def joint_tour_frequency(
         trace_label=trace_label,
         trace_choice_name="joint_tour_frequency",
         estimator=estimator,
+        fastmath=model_settings.sharrow_fastmath,
     )
 
     # convert indexes to alternative names

activitysim/abm/models/joint_tour_frequency_composition.py

@@ -136,6 +136,7 @@ def joint_tour_frequency_composition(
         trace_choice_name=trace_label,
         estimator=estimator,
         explicit_chunk_size=0,
+        fastmath=model_settings.sharrow_fastmath,
     )
 
     if estimator:

activitysim/abm/models/joint_tour_participation.py

@@ -418,6 +418,7 @@ def joint_tour_participation(
         trace_choice_name="participation",
         custom_chooser=participants_chooser,
         estimator=estimator,
+        fastmath=model_settings.sharrow_fastmath,
     )
 
     # choice is boolean (participate or not)

activitysim/abm/models/mandatory_tour_frequency.py

@@ -134,6 +134,7 @@ def mandatory_tour_frequency(
         trace_label=trace_label,
         trace_choice_name="mandatory_tour_frequency",
         estimator=estimator,
+        fastmath=model_settings.sharrow_fastmath,
     )
 
     # convert indexes to alternative names

activitysim/abm/models/non_mandatory_tour_frequency.py

@@ -312,6 +312,7 @@ def non_mandatory_tour_frequency(
             trace_choice_name="non_mandatory_tour_frequency",
             estimator=estimator,
             explicit_chunk_size=model_settings.explicit_chunk,
+            fastmath=model_settings.sharrow_fastmath,
         )
 
         if estimator:

activitysim/abm/models/school_escorting.py

@@ -395,6 +395,22 @@ class SchoolEscortSettings(BaseLogitComponentSettings):
     ```
     """
 
+    sharrow_fastmath: bool = False
+    """Use fastmath when evaluating this component with sharrow.
+
+    The fastmath option can be used to speed up the evaluation of expressions in
+    this component's spec files, but it does so by making some simplifying
+    assumptions about the math, e.g. that neither inputs nor outputs of any
+    computations are NaN or Inf.  This can lead to errors when the assumptions
+    are violated.  If running in sharrow test mode generates errors, try turning
+    this setting off.
+
+    The default for most components is True, but due to the prevalence of NaN
+    values in the school escorting model, the default is False for this
+    component. It can be turned on if desired, but it is recommended to test
+    carefully that results are as expected.
+    """
+
     SIMULATE_CHOOSER_COLUMNS: list[str] | None = None
 
     SPEC: None = None
@@ -550,6 +566,7 @@ def school_escorting(
             trace_choice_name="school_escorting_" + stage,
             estimator=estimator,
             explicit_chunk_size=model_settings.explicit_chunk,
+            fastmath=model_settings.sharrow_fastmath,
         )
 
         if estimator:

activitysim/abm/models/stop_frequency.py

@@ -214,6 +214,7 @@ def stop_frequency(
             trace_label=tracing.extend_trace_label(trace_label, segment_name),
             trace_choice_name="stops",
             estimator=estimator,
+            fastmath=model_settings.sharrow_fastmath,
         )
 
         # convert indexes to alternative names

activitysim/abm/models/telecommute_frequency.py

@@ -99,6 +99,7 @@ def telecommute_frequency(
         trace_label=trace_label,
         trace_choice_name="telecommute_frequency",
         estimator=estimator,
+        fastmath=model_settings.sharrow_fastmath,
     )
 
     choices = pd.Series(model_spec.columns[choices.values], index=choices.index)

activitysim/abm/models/transit_pass_ownership.py

@@ -93,6 +93,7 @@ def transit_pass_ownership(
         trace_label=trace_label,
         trace_choice_name="transit_pass_ownership",
         estimator=estimator,
+        fastmath=model_settings.sharrow_fastmath,
     )
 
     if estimator:

activitysim/abm/models/trip_mode_choice.py

@@ -277,6 +277,7 @@ def trip_mode_choice(
             trace_label=segment_trace_label,
             trace_choice_name="trip_mode_choice",
             estimator=estimator,
+            fastmath=model_settings.sharrow_fastmath,
         )
 
         if state.settings.trace_hh_id:

activitysim/abm/models/util/mode.py

@@ -34,6 +34,7 @@ def mode_choice_simulate(
     trace_choice_name,
     trace_column_names=None,
     estimator: Optional[Estimator] = None,
+    fastmath: bool = True,
 ):
     """
     common method for  both tour_mode_choice and trip_mode_choice
@@ -51,6 +52,7 @@ def mode_choice_simulate(
     trace_label
     trace_choice_name
     estimator
+    fastmath
 
     Returns
     -------
@@ -70,6 +72,7 @@ def mode_choice_simulate(
         trace_choice_name=trace_choice_name,
         estimator=estimator,
         trace_column_names=trace_column_names,
+        fastmath=fastmath,
     )
 
     # for consistency, always return dataframe, whether or not logsums were requested
@@ -170,6 +173,7 @@ def run_tour_mode_choice_simulate(
         trace_choice_name=trace_choice_name,
         trace_column_names=trace_column_names,
         estimator=estimator,
+        fastmath=model_settings.sharrow_fastmath,
     )
 
     return choices

activitysim/abm/models/vehicle_allocation.py

@@ -247,6 +247,7 @@ def vehicle_allocation(
             trace_label=trace_label,
             trace_choice_name="vehicle_allocation",
             estimator=estimator,
+            fastmath=model_settings.sharrow_fastmath,
         )
 
         # matching alt names to choices

activitysim/abm/models/vehicle_type_choice.py

@@ -485,6 +485,7 @@ def iterate_vehicle_type_choice(
                 trace_choice_name="vehicle_type",
                 estimator=estimator,
                 explicit_chunk_size=model_settings.explicit_chunk,
+                fastmath=model_settings.sharrow_fastmath,
             )
 
         # otherwise, "simple simulation" should suffice, with a model spec that enumerates
@@ -500,6 +501,7 @@ def iterate_vehicle_type_choice(
                 trace_label=trace_label,
                 trace_choice_name="vehicle_type",
                 estimator=estimator,
+                fastmath=model_settings.sharrow_fastmath,
             )
         else:
             raise NotImplementedError(simulation_type)

activitysim/abm/models/work_from_home.py

@@ -152,6 +152,7 @@ def work_from_home(
             trace_label=trace_label,
             trace_choice_name="work_from_home",
             estimator=estimator,
+            fastmath=model_settings.sharrow_fastmath,
         )
 
         if iterations_target_percent is not None:

activitysim/core/configuration/logit.py

@@ -78,6 +78,17 @@ class BaseLogitComponentSettings(PydanticReadable):
     sharrow_skip: bool = False
     """Skip sharrow when evaluating this component."""
 
+    sharrow_fastmath: bool = True
+    """Use fastmath when evaluating this component with sharrow.
+
+    The fastmath option can be used to speed up the evaluation of expressions in
+    this component's spec files, but it does so by making some simplifying
+    assumptions about the math, e.g. that neither inputs nor outputs of any
+    computations are NaN or Inf.  This can lead to errors when the assumptions
+    are violated.  If running in sharrow test mode generates errors, try turning
+    this setting off.
+    """
+
 
 class LogitComponentSettings(BaseLogitComponentSettings):
     """

activitysim/core/flow.py

@@ -142,6 +142,7 @@ def get_flow(
     choosers=None,
     interacts=None,
     zone_layer=None,
+    fastmath=True,
 ):
     extra_vars = only_simple(local_d)
     orig_col_name = local_d.get("orig_col_name", None)
@@ -184,6 +185,7 @@ def get_flow(
         zone_layer=zone_layer,
         aux_vars=aux_vars,
         primary_origin_col_name=primary_origin_col_name,
+        fastmath=fastmath,
     )
     flow.tree.aux_vars = aux_vars
     return flow
@@ -465,6 +467,7 @@ def new_flow(
     zone_layer=None,
     aux_vars=None,
     primary_origin_col_name=None,
+    fastmath=True,
 ):
     """
     Setup a new sharrow flow.
@@ -516,6 +519,11 @@ def new_flow(
     aux_vars : Mapping
         Extra values that are available to expressions and which are written
         only by reference into compiled code (and thus can be changed later).
+    fastmath : bool, default True
+        Whether to use fastmath in the numba compiled code. Leaving this option
+        set to "True" is generally a good idea to maximize performance, but it
+        can be turned off if some expressions are not compatible with fastmath,
+        i.e. when they include "NaN" or "Inf" values.
 
     Returns
     -------
@@ -700,6 +708,7 @@ def _apply_filter(_dataset, renames: list):
             extra_hash_data=extra_hash_data,
             hashing_level=0,
             boundscheck=False,
+            fastmath=fastmath,
         )
 
 
@@ -750,6 +759,7 @@ def apply_flow(
     required=False,
     interacts=None,
     zone_layer=None,
+    fastmath=True,
 ):
     """
     Apply a sharrow flow.
@@ -779,6 +789,11 @@ def apply_flow(
         Specify which zone layer of the skims is to be used.  You cannot use the
         'maz' zone layer in a one-zone model, but you can use the 'taz' layer in
         a two- or three-zone model (e.g. for destination pre-sampling).
+    fastmath : bool, default True
+        Whether to use fastmath in the numba compiled code. Leaving this option
+        set to "True" is generally a good idea to maximize performance, but it
+        can be turned off if some expressions are not compatible with fastmath,
+        i.e. when they include "NaN" or "Inf" values.
 
     Returns
     -------
@@ -807,6 +822,7 @@ def apply_flow(
                 choosers=choosers,
                 interacts=interacts,
                 zone_layer=zone_layer,
+                fastmath=fastmath,
             )
         except ValueError as err:
             if "unable to rewrite" in str(err):

activitysim/core/interaction_simulate.py

@@ -32,6 +32,7 @@ def eval_interaction_utilities(
     log_alt_losers=False,
     extra_data=None,
     zone_layer=None,
+    fastmath=True,
 ):
     """
     Compute the utilities for a single-alternative spec evaluated in the context of df
@@ -184,6 +185,7 @@ def replace_in_index_level(mi, level, *repls):
                 trace_label,
                 interacts=extra_data,
                 zone_layer=zone_layer,
+                fastmath=fastmath,
             )
             if sh_util is not None:
                 chunk_sizer.log_df(trace_label, "sh_util", sh_util)
@@ -566,6 +568,7 @@ def to_series(x):
                     retrace_eval_parts_ = pd.concat(retrace_eval_parts, axis=1)
 
                     re_sh_flow_load = sh_flow.load(
+                        sh_tree,
                         dtype=np.float32,
                     )
                     re_sh_flow_load_ = re_sh_flow_load[re_trace]
@@ -606,6 +609,7 @@ def _interaction_simulate(
     log_alt_losers=False,
     estimator=None,
     chunk_sizer=None,
+    fastmath: bool = True,
 ):
     """
     Run a MNL simulation in the situation in which alternatives must
@@ -721,6 +725,7 @@ def _interaction_simulate(
             estimator=estimator,
             log_alt_losers=log_alt_losers,
             extra_data=alternatives,
+            fastmath=fastmath,
         )
 
         # set this index here as this is how later code extracts the chosen alt id's
@@ -785,6 +790,7 @@ def _interaction_simulate(
             trace_rows,
             estimator=estimator,
             log_alt_losers=log_alt_losers,
+            fastmath=fastmath,
         )
         chunk_sizer.log_df(trace_label, "interaction_utilities", interaction_utilities)
         # mem.trace_memory_info(f"{trace_label}.init interaction_utilities", force_garbage_collect=True)
@@ -892,6 +898,7 @@ def interaction_simulate(
     trace_choice_name=None,
     estimator=None,
     explicit_chunk_size=0,
+    fastmath: bool = True,
 ):
     """
     Run a simulation in the situation in which alternatives must
@@ -969,6 +976,7 @@ def interaction_simulate(
             log_alt_losers=log_alt_losers,
             estimator=estimator,
             chunk_sizer=chunk_sizer,
+            fastmath=fastmath,
         )
 
         result_list.append(choices)