Improve memory use in area_potential_report rule

Author: sjpfenninger

config/config.yaml

@@ -6,7 +6,7 @@ tiny_files: False
 #   A good option is "epsg:8857" (WGS 84 / Equal Earth Greenwich) for global coverage
 buffer_crs: "epsg:8857"
 
-split_by: country_id
+split_by: country_id  # likely country_id or shape_id
 
 land_cover_types:
   POST_FLOODING: FARM

tests/integration/test_config.yaml

@@ -100,3 +100,5 @@ module_area_potentials:
         protected: 0
       shapes_buffer:
         land: 10000  # meters
+
+  overrides: {}

workflow/rules/process.smk

@@ -98,6 +98,23 @@ rule aggregate_area_potential:
         """
 
 
+rule plot_aggregated_area_potential:
+    message:
+        "Plot aggregated area potential for the tech {wildcards.tech} in {wildcards.shape}."
+    input:
+        rules.aggregate_area_potential.output.aggregated_area_potential,
+    output:
+        report(
+            "results/{shape}/area_potential_{tech}.png", category="area_potential_plot"
+        ),
+    log:
+        "logs/{shape}/plot_aggregated_area_potential_{tech}.log",
+    conda:
+        "../envs/default.yaml"
+    script:
+        "../scripts/tif_to_png.py"
+
+
 rule area_potential_report:
     message:
         "Generate an overview report of the area potential for all techs in shapes {wildcards.shape}."
@@ -107,16 +124,16 @@ rule area_potential_report:
             "results/{{shape}}/area_potential_{tech}.tif",
             tech=config["techs"].keys(),
         ),
+        area_potential_plots=expand(
+            "results/{{shape}}/area_potential_{tech}.png",
+            tech=config["techs"].keys(),
+        ),
     output:
         csv="results/{shape}/area_potential_report.csv",
         html=report(
             "results/{shape}/area_potential_report.html",
             category="area_potential_report_table",
         ),
-        png=report(
-            "results/{shape}/area_potential_report.png",
-            category="area_potential_report_image",
-        ),
     log:
         "logs/{shape}/area_potential_report.log",
     conda:

workflow/scripts/report.py

@@ -3,42 +3,39 @@
 import geopandas as gpd
 import pandas as pd
 import rioxarray as rxr
-import script_utils
 import xarray as xr
 from resample import _rasterize_regions
 
 
-def report(shapes, area_potentials, csv_path, html_path, png_path):
+def report(shapes, area_potentials, csv_path, html_path):
     """Generate a report summarizing area potentials for different technologies."""
     shapes = gpd.read_parquet(shapes)
 
-    ds_inputs = xr.Dataset()
-
-    # Collect the area potentials from the input files
-    for area_potential in area_potentials:
-        ds_inputs[area_potential] = rxr.open_rasterio(
-            area_potential, mask_and_scale=True
-        )
-
-    ds_inputs["regions"] = (
-        ("y", "x"),
-        _rasterize_regions(shapes, ds_inputs[area_potential]),
+    print("Generating reference raster and rasterizing regions...")
+    reference_raster = rxr.open_rasterio(area_potentials[0])
+    regions = xr.DataArray(
+        _rasterize_regions(shapes, reference_raster),
+        dims=("y", "x"),
+        coords={"y": reference_raster.y, "x": reference_raster.x},
     )
+    # regions = xr.DataArray(("y", "x"), _rasterize_regions(shapes, reference_raster))
+    del reference_raster
 
-    script_utils.plot_all_dataset_variables(
-        ds_inputs, ncols=2, savefig=png_path, categorical_vars=["regions"]
-    )
-
-    ds_inputs = ds_inputs.squeeze().drop_vars(["band", "spatial_ref"])
-
+    # Collect the area potentials from the input files
     # Group the area potentials by regions, sum them up, and collect the resulting Series
     # into a DataFrame, where each column corresponds to a technology's area potential,
     # and the index corresponds to the regions.
     dataframes = []
-    for area_potential in area_potentials:
-        dataframes.append(
-            ds_inputs[area_potential].groupby(ds_inputs["regions"]).sum().to_pandas()
+    for area_potential_file in area_potentials:
+        print(f"Processing area potential file: {area_potential_file}")
+        da_area_potential = (
+            rxr.open_rasterio(area_potential_file, mask_and_scale=True)
+            .squeeze()
+            .drop_vars(["band", "spatial_ref"])
         )
+        dataframes.append(da_area_potential.groupby(regions).sum().to_pandas())
+        del da_area_potential
+
     df = pd.concat(dataframes, axis=1)
 
     # Add metadata columns from shapes in front of the data columns
@@ -64,5 +61,4 @@ def report(shapes, area_potentials, csv_path, html_path, png_path):
         snakemake.input.area_potentials,
         snakemake.output.csv,
         snakemake.output.html,
-        snakemake.output.png,
     )

workflow/scripts/resample.py

@@ -289,26 +289,6 @@ def resample_inputs(
     resampled.to_netcdf(output_path, encoding=netcdf4_encoding)
 
     print("Saving image to plot path:", plot_path)
-    # If needed, resample `resampled` to fit within a maximum of `max_pixels` pixels
-    max_pixels = 5000000
-    total_pixels = resampled.sizes["y"] * resampled.sizes["x"]
-    if total_pixels > max_pixels:
-        # Calculate the new resolution to fit within the max_pixels limit
-        resolution_multiplier = 1 / math.sqrt(total_pixels / max_pixels)
-        new_y_size = int(resampled.sizes["y"] * resolution_multiplier)
-        new_x_size = int(resampled.sizes["x"] * resolution_multiplier)
-        print(
-            f"Resampling old size {resampled.sizes['y']} x {resampled.sizes['x']} "
-            f"to new size: {new_y_size} x {new_x_size} "
-            f"to fit within {max_pixels} pixels."
-        )
-
-        resampled = resampled.coarsen(
-            x=round(resampled.sizes["x"] / new_x_size),
-            y=round(resampled.sizes["y"] / new_y_size),
-            boundary="trim",
-        ).mean()
-
     script_utils.plot_all_dataset_variables(resampled, ncols=3, savefig=plot_path)
 
 

workflow/scripts/script_utils.py

@@ -29,6 +29,26 @@ def random_categorical_cmap(n, base_cmap="tab20", seed=42):
 
 def plot_all_dataset_variables(ds, ncols=2, savefig=None, categorical_vars=[]):
     """Plot all variables in an xarray dataset on a grid of plots."""
+    # If needed, resample `ds` to fit within a maximum of `max_pixels` pixels
+    max_pixels = 5000000
+    total_pixels = ds.sizes["y"] * ds.sizes["x"]
+    if total_pixels > max_pixels:
+        # Calculate the new resolution to fit within the max_pixels limit
+        resolution_multiplier = 1 / math.sqrt(total_pixels / max_pixels)
+        new_y_size = int(ds.sizes["y"] * resolution_multiplier)
+        new_x_size = int(ds.sizes["x"] * resolution_multiplier)
+        print(
+            f"Resampling old size {ds.sizes['y']} x {ds.sizes['x']} "
+            f"to new size: {new_y_size} x {new_x_size} "
+            f"to fit within {max_pixels} pixels."
+        )
+
+        ds = ds.coarsen(
+            x=round(ds.sizes["x"] / new_x_size),
+            y=round(ds.sizes["y"] / new_y_size),
+            boundary="trim",
+        ).mean()
+
     # Drop dimensionless variables
     ds = ds.drop_vars(lambda x: [v for v, da in x.variables.items() if not da.ndim])
 

workflow/scripts/tif_to_png.py

@@ -0,0 +1,18 @@
+"""This script plots a TIF file to PNG format."""
+
+import rioxarray as rxr
+from script_utils import plot_all_dataset_variables
+
+
+def tif_to_png(tif_file_in, png_file_out):
+    """Convert a TIF file to PNG format."""
+    ds = rxr.open_rasterio(tif_file_in, mask_and_scale=True).to_dataset(
+        name=tif_file_in
+    )
+    plot_all_dataset_variables(
+        ds, ncols=2, savefig=png_file_out, categorical_vars=["regions"]
+    )
+
+
+if __name__ == "__main__":
+    tif_to_png(snakemake.input[0], snakemake.output[0])