Speed Up Nearest Zone Calculation in Disaggregate Accessibilities (ActivitySim#1031)

* optimize nearest_zone from skims function

* increase speed of skim dict _lookup mapping function

* blacken

* handle empty inputs for nearest zone calc

* maz span not dependent on numbering

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

* map based on index instead of number

* batching skim lookup for disagg access nearest zone

* Revert "Merge branch 'disagg_access_nearest_zone_speed_up' of https://github.com/RSGInc/activitysim into disagg_access_nearest_zone_speed_up"

This reverts commit a9b7f69, reversing
changes made to 66ae28c.

* Revert "map based on index instead of number"

This reverts commit 66ae28c.

* Reapply "map based on index instead of number"

This reverts commit f1e135e.

* Reapply "Merge branch 'disagg_access_nearest_zone_speed_up' of https://github.com/RSGInc/activitysim into disagg_access_nearest_zone_speed_up"

This reverts commit 6c6367f.

---------

Co-authored-by: Jeffrey Newman <jeff@driftless.xyz>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: 3 people

activitysim/abm/tables/disaggregate_accessibility.py

@@ -27,13 +27,41 @@ def find_nearest_accessibility_zone(
     def weighted_average(df, values, weights):
         return df[values].T.dot(df[weights]) / df[weights].sum()
 
-    def nearest_skim(oz, zones):
-        # need to pass equal # of origins and destinations to skim_dict
-        orig_zones = np.full(shape=len(zones), fill_value=oz, dtype=int)
-        return (
-            oz,
-            zones[np.argmin(skim_dict.lookup(orig_zones, zones, "DIST"))],
-        )
+    def find_nearest_zones_via_skims(origin_zones, dest_zones, skim_dict):
+        """
+        Vectorized lookup to find nearest destination zone for each origin zone.
+        Performs a single batched skim lookup instead of one per origin zone.
+        """
+        origin_zones = np.asarray(origin_zones)
+        dest_zones = np.asarray(dest_zones)
+        n_origins = len(origin_zones)
+        n_dests = len(dest_zones)
+
+        # handle empty input case
+        if n_origins == 0 or n_dests == 0:
+            return []
+
+        # Process in batches to avoid allocating arrays of size n_origins * n_dests
+        max_pairs = 5_000_000
+        batch_size = max(1, int(max_pairs // n_dests))
+
+        results = []
+        for start in range(0, n_origins, batch_size):
+            batch_orig = origin_zones[start : start + batch_size]
+
+            # create all origin-destination pairs for this batch
+            all_orig = np.repeat(batch_orig, n_dests)
+            all_dest = np.tile(dest_zones, len(batch_orig))
+
+            # single skim lookup for all pairs in the batch
+            all_dists = skim_dict.lookup(all_orig, all_dest, "DIST")
+            dist_matrix = np.asarray(all_dists).reshape(len(batch_orig), n_dests)
+
+            # find the index of the nearest destination zone for each origin in the batch
+            nearest_indices = np.argmin(dist_matrix, axis=1)
+            results.extend(zip(batch_orig, dest_zones[nearest_indices]))
+
+        return results
 
     def nearest_node(oz, zones_df):
         _idx = util.nearest_node_index(_centroids.loc[oz].XY, zones_df.to_list())
@@ -70,7 +98,10 @@ def nearest_node(oz, zones_df):
 
     else:
         skim_dict = state.get_injectable("skim_dict")
-        nearest = [nearest_skim(Oz, accessibility_zones) for Oz in unmatched_zones]
+        # Vectorized lookup: single skim call for all origin-destination pairs
+        nearest = find_nearest_zones_via_skims(
+            unmatched_zones, accessibility_zones, skim_dict
+        )
 
     # Add the nearest zones to the matched zones
     matched = [(x, x) for x in matched_zones]

activitysim/core/skim_dictionary.py

@@ -39,6 +39,7 @@ class OffsetMapper(object):
 
     def __init__(self, offset_int=None, offset_list=None, offset_series=None):
         self.offset_int = self.offset_series = None
+        self._offset_array = None  # numpy array for fast O(1) lookups
 
         assert (offset_int is not None) + (offset_list is not None) + (
             offset_series is not None
@@ -72,6 +73,28 @@ def set_offset_series(self, offset_series):
         self.offset_series = offset_series
         self.offset_int = None
 
+        # Build numpy lookup array for fast O(1) mapping
+        # This replaces slow pandas Series.map() with direct numpy indexing
+        index_vals = offset_series.index.values
+        if len(index_vals) > 0:
+            min_zone = int(index_vals.min())
+            max_zone = int(index_vals.max())
+            span = max_zone - min_zone + 1
+            # Only build array if zone IDs are non-negative and range is reasonable
+            # (avoid huge arrays for sparse zone IDs)
+            if min_zone >= 0 and span <= len(index_vals) * 10:
+                self._offset_array_min_zone = min_zone
+                self._offset_array = np.full(span, NOT_IN_SKIM_ZONE_ID, dtype=np.int32)
+                self._offset_array[
+                    index_vals.astype(int) - min_zone
+                ] = offset_series.values.astype(np.int32)
+            else:
+                self._offset_array = None
+                self._offset_array_min_zone = 0
+        else:
+            self._offset_array = None
+            self._offset_array_min_zone = 0
+
     def set_offset_list(self, offset_list):
         """
         Convenience method to set offset_series using an integer list the same size as target skim dimension
@@ -110,6 +133,7 @@ def set_offset_int(self, offset_int):
 
         self.offset_int = int(offset_int)
         self.offset_series = None
+        self._offset_array = None  # not needed for simple int offset
 
     def map(self, zone_ids):
         """
@@ -124,7 +148,24 @@ def map(self, zone_ids):
         offsets : numpy array of int
         """
 
-        if self.offset_series is not None:
+        if self._offset_array is not None:
+            # Fast path: use numpy array indexing (O(1) per element)
+            zone_ids = np.asanyarray(zone_ids).astype(int)
+            min_zone = getattr(self, "_offset_array_min_zone", 0)
+            idx = zone_ids - min_zone
+            # Clip to valid range to avoid index errors, then mark out-of-range as NOT_IN_SKIM
+            max_valid = len(self._offset_array) - 1
+            valid_mask = (idx >= 0) & (idx <= max_valid)
+            # Use clip to safely index, then apply mask
+            clipped_idx = np.clip(idx, 0, max_valid)
+            offsets = np.where(
+                valid_mask,
+                self._offset_array[clipped_idx],
+                NOT_IN_SKIM_ZONE_ID,
+            )
+            return offsets
+
+        elif self.offset_series is not None:
             assert self.offset_int is None
             assert isinstance(self.offset_series, pd.Series)