From 13a3f67489bc5b4d17ab3733d3c36c2a8dba82dd Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Tue, 22 Oct 2024 11:30:12 +0200
Subject: [PATCH 01/18] Auto-adjust cooling technology constraints

Move from hard-coded constraints to automatically set, to avoid
infeasibilities.
---
 .../model/water/data/water_for_ppl.py         | 145 +++++++++++++-----
 1 file changed, 110 insertions(+), 35 deletions(-)

diff --git a/message_ix_models/model/water/data/water_for_ppl.py b/message_ix_models/model/water/data/water_for_ppl.py
index 0664d40166..a5d412fc67 100644
--- a/message_ix_models/model/water/data/water_for_ppl.py
+++ b/message_ix_models/model/water/data/water_for_ppl.py
@@ -26,16 +26,22 @@ def missing_tech(x: pd.Series) -> pd.Series:
     data_dic = {
         "geo_hpl": 1 / 0.850,
         "geo_ppl": 1 / 0.385,
+        "gas_hpl": 1 / 0.3,
+        "foil_hpl": 1 / 0.25,
         "nuc_hc": 1 / 0.326,
         "nuc_lc": 1 / 0.326,
         "solar_th_ppl": 1 / 0.385,
     }
 
-    if data_dic.get(x["technology"]):
+    if pd.notna(x["technology"]) and x["technology"] in data_dic:
+        value = data_dic.get(x["technology"])
+        if x["value"] < 1:
+            value = max(x["value"], value)
+        # for backwards compatibility
         if x["level"] == "cooling":
-            return pd.Series((data_dic.get(x["technology"]), "dummy_supply"))
+            return pd.Series({"value": value, "level": "dummy_supply"})
         else:
-            return pd.Series((data_dic.get(x["technology"]), x["level"]))
+            return pd.Series({"value": value, "level": x["level"]})
     else:
         return pd.Series({"value": x["value"], "level": x["level"]})
 
@@ -171,6 +177,87 @@ def hist_cap(x: pd.Series, context: "Context", hold_cost: pd.DataFrame) -> list:
     ]
 
 
+def relax_growth_constraint(ref_hist, scen, cooling_df, g_lo, constraint_type):
+    """
+    Function to check if the parent technologies are shut down and require relaxing the growth constraint.
+
+    Parameters:
+    ref_hist (pd.DataFrame): Historical data.
+    scen (Scenario): Scenario object to retrieve parameters.
+    cooling_df (pd.DataFrame): DataFrame containing cooling technologies and their parent technologies.
+    g_lo (pd.DataFrame): DataFrame containing growth constraints.
+    constraint_type (str): Type of constraint to check ("activity" or "capacity").
+
+    Returns:
+    pd.DataFrame: Updated g_lo DataFrame with relaxed growth constraints.
+    """
+    if constraint_type == "activity":
+        year_type = "year_act"
+        bound_param = "bound_activity_up"
+    elif constraint_type == "new_capacity":
+        year_type = "year_vtg"
+        bound_param = "bound_new_capacity_up"
+    else:
+        raise ValueError(
+            "Invalid constraint_type. Must be 'activity' or 'new_capacity'."
+        )
+
+    # Group by all variables of ref_hist apart from year_type and hist_value_col and only keep rows with max year_type
+    max_year_hist = (
+        ref_hist.loc[ref_hist.groupby(["node_loc", "technology"])[year_type].idxmax()]
+        .drop(columns="unit")
+        .rename(columns={year_type: "hist_year", "value": "hist_value"})
+    )
+
+    # Step 2: Check for bound_activity_up or bound_new_capacity_up conditions
+    bound_up_pare = scen.par(bound_param, {"technology": cooling_df["parent_tech"]})
+    # Get a set with unique year_type values and order them
+    years = np.sort(bound_up_pare[year_type].unique())
+
+    # In max_year_hist add the next year from years corresponding to the hist_year columns
+    max_year_hist["next_year"] = max_year_hist["hist_year"].apply(
+        lambda x: years[years > x][0]
+    )
+
+    # Merge the max_year_hist with bound_up_pare
+    bound_up = pd.merge(bound_up_pare, max_year_hist, how="left")
+    # Look at years just after the historical year
+    bound_up1 = bound_up[bound_up[year_type] == bound_up["next_year"]]
+    # Categories that might break the growth constraints
+    bound_up1 = bound_up1[bound_up1["value"] < 0.9 * bound_up1["hist_value"]]
+    # not look ad sudden contraints after sthe starting year
+    bound_up = bound_up.sort_values(by=["node_loc", "technology", year_type])
+    # Check if value for a year is greater than the value of the next year (after sorting)
+    bound_up["next_value"] = bound_up.groupby(["node_loc", "technology"])[
+        "value"
+    ].shift(-1)
+    bound_up2 = bound_up[bound_up["value"] > 0.9 * bound_up["next_value"]]
+    bound_up2 = bound_up2.drop(columns=["next_value"])
+    # combine bound 1 and 2
+    combined_bound = (
+        pd.concat([bound_up1, bound_up2]).drop_duplicates().reset_index(drop=True)
+    )
+    # Keep only node_loc, technology, and year_type
+    combined_bound = combined_bound[["node_loc", "technology", year_type]]
+    # Add columns with value "remove" to be able to use make_matched_dfs
+    combined_bound["rem"] = "remove"
+    combined_bound.rename(columns={"technology": "parent_tech"}, inplace=True)
+
+    # map_par tec to parent tec
+    map_parent = cooling_df[["technology_name", "parent_tech"]]
+    map_parent.rename(columns={"technology_name": "technology"}, inplace=True)
+    # expand bound_up to all cooling technologies in map_parent
+    combined_bound = pd.merge(combined_bound, map_parent, how="left")
+
+    # Merge to g_lo to be able to remove the technologies
+    g_lo = pd.merge(g_lo, combined_bound, how="left")
+    g_lo = g_lo[g_lo["rem"] != "remove"]
+    # Remove column rem and parent_tech
+    g_lo = g_lo.drop(columns=["rem", "parent_tech"])
+
+    return g_lo
+
+
 # water & electricity for cooling technologies
 @minimum_version("message_ix 3.7")
 def cool_tech(context: "Context") -> dict[str, pd.DataFrame]:
@@ -257,11 +344,9 @@ def cool_tech(context: "Context") -> dict[str, pd.DataFrame]:
     )
     # cooling fraction = H_cool = Hi - 1 - Hi*(h_fg)
     # where h_fg (flue gasses losses) = 0.1
-    ref_input["cooling_fraction"] = ref_input["value"] * 0.9 - 1
+    # ref_input["cooling_fraction"] = ref_input["value"] * 0.9 - 1 # probably obsolete
 
-    ref_input[["value", "level"]] = ref_input[["technology", "value", "level"]].apply(
-        missing_tech, axis=1
-    )[["value", "level"]]
+    ref_input[["value", "level"]] = ref_input.apply(missing_tech, axis=1)
 
     # Combines the input df of parent_tech with water withdrawal data
     input_cool = (
@@ -276,18 +361,22 @@ def cool_tech(context: "Context") -> dict[str, pd.DataFrame]:
     # Convert year values into integers to be compatibel for model
     input_cool.year_vtg = input_cool.year_vtg.astype(int)
     input_cool.year_act = input_cool.year_act.astype(int)
-    # Drops extra technologies from the data
+    # Drops extra technologies from the data. backwards compatibility
     input_cool = input_cool[
         (input_cool["level"] != "water_supply") & (input_cool["level"] != "cooling")
     ]
-
+    # heat plants need no cooling
     input_cool = input_cool[
         ~input_cool["technology_name"].str.contains("hpl", na=False)
     ]
-    input_cool = input_cool[
-        (input_cool["node_loc"] != f"{context.regions}_GLB")
-        & (input_cool["node_origin"] != f"{context.regions}_GLB")
-    ]
+    # Swap node_loc if node_loc equals "{context.regions}_GLB"
+    input_cool.loc[input_cool["node_loc"] == f"{context.regions}_GLB", "node_loc"] = (
+        input_cool["node_origin"]
+    )
+    # Swap node_origin if node_origin equals "{context.regions}_GLB"
+    input_cool.loc[
+        input_cool["node_origin"] == f"{context.regions}_GLB", "node_origin"
+    ] = input_cool["node_loc"]
 
     input_cool["cooling_fraction"] = input_cool.apply(cooling_fr, axis=1)
 
@@ -347,7 +436,7 @@ def cool_tech(context: "Context") -> dict[str, pd.DataFrame]:
     con1 = input_cool["technology_name"].str.endswith("ot_saline", na=False)
     con2 = input_cool["technology_name"].str.endswith("air", na=False)
     icmse_df = input_cool[(~con1) & (~con2)]
-
+    # electricity inputs
     inp = make_df(
         "input",
         node_loc=electr["node_loc"],
@@ -747,14 +836,11 @@ def cool_tech(context: "Context") -> dict[str, pd.DataFrame]:
         unit="%",
         time="year",
     ).pipe(broadcast, year_act=info.Y, node_loc=node_region)
-    # Alligining certain technologies with growth constriants
-    g_lo.loc[g_lo["technology"].str.contains("bio_ppl|loil_ppl"), "value"] = -0.5
-    g_lo.loc[g_lo["technology"].str.contains("coal_ppl_u|coal_ppl"), "value"] = -0.5
-    g_lo.loc[
-        (g_lo["technology"].str.contains("coal_ppl_u|coal_ppl"))
-        & (g_lo["node_loc"].str.contains("CPA|PAS")),
-        "value",
-    ] = -1
+
+    # relax growth constraints for activity jumps
+    g_lo = relax_growth_constraint(ref_hist_act, scen, cooling_df, g_lo, "activity")
+    # relax growth constraints for capacity jumps
+    g_lo = relax_growth_constraint(ref_hist_cap, scen, cooling_df, g_lo, "new_capacity")
     results["growth_activity_lo"] = g_lo
 
     # growth activity up on saline water
@@ -769,18 +855,6 @@ def cool_tech(context: "Context") -> dict[str, pd.DataFrame]:
     ).pipe(broadcast, year_act=info.Y, node_loc=node_region)
     results["growth_activity_up"] = g_up
 
-    # # adding initial activity
-    # in_lo = h_act.copy()
-    # in_lo.drop(columns='mode', inplace=True)
-    # in_lo = in_lo[in_lo['year_act'] == 2015]
-    # in_lo_1 = make_df('initial_activity_lo',
-    #                   node_loc=in_lo['node_loc'],
-    #                   technology=in_lo['technology'],
-    #                   time='year',
-    #                   value=in_lo['value'],
-    #                   unit='GWa').pipe(broadcast, year_act=[2015, 2020])
-    # results['initial_activity_lo'] = in_lo_1
-
     return results
 
 
@@ -839,7 +913,8 @@ def non_cooling_tec(context: "Context") -> dict[str, pd.DataFrame]:
 
     n_cool_df = scen.par("output", {"technology": non_cool_tech})
     n_cool_df = n_cool_df[
-        (n_cool_df["node_loc"] != "R11_GLB") & (n_cool_df["node_dest"] != "R11_GLB")
+        (n_cool_df["node_loc"] != f"{context.regions}_GLB")
+        & (n_cool_df["node_dest"] != f"{context.regions}_GLB")
     ]
     n_cool_df_merge = pd.merge(n_cool_df, non_cool_df, on="technology", how="right")
     n_cool_df_merge.dropna(inplace=True)

From cdc15a12c9ba8cafee41aecc31b4a63cbc57308a Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Tue, 22 Oct 2024 11:31:34 +0200
Subject: [PATCH 02/18] Remove all dependencies from R11

This was from another older branch.
---
 message_ix_models/model/water/data/water_supply.py | 6 ++++--
 message_ix_models/model/water/report.py            | 2 +-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/message_ix_models/model/water/data/water_supply.py b/message_ix_models/model/water/data/water_supply.py
index 471513c359..5173161514 100644
--- a/message_ix_models/model/water/data/water_supply.py
+++ b/message_ix_models/model/water/data/water_supply.py
@@ -54,7 +54,7 @@ def map_basin_region_wat(context: "Context") -> pd.DataFrame:
         df_sw["MSGREG"] = (
             context.map_ISO_c[context.regions]
             if context.type_reg == "country"
-            else f"{context.regions}_" + df_sw["BCU_name"].str[-3:]
+            else f"{context.regions}_" + df_sw["BCU_name"].str.split("|").str[-1]
         )
 
         df_sw = df_sw.set_index(["MSGREG", "BCU_name"])
@@ -97,7 +97,9 @@ def map_basin_region_wat(context: "Context") -> pd.DataFrame:
         if context.type_reg == "country":
             df_sw["MSGREG"] = context.map_ISO_c[context.regions]
         else:
-            df_sw["MSGREG"] = f"{context.regions}_" + df_sw["BCU_name"].str[-3:]
+            df_sw["MSGREG"] = (
+                f"{context.regions}_" + df_sw["BCU_name"].str.split("|").str[-1]
+            )
 
         df_sw = df_sw.set_index(["MSGREG", "BCU_name"])
         df_sw.drop(columns="Unnamed: 0", inplace=True)
diff --git a/message_ix_models/model/water/report.py b/message_ix_models/model/water/report.py
index 39ee421ee3..15e9b3910d 100644
--- a/message_ix_models/model/water/report.py
+++ b/message_ix_models/model/water/report.py
@@ -1113,7 +1113,7 @@ def report(sc: Scenario, reg: str, sdgs: bool = False) -> None:
     for ur in ["urban", "rural"]:
         # CHANGE TO URBAN AND RURAL POP
         pop_tot = sc.timeseries(variable=("Population|" + ur.capitalize()))
-        pop_tot = pop_tot[-(pop_tot.region == "GLB region (R11)")]
+        pop_tot = pop_tot[-(pop_tot.region == "GLB region (R11)")]  # ONLY R11!!!
         pop_reg = np.unique(pop_tot["region"])
         # need to change names
         reg_map = mp2.regions()

From 75648546b2be4917f21675938a48872995dd158a Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Tue, 22 Oct 2024 12:04:46 +0200
Subject: [PATCH 03/18] Add new test_relax_growth_constraint

---
 .../model/water/data/water_for_ppl.py         |   2 +
 .../model/water/data/test_water_for_ppl.py    | 106 +++++++++++++++++-
 2 files changed, 107 insertions(+), 1 deletion(-)

diff --git a/message_ix_models/model/water/data/water_for_ppl.py b/message_ix_models/model/water/data/water_for_ppl.py
index a5d412fc67..5a245801ea 100644
--- a/message_ix_models/model/water/data/water_for_ppl.py
+++ b/message_ix_models/model/water/data/water_for_ppl.py
@@ -248,6 +248,8 @@ def relax_growth_constraint(ref_hist, scen, cooling_df, g_lo, constraint_type):
     map_parent.rename(columns={"technology_name": "technology"}, inplace=True)
     # expand bound_up to all cooling technologies in map_parent
     combined_bound = pd.merge(combined_bound, map_parent, how="left")
+    # rename tear_type to year_act, because g_lo use it
+    combined_bound.rename(columns={year_type: "year_act"}, inplace=True)
 
     # Merge to g_lo to be able to remove the technologies
     g_lo = pd.merge(g_lo, combined_bound, how="left")
diff --git a/message_ix_models/tests/model/water/data/test_water_for_ppl.py b/message_ix_models/tests/model/water/data/test_water_for_ppl.py
index 869cff4064..f6af40668b 100644
--- a/message_ix_models/tests/model/water/data/test_water_for_ppl.py
+++ b/message_ix_models/tests/model/water/data/test_water_for_ppl.py
@@ -5,7 +5,11 @@
 from message_ix_models import ScenarioInfo
 
 # from message_ix_models.model.structure import get_codes
-from message_ix_models.model.water.data.water_for_ppl import cool_tech, non_cooling_tec
+from message_ix_models.model.water.data.water_for_ppl import (
+    cool_tech,
+    non_cooling_tec,
+    relax_growth_constraint,
+)
 
 
 @cool_tech.minimum_version
@@ -201,3 +205,103 @@ def test_non_cooling_tec(request, test_context):
             "year_act",
         ]
     )
+
+
+# Mock function for scen.par
+class MockScenario:
+    def par(self, param, filters):
+        if param == "bound_activity_up":
+            return pd.DataFrame(
+                {
+                    "node_loc": ["R12_AFR", "R12_AFR", "R12_AFR"],
+                    "technology": ["coal_ppl", "coal_ppl", "coal_ppl"],
+                    "year_act": [2030, 2040, 2050],
+                    "value": [30, 15, 0],
+                }
+            )
+        elif param == "bound_new_capacity_up":
+            return pd.DataFrame(
+                {
+                    "node_loc": ["R12_AFR", "R12_AFR", "R12_AFR"],
+                    "technology": ["coal_ppl", "coal_ppl", "coal_ppl"],
+                    "year_vtg": [2030, 2040, 2050],
+                    "value": [30, 15, 0],
+                }
+            )
+        return pd.DataFrame()
+
+
+@pytest.mark.parametrize("constraint_type", ["activity", "new_capacity"])
+def test_relax_growth_constraint(constraint_type):
+    # Sample data for g_lo
+    if constraint_type == "activity":
+        year_type = "year_act"
+    elif constraint_type == "new_capacity":
+        year_type = "year_vtg"
+    else:
+        raise ValueError(
+            "Invalid constraint_type. Must be 'activity' or 'new_capacity'."
+        )
+
+    g_lo = pd.DataFrame(
+        {
+            "node_loc": ["R12_AFR", "R12_AFR", "R12_AFR", "R12_AFR"],
+            "technology": [
+                "coal_ppl__ot_fresh",
+                "coal_ppl__ot_fresh",
+                "coal_ppl__ot_fresh",
+                "gas_ppl__ot_fresh",
+            ],
+            "year_act": [2030, 2040, 2050, 2030],
+            "time": ["year", "year", "year", "year"],
+            "value": [-0.05, -0.05, -0.05, -0.05],
+            "unit": ["%", "%", "%", "%"],
+        }
+    )
+
+    # Sample data for ref_hist
+    ref_hist = pd.DataFrame(
+        {
+            "node_loc": ["R12_AFR", "R12_AFR", "R12_AFR"],
+            "technology": ["coal_ppl", "coal_ppl", "coal_ppl"],
+            year_type: [2015, 2020, 2025],
+            "time": ["year", "year", "year"],
+            "value": [30, 50, 80],
+            "unit": ["GWa", "GWa", "GWa"],
+        }
+    )
+
+    # Sample data for cooling_df
+    cooling_df = pd.DataFrame(
+        {
+            "technology_name": [
+                "coal_ppl__ot_fresh",
+                "coal_ppl__ot_fresh",
+                "coal_ppl__ot_fresh",
+            ],
+            "parent_tech": ["coal_ppl", "coal_ppl", "coal_ppl"],
+        }
+    )
+
+    # Instantiate mock scenario
+    scen = MockScenario()
+
+    # Call the function with mock data
+    result = relax_growth_constraint(ref_hist, scen, cooling_df, g_lo, constraint_type)
+    # reset_index to make the comparison easier
+    result = result.reset_index(drop=True)
+
+    # Expected result (this is an example, you should adjust this based on expected behavior)
+    expected_result = pd.DataFrame(
+        {
+            "node_loc": ["R12_AFR", "R12_AFR"],
+            "technology": ["coal_ppl__ot_fresh", "gas_ppl__ot_fresh"],
+            "year_act": [2050, 2030],
+            "time": ["year", "year"],
+            "value": [-0.05, -0.05],
+            "unit": ["%", "%"],
+        }
+    )
+
+    # Assert that the result matches the expected DataFrame
+    pd.testing.assert_frame_equal(result, expected_result)

From c88ec39a29034d1bef369e0b24754e8d79ac29c2 Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Tue, 22 Oct 2024 14:03:29 +0200
Subject: [PATCH 04/18] Shorten comments for Ruff

---
 message_ix_models/model/water/data/water_for_ppl.py  | 12 +++++++-----
 .../tests/model/water/data/test_water_for_ppl.py     |  2 +-
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/message_ix_models/model/water/data/water_for_ppl.py b/message_ix_models/model/water/data/water_for_ppl.py
index 5a245801ea..36fa4b6858 100644
--- a/message_ix_models/model/water/data/water_for_ppl.py
+++ b/message_ix_models/model/water/data/water_for_ppl.py
@@ -179,12 +179,14 @@ def hist_cap(x: pd.Series, context: "Context", hold_cost: pd.DataFrame) -> list:
 
 def relax_growth_constraint(ref_hist, scen, cooling_df, g_lo, constraint_type):
     """
-    Function to check if the parent technologies are shut down and require relaxing the growth constraint.
+    Function to check if the parent technologies are shut down and require
+    relaxing the growth constraint.
 
     Parameters:
     ref_hist (pd.DataFrame): Historical data.
     scen (Scenario): Scenario object to retrieve parameters.
-    cooling_df (pd.DataFrame): DataFrame containing cooling technologies and their parent technologies.
+    cooling_df (pd.DataFrame): DataFrame containing cooling technologies and
+    their parent technologies.
     g_lo (pd.DataFrame): DataFrame containing growth constraints.
     constraint_type (str): Type of constraint to check ("activity" or "capacity").
 
@@ -202,7 +204,7 @@ def relax_growth_constraint(ref_hist, scen, cooling_df, g_lo, constraint_type):
             "Invalid constraint_type. Must be 'activity' or 'new_capacity'."
         )
 
-    # Group by all variables of ref_hist apart from year_type and hist_value_col and only keep rows with max year_type
+    # keep rows with max year_type
     max_year_hist = (
         ref_hist.loc[ref_hist.groupby(["node_loc", "technology"])[year_type].idxmax()]
         .drop(columns="unit")
@@ -214,7 +216,7 @@ def relax_growth_constraint(ref_hist, scen, cooling_df, g_lo, constraint_type):
     # Get a set with unique year_type values and order them
     years = np.sort(bound_up_pare[year_type].unique())
 
-    # In max_year_hist add the next year from years corresponding to the hist_year columns
+    # In max_year_hist add the next year from years matching the hist_year columns
     max_year_hist["next_year"] = max_year_hist["hist_year"].apply(
         lambda x: years[years > x][0]
     )
@@ -227,7 +229,7 @@ def relax_growth_constraint(ref_hist, scen, cooling_df, g_lo, constraint_type):
     bound_up1 = bound_up1[bound_up1["value"] < 0.9 * bound_up1["hist_value"]]
     # not look ad sudden contraints after sthe starting year
     bound_up = bound_up.sort_values(by=["node_loc", "technology", year_type])
-    # Check if value for a year is greater than the value of the next year (after sorting)
+    # Check if value for a year is greater than the value of the next year
     bound_up["next_value"] = bound_up.groupby(["node_loc", "technology"])[
         "value"
     ].shift(-1)
diff --git a/message_ix_models/tests/model/water/data/test_water_for_ppl.py b/message_ix_models/tests/model/water/data/test_water_for_ppl.py
index f6af40668b..d11afc9df5 100644
--- a/message_ix_models/tests/model/water/data/test_water_for_ppl.py
+++ b/message_ix_models/tests/model/water/data/test_water_for_ppl.py
@@ -291,7 +291,7 @@ def test_relax_growth_constraint(constraint_type):
     # reset_index to make the comparison easier
     result = result.reset_index(drop=True)
 
-    # Expected result (this is an example, you should adjust this based on expected behavior)
+    # Expected result
     expected_result = pd.DataFrame(
         {
             "node_loc": ["R12_AFR", "R12_AFR"],

From 979848b8106c02624db51e3cdcb96bfd997801e5 Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Wed, 23 Oct 2024 10:55:29 +0200
Subject: [PATCH 05/18] Improve docstrings and code with Literal

---
 .../model/water/data/water_for_ppl.py         | 61 +++++++++++--------
 1 file changed, 36 insertions(+), 25 deletions(-)

diff --git a/message_ix_models/model/water/data/water_for_ppl.py b/message_ix_models/model/water/data/water_for_ppl.py
index 36fa4b6858..1af7bb0742 100644
--- a/message_ix_models/model/water/data/water_for_ppl.py
+++ b/message_ix_models/model/water/data/water_for_ppl.py
@@ -1,6 +1,6 @@
 """Prepare data for water use for cooling & energy technologies."""
 
-from typing import Any
+from typing import Any, Literal, Union
 
 import numpy as np
 import pandas as pd
@@ -177,32 +177,43 @@ def hist_cap(x: pd.Series, context: "Context", hold_cost: pd.DataFrame) -> list:
     ]
 
 
-def relax_growth_constraint(ref_hist, scen, cooling_df, g_lo, constraint_type):
+def relax_growth_constraint(
+    ref_hist: pd.DataFrame,
+    scen,
+    cooling_df: pd.DataFrame,
+    g_lo: pd.DataFrame,
+    constraint_type: Literal[Union["activity", "new_capacity"]],
+) -> pd.DataFrame:
     """
-    Function to check if the parent technologies are shut down and require
-    relaxing the growth constraint.
-
-    Parameters:
-    ref_hist (pd.DataFrame): Historical data.
-    scen (Scenario): Scenario object to retrieve parameters.
-    cooling_df (pd.DataFrame): DataFrame containing cooling technologies and
-    their parent technologies.
-    g_lo (pd.DataFrame): DataFrame containing growth constraints.
-    constraint_type (str): Type of constraint to check ("activity" or "capacity").
-
-    Returns:
-    pd.DataFrame: Updated g_lo DataFrame with relaxed growth constraints.
+    Checks if the parent technologies are shut down and require relaxing
+    the growth constraint.
+
+    Parameters
+    ----------
+    ref_hist : pd.DataFrame
+        Historical data in the reference scenario.
+    scen : Scenario
+        Scenario object to retrieve necessary parameters.
+    cooling_df : pd.DataFrame
+        DataFrame containing information on cooling technologies and their
+        parent technologies.
+    g_lo : pd.DataFrame
+        DataFrame containing growth constraints for each technology.
+    constraint_type : {"activity", "new_capacity"}
+        Type of constraint to check, either "activity" for operational limits or
+        "new_capacity" for capacity expansion limits.
+
+    Returns
+    -------
+    pd.DataFrame
+        Updated `g_lo` DataFrame with relaxed growth constraints.
     """
-    if constraint_type == "activity":
-        year_type = "year_act"
-        bound_param = "bound_activity_up"
-    elif constraint_type == "new_capacity":
-        year_type = "year_vtg"
-        bound_param = "bound_new_capacity_up"
-    else:
-        raise ValueError(
-            "Invalid constraint_type. Must be 'activity' or 'new_capacity'."
-        )
+    year_type = "year_act" if constraint_type == "activity" else "year_vtg"
+    bound_param = (
+        "bound_activity_up"
+        if constraint_type == "activity"
+        else "bound_new_capacity_up"
+    )
 
     # keep rows with max year_type
     max_year_hist = (

From a1e38d6a246607420b84203841f7afa3067ae0f6 Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Wed, 23 Oct 2024 10:56:18 +0200
Subject: [PATCH 06/18] Improve Mock and mark.para in test function

---
 .../model/water/data/test_water_for_ppl.py    | 55 ++++++++-----------
 1 file changed, 23 insertions(+), 32 deletions(-)

diff --git a/message_ix_models/tests/model/water/data/test_water_for_ppl.py b/message_ix_models/tests/model/water/data/test_water_for_ppl.py
index d11afc9df5..b4aff50637 100644
--- a/message_ix_models/tests/model/water/data/test_water_for_ppl.py
+++ b/message_ix_models/tests/model/water/data/test_water_for_ppl.py
@@ -1,3 +1,5 @@
+from typing import Literal
+
 import pandas as pd
 import pytest
 from message_ix import Scenario
@@ -209,40 +211,29 @@ def test_non_cooling_tec(request, test_context):
 
 # Mock function for scen.par
 class MockScenario:
-    def par(self, param, filters):
-        if param == "bound_activity_up":
-            return pd.DataFrame(
-                {
-                    "node_loc": ["R12_AFR", "R12_AFR", "R12_AFR"],
-                    "technology": ["coal_ppl", "coal_ppl", "coal_ppl"],
-                    "year_act": [2030, 2040, 2050],
-                    "value": [30, 15, 0],
-                }
-            )
-        elif param == "bound_new_capacity_up":
-            return pd.DataFrame(
-                {
-                    "node_loc": ["R12_AFR", "R12_AFR", "R12_AFR"],
-                    "technology": ["coal_ppl", "coal_ppl", "coal_ppl"],
-                    "year_vtg": [2030, 2040, 2050],
-                    "value": [30, 15, 0],
-                }
-            )
-        return pd.DataFrame()
-
-
-@pytest.mark.parametrize("constraint_type", ["activity", "new_capacity"])
-def test_relax_growth_constraint(constraint_type):
-    # Sample data for g_lo
-    if constraint_type == "activity":
-        year_type = "year_act"
-    elif constraint_type == "new_capacity":
-        year_type = "year_vtg"
-    else:
-        raise ValueError(
-            "Invalid constraint_type. Must be 'activity' or 'new_capacity'."
+    def par(
+        self,
+        param: Literal["bound_activity_up", "bound_new_capacity_up"],
+        filters: dict,
+    ) -> pd.DataFrame:
+        year_type = "year_act" if param == "bound_activity_up" else "year_vtg"
+
+        return pd.DataFrame(
+            {
+                "node_loc": ["R12_AFR", "R12_AFR", "R12_AFR"],
+                "technology": ["coal_ppl", "coal_ppl", "coal_ppl"],
+                year_type: [2030, 2040, 2050],
+                "value": [30, 15, 0],
+            }
         )
 
+
+@pytest.mark.parametrize(
+    "constraint_type, year_type",
+    [("activity", "year_act"), ("new_capacity", "year_vtg")],
+)
+def test_relax_growth_constraint(constraint_type, year_type):
+    # Sample data for g_lo
     g_lo = pd.DataFrame(
         {
             "node_loc": ["R12_AFR", "R12_AFR", "R12_AFR", "R12_AFR"],

From 57de933213cd65846e43c807ac3577b563e1a10e Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Wed, 23 Oct 2024 10:57:03 +0200
Subject: [PATCH 07/18] Clarify report comment; recode water supply fix

---
 message_ix_models/model/water/data/water_supply.py | 11 +++++------
 message_ix_models/model/water/report.py            |  3 ++-
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/message_ix_models/model/water/data/water_supply.py b/message_ix_models/model/water/data/water_supply.py
index 5173161514..f135487905 100644
--- a/message_ix_models/model/water/data/water_supply.py
+++ b/message_ix_models/model/water/data/water_supply.py
@@ -94,12 +94,11 @@ def map_basin_region_wat(context: "Context") -> pd.DataFrame:
         # Reading data, the data is spatially and temporally aggregated from GHMs
         df_sw["BCU_name"] = df_x["BCU_name"]
 
-        if context.type_reg == "country":
-            df_sw["MSGREG"] = context.map_ISO_c[context.regions]
-        else:
-            df_sw["MSGREG"] = (
-                f"{context.regions}_" + df_sw["BCU_name"].str.split("|").str[-1]
-            )
+        df_sw["MSGREG"] = (
+            context.map_ISO_c[context.regions]
+            if context.type_reg == "country"
+            else f"{context.regions}_" + df_sw["BCU_name"].str.split("|").str[-1]
+        )
 
         df_sw = df_sw.set_index(["MSGREG", "BCU_name"])
         df_sw.drop(columns="Unnamed: 0", inplace=True)
diff --git a/message_ix_models/model/water/report.py b/message_ix_models/model/water/report.py
index 15e9b3910d..219cb7cbd4 100644
--- a/message_ix_models/model/water/report.py
+++ b/message_ix_models/model/water/report.py
@@ -1113,7 +1113,8 @@ def report(sc: Scenario, reg: str, sdgs: bool = False) -> None:
     for ur in ["urban", "rural"]:
         # CHANGE TO URBAN AND RURAL POP
         pop_tot = sc.timeseries(variable=("Population|" + ur.capitalize()))
-        pop_tot = pop_tot[-(pop_tot.region == "GLB region (R11)")]  # ONLY R11!!!
+        # ONLY R11!!! Need to fix when updating the reporting to work with any region
+        pop_tot = pop_tot[-(pop_tot.region == "GLB region (R11)")]
         pop_reg = np.unique(pop_tot["region"])
         # need to change names
         reg_map = mp2.regions()

From acab52cf3f9a83dd581e5a9e28cb7dc5e941764f Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Wed, 23 Oct 2024 11:58:02 +0200
Subject: [PATCH 08/18] Fix breaking test

---
 message_ix_models/model/water/data/water_for_ppl.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/message_ix_models/model/water/data/water_for_ppl.py b/message_ix_models/model/water/data/water_for_ppl.py
index 1af7bb0742..e47a4154db 100644
--- a/message_ix_models/model/water/data/water_for_ppl.py
+++ b/message_ix_models/model/water/data/water_for_ppl.py
@@ -229,12 +229,13 @@ def relax_growth_constraint(
 
     # In max_year_hist add the next year from years matching the hist_year columns
     max_year_hist["next_year"] = max_year_hist["hist_year"].apply(
-        lambda x: years[years > x][0]
+        lambda x: years[years > x][0] if any(years > x) else None
     )
 
     # Merge the max_year_hist with bound_up_pare
     bound_up = pd.merge(bound_up_pare, max_year_hist, how="left")
-    # Look at years just after the historical year
+    # subset of first year after the historical
+    # if next_year = None (single year test case) bound_up1 is simply empty
     bound_up1 = bound_up[bound_up[year_type] == bound_up["next_year"]]
     # Categories that might break the growth constraints
     bound_up1 = bound_up1[bound_up1["value"] < 0.9 * bound_up1["hist_value"]]

From 9abd3e374d45d1b69744c6c8178970325a43430e Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Wed, 23 Oct 2024 11:58:27 +0200
Subject: [PATCH 09/18] Add set as default for the cooling function

---
 message_ix_models/model/water/cli.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/message_ix_models/model/water/cli.py b/message_ix_models/model/water/cli.py
index 5b658e1771..2b345448b4 100644
--- a/message_ix_models/model/water/cli.py
+++ b/message_ix_models/model/water/cli.py
@@ -267,6 +267,9 @@ def cooling(context, regions, rcps, rels):
     # Build
     build(context, scen)
 
+    # Set scenario as default
+    scen.set_as_default()
+
     # Solve
     scen.solve(solve_options={"lpmethod": "4"}, case=caseName)
 

From 04811476b555d157f86109d160dee051071fed3b Mon Sep 17 00:00:00 2001
From: Fridolin Glatter <glatter@iiasa.ac.at>
Date: Fri, 25 Oct 2024 09:30:23 +0200
Subject: [PATCH 10/18] Add #242 to doc/whatsnew

---
 doc/whatsnew.rst | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/doc/whatsnew.rst b/doc/whatsnew.rst
index 6a2f4ee2f9..2c3b4d9e24 100644
--- a/doc/whatsnew.rst
+++ b/doc/whatsnew.rst
@@ -4,7 +4,8 @@ What's new
 Next release
 ============
 
-- Fix the nexus/cooling function and adding test for checking some input data (:pull:`236`).
+- Make setup of constraints for cooling technologies flexible (:pull:`242`). 
+- Fix the nexus/cooling function and add test for checking some input data (:pull:`236`).
 - Add :doc:`/project/circeular` project code and documentation (:pull:`232`).
 - Update water availability data and major code editing to allow a new test suite for the water module (:pull:`106`).
 - Expand :doc:`repro` with sections on :ref:`repro-doc` and :ref:`versioning`, including :ref:`a list of external model names and ‘versions’ <model-names>` like “MESSAGEix-GLOBIOM 2.0” (:issue:`224`, :pull:`226`).

From 0d0ef4bafdd32c303fd313244c62876995511ff7 Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Mon, 28 Oct 2024 14:40:25 +0100
Subject: [PATCH 11/18] Add cooling to new csp solar tehnologies

---
 .../cool_techs_country_share.csv              |  4 +-
 .../cool_techs_region_share.csv               |  4 +-
 .../cool_techs_region_share_R11.csv           |  4 +-
 .../cooltech_cost_and_shares_country.csv      |  4 +-
 .../cooltech_cost_and_shares_ssp_msg.csv      |  4 +-
 .../cooltech_cost_and_shares_ssp_msg_R11.csv  |  4 +-
 message_ix_models/data/water/set.yaml         | 63 +++++++++++-
 message_ix_models/data/water/technology.yaml  | 98 +++++++++++++++++++
 8 files changed, 170 insertions(+), 15 deletions(-)

diff --git a/message_ix_models/data/water/ppl_cooling_tech/cool_techs_country_share.csv b/message_ix_models/data/water/ppl_cooling_tech/cool_techs_country_share.csv
index bfeb8f2995..a13e09acc2 100644
--- a/message_ix_models/data/water/ppl_cooling_tech/cool_techs_country_share.csv
+++ b/message_ix_models/data/water/ppl_cooling_tech/cool_techs_country_share.csv
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:669c1dd152ff47b2b203a29decd8ded056a6ca76d7823af8665c29d274e68625
-size 31012
+oid sha256:7b20ea1a29b19485f529db6bd25eb98495b82db58a39826cf27da5bb7e2be15b
+size 26173
diff --git a/message_ix_models/data/water/ppl_cooling_tech/cool_techs_region_share.csv b/message_ix_models/data/water/ppl_cooling_tech/cool_techs_region_share.csv
index c26d0538ef..822b3f64f1 100644
--- a/message_ix_models/data/water/ppl_cooling_tech/cool_techs_region_share.csv
+++ b/message_ix_models/data/water/ppl_cooling_tech/cool_techs_region_share.csv
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f03cf7b25a4eb15af32252a5498f7534ab1298a540b9e4041905084b4054b9e3
-size 7094
+oid sha256:e234d4a8238dc36e38b7836e74b52c54522c6d11bb5fc0441d5bfea616fbc698
+size 5154
diff --git a/message_ix_models/data/water/ppl_cooling_tech/cool_techs_region_share_R11.csv b/message_ix_models/data/water/ppl_cooling_tech/cool_techs_region_share_R11.csv
index a9ff1fa298..91620e6f80 100644
--- a/message_ix_models/data/water/ppl_cooling_tech/cool_techs_region_share_R11.csv
+++ b/message_ix_models/data/water/ppl_cooling_tech/cool_techs_region_share_R11.csv
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f4b53fd2df77644122011041ba958ff0fd84646bdfb81a4bd918ee10e6a8dd15
-size 7138
+oid sha256:9137797fc3b6e1f67e3b5224014b8bd7c82fac39eca5d006232c766dff9af446
+size 5198
diff --git a/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_country.csv b/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_country.csv
index bf9beaaaf6..1053944064 100644
--- a/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_country.csv
+++ b/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_country.csv
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:71d69e10931793d6a410226ea2098530cc568f80f4dec5e8587c9f7d5aa8d69e
-size 50969
+oid sha256:72ee20f7102b8f71ab9fed966b042d2bc59d74d0c2d404dd5ca469d180d09f77
+size 52407
diff --git a/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_ssp_msg.csv b/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_ssp_msg.csv
index f8af3242c5..b6a866da71 100644
--- a/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_ssp_msg.csv
+++ b/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_ssp_msg.csv
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:024f4981a4232caf1c206f6c2206b118d4230164a4e4dd9dafa21aa6887f0c04
-size 11413
+oid sha256:500138012d5bf33da1ba44745053cada82ce594a091a286887f3c591b9513a4d
+size 12039
diff --git a/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_ssp_msg_R11.csv b/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_ssp_msg_R11.csv
index 2668b849f7..1d20334097 100644
--- a/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_ssp_msg_R11.csv
+++ b/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_ssp_msg_R11.csv
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ed4e738f2dc9667ea9faa1f5420c78b52007de701aa96f29f44cce18c658c404
-size 11457
+oid sha256:9a9420d1197e1d5856441b4cf211dd360ea8c7432490e8e568266e186c01f8fa
+size 12083
diff --git a/message_ix_models/data/water/set.yaml b/message_ix_models/data/water/set.yaml
index 913cd0f160..edd8e36fd6 100644
--- a/message_ix_models/data/water/set.yaml
+++ b/message_ix_models/data/water/set.yaml
@@ -50,6 +50,8 @@ nexus:
     - cooling__loil_ppl
     - cooling__nuc_hc
     - cooling__solar_th_ppl
+    - cooling__csp_sm1_ppl
+    - cooling__csp_sm3_ppl
 
     add:
     - electr
@@ -210,9 +212,16 @@ nexus:
     - nuc_hc__ot_saline
     - solar_th_ppl__ot_fresh
     - solar_th_ppl__cl_fresh
-    - solar_th_ppl__cl_fresh
     - solar_th_ppl__ot_saline
     - solar_th_ppl__air
+    - csp_sm1_ppl__ot_fresh
+    - csp_sm1_ppl__cl_fresh
+    - csp_sm1_ppl__ot_saline
+    - csp_sm1_ppl__air
+    - csp_sm3_ppl__ot_fresh
+    - csp_sm3_ppl__cl_fresh
+    - csp_sm3_ppl__ot_saline
+    - csp_sm3_ppl__air
     - extract__freshwater_supply
     - extract__freshwater_instream
     - extract__saline_supply
@@ -318,6 +327,14 @@ nexus:
     - solar_th_ppl__cl_fresh
     - solar_th_ppl__ot_saline
     - solar_th_ppl__air
+    - csp_sm1_ppl__ot_fresh
+    - csp_sm1_ppl__cl_fresh
+    - csp_sm1_ppl__ot_saline
+    - csp_sm1_ppl__air
+    - csp_sm3_ppl__ot_fresh
+    - csp_sm3_ppl__cl_fresh
+    - csp_sm3_ppl__ot_saline
+    - csp_sm3_ppl__air
 
   # cat_addon includes cooling technology addons for parent technologies
   type_addon:
@@ -345,6 +362,8 @@ nexus:
     - cooling__nuc_hc
     - cooling__nuc_lc
     - cooling__solar_th_ppl
+    - cooling__csp_sm1_ppl
+    - cooling__csp_sm3_ppl
 
   map_tec_addon:
     add:
@@ -370,7 +389,8 @@ nexus:
     - [loil_ppl, cooling__loil_ppl]
     - [nuc_hc, cooling__nuc_hc]
     - [nuc_lc, cooling__nuc_lc]
-    # - [solar_th_ppl, cooling__solar_th_ppl]
+    - [csp_sm1_ppl, cooling__csp_sm1_ppl]
+    - [csp_sm3_ppl, cooling__csp_sm3_ppl]
 
 
   cat_addon:
@@ -468,6 +488,14 @@ nexus:
     - [cooling__solar_th_ppl, solar_th_ppl__cl_fresh]
     - [cooling__solar_th_ppl, solar_th_ppl__ot_saline]
     - [cooling__solar_th_ppl, solar_th_ppl__air]
+    - [cooling__csp_sm1_ppl, csp_sm1_ppl__ot_fresh]
+    - [cooling__csp_sm1_ppl, csp_sm1_ppl__cl_fresh]
+    - [cooling__csp_sm1_ppl, csp_sm1_ppl__ot_saline]
+    - [cooling__csp_sm1_ppl, csp_sm1_ppl__air]
+    - [cooling__csp_sm3_ppl, csp_sm3_ppl__ot_fresh]
+    - [cooling__csp_sm3_ppl, csp_sm3_ppl__cl_fresh]
+    - [cooling__csp_sm3_ppl, csp_sm3_ppl__ot_saline]
+    - [cooling__csp_sm3_ppl, csp_sm3_ppl__air]
 
   type_tec:
     add:
@@ -611,6 +639,8 @@ cooling:
     - cooling__loil_ppl
     - cooling__nuc_hc
     - cooling__solar_th_ppl
+    - cooling__csp_sm1_ppl
+    - cooling__csp_sm3_ppl
 
     add:
     - electr
@@ -747,9 +777,16 @@ cooling:
     - nuc_hc__ot_saline
     - solar_th_ppl__ot_fresh
     - solar_th_ppl__cl_fresh
-    - solar_th_ppl__cl_fresh
     - solar_th_ppl__ot_saline
     - solar_th_ppl__air
+    - csp_sm1_ppl__ot_fresh
+    - csp_sm1_ppl__cl_fresh
+    - csp_sm1_ppl__ot_saline
+    - csp_sm1_ppl__air
+    - csp_sm3_ppl__ot_fresh
+    - csp_sm3_ppl__cl_fresh
+    - csp_sm3_ppl__ot_saline
+    - csp_sm3_ppl__air
     - extract__upstream_landuse
     - extract__saline_supply
     - extract__freshwater_supply
@@ -856,6 +893,14 @@ cooling:
     - solar_th_ppl__cl_fresh
     - solar_th_ppl__ot_saline
     - solar_th_ppl__air
+    - csp_sm1_ppl__ot_fresh
+    - csp_sm1_ppl__cl_fresh
+    - csp_sm1_ppl__ot_saline
+    - csp_sm1_ppl__air
+    - csp_sm3_ppl__ot_fresh
+    - csp_sm3_ppl__cl_fresh
+    - csp_sm3_ppl__ot_saline
+    - csp_sm3_ppl__air
 
   # cat_addon includes cooling technology addons for parent technologies
   type_addon:
@@ -883,6 +928,8 @@ cooling:
     - cooling__nuc_hc
     - cooling__nuc_lc
     - cooling__solar_th_ppl
+    - cooling__csp_sm1_ppl
+    - cooling__csp_sm3_ppl
 
   map_tec_addon:
     add:
@@ -909,6 +956,8 @@ cooling:
     - [nuc_hc, cooling__nuc_hc]
     - [nuc_lc, cooling__nuc_lc]
     # - [solar_th_ppl, cooling__solar_th_ppl]
+    - [csp_sm1_ppl, cooling__csp_sm1_ppl]
+    - [csp_sm3_ppl, cooling__csp_sm3_ppl]
 
 
   cat_addon:
@@ -1006,6 +1055,14 @@ cooling:
     - [cooling__solar_th_ppl, solar_th_ppl__cl_fresh]
     - [cooling__solar_th_ppl, solar_th_ppl__ot_saline]
     - [cooling__solar_th_ppl, solar_th_ppl__air]
+    - [cooling__csp_sm1_ppl, csp_sm1_ppl__ot_fresh]
+    - [cooling__csp_sm1_ppl, csp_sm1_ppl__cl_fresh]
+    - [cooling__csp_sm1_ppl, csp_sm1_ppl__ot_saline]
+    - [cooling__csp_sm1_ppl, csp_sm1_ppl__air]
+    - [cooling__csp_sm3_ppl, csp_sm3_ppl__ot_fresh]
+    - [cooling__csp_sm3_ppl, csp_sm3_ppl__cl_fresh]
+    - [cooling__csp_sm3_ppl, csp_sm3_ppl__ot_saline]
+    - [cooling__csp_sm3_ppl, csp_sm3_ppl__air]
 
 
   unit:
diff --git a/message_ix_models/data/water/technology.yaml b/message_ix_models/data/water/technology.yaml
index 03af31c718..6b153f9090 100644
--- a/message_ix_models/data/water/technology.yaml
+++ b/message_ix_models/data/water/technology.yaml
@@ -576,6 +576,54 @@ nexus:
       technology using parasitic electricity
     input: {commodity: electr}
 
+  csp_sm1_ppl__ot_fresh:
+    description: >-
+      Solar thermal power plant with storage cooling by once through cooling
+      technology using freshwater supply
+    input: {commodity: freshwater_supply}
+
+  csp_sm1_ppl__cl_fresh:
+    description: >-
+      Solar thermal power plant with storage cooling by closed loop cooling
+      technology using freshwater supply & parasitic electricity
+    input: {commodity: freshwater_supply,electr}
+
+  csp_sm1_ppl__ot_saline:
+    description: >-
+      Solar thermal power plant with storage cooling by once through cooling
+      technology using freshwater supply
+    input: {commodity: saline_supply_ppl}
+
+  csp_sm1_ppl__air:
+    description: >-
+      Solar thermal power plant with storage cooling by once through cooling
+      technology using parasitic electricity
+    input: {commodity: electr}
+
+  csp_sm3_ppl__ot_fresh:
+    description: >-
+      Solar thermal power plant with storage cooling by once through cooling
+      technology using freshwater supply
+    input: {commodity: freshwater_supply}
+
+  csp_sm3_ppl__cl_fresh:
+    description: >-
+      Solar thermal power plant with storage cooling by closed loop cooling
+      technology using freshwater supply & parasitic electricity
+    input: {commodity: freshwater_supply,electr}
+
+  csp_sm3_ppl__ot_saline:
+    description: >-
+      Solar thermal power plant with storage cooling by once through cooling
+      technology using freshwater supply
+    input: {commodity: saline_supply_ppl}
+
+  csp_sm3_ppl__air:
+    description: >-
+      Solar thermal power plant with storage cooling by once through cooling
+      technology using parasitic electricity
+    input: {commodity: electr}
+
   urban_t_d:
     description: >-
       Urban water transmission & distribution
@@ -1324,3 +1372,53 @@ cooling:
       technology using parasitic electricity
     input: {commodity: electr}
 
+  csp_sm1_ppl__ot_fresh:
+    description: >-
+      Solar thermal power plant with storage cooling by once through cooling
+      technology using freshwater supply
+    input: {commodity: freshwater_supply}
+
+  csp_sm1_ppl__cl_fresh:
+    description: >-
+      Solar thermal power plant with storage cooling by closed loop cooling
+      technology using freshwater supply & parasitic electricity
+    input: {commodity: freshwater_supply,electr}
+
+  csp_sm1_ppl__ot_saline:
+    description: >-
+      Solar thermal power plant with storage cooling by once through cooling
+      technology using freshwater supply
+    input: {commodity: saline_supply_ppl}
+
+  csp_sm1_ppl__air:
+    description: >-
+      Solar thermal power plant with storage cooling by once through cooling
+      technology using parasitic electricity
+    input: {commodity: electr}
+
+  csp_sm3_ppl__ot_fresh:
+    description: >-
+      Solar thermal power plant with storage cooling by once through cooling
+      technology using freshwater supply
+    input: {commodity: freshwater_supply}
+
+  csp_sm3_ppl__cl_fresh:
+    description: >-
+      Solar thermal power plant with storage cooling by closed loop cooling
+      technology using freshwater supply & parasitic electricity
+    input: {commodity: freshwater_supply,electr}
+
+  csp_sm3_ppl__ot_saline:
+    description: >-
+      Solar thermal power plant with storage cooling by once through cooling
+      technology using freshwater supply
+    input: {commodity: saline_supply_ppl}
+
+  csp_sm3_ppl__air:
+    description: >-
+      Solar thermal power plant with storage cooling by once through cooling
+      technology using parasitic electricity
+    input: {commodity: electr}
+
+  
+

From e9c2516a99d2df4eca926091cdf4f33c4bd585bf Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Mon, 28 Oct 2024 14:40:48 +0100
Subject: [PATCH 12/18] Remove obsolete line on cooling factor

---
 message_ix_models/model/water/data/water_for_ppl.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/message_ix_models/model/water/data/water_for_ppl.py b/message_ix_models/model/water/data/water_for_ppl.py
index e47a4154db..7c101fe89d 100644
--- a/message_ix_models/model/water/data/water_for_ppl.py
+++ b/message_ix_models/model/water/data/water_for_ppl.py
@@ -358,9 +358,6 @@ def cool_tech(context: "Context") -> dict[str, pd.DataFrame]:
     ref_hist_cap: pd.DataFrame = scen.par(
         "historical_new_capacity", {"technology": cooling_df["parent_tech"]}
     )
-    # cooling fraction = H_cool = Hi - 1 - Hi*(h_fg)
-    # where h_fg (flue gasses losses) = 0.1
-    # ref_input["cooling_fraction"] = ref_input["value"] * 0.9 - 1 # probably obsolete
 
     ref_input[["value", "level"]] = ref_input.apply(missing_tech, axis=1)
 

From fcf31ef7b5a5c70a50ad0e3c0a459753b0d5409b Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Mon, 28 Oct 2024 15:42:50 +0100
Subject: [PATCH 13/18] Add R12 files lost in previous update

---
 .../water/ppl_cooling_tech/cool_techs_region_share_R12.csv    | 4 ++--
 .../ppl_cooling_tech/cooltech_cost_and_shares_ssp_msg_R12.csv | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/message_ix_models/data/water/ppl_cooling_tech/cool_techs_region_share_R12.csv b/message_ix_models/data/water/ppl_cooling_tech/cool_techs_region_share_R12.csv
index 07410e91db..53d44cf424 100644
--- a/message_ix_models/data/water/ppl_cooling_tech/cool_techs_region_share_R12.csv
+++ b/message_ix_models/data/water/ppl_cooling_tech/cool_techs_region_share_R12.csv
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e54b2a38792540e3e83e0ae4812286137235ab0f9fb7c9a26f70f6d9bb910339
-size 7528
+oid sha256:64f8064f5971680fd928e35df10c13ba24b202f3328a8686e61536ba69586fcb
+size 5492
diff --git a/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_ssp_msg_R12.csv b/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_ssp_msg_R12.csv
index d6ebd412de..7c3422848e 100644
--- a/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_ssp_msg_R12.csv
+++ b/message_ix_models/data/water/ppl_cooling_tech/cooltech_cost_and_shares_ssp_msg_R12.csv
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2962d1cfbae328fdc027823e2de151020500047d01f3ab3ba855b95bd3e80ded
-size 11903
+oid sha256:9639f5b03a8deaff90f211cb146ffbf0cdf83ab3bf88927ab738bc0bcd8936a4
+size 12545

From c78e01666d1d9c1dffc5196e4a29a9ec496fb730 Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Wed, 30 Oct 2024 09:54:32 +0100
Subject: [PATCH 14/18] Remove lower growth constraint on cooling tech and add
 upper one

---
 .../model/water/data/water_for_ppl.py         | 52 +++++++++----------
 1 file changed, 24 insertions(+), 28 deletions(-)

diff --git a/message_ix_models/model/water/data/water_for_ppl.py b/message_ix_models/model/water/data/water_for_ppl.py
index 7c101fe89d..37295e70f2 100644
--- a/message_ix_models/model/water/data/water_for_ppl.py
+++ b/message_ix_models/model/water/data/water_for_ppl.py
@@ -182,7 +182,7 @@ def relax_growth_constraint(
     scen,
     cooling_df: pd.DataFrame,
     g_lo: pd.DataFrame,
-    constraint_type: Literal[Union["activity", "new_capacity"]],
+    constraint_type: Literal[Union["activity", "new_capacity", "total_capacity"]],
 ) -> pd.DataFrame:
     """
     Checks if the parent technologies are shut down and require relaxing
@@ -208,18 +208,22 @@ def relax_growth_constraint(
     pd.DataFrame
         Updated `g_lo` DataFrame with relaxed growth constraints.
     """
-    year_type = "year_act" if constraint_type == "activity" else "year_vtg"
+    year_type = "year_vtg" if constraint_type == "new_capacity" else "year_act"
+    year_hist = "year_act" if constraint_type == "activity" else "year_vtg"
+    print(year_type)
     bound_param = (
-        "bound_activity_up"
+        "bound_activity_lo"
         if constraint_type == "activity"
-        else "bound_new_capacity_up"
+        else "bound_new_capacity_lo"
+        if constraint_type == "new_capacity"
+        else "bound_total_capacity_lo"
     )
-
+    print(bound_param)
     # keep rows with max year_type
     max_year_hist = (
-        ref_hist.loc[ref_hist.groupby(["node_loc", "technology"])[year_type].idxmax()]
+        ref_hist.loc[ref_hist.groupby(["node_loc", "technology"])[year_hist].idxmax()]
         .drop(columns="unit")
-        .rename(columns={year_type: "hist_year", "value": "hist_value"})
+        .rename(columns={year_hist: "hist_year", "value": "hist_value"})
     )
 
     # Step 2: Check for bound_activity_up or bound_new_capacity_up conditions
@@ -238,14 +242,14 @@ def relax_growth_constraint(
     # if next_year = None (single year test case) bound_up1 is simply empty
     bound_up1 = bound_up[bound_up[year_type] == bound_up["next_year"]]
     # Categories that might break the growth constraints
-    bound_up1 = bound_up1[bound_up1["value"] < 0.9 * bound_up1["hist_value"]]
+    bound_up1 = bound_up1[bound_up1["value"] > 0.9 * bound_up1["hist_value"]]
     # not look ad sudden contraints after sthe starting year
     bound_up = bound_up.sort_values(by=["node_loc", "technology", year_type])
     # Check if value for a year is greater than the value of the next year
     bound_up["next_value"] = bound_up.groupby(["node_loc", "technology"])[
         "value"
     ].shift(-1)
-    bound_up2 = bound_up[bound_up["value"] > 0.9 * bound_up["next_value"]]
+    bound_up2 = bound_up[bound_up["value"] < 0.9 * bound_up["next_value"]]
     bound_up2 = bound_up2.drop(columns=["next_value"])
     # combine bound 1 and 2
     combined_bound = (
@@ -841,31 +845,23 @@ def cool_tech(context: "Context") -> dict[str, pd.DataFrame]:
     results["capacity_factor"] = df
     # results = {par_name: pd.concat(dfs) for par_name, dfs in results.items()}
 
-    # growth activity low to allow the cooling techs to be operational
-    g_lo = make_df(
-        "growth_activity_lo",
-        technology=inp["technology"].drop_duplicates(),
-        value=-0.05,
-        unit="%",
-        time="year",
-    ).pipe(broadcast, year_act=info.Y, node_loc=node_region)
-
-    # relax growth constraints for activity jumps
-    g_lo = relax_growth_constraint(ref_hist_act, scen, cooling_df, g_lo, "activity")
-    # relax growth constraints for capacity jumps
-    g_lo = relax_growth_constraint(ref_hist_cap, scen, cooling_df, g_lo, "new_capacity")
-    results["growth_activity_lo"] = g_lo
-
-    # growth activity up on saline water
-    inp_saline = inp[inp["technology"].str.endswith("ot_saline")]
-
+    # growth activity up to avoid sudden switch in the cooling techs
     g_up = make_df(
         "growth_activity_up",
-        technology=inp_saline["technology"].drop_duplicates(),
+        technology=inp["technology"].drop_duplicates(),
         value=0.05,
         unit="%",
         time="year",
     ).pipe(broadcast, year_act=info.Y, node_loc=node_region)
+
+    # relax growth constraints for activity jumps of parent technologies
+    g_up = relax_growth_constraint(ref_hist_act, scen, cooling_df, g_up, "activity")
+    # relax growth constraints for capacity jumps of parent technologies
+    g_up = relax_growth_constraint(ref_hist_cap, scen, cooling_df, g_up, "new_capacity")
+    g_up = relax_growth_constraint(
+        ref_hist_cap, scen, cooling_df, g_up, "total_capacity"
+    )
+
     results["growth_activity_up"] = g_up
 
     return results

From 8fb76d5b68e1df1ce52800e340421718f5591221 Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Wed, 30 Oct 2024 09:54:57 +0100
Subject: [PATCH 15/18] Update function test for upper growth constraints

---
 .../tests/model/water/data/test_water_for_ppl.py          | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/message_ix_models/tests/model/water/data/test_water_for_ppl.py b/message_ix_models/tests/model/water/data/test_water_for_ppl.py
index b4aff50637..dace7352e5 100644
--- a/message_ix_models/tests/model/water/data/test_water_for_ppl.py
+++ b/message_ix_models/tests/model/water/data/test_water_for_ppl.py
@@ -213,17 +213,17 @@ def test_non_cooling_tec(request, test_context):
 class MockScenario:
     def par(
         self,
-        param: Literal["bound_activity_up", "bound_new_capacity_up"],
+        param: Literal["bound_activity_lo", "bound_new_capacity_lo"],
         filters: dict,
     ) -> pd.DataFrame:
-        year_type = "year_act" if param == "bound_activity_up" else "year_vtg"
+        year_type = "year_act" if param == "bound_activity_lo" else "year_vtg"
 
         return pd.DataFrame(
             {
                 "node_loc": ["R12_AFR", "R12_AFR", "R12_AFR"],
                 "technology": ["coal_ppl", "coal_ppl", "coal_ppl"],
                 year_type: [2030, 2040, 2050],
-                "value": [30, 15, 0],
+                "value": [15, 30, 32],
             }
         )
 
@@ -257,7 +257,7 @@ def test_relax_growth_constraint(constraint_type, year_type):
             "technology": ["coal_ppl", "coal_ppl", "coal_ppl"],
             year_type: [2015, 2020, 2025],
             "time": ["year", "year", "year"],
-            "value": [30, 50, 80],
+            "value": [1, 2, 3],
             "unit": ["GWa", "GWa", "GWa"],
         }
     )

From 46d928e0cfba11538def3e78eeacbdca275f8a27 Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Wed, 30 Oct 2024 09:55:14 +0100
Subject: [PATCH 16/18] Change wrong water supply level

---
 message_ix_models/model/water/data/water_supply.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/message_ix_models/model/water/data/water_supply.py b/message_ix_models/model/water/data/water_supply.py
index f135487905..065244072c 100644
--- a/message_ix_models/model/water/data/water_supply.py
+++ b/message_ix_models/model/water/data/water_supply.py
@@ -247,7 +247,7 @@ def add_water_supply(context: "Context") -> dict[str, pd.DataFrame]:
                     unit="km3",
                     year_vtg=year_wat,
                     year_act=year_wat,
-                    level="water_supply",
+                    level="saline_supply",
                     commodity="saline_ppl",
                     mode="M1",
                     time="year",

From abba7be27a659b16efca617da00cc0af835ab4cb Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Wed, 30 Oct 2024 10:56:31 +0100
Subject: [PATCH 17/18] Fix test parameter and numbers

---
 .../tests/model/water/data/test_water_for_ppl.py            | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/message_ix_models/tests/model/water/data/test_water_for_ppl.py b/message_ix_models/tests/model/water/data/test_water_for_ppl.py
index dace7352e5..d7e92f8372 100644
--- a/message_ix_models/tests/model/water/data/test_water_for_ppl.py
+++ b/message_ix_models/tests/model/water/data/test_water_for_ppl.py
@@ -223,7 +223,7 @@ def par(
                 "node_loc": ["R12_AFR", "R12_AFR", "R12_AFR"],
                 "technology": ["coal_ppl", "coal_ppl", "coal_ppl"],
                 year_type: [2030, 2040, 2050],
-                "value": [15, 30, 32],
+                "value": [15, 150, 2000],
             }
         )
 
@@ -234,7 +234,7 @@ def par(
 )
 def test_relax_growth_constraint(constraint_type, year_type):
     # Sample data for g_lo
-    g_lo = pd.DataFrame(
+    g_up = pd.DataFrame(
         {
             "node_loc": ["R12_AFR", "R12_AFR", "R12_AFR", "R12_AFR"],
             "technology": [
@@ -278,7 +278,7 @@ def test_relax_growth_constraint(constraint_type, year_type):
     scen = MockScenario()
 
     # Call the function with mock data
-    result = relax_growth_constraint(ref_hist, scen, cooling_df, g_lo, constraint_type)
+    result = relax_growth_constraint(ref_hist, scen, cooling_df, g_up, constraint_type)
     # reset_index to make the comparison easier
     result = result.reset_index(drop=True)
 

From eb7aa1811e8697ed1538f40a1a24dd7015a26c91 Mon Sep 17 00:00:00 2001
From: adrivinca <vinca@iiasa.ac.at>
Date: Mon, 4 Nov 2024 10:24:00 +0100
Subject: [PATCH 18/18] Add warning note in doc and expand whatsnew

---
 doc/water/index.rst | 4 ++++
 doc/whatsnew.rst    | 2 +-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/doc/water/index.rst b/doc/water/index.rst
index 2de4ed9b69..925ef0f17e 100644
--- a/doc/water/index.rst
+++ b/doc/water/index.rst
@@ -1,6 +1,10 @@
 MESSAGEix-Nexus (:mod:`.model.water`)
 *************************************
 
+.. note:: The current version of :mod:`.model.water` is configured to work with scenarios that have technology set elements including ``csp_sm1_ppl`` and ``csp_sm3_ppl`` instead of ``solar_th_ppl``.
+   Previous versions (in :mod:`message_ix_models` v2024.8.6 and earlier) were configured to work with scenarios including the technology ``solar_th_ppl``.
+   See further discussion at :pull:`242`.
+
 :mod:`message_ix_models.model.water` adds water usage and demand related representation to the MESSAGEix-GLOBIOM global model.
 The resulting model is referred to as **“MESSAGEix-Nexus”**.
 This work extends the water sector linkage described by Parkinson et al. (2019) :cite:`parkinson-2019`.
diff --git a/doc/whatsnew.rst b/doc/whatsnew.rst
index 2c3b4d9e24..ad8e3db8d9 100644
--- a/doc/whatsnew.rst
+++ b/doc/whatsnew.rst
@@ -4,7 +4,7 @@ What's new
 Next release
 ============
 
-- Make setup of constraints for cooling technologies flexible (:pull:`242`). 
+- Make setup of constraints for cooling technologies flexible and update solar csp tech. name (:pull:`242`). 
 - Fix the nexus/cooling function and add test for checking some input data (:pull:`236`).
 - Add :doc:`/project/circeular` project code and documentation (:pull:`232`).
 - Update water availability data and major code editing to allow a new test suite for the water module (:pull:`106`).