Add config for n_processes, remove dataclass decorators

scripts/3.3_assign_facility_all.py

@@ -91,6 +91,7 @@ def main(config_file):
         gdf_facility_type_col="activities",
         gdf_sample_col="floor_area",
         neighboring_zones=zone_neighbors,
+        n_processes=config.parameters.n_processes,
     )
 
     # Map the activity_id and activity_geometry to the activity_chains_home_df DataFrame
@@ -113,6 +114,7 @@ def main(config_file):
         gdf_facility_type_col="activities",
         gdf_sample_col="floor_area",
         neighboring_zones=zone_neighbors,
+        n_processes=config.parameters.n_processes,
     )
 
     # Map the activity_id and activity_geometry to the activity_chains_df DataFrame
@@ -156,6 +158,7 @@ def main(config_file):
         gdf_sample_col="floor_area",
         neighboring_zones=zone_neighbors,
         fallback_type="education",
+        n_processes=config.parameters.n_processes,
     )
 
     logger.info(f"Shape of activity chains edu: {activity_chains_edu.shape}")
@@ -196,6 +199,7 @@ def main(config_file):
         gdf_sample_col="floor_area",
         neighboring_zones=zone_neighbors,
         fallback_to_random=True,
+        n_processes=config.parameters.n_processes,
     )
 
     # Map the activity_id and activity_geometry to the activity_chains_home_df DataFrame

src/acbm/assigning/select_facility.py

@@ -174,6 +174,7 @@ def select_facility(
     neighboring_zones: Optional[dict] = None,
     fallback_type: Optional[str] = None,
     fallback_to_random: bool = False,
+    n_processes: int = max(int(cpu_count() * 0.75), 1),
 ) -> dict[str, Tuple[str, Point] | Tuple[float, float]]:
     """
     Select facilities for each row in the DataFrame based on the provided logic.
@@ -205,8 +206,7 @@ def select_facility(
         keys with selected facility ID and facility ID's geometry, or (np.nan, np.nan)
     """
     # TODO: check if this is deterministic for a given seed (or pass seed to pool)
-    # TODO: update to be configurable
-    with Pool(max(int(cpu_count() * 0.75), 1)) as p:
+    with Pool(n_processes) as p:
         # Set to a large enough chunk size so that each process
         # has a sufficiently large amount of processing to do.
         chunk_size = 16_000

src/acbm/config.py

@@ -3,6 +3,7 @@
 from dataclasses import dataclass
 from hashlib import sha256
 from logging import Logger
+from multiprocessing import cpu_count
 from pathlib import Path
 from typing import Tuple
 
@@ -33,6 +34,11 @@ class Parameters(BaseModel):
     tolerance_edu: float | None = None
     common_household_day: bool = True
     part_time_work_prob: float = 0.7
+    n_processes = Field(
+        description="Number of processes to use in multiprocessing",
+        default=max(int(cpu_count() * 0.75), 1),
+        exclude=True,
+    )
 
 
 @dataclass(frozen=True)