mosaicing, upsampling, other raking

Author: collijk

src/rra_population_model/cli_options.py

@@ -126,7 +126,7 @@ def with_block_key() -> ClickOption[_P, _T]:
         "-b",
         type=click.STRING,
         required=True,
-        help="Block key of block to run.",
+        help="Block key to run.",
     )
 
 

src/rra_population_model/data.py

@@ -333,7 +333,7 @@ def itu_mask_path(self, iso3: str) -> Path:
         return self.itu_masks / f"{iso3}.tif"
 
     def list_itu_iso3s(self) -> list[str]:
-        return [f.stem for f in self.itu_masks.glob("*")]
+        return [f.stem for f in self.itu_masks.glob("*.tif")]
 
     def load_itu_mask(self, iso3: str) -> rt.RasterArray:
         path = self.itu_mask_path(iso3)
@@ -532,20 +532,15 @@ def save_tile_training_data(
             touch(raster_path, clobber=True)
             save_raster(raster, raster_path)
 
-    def save_summary_training_data(
+    def save_summary_people_per_structure(
         self,
+        data: pd.DataFrame,
         resolution: str,
-        people_per_structure: gpd.GeoDataFrame,
-        pixel_area_weights: gpd.GeoDataFrame,
     ) -> None:
         root = self.training_data_root(resolution)
-        pps_path = root / "people_per_structure.parquet"
-        touch(pps_path, clobber=True)
-        people_per_structure.to_parquet(pps_path)
-
-        paw_path = root / "pixel_area_weights.parquet"
-        touch(paw_path, clobber=True)
-        pixel_area_weights.to_parquet(paw_path)
+        path = root / f"people_per_structure.parquet"
+        touch(path, clobber=True)
+        data.to_parquet(path)
 
     def load_people_per_structure(
         self, resolution: str, tile_key: str | None = None
@@ -725,6 +720,15 @@ def raked_prediction_path(
             / f"{block_key}.tif"
         )
 
+    def list_raked_prediction_time_points(
+        self, resolution: str, version: str
+    ) -> list[str]:
+        return [
+            p.name
+            for p in self.raked_predictions_root(resolution, version).iterdir()
+            if p.is_dir()
+        ]
+
     def save_raked_prediction(
         self,
         raster: rt.RasterArray,
@@ -759,6 +763,26 @@ def compiled_prediction_path(
             / f"{group_key}.tif"
         )
 
+    def compiled_prediction_vrt_path(
+        self, time_point: str, model_spec: "ModelSpecification"
+    ) -> Path:
+        resolution = model_spec.resolution
+        version = model_spec.model_version
+        return (
+            self.compiled_predictions_root(resolution, version)
+            / time_point
+            / "index.vrt"
+        )
+
+    def list_compiled_prediction_time_points(
+        self, resolution: str, version: str
+    ) -> list[str]:
+        return [
+            p.name
+            for p in self.compiled_predictions_root(resolution, version).iterdir()
+            if p.is_dir()
+        ]
+
     def save_compiled_prediction(
         self,
         raster: rt.RasterArray,

src/rra_population_model/model/inference/runner.py

@@ -66,19 +66,25 @@ def inference_main(
 
     modeling_frame = pm_data.load_modeling_frame(model_spec.resolution)
     block_keys = modeling_frame.block_key.unique().tolist()
+    block_keys = [
+        block_key
+        for block_key in block_keys
+        if not pm_data.raw_prediction_path(block_key, time_point, model_spec).exists()
+    ]
 
     datamodule = InferenceDataModule(
         model_spec.model_dump(),
         block_keys,
         time_point,
-        num_workers=0,
+        num_workers=4,
     )
     pred_writer = CustomWriter(
         pm_data, model.specification, time_point, write_interval="batch"
     )
     trainer = Trainer(
         callbacks=[pred_writer],
         enable_progress_bar=progress_bar,
+        devices=2,
     )
     trainer.predict(model, datamodule, return_predictions=False)
 

src/rra_population_model/model_prep/training_data/metadata.py

@@ -54,7 +54,17 @@ def get_training_metadata(
         # for at least one tile it causes a GEOS exception.
         # This should be investigated, but just going with something
         # that appears to work for now.
-        full_shape = intersecting_admins.union_all()
+        try:
+            full_shape = intersecting_admins.union_all()
+        except shapely.errors.GEOSException:
+            # Last ditch effort, actually move around the boundaries a little
+            buffer_size = 0.01 # 1 cm
+            full_shape = (
+                intersecting_admins
+                .buffer(buffer_size)
+                .buffer(-buffer_size)
+                .union_all()
+            )
 
     overlaps = model_frame.intersects(full_shape)
     neighborhood_keys = model_frame[overlaps].tile_key.tolist()

src/rra_population_model/model_prep/training_data/runner.py

@@ -158,19 +158,13 @@ def training_data(
             "project": "proj_rapidresponse",
         },
         max_attempts=5,
+        log_root=pm_data.log_dir("model_prep_training_data"),
     )
 
     if status != "D":
         msg = f"Workflow failed with status {status}."
         raise RuntimeError(msg)
 
     print("Building summary datasets.")
-    people_per_structure, pixel_area_weight = utils.build_summary_data(
-        pm_data, resolution
-    )
-
-    pm_data.save_summary_training_data(
-        resolution,
-        people_per_structure,
-        pixel_area_weight,
-    )
+    people_per_structure = utils.build_summary_people_per_structure(resolution, pm_data)
+    pm_data.save_summary_people_per_structure(people_per_structure, resolution)

src/rra_population_model/model_prep/training_data/utils.py

@@ -85,22 +85,17 @@ def build_arg_list(
     return to_run
 
 
-def build_summary_data(
+def build_summary_people_per_structure(
     pm_data: PopulationModelData,
     resolution: str,
-) -> tuple[gpd.GeoDataFrame, gpd.GeoDataFrame]:
+) -> pd.DataFrame:
     tile_dirs = list(pm_data.tile_training_data_root(resolution).iterdir())
-    model_gdfs = []
-    for tile_dir in tqdm.tqdm(tile_dirs):
-        tile_pps = pm_data.load_people_per_structure(resolution, tile_dir.name)
-        tile_paw = pm_data.load_pixel_area_weights(resolution, tile_dir.name)
-        model_gdfs.append((tile_pps, tile_paw))
+    data = pd.concat([
+        pm_data.load_people_per_structure(resolution, tile_dir.name)
+        for tile_dir in tqdm.tqdm(tile_dirs)
+    ], ignore_index=True)
+    return data
 
-    pps, paw = zip(*model_gdfs, strict=False)
-
-    people_per_structure_gdf = pd.concat(pps, ignore_index=True)
-    pixel_area_weight_gdf = pd.concat(paw, ignore_index=True)
-    return people_per_structure_gdf, pixel_area_weight_gdf
 
 
 def safe_divide(

src/rra_population_model/postprocess/__init__.py

@@ -1,7 +1,6 @@
-from rra_population_model.postprocess.compile_results import (
-    compile_results,
-    compile_results_task,
-    upsample_task,
+from rra_population_model.postprocess.mosaic.runner import (
+    mosaic,
+    mosaic_task,
 )
 from rra_population_model.postprocess.rake.runner import (
     rake,
@@ -15,18 +14,23 @@
     raking_factors,
     raking_factors_task,
 )
+from rra_population_model.postprocess.upsample.runner import (
+    upsample,
+    upsample_task,
+)
 
 RUNNERS = {
     "raking_factors": raking_factors,
     "rake": rake,
     "rake_itu": rake_itu,
-    "compile": compile_results,
+    "mosaic": mosaic,
+    "upsample": upsample,
 }
 
 TASK_RUNNERS = {
     "raking_factors": raking_factors_task,
     "rake": rake_task,
     "rake_itu": rake_itu_task,
-    "compile": compile_results_task,
+    "mosaic": mosaic_task,
     "upsample": upsample_task,
 }