Bring in aggregation code

Author: collijk

src/climate_data/aggregate/hierarchy.py

@@ -96,9 +96,15 @@ def hierarchy_main(
             draw_results.append(draw_df)
 
             pbar.update()
+        draw_df = (
+            pd.concat(draw_results, ignore_index=True)
+            .groupby(["location_id", "year_id"])
+            .sum()
+            .reset_index()
+        )
 
         agg_df = utils.aggregate_climate_to_hierarchy(
-            pd.concat(draw_results, ignore_index=True),
+            draw_df,
             hierarchy_df,
         ).set_index(["location_id", "year_id"])
         all_results.append(agg_df["value"].rename(draw))
@@ -212,8 +218,8 @@ def hierarchy(
         task_resources={
             "queue": queue,
             "cores": 1,
-            "memory": "16G",
-            "runtime": "30m",
+            "memory": "50G",
+            "runtime": "200m",
             "project": "proj_rapidresponse",
         },
         log_root=ca_data.log_dir("aggregate_hierarchy"),

src/climate_data/aggregate/pixel.py

@@ -38,7 +38,7 @@ def pixel_main(
         hierarchy, block_key, pm_data
     )
 
-    years = cdc.ALL_YEARS
+    years = [int(y) for y in cdc.ALL_YEARS]
     measures = cdc.AGGREGATION_MEASURES
     scenarios = cdc.AGGREGATION_SCENARIOS
 
@@ -67,7 +67,12 @@ def pixel_main(
 
             # Pull out and rasterize the climate data for the current year
             clim_arr = (
-                to_raster(ds.sel(year=year)["value"], no_data_value=np.nan)  # noqa: SLF001
+                to_raster(  # noqa: SLF001
+                    ds.sel(year=year)["value"],
+                    no_data_value=np.nan,
+                    lat_col="latitude",
+                    lon_col="longitude",
+                )
                 .resample_to(pop_raster, "nearest")
                 .astype(np.float32)
                 ._ndarray
@@ -164,14 +169,15 @@ def pixel(
     block_keys = modeling_frame["block_key"].unique().tolist()
     block_keys = clio.convert_choice(block_key, block_keys)
 
+    print("Checking for existing results")
     jobs = []
-    for h, b, d in itertools.product(hierarchy, block_keys, draw):
+    hbd = list(itertools.product(hierarchy, block_keys, draw))
+    for h, b, d in tqdm.tqdm(hbd):
         if not ca_data.raw_results_path(agg_version, h, b, d).exists():
             jobs.append((h, b, d))
     jobs = list(set(jobs))
 
     print(f"Running {len(jobs)} jobs")
-
     jobmon.run_parallel(
         runner="cdtask aggregate",
         task_name="pixel",

src/climate_data/aggregate/utils.py

@@ -13,6 +13,10 @@
     PopulationModelData,
 )
 
+MAX_BOUNDS = {
+    "ESRI:54034": (-20037508.34, 20037508.34, -6363885.33, 6363885.33),
+}
+
 
 def build_location_masks(
     hierarchy: str,
@@ -136,11 +140,32 @@ def get_bbox(raster: rt.RasterArray, crs: str | None = None) -> shapely.Polygon:
     shapely.Polybon
         The bounding box of the raster in the CRS specified by the crs parameter.
     """
+    if raster.crs not in MAX_BOUNDS:
+        msg = f"Unsupported CRS: {raster.crs}"
+        raise ValueError(msg)
+
+    xmin_clip, xmax_clip, ymin_clip, ymax_clip = MAX_BOUNDS[raster.crs]
     xmin, xmax, ymin, ymax = raster.bounds
+
+    xmin = np.clip(xmin, xmin_clip, xmax_clip)
+    xmax = np.clip(xmax, xmin_clip, xmax_clip)
+    ymin = np.clip(ymin, ymin_clip, ymax_clip)
+    ymax = np.clip(ymax, ymin_clip, ymax_clip)
+
     bbox = gpd.GeoSeries([shapely.box(xmin, ymin, xmax, ymax)], crs=raster.crs)
-    if crs is not None:
-        bbox = bbox.to_crs(crs)
-    return cast(shapely.Polygon, bbox.iloc[0])
+    out_bbox = bbox.to_crs(crs) if crs is not None else bbox.copy()
+
+    # Check that our transformation didn't do something weird
+    # (e.g. artificially clip the bounds or have the bounds extend over the
+    # antimeridian)
+    check_bbox = out_bbox.to_crs(raster.crs)
+    area_change = (np.abs(bbox.area - check_bbox.area) / bbox.area).iloc[0]
+    tolerance = 1e-6
+    if area_change > tolerance:
+        msg = f"Area change: {area_change}"
+        raise ValueError(msg)
+
+    return cast(shapely.Polygon, out_bbox.iloc[0])
 
 
 def aggregate_climate_to_hierarchy(
@@ -173,21 +198,17 @@ def aggregate_climate_to_hierarchy(
         subset["parent_id"] = parent_map
 
         parent_values = (
-            subset.groupby(["year_id", "scenario", "parent_id"])[
-                ["weighted_climate", "population"]
-            ]
+            subset.groupby(["year_id", "parent_id"])[["weighted_climate", "population"]]
             .sum()
             .reset_index()
             .rename(columns={"parent_id": "location_id"})
             .set_index("location_id")
         )
-        parent_values["value"] = (
-            parent_values.weighted_climate / parent_values.population
-        )
         results = pd.concat([results, parent_values])
     results = (
         results.reset_index()
         .sort_values(["location_id", "year_id"])
         .reset_index(drop=True)
     )
+    parent_values["value"] = parent_values.weighted_climate / parent_values.population
     return results