Merge pull request #698 from cal-adapt/fix/warming-level-years

Fix/warming level years

Author: neilSchroeder

.github/workflows/ci-not-main.yml

@@ -28,12 +28,8 @@ jobs:
       matrix:
         python-version: ["3.12", "3.13"]
       fail-fast: false
-    defaults:
-      on:
-        push:
-          branches-ignore:
-            - main
-          python-version: ${{ matrix.python-version }}
+    steps:
+      - uses: actions/checkout@v3
       - name: Install uv
         run: |
           python -m pip install --upgrade pip

climakitae/new_core/param_validation/warming_param_validator.py

@@ -61,6 +61,10 @@ def validate_warming_level_param(
         logger.warning(msg)
         return False
 
+    # check that catalog is "cadcat"
+    if not _check_catalog(query):
+        return False
+
     # validate query
     if not _check_query(query):
         return False
@@ -71,6 +75,42 @@ def validate_warming_level_param(
     return True
 
 
+def _check_catalog(query: Any) -> bool:
+    """
+    Validates that the catalog is "cadcat" (warming levels only supported for cadcat).
+
+    Parameters
+    ----------
+    query : Any
+        The query dictionary that may contain a 'catalog' key.
+
+    Returns
+    -------
+    bool
+        True if the catalog is "cadcat", False otherwise.
+    """
+    logger.debug("_check_catalog called with query: %s", query)
+
+    if not isinstance(query, dict):
+        return False
+
+    catalog = query.get("catalog", UNSET)
+    if catalog is UNSET:
+        # If catalog is not specified, assume it will default to cadcat
+        return True
+
+    if catalog != "cadcat":
+        msg = (
+            f"Warming level processor is not supported for '{catalog}' catalog. "
+            "Warming levels are only available for 'cadcat' (Cal-Adapt climate data). "
+            "Please use time_slice processor instead to specify a time range."
+        )
+        logger.warning(msg)
+        return False
+
+    return True
+
+
 def _check_input_types(
     value: dict[str, Any],
 ) -> bool:

climakitae/new_core/processors/concatenate.py

@@ -3,6 +3,7 @@
 import logging
 from typing import Any, Dict, Iterable, List, Union
 
+import numpy as np
 import xarray as xr
 
 from climakitae.core.constants import _NEW_ATTRS_KEY, CATALOG_REN_ENERGY_GEN, UNSET
@@ -60,8 +61,6 @@ def __init__(self, value: str = "time"):
         self.name = "concat"
         self.catalog = None
         self.needs_catalog = True
-        # Log initialization
-        logger.debug("Concat processor initialized with dim_name=%s", self.dim_name)
 
     def execute(
         self,
@@ -91,17 +90,8 @@ def execute(
             A single dataset with concatenated data.
 
         """
-        logger.debug(
-            "Concat.execute called with dim_name=%s result_type=%s",
-            self.dim_name,
-            type(result).__name__,
-        )
-
         if isinstance(result, (xr.Dataset, xr.DataArray)):
             # If we receive a single dataset, just return it
-            logger.debug(
-                "Concat.execute received single dataset/dataarray, returning as-is"
-            )
             return result
 
         # Route to appropriate concat method based on catalog
@@ -138,7 +128,6 @@ def _execute_hdp_concat(
             A single dataset with concatenated station data.
 
         """
-        logger.debug("Using HDP concatenation logic")
         datasets_to_concat = []
         station_ids = []
 
@@ -173,6 +162,11 @@ def _execute_hdp_concat(
             join="outer",
         )
 
+        # Drop "simulation" coordinate if it exists (distinct from "sim" dimension)
+        if "simulation" in concatenated.coords:
+            concatenated = concatenated.drop_vars("simulation")
+            logger.debug("Dropped 'simulation' coordinate from concatenated dataset.")
+
         logger.info("Concatenated HDP datasets along station_id dimension.")
         self.update_context(context, station_ids)
         return concatenated
@@ -203,15 +197,12 @@ def _execute_gridded_concat(
             A single dataset with concatenated data.
 
         """
-        logger.debug("Using gridded data concatenation logic")
-
         # Special handling for time dimension concatenation
         if self.dim_name == "time" and isinstance(result, dict):
             # Handle time domain extension for dictionaries
             result = extend_time_domain(result)  # type: ignore
             # After extending time domain, switch to standard sim concatenation
             self.dim_name = "sim"
-            logger.info("Time-domain extension applied; switching concat dim to 'sim'")
 
         datasets_to_concat = []
         concat_attrs = (
@@ -390,6 +381,66 @@ def _execute_gridded_concat(
 
         logger.info("Concatenated datasets along '%s' dimension.", self.dim_name)
 
+        # Drop "simulation" coordinate if it exists (distinct from "sim" dimension)
+        if "simulation" in concatenated.coords:
+            concatenated = concatenated.drop_vars("simulation")
+            logger.debug("Dropped 'simulation' coordinate from concatenated dataset.")
+
+        # Fix centered_year coordinate if it was assigned per-simulation in warming_level processor
+        # After concatenating along 'sim', we need to reconstruct it as 2D (sim, warming_level)
+        if (
+            "_sim_centered_years" in context
+            and "sim" in concatenated.dims
+            and "warming_level" in concatenated.dims
+        ):
+            sim_centered_years = context.get("_sim_centered_years", {})
+
+            if sim_centered_years:
+                sim_names = list(concatenated.sim.values)
+                warming_levels = list(concatenated.warming_level.values)
+
+                # Build 2D centered_year array indexed by (sim, warming_level)
+                cy_2d = np.full((len(sim_names), len(warming_levels)), np.nan)
+
+                for i, sim_name in enumerate(sim_names):
+                    # Try to find matching key in sim_centered_years
+                    # sim_name might be like 'WRF_UCLA_EC-Earth3_ssp370_day_d03_r1i1p1f1'
+                    # while sim_centered_years keys are like 'WRF.UCLA.EC-Earth3.ssp370.day.d03.r1i1p1f1'
+                    matching_key = None
+                    sim_name_dots = sim_name.replace("_", ".")
+
+                    for key in sim_centered_years.keys():
+                        if key.replace(".", "_") == sim_name or sim_name_dots == key:
+                            matching_key = key
+                            break
+
+                    if matching_key and matching_key in sim_centered_years:
+                        cy_values = sim_centered_years[matching_key]
+                        # Fill in the values for this simulation
+                        for j, cy_val in enumerate(cy_values):
+                            if j < len(warming_levels):
+                                cy_2d[i, j] = cy_val
+                        logger.debug(
+                            "Assigned centered_years for sim %s (key=%s): %s",
+                            sim_name,
+                            matching_key,
+                            cy_values,
+                        )
+                    else:
+                        logger.debug(
+                            "Could not find matching key for sim %s in sim_centered_years",
+                            sim_name,
+                        )
+
+                # Assign as 2D coordinate
+                concatenated.coords["centered_year"] = (["sim", "warming_level"], cy_2d)
+                logger.debug(
+                    "Successfully created 2D centered_year coordinate with shape %s",
+                    cy_2d.shape,
+                )
+            else:
+                logger.debug("sim_centered_years is empty in context")
+
         self.update_context(context, attr_ids)
 
         # Set resolution attribute if available

climakitae/new_core/processors/warming_level.py

@@ -165,19 +165,33 @@ def execute(
                 # If not, assume no member_id is present
                 member_ids.append(None)
 
+        logger.debug(
+            "execute: Extracted member_ids=%s for keys=%s", member_ids, list(ret.keys())
+        )
+
         # get center years for each key for each warming level
         center_years = self.get_center_years(member_ids, ret.keys())
+        logger.debug(
+            "execute: Received center_years dict with %d keys",
+            len(center_years),
+        )
+
         retkeys = list(ret.keys())
 
         for key in retkeys:
             if key not in center_years:
                 del ret[key]
 
+        # Build a mapping of simulation keys to their center years for proper coordinate assignment
+        sim_centered_years = {}
+
         for key, years in center_years.items():
             if not years:
                 continue
 
             slices = []
+            valid_wls = []  # Track which warming levels were successfully processed
+            valid_center_years = []  # Track center years for valid warming levels
 
             for year, wl in zip(years, self.warming_levels):
                 if year is None or pd.isna(year):
@@ -227,13 +241,20 @@ def execute(
                 # Expand dimensions to include warming_level as a new dimension
                 da_slice = da_slice.expand_dims({"warming_level": [wl]})
 
-                # Add simulation and centered_year coordinates
-                da_slice = da_slice.assign_coords(
-                    simulation=key,
-                    centered_year=(["warming_level"], [center_year]),
+                # Add simulation coordinate (but NOT centered_year yet)
+                da_slice = da_slice.assign_coords(simulation=key)
+
+                logger.debug(
+                    "execute: key=%s, wl=%s -> assigned coord simulation=%s",
+                    key,
+                    wl,
+                    key,
                 )
 
                 slices.append(da_slice)
+                valid_wls.append(wl)
+                valid_center_years.append(center_year)
+
             # After processing all warming levels, check if we have any valid slices.
             if not slices:
                 msg = f"No valid slices found for {key}. Ensure the warming level times table is correctly configured."
@@ -247,6 +268,18 @@ def execute(
                 slices, dim="warming_level", join="outer", fill_value=np.nan
             )
 
+            # Store center years for this simulation key
+            # DO NOT assign as coordinate here - it will be broadcast incorrectly
+            # Store in sim_centered_years dict to be reconstructed in concatenate processor
+            sim_centered_years[key] = valid_center_years
+
+        # Store center years in context for reconstruction after concatenation
+        context["_sim_centered_years"] = sim_centered_years
+        logger.debug(
+            "execute: Stored %d simulation center year mappings in context",
+            len(sim_centered_years),
+        )
+
         self.update_context(context)
         return ret
 
@@ -351,9 +384,17 @@ def get_center_years(
         """
         center_years = {}
 
+        logger.debug(
+            "get_center_years: Computing center years for %d keys",
+            len(list(keys)),
+        )
+
         # cols are key.join("_"), values are warming levels, index is time
         for key, member_id in zip(keys, member_ids):
             if member_id is None:
+                logger.debug(
+                    "get_center_years: Skipping key=%s (member_id is None)", key
+                )
                 continue  # Skip if member_id is None
 
             # build the key for the warming level table
@@ -363,24 +404,54 @@ def get_center_years(
                 center_years[key] = []
 
             for wl in self.warming_levels:
+                logger.debug(
+                    "get_center_years: Processing warming_level=%s for key=%s",
+                    wl,
+                    key,
+                )
+
                 if str(wl) not in self.warming_level_times.columns:
+                    logger.debug(
+                        "get_center_years: wl=%s not in warming_level_times.columns, using warming_level_times_idx",
+                        wl,
+                    )
                     # warming level is not an integer value, so we need to try to find
                     # the first year that the simulation crosses the given warming level
                     wl_table_key = f"{key_list[2]}_{member_id}_{key_list[3]}"
+                    logger.debug(
+                        "get_center_years: Looking for wl_table_key=%s in warming_level_times_idx",
+                        wl_table_key,
+                    )
+
                     if wl_table_key not in self.warming_level_times_idx.columns:
                         logger.warning(
-                            f"\n\nWarming level table does not contain data for {wl_table_key}. "
-                            "\nEnsure the warming level times table is correctly configured."
+                            f"\n\nWarning: Warming level table does not contain data for {wl_table_key}. "
+                            "\nEnsure the warming level times table is correctly configured. "
+                            f"Available columns: {list(self.warming_level_times_idx.columns)[:5]}..."
                         )
                         center_years[key].append(np.nan)
                         continue  # exit warming levels loop
 
                     # get the FIRST index value for this key
                     mask = (self.warming_level_times_idx[wl_table_key] >= wl).dropna()
+                    logger.debug(
+                        "get_center_years: mask for %s >= %s has %d True values",
+                        wl_table_key,
+                        wl,
+                        mask.sum(),
+                    )
+
                     if mask.any():
-                        center_years[key].append(
-                            mask.idxmax()  # Get the first occurrence of the warming level
+                        center_year = (
+                            mask.idxmax()
+                        )  # Get the first occurrence of the warming level
+                        logger.debug(
+                            "get_center_years: Found center_year=%s for wl_table_key=%s, wl=%s",
+                            center_year,
+                            wl_table_key,
+                            wl,
                         )
+                        center_years[key].append(center_year)
                     else:
                         max_valid_wl = (
                             self.warming_level_times.loc[
@@ -395,11 +466,34 @@ def get_center_years(
                         )
                         center_years[key].append(np.nan)
                 else:
+                    logger.debug(
+                        "get_center_years: wl=%s found in warming_level_times.columns",
+                        wl,
+                    )
                     # this is a common warming level, so we can just get the year
                     tuple_key = (key_list[2], member_id, key_list[3])
+                    logger.debug(
+                        "get_center_years: Looking up tuple_key=%s for wl=%s",
+                        tuple_key,
+                        wl,
+                    )
+
                     center_time = pd.to_datetime(
                         self.warming_level_times.loc[tuple_key, str(wl)]
                     )
-                    center_years[key].append(center_time.year)
+                    center_year = center_time.year
+                    logger.debug(
+                        "get_center_years: Found center_year=%s from warming_level_times for tuple_key=%s, wl=%s",
+                        center_year,
+                        tuple_key,
+                        wl,
+                    )
+                    center_years[key].append(center_year)
+
+        logger.debug(
+            "get_center_years: Returning center_years with %d keys", len(center_years)
+        )
+        for k, v in center_years.items():
+            logger.debug("get_center_years: key=%s -> center_years=%s", k, v)
 
         return center_years