Zero out rather than NaN areas with no potential

Author: sjpfenninger

workflow/scripts/_script_utils.py

@@ -27,6 +27,15 @@ def random_categorical_cmap(n, base_cmap="tab20", seed=42):
     return mcolors.ListedColormap(colors)
 
 
+def plot_with_zero_separate(ax, da, cmap="viridis", zero_color="#e0e0e0"):
+    """Plot data array with zero values in a separate color."""
+    da.where(da != 0).plot(ax=ax, cmap=cmap)
+    da.where(da == 0).plot(
+        ax=ax, cmap=mcolors.ListedColormap([zero_color]), add_colorbar=False
+    )
+    return ax
+
+
 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
@@ -63,7 +72,8 @@ def plot_all_dataset_variables(ds, ncols=2, savefig=None, categorical_vars=[]):
             cmap = random_categorical_cmap(len(ds[var].values))
         else:
             cmap = "viridis"
-        ds[var].plot(ax=axes[i], cmap=cmap)
+
+        plot_with_zero_separate(ax=axes[i], da=ds[var], cmap=cmap)
         axes[i].set_title(var)
         # We want to save rasterized images also for e.g. PDF output
         # Any actor with a zorder below the value given here is rasterized

workflow/scripts/area_potential.py

@@ -7,6 +7,7 @@
 import matplotlib.pyplot as plt
 import xarray as xr
 import yaml
+from _script_utils import plot_with_zero_separate
 
 
 @click.command()
@@ -55,22 +56,23 @@ def get_area_potential(
     # Start with the configured pixel area as a base
     potential_da = ds[config["initial_area"]].squeeze(drop=True)  # Drop `band`
 
-    # Drop pixels from binary layers with share 0 from potential_da
+    # Zero out pixels from binary layers with share 0 from potential_da
     binary_layers = config.get("binary_layers", {})
     zero_binary_layers = [layer for layer, value in binary_layers.items() if value == 0]
     for layer in zero_binary_layers:
         if layer in ds:
-            potential_da = potential_da.where(~(ds[layer] > 0))
+            potential_da = potential_da.where(~(ds[layer] > 0), other=0)
         else:
             print(f"Warning: Layer '{layer}' not found in dataset. Skipping.")
 
-    # Apply the continuous_layers criteria to drop additional pixels
+    # Apply the continuous_layers criteria to zero out additional pixels
     continuous_layers = config.get("continuous_layers", {})
     for layer, layer_config in continuous_layers.items():
         if layer in ds:
             # Apply the min-max criteria
             potential_da = potential_da.where(
-                (ds[layer] <= layer_config["max"]) & (ds[layer] >= layer_config["min"])
+                (ds[layer] <= layer_config["max"]) & (ds[layer] >= layer_config["min"]),
+                other=0,
             )
             # If a share is defined, multiply the pixel area by the share
             if "share" in layer_config:
@@ -111,7 +113,8 @@ def get_area_potential(
     potential_da = potential_da.transpose("band", "y", "x")
     potential_da.rio.write_crs(ds.rio.crs, inplace=True)
 
-    potential_da.plot()
+    fig, ax = plt.subplots(1, 1)
+    ax = plot_with_zero_separate(ax=ax, da=potential_da)
     plt.savefig(plot_path, bbox_inches="tight")
 
     # Fill NaN with a nodata value only after plotting