openforis-whisp 2.0.0b2__py3-none-any.whl → 3.0.0a1__py3-none-any.whl

openforis_whisp/datasets.py

@@ -7,8 +7,10 @@
 # b) a prefix of "nXX_" if it is national/sub-national dataset (where XX is replaced by that country code), or a prefix of 'g_' if it covers more than one country.
 # c) a name for your image, defined by ".rename('add_your_image_name_here')". This becomes the column header in the output table.
 
-# NB for all  the above you will need to be running the package in editable mode for these local changes to take effect.
-# Editable mode runs the package locally and thus changes to any files are reflected immediately.
+# Tips:
+#  -Avoid getInfo() and for loops to speed up processing by keeping everything in the Earth Engine API.
+#  -For all the above you will need to be running the package in editable mode for these local changes to take effect.
+#   Editable mode runs the package locally and thus changes to any files are reflected immediately.
 
 import ee
 
@@ -24,6 +26,11 @@ from datetime import datetime
 # defining here instead of importing from config_runtime, to allow functioning as more of a standalone script
 geometry_area_column = "Area"
 
+# Calculate current year once at module load time (not in functions)
+# This avoids repeated datetime calls and potential .getInfo() calls
+CURRENT_YEAR = datetime.now().year
+CURRENT_YEAR_2DIGIT = CURRENT_YEAR % 100  # Last two digits for RADD datasets
+
 import inspect
 
 import logging
@@ -49,13 +56,13 @@ def g_esa_worldcover_trees_prep():
     esa_worldcover_trees_2020 = esa_worldcover_2020_raw.eq(95).Or(
         esa_worldcover_2020_raw.eq(10)
     )  # get trees and mnangroves
-    return esa_worldcover_trees_2020.rename("ESA_TC_2020")
+    return esa_worldcover_trees_2020.rename("ESA_TC_2020").selfMask()
 
 
 # EUFO_2020
 def g_jrc_gfc_2020_prep():
     jrc_gfc2020_raw = ee.ImageCollection("JRC/GFC2020/V2")
-    return jrc_gfc2020_raw.mosaic().rename("EUFO_2020")
+    return jrc_gfc2020_raw.mosaic().rename("EUFO_2020").selfMask()
 
 
 # GFC_TC_2020
@@ -64,7 +71,7 @@ def g_glad_gfc_10pc_prep():
     gfc_treecover2000 = gfc.select(["treecover2000"])
     gfc_loss2001_2020 = gfc.select(["lossyear"]).lte(20)
     gfc_treecover2020 = gfc_treecover2000.where(gfc_loss2001_2020.eq(1), 0)
-    return gfc_treecover2020.gt(10).rename("GFC_TC_2020")
+    return gfc_treecover2020.gt(10).rename("GFC_TC_2020").selfMask()
 
 
 # GLAD_Primary
@@ -77,8 +84,10 @@ def g_glad_pht_prep():
     )
     gfc = ee.Image("UMD/hansen/global_forest_change_2024_v1_12")
     gfc_loss2001_2020 = gfc.select(["lossyear"]).lte(20)
-    return primary_ht_forests2001.where(gfc_loss2001_2020.eq(1), 0).rename(
-        "GLAD_Primary"
+    return (
+        primary_ht_forests2001.where(gfc_loss2001_2020.eq(1), 0)
+        .rename("GLAD_Primary")
+        .selfMask()
     )
 
 
@@ -90,7 +99,7 @@ def g_jrc_tmf_undisturbed_prep():
         .mosaic()
         .eq(1)
     )  # update from https://github.com/forestdatapartnership/whisp/issues/42
-    return TMF_undist_2020.rename("TMF_undist")
+    return TMF_undist_2020.rename("TMF_undist").selfMask()
 
 
 # Forest Persistence FDaP
@@ -99,7 +108,7 @@ def g_fdap_forest_prep():
         "projects/forestdatapartnership/assets/community_forests/ForestPersistence_2020"
     )
     fdap_forest = fdap_forest_raw.gt(0.75)
-    return fdap_forest.rename("Forest_FDaP")
+    return fdap_forest.rename("Forest_FDaP").selfMask()
 
 
 #########################primary forest
@@ -107,27 +116,27 @@ def g_fdap_forest_prep():
 def g_gft_primary_prep():
     gft_raw = ee.ImageCollection("JRC/GFC2020_subtypes/V0").mosaic()
     gft_primary = gft_raw.eq(10)
-    return gft_primary.rename("GFT_primary")
+    return gft_primary.rename("GFT_primary").selfMask()
 
 
 # Intact Forest Landscape 2020
 def g_ifl_2020_prep():
     IFL_2020 = ee.Image("users/potapovpeter/IFL_2020")
-    return IFL_2020.rename("IFL_2020")
+    return IFL_2020.rename("IFL_2020").selfMask()
 
 
 # European Primary Forest Dataset
 def g_epfd_prep():
     EPFD = ee.FeatureCollection("HU_BERLIN/EPFD/V2/polygons")
     EPFD_binary = ee.Image().paint(EPFD, 1)
-    return EPFD_binary.rename("European_Primary_Forest")
+    return EPFD_binary.rename("European_Primary_Forest").selfMask()
 
 
 # EUFO JRC Global forest type - naturally regenerating planted/plantation forests
 def g_gft_nat_reg_prep():
     gft_raw = ee.ImageCollection("JRC/GFC2020_subtypes/V0").mosaic()
     gft_nat_reg = gft_raw.eq(1)
-    return gft_nat_reg.rename("GFT_naturally_regenerating")
+    return gft_nat_reg.rename("GFT_naturally_regenerating").selfMask()
 
 
 #########################planted and plantation forests
@@ -136,13 +145,13 @@ def g_gft_nat_reg_prep():
 def g_gft_plantation_prep():
     gft_raw = ee.ImageCollection("JRC/GFC2020_subtypes/V0").mosaic()
     gft_plantation = gft_raw.eq(20)
-    return gft_plantation.rename("GFT_planted_plantation")
+    return gft_plantation.rename("GFT_planted_plantation").selfMask()
 
 
 def g_iiasa_planted_prep():
     iiasa = ee.Image("projects/sat-io/open-datasets/GFM/FML_v3-2")
     iiasa_PL = iiasa.eq(31).Or(iiasa.eq(32))
-    return iiasa_PL.rename("IIASA_planted_plantation")
+    return iiasa_PL.rename("IIASA_planted_plantation").selfMask()
 
 
 #########################TMF regrowth in 2023
@@ -151,7 +160,7 @@ def g_tmf_regrowth_prep():
     TMF_AC = ee.ImageCollection("projects/JRC/TMF/v1_2024/AnnualChanges").mosaic()
     TMF_AC_2023 = TMF_AC.select("Dec2023")
     Regrowth_TMF = TMF_AC_2023.eq(4)
-    return Regrowth_TMF.rename("TMF_regrowth_2023")
+    return Regrowth_TMF.rename("TMF_regrowth_2023").selfMask()
 
 
 ############tree crops
@@ -168,7 +177,7 @@ def g_jrc_tmf_plantation_prep():
     plantation_2020 = plantation.where(
         deforestation_year.gte(2021), 0
     )  # update from https://github.com/forestdatapartnership/whisp/issues/42
-    return plantation_2020.rename("TMF_plant")
+    return plantation_2020.rename("TMF_plant").selfMask()
 
 
 # # Oil_palm_Descals
@@ -181,20 +190,22 @@ def g_creaf_descals_palm_prep():
         )
         .mosaic()
         .select("minNBR_date")
-    )
+    ).selfMask()
 
     # Calculate the year of plantation and select all below and including 2020
     oil_palm_plantation_year = img.divide(365).add(1970).floor().lte(2020)
 
     # Create a mask for plantations in the year 2020 or earlier
-    plantation_2020 = oil_palm_plantation_year.lte(2020).selfMask()
-    return plantation_2020.rename("Oil_palm_Descals")
+    plantation_2020 = oil_palm_plantation_year.lte(2020)
+    return plantation_2020.rename("Oil_palm_Descals").selfMask()
 
 
 # Cocoa_ETH
 def g_eth_kalischek_cocoa_prep():
-    return ee.Image("projects/ee-nk-cocoa/assets/cocoa_map_threshold_065").rename(
-        "Cocoa_ETH"
+    return (
+        ee.Image("projects/ee-nk-cocoa/assets/cocoa_map_threshold_065")
+        .rename("Cocoa_ETH")
+        .selfMask()
     )
 
 
@@ -212,7 +223,7 @@ def g_fdap_palm_prep():
         .mosaic()
         .gt(0.88)  # Precision and recall ~78% at 0.88 threshold.
     )
-    return fdap_palm.rename("Oil_palm_FDaP")
+    return fdap_palm.rename("Oil_palm_FDaP").selfMask()
 
 
 def g_fdap_palm_2023_prep():
@@ -224,7 +235,7 @@ def g_fdap_palm_2023_prep():
         .mosaic()
         .gt(0.88)  # Precision and recall ~78% at 0.88 threshold.
     )
-    return fdap_palm.rename("Oil_palm_2023_FDaP")
+    return fdap_palm.rename("Oil_palm_2023_FDaP").selfMask()
 
 
 # Cocoa FDaP
@@ -237,7 +248,7 @@ def g_fdap_cocoa_prep():
         .mosaic()
         .gt(0.96)  # Precision and recall ~87% 0.96 threshold.
     )
-    return fdap_cocoa.rename("Cocoa_FDaP")
+    return fdap_cocoa.rename("Cocoa_FDaP").selfMask()
 
 
 def g_fdap_cocoa_2023_prep():
@@ -249,7 +260,7 @@ def g_fdap_cocoa_2023_prep():
         .mosaic()
         .gt(0.96)  # Precision and recall ~87% 0.96 threshold.
     )
-    return fdap_cocoa.rename("Cocoa_2023_FDaP")
+    return fdap_cocoa.rename("Cocoa_2023_FDaP").selfMask()
 
 
 # Rubber FDaP
@@ -262,7 +273,7 @@ def g_fdap_rubber_prep():
         .mosaic()
         .gt(0.59)  # Precision and recall ~80% 0.59 threshold.
     )
-    return fdap_rubber.rename("Rubber_FDaP")
+    return fdap_rubber.rename("Rubber_FDaP").selfMask()
 
 
 def g_fdap_rubber_2023_prep():
@@ -274,7 +285,7 @@ def g_fdap_rubber_2023_prep():
         .mosaic()
         .gt(0.59)  # Threshold for Rubber
     )
-    return fdap_rubber.rename("Rubber_2023_FDaP")
+    return fdap_rubber.rename("Rubber_2023_FDaP").selfMask()
 
 
 # # Coffee FDaP
@@ -291,7 +302,7 @@ def g_fdap_coffee_2020_prep():
         .gt(0.99)  # Precision and recall ~54% 0.99 threshold.
     )
 
-    return coffee_2020.rename("Coffee_FDaP")
+    return coffee_2020.rename("Coffee_FDaP").selfMask()
 
 
 def g_fdap_coffee_2023_prep():
@@ -306,7 +317,7 @@ def g_fdap_coffee_2023_prep():
         .mosaic()
         .gt(0.99)  # Precision and recall ~54% 0.99 threshold.
     )
-    return coffee_2023.rename("Coffee_FDaP_2023")
+    return coffee_2023.rename("Coffee_FDaP_2023").selfMask()
 
 
 # Rubber_RBGE  - from Royal Botanical Gardens of Edinburgh (RBGE) NB for 2021
@@ -315,14 +326,16 @@ def g_rbge_rubber_prep():
         ee.Image(
             "users/wangyxtina/MapRubberPaper/rRubber10m202122_perc1585DifESAdist5pxPF"
         )
-        .unmask()
         .rename("Rubber_RBGE")
+        .selfMask()
     )
 
 
 # soy 2020 South America
 def g_soy_song_2020_prep():
-    return ee.Image("projects/glad/soy_annual_SA/2020").unmask().rename("Soy_Song_2020")
+    return (
+        ee.Image("projects/glad/soy_annual_SA/2020").rename("Soy_Song_2020").selfMask()
+    )
 
 
 ##############
@@ -336,7 +349,7 @@ def g_esri_2023_tc_prep():
     esri_lulc10_TC = (
         esri_lulc10_raw.filterDate("2023-01-01", "2023-12-31").mosaic().eq(2)
     )
-    return esri_lulc10_TC.rename("ESRI_2023_TC")
+    return esri_lulc10_TC.rename("ESRI_2023_TC").selfMask()
 
 
 # ESRI 2023 - Crop
@@ -353,7 +366,7 @@ def g_esri_2020_2023_crop_prep():
 
     newCrop = esri_lulc10_crop_2023.And(esri_lulc10_crop_2020.Not())
 
-    return newCrop.rename("ESRI_crop_gain_2020_2023")
+    return newCrop.rename("ESRI_crop_gain_2020_2023").selfMask()
 
 
 #### disturbances by year
@@ -363,23 +376,13 @@ def g_radd_year_prep():
     from datetime import datetime
 
     radd = ee.ImageCollection("projects/radar-wur/raddalert/v1")
-
     radd_date = (
         radd.filterMetadata("layer", "contains", "alert").select("Date").mosaic()
     )
-    # date of avaialbility
-    start_year = 19  ## (starts 2019 in Africa, then 2020 for S America and Asia: https://data.globalforestwatch.org/datasets/gfw::deforestation-alerts-radd/about
-
-    current_year = (
-        datetime.now().year
-        % 100
-        # NB the % 100 part gets last two digits needed
-    )
+    start_year = 19
+    current_year = datetime.now().year % 100
 
-    img_stack = None
-    # Generate an image based on GFC with one band of forest tree loss per year from 2001 to <current year>
-    for year in range(start_year, current_year + 1):
-        # gfc_loss_year = gfc.select(['lossyear']).eq(i).And(gfc.select(['treecover2000']).gt(10)) # use any definition of loss
+    def make_band(year, img_stack):
         start = year * 1000
         end = year * 1000 + 365
         radd_year = (
@@ -387,13 +390,39 @@ def g_radd_year_prep():
             .updateMask(radd_date.lte(end))
             .gt(0)
             .rename("RADD_year_" + "20" + str(year))
+            .selfMask()
+        )
+        return ee.Image(img_stack).addBands(radd_year)
+
+    years = ee.List.sequence(start_year, current_year)
+    first_year = ee.Number(years.get(0))
+    start = first_year.multiply(1000)
+    end = first_year.multiply(1000).add(365)
+    band_name = ee.String("RADD_year_").cat("20").cat(first_year.format("%02d"))
+    first_band = (
+        radd_date.updateMask(radd_date.gte(start))
+        .updateMask(radd_date.lte(end))
+        .gt(0)
+        .rename(band_name)
+        .selfMask()
+    )
+
+    def make_band(year, img_stack):
+        year_num = ee.Number(year)
+        start = year_num.multiply(1000)
+        end = year_num.multiply(1000).add(365)
+        band_name = ee.String("RADD_year_").cat("20").cat(year_num.format("%02d"))
+        radd_year = (
+            radd_date.updateMask(radd_date.gte(start))
+            .updateMask(radd_date.lte(end))
+            .gt(0)
+            .rename(band_name)
+            .selfMask()
         )
+        return ee.Image(img_stack).addBands(radd_year)
 
-        if img_stack is None:
-            img_stack = radd_year
-        else:
-            img_stack = img_stack.addBands(radd_year)
-    return img_stack
+    img_stack = years.slice(1).iterate(make_band, first_band)
+    return ee.Image(img_stack)
 
 
 # TMF_def_2000 to TMF_def_2023
@@ -403,7 +432,9 @@ def g_tmf_def_per_year_prep():
     img_stack = None
     # Generate an image based on GFC with one band of forest tree loss per year from 2001 to 2022
     for i in range(0, 24 + 1):
-        tmf_def_year = tmf_def.eq(2000 + i).rename("TMF_def_" + str(2000 + i))
+        year_num = ee.Number(2000 + i)
+        band_name = ee.String("TMF_def_").cat(year_num.format("%d"))
+        tmf_def_year = tmf_def.eq(year_num).rename(band_name).selfMask()
         if img_stack is None:
             img_stack = tmf_def_year
         else:
@@ -418,7 +449,9 @@ def g_tmf_deg_per_year_prep():
     img_stack = None
     # Generate an image based on GFC with one band of forest tree loss per year from 2001 to 2022
     for i in range(0, 24 + 1):
-        tmf_def_year = tmf_def.eq(2000 + i).rename("TMF_deg_" + str(2000 + i))
+        year_num = ee.Number(2000 + i)
+        band_name = ee.String("TMF_deg_").cat(year_num.format("%d"))
+        tmf_def_year = tmf_def.eq(year_num).rename(band_name).selfMask()
         if img_stack is None:
             img_stack = tmf_def_year
         else:
@@ -433,10 +466,12 @@ def g_glad_gfc_loss_per_year_prep():
     img_stack = None
     # Generate an image based on GFC with one band of forest tree loss per year from 2001 to 2022
     for i in range(1, 24 + 1):
+        year_num = ee.Number(2000 + i)
+        band_name = ee.String("GFC_loss_year_").cat(year_num.format("%d"))
         gfc_loss_year = (
             gfc.select(["lossyear"]).eq(i).And(gfc.select(["treecover2000"]).gt(10))
         )
-        gfc_loss_year = gfc_loss_year.rename("GFC_loss_year_" + str(2000 + i))
+        gfc_loss_year = gfc_loss_year.rename(band_name).selfMask()
         if img_stack is None:
             img_stack = gfc_loss_year
         else:
@@ -457,6 +492,8 @@ def g_modis_fire_prep():
     img_stack = None
 
     for year in range(start_year, end_year + 1):
+        year_num = ee.Number(year)
+        band_name = ee.String("MODIS_fire_").cat(year_num.format("%d"))
         date_st = f"{year}-01-01"
         date_ed = f"{year}-12-31"
         modis_year = (
@@ -464,7 +501,8 @@ def g_modis_fire_prep():
             .mosaic()
             .select(["BurnDate"])
             .gte(0)
-            .rename(f"MODIS_fire_{year}")
+            .rename(band_name)
+            .selfMask()
         )
         img_stack = modis_year if img_stack is None else img_stack.addBands(modis_year)
 
@@ -484,6 +522,8 @@ def g_esa_fire_prep():
     img_stack = None
 
     for year in range(start_year, end_year + 1):
+        year_num = ee.Number(year)
+        band_name = ee.String("ESA_fire_").cat(year_num.format("%d"))
         date_st = f"{year}-01-01"
         date_ed = f"{year}-12-31"
         esa_year = (
@@ -491,107 +531,18 @@ def g_esa_fire_prep():
             .mosaic()
             .select(["BurnDate"])
             .gte(0)
-            .rename(f"ESA_fire_{year}")
+            .rename(band_name)
+            .selfMask()
         )
         img_stack = esa_year if img_stack is None else img_stack.addBands(esa_year)
 
     return img_stack
 
 
-# # DIST_alert_2024 to DIST_alert_< current year >
-# # Notes:
-# # 1) so far only available for 2024 onwards in GEE
-# # TO DO - see if gee asset for pre 2020-2024 is available from GLAD team, else download from nasa and put in Whisp assets
-# # 2) masked alerts (as dist alerts are for all vegetation) to JRC EUFO 2020 layer, as close to EUDR definition
-# # TO DO - ask opinions on if others (such as treecover data from GLAD team) should be used instead
-
-
-# def glad_dist_year_prep():
-
-#     # Load the vegetation disturbance collections
-
-#     #  Vegetation disturbance status (0-8, class flag, 8-bit)
-#     VEGDISTSTATUS = ee.ImageCollection(
-#         "projects/glad/HLSDIST/current/VEG-DIST-STATUS"
-#     ).mosaic()
-#     # Initial vegetation disturbance date (>0: days since 2020-12-31, 16-bit)
-#     VEGDISTDATE = ee.ImageCollection(
-#         "projects/glad/HLSDIST/current/VEG-DIST-DATE"
-#     ).mosaic()
-
-#     # NB relies on initial date of disturbance - consider if last date needed? : VEGLASTDATE = ee.ImageCollection("projects/glad/HLSDIST/current/VEG-LAST-DATE").mosaic(); # Last assessed observation date (≥1, days, 16-bit)
-
-#     # Key for high-confidence alerts (values 3, 6, 7, 8)
-#     high_conf_values = [3, 6, 7, 8]
-#     # where:
-#     # 3 = <50% loss, high confidence, ongoing
-#     # 6 = ≥50% loss, high confidence, ongoing
-#     # 7 = <50% loss, high confidence, finished
-#     # 8 = ≥50% loss, high confidence, finished
-#     # Note could use <50% loss (i.e. only 6 and 7) for if want to be more strict
-
-#     # Create high-confidence mask
-#     dist_high_conf = VEGDISTSTATUS.remap(
-#         high_conf_values, [1] * len(high_conf_values), 0
-#     )
-
-#     # Determine start year and current year dynamically
-#     start_year = 2024  # Set the first year of interest
-#     current_year = datetime.now().year
-
-#     # Calculate days since December 31, 2020 for start and end dates (server-side)
-#     start_of_2020 = ee.Date("2020-12-31").millis().divide(86400000).int()
-
-#     # Create a list to hold the yearly images
-#     yearly_images = []
-
-#     for year in range(start_year, current_year + 1):
-#         start_of_year = (
-#             ee.Date(f"{year}-01-01")
-#             .millis()
-#             .divide(86400000)
-#             .int()
-#             .subtract(start_of_2020)
-#         )
-#         start_of_next_year = (
-#             ee.Date(f"{year + 1}-01-01")
-#             .millis()
-#             .divide(86400000)
-#             .int()
-#             .subtract(start_of_2020)
-#         )
-
-#         # Filter VEG-DIST-DATE for the selected year
-#         dist_year = VEGDISTDATE.gte(start_of_year).And(
-#             VEGDISTDATE.lt(start_of_next_year)
-#         )
-
-#         # Apply high-confidence mask and rename the band
-#         high_conf_year = dist_year.updateMask(dist_high_conf).rename(
-#             f"DIST_year_{year}"
-#         )
-
-#         # Append the year's data to the list
-#         yearly_images.append(high_conf_year)
-
-#     # Combine all yearly images into a single image
-#     img_stack = ee.Image.cat(yearly_images)
-
-#     # Rename the bands correctly
-#     band_names = [f"DIST_year_{year}" for year in range(start_year, current_year + 1)]
-#     img_stack = img_stack.select(img_stack.bandNames(), band_names)
-
-#     return img_stack.updateMask(
-#         jrc_gfc_2020_prep()
-#     )  # mask yearly dist alerts to forest cover in 2020
-
-
 #### disturbances combined (split into before and after 2020)
 
 # RADD_after_2020
 def g_radd_after_2020_prep():
-    from datetime import datetime
-
     radd = ee.ImageCollection("projects/radar-wur/raddalert/v1")
 
     radd_date = (
@@ -600,9 +551,8 @@ def g_radd_after_2020_prep():
     # date of avaialbility
     start_year = 21  ## (starts 2019 in Africa, then 2020 for S America and Asia: https://data.globalforestwatch.org/datasets/gfw::deforestation-alerts-radd/about)
 
-    current_year = (
-        datetime.now().year % 100
-    )  # NB the % 100 part gets last two digits needed
+    # Use pre-calculated current year (avoids repeated datetime calls)
+    current_year = CURRENT_YEAR_2DIGIT
     start = start_year * 1000
     end = current_year * 1000 + 365
     return (
@@ -610,13 +560,11 @@ def g_radd_after_2020_prep():
         .updateMask(radd_date.lte(end))
         .gt(0)
         .rename("RADD_after_2020")
-    )
+    ).selfMask()
 
 
 # RADD_before_2020
 def g_radd_before_2020_prep():
-    from datetime import datetime
-
     radd = ee.ImageCollection("projects/radar-wur/raddalert/v1")
 
     radd_date = (
@@ -625,8 +573,6 @@ def g_radd_before_2020_prep():
     # date of avaialbility
     start_year = 19  ## (starts 2019 in Africa, then 2020 for S America and Asia: https://data.globalforestwatch.org/datasets/gfw::deforestation-alerts-radd/about)
 
-    # current_year = datetime.now().year % 100 # NB the % 100 part gets last two digits needed
-
     start = start_year * 1000
     end = 20 * 1000 + 365
     return (
@@ -634,7 +580,7 @@ def g_radd_before_2020_prep():
         .updateMask(radd_date.lte(end))
         .gt(0)
         .rename("RADD_before_2020")
-    )
+    ).selfMask()
 
 
 # # DIST_after_2020
@@ -662,25 +608,35 @@ def g_radd_before_2020_prep():
 # TMF_deg_before_2020
 def g_tmf_deg_before_2020_prep():
     tmf_deg = ee.ImageCollection("projects/JRC/TMF/v1_2024/DegradationYear").mosaic()
-    return (tmf_deg.lte(2020)).And(tmf_deg.gte(2000)).rename("TMF_deg_before_2020")
+    return (
+        (tmf_deg.lte(2020))
+        .And(tmf_deg.gte(2000))
+        .rename("TMF_deg_before_2020")
+        .selfMask()
+    )
 
 
 # TMF_deg_after_2020
 def g_tmf_deg_after_2020_prep():
     tmf_deg = ee.ImageCollection("projects/JRC/TMF/v1_2024/DegradationYear").mosaic()
-    return tmf_deg.gt(2020).rename("TMF_deg_after_2020")
+    return tmf_deg.gt(2020).rename("TMF_deg_after_2020").selfMask()
 
 
 # tmf_def_before_2020
 def g_tmf_def_before_2020_prep():
     tmf_def = ee.ImageCollection("projects/JRC/TMF/v1_2024/DeforestationYear").mosaic()
-    return (tmf_def.lte(2020)).And(tmf_def.gte(2000)).rename("TMF_def_before_2020")
+    return (
+        (tmf_def.lte(2020))
+        .And(tmf_def.gte(2000))
+        .rename("TMF_def_before_2020")
+        .selfMask()
+    )
 
 
 # tmf_def_after_2020
 def g_tmf_def_after_2020_prep():
     tmf_def = ee.ImageCollection("projects/JRC/TMF/v1_2024/DeforestationYear").mosaic()
-    return tmf_def.gt(2020).rename("TMF_def_after_2020")
+    return tmf_def.gt(2020).rename("TMF_def_after_2020").selfMask()
 
 
 # GFC_loss_before_2020 (loss within 10 percent cover; includes 2020; correct for version 11)
@@ -690,7 +646,7 @@ def g_glad_gfc_loss_before_2020_prep():
     gfc_loss = (
         gfc.select(["lossyear"]).lte(20).And(gfc.select(["treecover2000"]).gt(10))
     )
-    return gfc_loss.rename("GFC_loss_before_2020")
+    return gfc_loss.rename("GFC_loss_before_2020").selfMask()
 
 
 # GFC_loss_after_2020 (loss within 10 percent cover; correct for version 11)
@@ -698,7 +654,7 @@ def g_glad_gfc_loss_after_2020_prep():
     # Load the Global Forest Change dataset
     gfc = ee.Image("UMD/hansen/global_forest_change_2024_v1_12")
     gfc_loss = gfc.select(["lossyear"]).gt(20).And(gfc.select(["treecover2000"]).gt(10))
-    return gfc_loss.rename("GFC_loss_after_2020")
+    return gfc_loss.rename("GFC_loss_after_2020").selfMask()
 
 
 # MODIS_fire_before_2020
@@ -714,14 +670,15 @@ def g_modis_fire_before_2020_prep():
         .select(["BurnDate"])
         .gte(0)
         .rename("MODIS_fire_before_2020")
-    )
+    ).selfMask()
 
 
 # MODIS_fire_after_2020
 def g_modis_fire_after_2020_prep():
     modis_fire = ee.ImageCollection("MODIS/061/MCD64A1")
     start_year = 2021
-    end_year = datetime.now().year
+    # Use pre-calculated current year (avoids repeated datetime calls)
+    end_year = CURRENT_YEAR - 1  # Use year - 1 to ensure data availability
     date_st = str(start_year) + "-01-01"
     date_ed = str(end_year) + "-12-31"
     return (
@@ -730,7 +687,7 @@ def g_modis_fire_after_2020_prep():
         .select(["BurnDate"])
         .gte(0)
         .rename("MODIS_fire_after_2020")
-    )
+    ).selfMask()
 
 
 # ESA_fire_before_2020
@@ -746,7 +703,7 @@ def g_esa_fire_before_2020_prep():
         .select(["BurnDate"])
         .gte(0)
         .rename("ESA_fire_before_2020")
-    )
+    ).selfMask()
 
 
 #########################logging concessions
@@ -792,7 +749,7 @@ def g_logging_concessions_before_2020_prep():
         ]
     ).mosaic()
 
-    return logging_concessions_binary.rename("GFW_logging_before_2020")
+    return logging_concessions_binary.rename("GFW_logging_before_2020").selfMask()
 
 
 #########################national datasets
@@ -810,7 +767,7 @@ def g_logging_concessions_before_2020_prep():
 def nbr_terraclass_amz20_primary_prep():
     tcamz20 = ee.Image("projects/ee-whisp/assets/NBR/terraclass_amz_2020")
     tcamz20_f = tcamz20.eq(1)
-    return tcamz20_f.rename("nBR_INPE_TC_primary_forest_Amazon_2020")
+    return tcamz20_f.rename("nBR_INPE_TC_primary_forest_Amazon_2020").selfMask()
 
 
 # [Official NFMS dataset] Brazilian Forest Service dataset on natural forest cover from PRODES and TerraClass data, base year 2022
@@ -824,7 +781,7 @@ def nbr_bfs_ptn_f20_prep():
     bfs_fptn20 = ee.FeatureCollection("projects/ee-whisp/assets/NBR/bfs_ptn_2020")
 
     bfs_fptn20_binary = ee.Image().paint(bfs_fptn20, 1)
-    return bfs_fptn20_binary.rename("nBR_BFS_primary_forest_Pantanal_2020")
+    return bfs_fptn20_binary.rename("nBR_BFS_primary_forest_Pantanal_2020").selfMask()
 
 
 # Caatinga - filtered with QGIS because the original geodatabase is too large to export as a shapefile (GEE accepted format)
@@ -832,35 +789,39 @@ def nbr_bfs_ptn_f20_prep():
 def nbr_bfs_caat_f20_prep():
     bfs_fcaat20 = ee.FeatureCollection("projects/ee-whisp/assets/NBR/bfs_caat_2020")
     bfs_fcaat20_binary = ee.Image().paint(bfs_fcaat20, 1)
-    return bfs_fcaat20_binary.rename("nBR_BFS_primary_forest_Caatinga_2020")
+    return bfs_fcaat20_binary.rename("nBR_BFS_primary_forest_Caatinga_2020").selfMask()
 
 
 # Atlantic Forest - filtered with QGIS because the original geodatabase is too large to export as a shapefile (GEE accepted format)
 def nbr_bfs_atlf_f20_prep():
     bfs_fatlf20 = ee.FeatureCollection("projects/ee-whisp/assets/NBR/bfs_atlf_2020")
     bfs_fatlf20_binary = ee.Image().paint(bfs_fatlf20, 1)
-    return bfs_fatlf20_binary.rename("nBR_BFS_primary_forest_AtlanticForest_2020")
+    return bfs_fatlf20_binary.rename(
+        "nBR_BFS_primary_forest_AtlanticForest_2020"
+    ).selfMask()
 
 
 # Pampa - filtered in QGIS to save some storage space
 def nbr_bfs_pmp_f20_prep():
     bfs_fpmp20 = ee.FeatureCollection("projects/ee-whisp/assets/NBR/bfs_pmp_2020")
     bfs_fpmp20_binary = ee.Image().paint(bfs_fpmp20, 1)
-    return bfs_fpmp20_binary.rename("nBR_BFS_primary_forest_Pampa_2020")
+    return bfs_fpmp20_binary.rename("nBR_BFS_primary_forest_Pampa_2020").selfMask()
 
 
 ##########################secondary forests###############################################
 def nbr_terraclass_amz20_secondary_prep():
     tcamz20 = ee.Image("projects/ee-whisp/assets/NBR/terraclass_amz_2020")
     tcamz20_f = tcamz20.eq(2)
-    return tcamz20_f.rename("nBR_INPE_TC_secondary_forest_Amazon_2020")
+    return tcamz20_f.rename("nBR_INPE_TC_secondary_forest_Amazon_2020").selfMask()
 
 
 # Cerrado - filtered with QGIS because the original geodatabase is too large to export as a shapefile (GEE accepted format)
 def nbr_bfs_cer_f20_prep():
     bfs_fcer20 = ee.FeatureCollection("projects/ee-whisp/assets/NBR/bfs_cerr_2020")
     bfs_fcer20_binary = ee.Image().paint(bfs_fcer20, 1)
-    return bfs_fcer20_binary.rename("nBR_BFS_primary_and_secondary_forest_Cerrado_2020")
+    return bfs_fcer20_binary.rename(
+        "nBR_BFS_primary_and_secondary_forest_Cerrado_2020"
+    ).selfMask()
 
 
 # %%
@@ -879,7 +840,9 @@ def nbr_mapbiomasc9_f20_prep():
         .Or(mapbiomasc9_20.eq(6))
         .Or(mapbiomasc9_20.eq(49))
     )
-    return mapbiomasc9_20_forest.rename("nBR_MapBiomas_col9_forest_Brazil_2020")
+    return mapbiomasc9_20_forest.rename(
+        "nBR_MapBiomas_col9_forest_Brazil_2020"
+    ).selfMask()
 
 
 # ### ########################NBR plantation forest in 2020:#######################################
@@ -890,7 +853,7 @@ def nbr_mapbiomasc9_f20_prep():
 def nbr_terraclass_amz20_silv_prep():
     tcamz20 = ee.Image("projects/ee-whisp/assets/NBR/terraclass_amz_2020")
     tcamz20_silviculture = tcamz20.eq(9)
-    return tcamz20_silviculture.rename("nBR_INPE_TCsilviculture_Amazon_2020")
+    return tcamz20_silviculture.rename("nBR_INPE_TCsilviculture_Amazon_2020").selfMask()
 
 
 # [Official NFMS dataset] INPE/EMBRAPA TerraClass land use/cover in the Cerrado biome, 2020
@@ -899,7 +862,9 @@ def nbr_terraclass_amz20_silv_prep():
 def nbr_terraclass_silv_cer20_prep():
     tccer20 = ee.Image("projects/ee-whisp/assets/NBR/terraclass_cer_2020")
     tccer20_silviculture = tccer20.eq(9)
-    return tccer20_silviculture.rename("nBR_INPE_TCsilviculture_Cerrado_2020")
+    return tccer20_silviculture.rename(
+        "nBR_INPE_TCsilviculture_Cerrado_2020"
+    ).selfMask()
 
 
 # [non-official dataset by MapBiomas multisector initiative]
@@ -913,7 +878,7 @@ def nbr_mapbiomasc9_silv20_prep():
     mapbiomasc9_20_silviculture = mapbiomasc9_20.eq(9)
     return mapbiomasc9_20_silviculture.rename(
         "nBR_MapBiomas_col9_silviculture_Brazil_2020"
-    )
+    ).selfMask()
 
 
 ################ ### NBR Disturbances before 2020:########################################
@@ -958,8 +923,9 @@ def nbr_prodes_before_2020_prep():
     prodes_before_20_mask = prodes.remap(
         prodes_before_20_dn, [1] * len(prodes_before_20_dn)
     )  # .eq(1)
-    prodes_before_20 = prodes_before_20_mask.selfMask()
-    return prodes_before_20.rename("nBR_PRODES_deforestation_Brazil_before_2020")
+    return prodes_before_20_mask.rename(
+        "nBR_PRODES_deforestation_Brazil_before_2020"
+    ).selfMask()
 
 
 ## Caution: 1) includes deforestation and conversion of other wooded land and grassland
@@ -984,7 +950,9 @@ def nbr_deter_amazon_before_2020_prep():
     ).filter(ee.Filter.lt("formatted_date", ee.Date("2020-12-31")))
 
     deter_deg_binary = ee.Image().paint(deter_deg, 1)
-    return deter_deg_binary.rename("nBR_DETER_forestdegradation_Amazon_before_2020")
+    return deter_deg_binary.rename(
+        "nBR_DETER_forestdegradation_Amazon_before_2020"
+    ).selfMask()
 
 
 ################ ### NBR Disturbances after 2020:########################################
@@ -1001,7 +969,9 @@ def nbr_prodes_after_2020_prep():
         prodes_after_20_dn, [1] * len(prodes_after_20_dn)
     )  # .eq(1)
     prodes_after_20 = prodes_after_20_mask.selfMask()
-    return prodes_after_20.rename("nBR_PRODES_deforestation_Brazil_after_2020")
+    return prodes_after_20.rename(
+        "nBR_PRODES_deforestation_Brazil_after_2020"
+    ).selfMask()
 
 
 # %%
@@ -1023,7 +993,9 @@ def nbr_deter_amazon_after_2020_prep():
     ).filter(ee.Filter.gt("formatted_date", ee.Date("2021-01-01")))
 
     deter_deg_binary = ee.Image().paint(deter_deg, 1)
-    return deter_deg_binary.rename("nBR_DETER_forestdegradation_Amazon_after_2020")
+    return deter_deg_binary.rename(
+        "nBR_DETER_forestdegradation_Amazon_after_2020"
+    ).selfMask()
 
 
 # ########################## NBR commodities - permanent/perennial crops in 2020:###############################
@@ -1037,7 +1009,7 @@ def nbr_terraclass_amz_cer20_pc_prep():
     tccer20 = ee.Image("projects/ee-whisp/assets/NBR/terraclass_cer_2020")
     tccer20_pc = tccer20.eq(12).Or(tccer20.eq(13))
     tc_pc = ee.ImageCollection([tcamz20_pc, tccer20_pc]).mosaic()
-    return tc_pc.rename("nBR_INPE_TCamz_cer_perennial_2020")
+    return tc_pc.rename("nBR_INPE_TCamz_cer_perennial_2020").selfMask()
 
 
 # [non-official dataset by MapBiomas multisector initiative]
@@ -1049,7 +1021,7 @@ def nbr_mapbiomasc9_cof_prep():
         "projects/mapbiomas-public/assets/brazil/lulc/collection9/mapbiomas_collection90_integration_v1"
     ).select("classification_2020")
     mapbiomasc9_20_coffee = mapbiomasc9_20.eq(46)
-    return mapbiomasc9_20_coffee.rename("nBR_MapBiomas_col9_coffee_2020")
+    return mapbiomasc9_20_coffee.rename("nBR_MapBiomas_col9_coffee_2020").selfMask()
 
 
 # [non-official dataset by MapBiomas multisector initiative]
@@ -1061,7 +1033,7 @@ def nbr_mapbiomasc9_po_prep():
         "projects/mapbiomas-public/assets/brazil/lulc/collection9/mapbiomas_collection90_integration_v1"
     ).select("classification_2020")
     mapbiomasc9_20_palm = mapbiomasc9_20.eq(35)
-    return mapbiomasc9_20_palm.rename("nBR_MapBiomas_col9_palmoil_2020")
+    return mapbiomasc9_20_palm.rename("nBR_MapBiomas_col9_palmoil_2020").selfMask()
 
 
 # [non-official dataset by MapBiomas multisector initiative]
@@ -1073,7 +1045,7 @@ def nbr_mapbiomasc9_pc_prep():
         "projects/mapbiomas-public/assets/brazil/lulc/collection9/mapbiomas_collection90_integration_v1"
     ).select("classification_2020")
     mapbiomasc9_20_pc = mapbiomasc9_20.eq(35).Or(mapbiomasc9_20.eq(46))
-    return mapbiomasc9_20_pc.rename("nBR_MapBiomas_col9_pc_2020")
+    return mapbiomasc9_20_pc.rename("nBR_MapBiomas_col9_pc_2020").selfMask()
 
 
 # ######################## NBR commodities - annual crops in 2020:##############################
@@ -1089,7 +1061,7 @@ def nbr_terraclass_amz_cer20_ac_prep():
     tccer20 = ee.Image("projects/ee-whisp/assets/NBR/terraclass_cer_2020")
     tccer20_ac = tccer20.eq(14).Or(tccer20.eq(15))
     tc_ac = ee.ImageCollection([tcamz20_ac, tccer20_ac]).mosaic()
-    return tc_ac.rename("nBR_INPE_TCamz_cer_annual_2020")
+    return tc_ac.rename("nBR_INPE_TCamz_cer_annual_2020").selfMask()
 
 
 # [non-official dataset by MapBiomas multisector initiative]
@@ -1101,7 +1073,7 @@ def nbr_mapbiomasc9_soy_prep():
         "projects/mapbiomas-public/assets/brazil/lulc/collection9/mapbiomas_collection90_integration_v1"
     ).select("classification_2020")
     mapbiomasc9_20_soy = mapbiomasc9_20.eq(39)
-    return mapbiomasc9_20_soy.rename("nBR_MapBiomas_col9_soy_2020")
+    return mapbiomasc9_20_soy.rename("nBR_MapBiomas_col9_soy_2020").selfMask()
 
 
 # [non-official dataset by MapBiomas multisector initiative]
@@ -1121,7 +1093,7 @@ def nbr_mapbiomasc9_ac_prep():
         .Or(mapbiomasc9_20.eq(40))
         .Or(mapbiomasc9_20.eq(62))
     )
-    return mapbiomasc9_20_ac.rename("nBR_MapBiomas_col9_annual_crops_2020")
+    return mapbiomasc9_20_ac.rename("nBR_MapBiomas_col9_annual_crops_2020").selfMask()
 
 
 # ################################### NBR commodities - pasture/livestock in 2020:##############################
@@ -1134,7 +1106,7 @@ def nbr_mapbiomasc9_ac_prep():
 def nbr_terraclass_amz20_pasture_prep():
     tcamz20 = ee.Image("projects/ee-whisp/assets/NBR/terraclass_amz_2020")
     tcamz20_pasture = tcamz20.eq(10).Or(tcamz20.eq(11))
-    return tcamz20_pasture.rename("nBR_INPE_TCamz_pasture_2020")
+    return tcamz20_pasture.rename("nBR_INPE_TCamz_pasture_2020").selfMask()
 
 
 # %%
@@ -1146,7 +1118,7 @@ def nbr_terraclass_amz20_pasture_prep():
 def nbr_terraclass_cer20_ac_prep():
     tccer20 = ee.Image("projects/ee-whisp/assets/NBR/terraclass_cer_2020")
     tccer20_pasture = tccer20.eq(11)
-    return tccer20_pasture.rename("nBR_INPE_TCcer_pasture_2020")
+    return tccer20_pasture.rename("nBR_INPE_TCcer_pasture_2020").selfMask()
 
 
 # %%
@@ -1159,7 +1131,7 @@ def nbr_mapbiomasc9_pasture_prep():
         "projects/mapbiomas-public/assets/brazil/lulc/collection9/mapbiomas_collection90_integration_v1"
     ).select("classification_2020")
     mapbiomasc9_20_pasture = mapbiomasc9_20.eq(15)
-    return mapbiomasc9_20_pasture.rename("nBR_MapBiomas_col9_pasture_2020")
+    return mapbiomasc9_20_pasture.rename("nBR_MapBiomas_col9_pasture_2020").selfMask()
 
 
 ###################################################################
@@ -1169,13 +1141,13 @@ def nbr_mapbiomasc9_pasture_prep():
 def nco_ideam_forest_2020_prep():
     ideam_forest_raw = ee.Image("projects/ee-whisp/assets/nCO/ideam_2020_geo")
     ideam_forest = ideam_forest_raw.eq(1)  # get forest class
-    return ideam_forest.rename("nCO_ideam_forest_2020")
+    return ideam_forest.rename("nCO_ideam_forest_2020").selfMask()
 
 
 def nco_ideam_eufo_commission_2020_prep():
     ideam_agroforest_raw = ee.Image("projects/ee-whisp/assets/nCO/ideam_2020_geo_EUFO")
     ideam_agroforest = ideam_agroforest_raw.eq(4)  # get forest class
-    return ideam_agroforest.rename("nCO_ideam_eufo_commission_2020")
+    return ideam_agroforest.rename("nCO_ideam_eufo_commission_2020").selfMask()
 
 
 # Cocoa_bnetd
@@ -1185,47 +1157,138 @@ def nci_ocs2020_prep():
         .select("classification")
         .eq(9)
         .rename("nCI_Cocoa_bnetd")
-    )  # cocoa from national land cover map for Côte d'Ivoire
+    ).selfMask()  # cocoa from national land cover map for Côte d'Ivoire
+
+
+# ============================================================================
+# CONTEXT BANDS (Administrative boundaries and water mask)
+# ============================================================================
+
+
+def g_gaul_admin_code():
+    """
+    GAUL 2024 Level 1 administrative boundary codes (500m resolution).
+    Used for spatial context and administrative aggregation.
+
+    Returns
+    -------
+    ee.Image
+        Image with admin codes renamed to 'admin_code' (as int32)
+    """
+    admin_image = ee.Image(
+        "projects/ee-andyarnellgee/assets/gaul_2024_level_1_code_500m"
+    )
+    # Cast to int32 to ensure integer GAUL codes, then rename
+    return admin_image.rename("admin_code")
+
+
+def g_water_mask_prep():
+    """
+    Water mask from JRC/USGS combined dataset.
+    Used to identify water bodies for downstream filtering and context.
+
+    Multiplied by pixel area to get water area in hectares.
+
+    Returns
+    -------
+    ee.Image
+        Binary water mask image renamed to In_waterbody (will be multiplied by pixel area)
+    """
+    from openforis_whisp.parameters.config_runtime import water_flag
+
+    water_mask_image = ee.Image("projects/ee-andyarnellgee/assets/water_mask_jrc_usgs")
+    return water_mask_image.selfMask().rename(water_flag)
 
 
 ###Combining datasets
 
 
-def combine_datasets(national_codes=None):
-    """Combines datasets into a single multiband image, with fallback if assets are missing."""
-    img_combined = ee.Image(1).rename(geometry_area_column)
+def combine_datasets(
+    national_codes=None,
+    validate_bands=False,
+    include_context_bands=True,
+    auto_recovery=False,
+):
+    """
+    Combines datasets into a single multiband image, with fallback if assets are missing.
+
+    Parameters
+    ----------
+    national_codes : list, optional
+        List of ISO2 country codes to include national datasets
+    validate_bands : bool, optional
+        If True, validates band names with a slow .getInfo() call (default: False)
+        Only enable for debugging. Normal operation relies on exception handling.
+    include_context_bands : bool, optional
+        If True (default), includes context bands (admin_code, water_flag) in the output.
+        Set to False when using stats.py implementations that compile datasets differently.
+    auto_recovery : bool, optional
+        If True (default), automatically enables validate_bands when an error is detected
+        during initial assembly. This allows graceful recovery from missing/broken datasets.
 
-    # Combine images directly
-    for img in [func() for func in list_functions(national_codes=national_codes)]:
+    Returns
+    -------
+    ee.Image
+        Combined multiband image with all datasets (and optionally context bands)
+    """
+    # Step 1: Combine all main dataset images
+    all_images = [ee.Image(1).rename(geometry_area_column)]
+    for func in list_functions(national_codes=national_codes):
         try:
-            img_combined = img_combined.addBands(img)
+            all_images.append(func())
         except ee.EEException as e:
-            # logger.error(f"Error adding image: {e}")
-            print(f"Error adding image: {e}")
-
-    try:
-        # Attempt to print band names to check for errors
-        # print(img_combined.bandNames().getInfo())
-        img_combined.bandNames().getInfo()
+            print(f"Error loading image: {e}")
 
-    except ee.EEException as e:
-        # logger.error(f"Error printing band names: {e}")
-        # logger.info("Running code for filtering to only valid datasets due to error in input")
-        print("using valid datasets filter due to error in input")
-        # Validate images
-        images_to_test = [
-            func() for func in list_functions(national_codes=national_codes)
-        ]
-        valid_imgs = keep_valid_images(images_to_test)  # Validate images
+    img_combined = ee.Image.cat(all_images)
 
-        # Retry combining images after validation
-        img_combined = ee.Image(1).rename(geometry_area_column)
-        for img in valid_imgs:
-            img_combined = img_combined.addBands(img)
+    # Step 2: Determine if validation needed
+    should_validate = validate_bands
+    if auto_recovery and not validate_bands:
+        try:
+            # Fast error detection: batch check main + context bands in one call
+            bands_to_check = [img_combined.bandNames().get(0)]
+            if include_context_bands:
+                admin_image = g_gaul_admin_code()
+                water_mask = g_water_mask_prep()
+                bands_to_check.extend(
+                    [admin_image.bandNames().get(0), water_mask.bandNames().get(0)]
+                )
+            ee.List(bands_to_check).getInfo()  # trigger error if any band is invalid
+        except ee.EEException as e:
+            print(f"Error detected, enabling recovery mode: {str(e)[:80]}...")
+            should_validate = True
 
+    # Step 3: Validate and recover if needed
+    if should_validate:
+        try:
+            img_combined.bandNames().getInfo()  # check all bands
+        except ee.EEException as e:
+            print("Using valid datasets filter due to error in validation")
+            valid_imgs = keep_valid_images(
+                [func() for func in list_functions(national_codes=national_codes)]
+            )
+            all_images_retry = [ee.Image(1).rename(geometry_area_column)]
+            all_images_retry.extend(valid_imgs)
+            img_combined = ee.Image.cat(all_images_retry)
+
+    # Step 4: Multiply main datasets by pixel area
     img_combined = img_combined.multiply(ee.Image.pixelArea())
-    print("Whisp multiband image compiled")
 
+    # Step 5: Add context bands (admin_code only - water mask is now in prep functions)
+    if include_context_bands:
+        for band_func, band_name in [
+            (g_gaul_admin_code, "admin_code"),
+            (g_water_mask_prep, "In_waterbody"),
+        ]:
+            try:
+                band_img = band_func()
+                if should_validate:
+                    band_img.bandNames().getInfo()
+                img_combined = img_combined.addBands(band_img)
+            except ee.EEException as e:
+                print(f"Warning: Could not add {band_name} band: {e}")
+
+    print("Whisp multiband image compiled")
     return img_combined
 
 
@@ -1245,9 +1308,12 @@ def combine_datasets(national_codes=None):
 def list_functions(national_codes=None):
     """
     Returns a list of functions that end with "_prep" and either:
-    - Start with "g_" (global/regional products)
+    - Start with "g_" (global/regional products, excluding context bands)
     - Start with any provided national code prefix (nXX_)
 
+    Context band functions (g_gaul_admin_code, g_water_mask_prep) are handled
+    separately and excluded from this list to avoid duplication.
+
     Args:
         national_codes: List of ISO2 country codes (without the 'n' prefix)
     """
@@ -1258,15 +1324,19 @@ def list_functions(national_codes=None):
     if national_codes is None:
         national_codes = []
 
+    # Context band functions that are handled separately
+    context_functions = {"g_gaul_admin_code", "g_water_mask_prep"}
+
     # Create prefixes list with proper formatting ('n' + code + '_')
     allowed_prefixes = ["g_"] + [f"n{code.lower()}_" for code in national_codes]
 
-    # Filter functions in a single pass
+    # Filter functions in a single pass, excluding context band functions
     functions = [
         func
         for name, func in inspect.getmembers(current_module, inspect.isfunction)
         if name.endswith("_prep")
         and any(name.startswith(prefix) for prefix in allowed_prefixes)
+        and name not in context_functions
     ]
 
     return functions
@@ -1350,3 +1420,6 @@ def combine_custom_bands(custom_images, custom_bands_info):
     custom_ee_image = custom_ee_image.multiply(ee.Image.pixelArea())
 
     return custom_ee_image  # Only return the image
+
+
+# %%