Speed Up Nearest Zone Calculation in Disaggregate Accessibilities

activitysim/abm/tables/disaggregate_accessibility.py

@@ -27,13 +27,35 @@ 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 []
+
+        # Create all origin-destination pairs in one go
+        # all_orig: [o1, o1, o1, ..., o2, o2, o2, ..., oN, oN, oN, ...]
+        # all_dest: [d1, d2, d3, ..., d1, d2, d3, ..., d1, d2, d3, ...]
+        all_orig = np.repeat(origin_zones, n_dests)
+        all_dest = np.tile(dest_zones, n_origins)
+
+        # Single skim lookup for all pairs
+        all_dists = skim_dict.lookup(all_orig, all_dest, "DIST")
+
+        # Reshape to (n_origins, n_dests) and find argmin per origin
+        dist_matrix = np.asarray(all_dists).reshape(n_origins, n_dests)
+        nearest_indices = np.argmin(dist_matrix, axis=1)
+
+        # Return list of (origin, nearest_dest) tuples
+        return list(zip(origin_zones, dest_zones[nearest_indices]))
 
     def nearest_node(oz, zones_df):
         _idx = util.nearest_node_index(_centroids.loc[oz].XY, zones_df.to_list())
@@ -70,7 +92,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,26 @@ 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())
+            # 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 (max_zone - min_zone + 1) <= len(index_vals) * 10:
+                self._offset_array = np.full(
+                    max_zone + 1, NOT_IN_SKIM_ZONE_ID, dtype=np.int32
+                )
+                self._offset_array[
+                    index_vals.astype(int)
+                ] = offset_series.values.astype(np.int32)
+            else:
+                self._offset_array = None
+        else:
+            self._offset_array = None
+
     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 +131,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 +146,22 @@ 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)
+            # 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 = (zone_ids >= 0) & (zone_ids <= max_valid)
+            # Use clip to safely index, then apply mask
+            clipped_ids = np.clip(zone_ids, 0, max_valid)
+            offsets = np.where(
+                valid_mask,
+                self._offset_array[clipped_ids],
+                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)