Merge pull request #9 from ClimateImpactLab/main_recipe_changes

Main recipe changes

Author: kitschwarz

dscim/menu/main_recipe.py

@@ -161,6 +161,10 @@ def __init__(
             "discount_factors",
             "uncollapsed_sccs",
             "scc",
+            "uncollapsed_discount_factors",
+            "uncollapsed_marginal_damages",
+            "global_consumption",
+            "global_consumption_no_pulse",
         ],
         **kwargs,
     ):
@@ -285,6 +289,7 @@ def damage_function():
         def scc():
             damage_function()
             self.global_consumption
+            self.global_consumption_no_pulse
             self.logger.info("Processing SCC calculation ...")
             if self.fit_type == "quantreg":
                 self.full_uncertainty_iqr
@@ -293,6 +298,8 @@ def scc():
                     self.stream_discount_factors
                     self.calculate_scc
                 self.uncollapsed_sccs
+                self.uncollapsed_marginal_damages
+                self.uncollapsed_discount_factors
 
         course_dict = {"damage_function": damage_function, "scc": scc}
 
@@ -1173,7 +1180,7 @@ def calculate_discount_factors(self, cons_pc):
             rhos = xr.DataArray(self.rho, coords=[cons_pc.year])
         else:
             # plug the unique rho in an array, and compute e^rho - 1
-            rhos = xr.ufuncs.expm1(xr.DataArray(self.rho, coords=[cons_pc.year]))
+            rhos = np.expm1(xr.DataArray(self.rho, coords=[cons_pc.year]))
 
         stream_rhos = np.divide(
             1, np.multiply.accumulate((rhos.values + 1), rhos.dims.index("year"))
@@ -1309,6 +1316,40 @@ def stream_discount_factors(self):
             fair_aggregation=self.fair_aggregation,
         )
 
+    @cachedproperty
+    @save("uncollapsed_discount_factors")
+    def uncollapsed_discount_factors(self):
+        pop = self.collapsed_pop.sum("region")
+        pop = pop.reindex(
+            year=range(pop.year.min().values, self.ext_end_year + 1),
+            method="ffill",
+        )
+        f = self.calculate_discount_factors(
+            self.global_consumption_no_pulse / pop
+        ).to_dataset(name="discount_factor")
+        for var in f.variables:
+            f[var].encoding.clear()
+
+        return f
+
+    @cachedproperty
+    @save("uncollapsed_marginal_damages")
+    def uncollapsed_marginal_damages(self):
+
+        md = (
+            (
+                (self.global_consumption_no_pulse - self.global_consumption_pulse)
+                * self.climate.conversion
+            )
+            .rename("marginal_damages")
+            .to_dataset()
+        )
+
+        for var in md.variables:
+            md[var].encoding.clear()
+
+        return md
+
     def ce(self, obj, dims):
         """Rechunk data appropriately and apply the certainty equivalence
         calculation. This is done in a loop to avoid memory crashes.

dscim/preprocessing/input_damages.py

@@ -14,7 +14,7 @@
 from pathlib import Path
 from itertools import product
 from functools import partial
-from p_tqdm import p_map
+from p_tqdm import p_map, p_umap
 from dscim.menu.simple_storage import EconVars
 from zarr.errors import GroupNotFoundError
 
@@ -220,13 +220,50 @@ def concatenate_labor_damages(
 
         # save out
         if format_file == "zarr":
-            concat_ds.to_zarr(f"{path_to_file}.zarr", consolidated=True, mode="w")
+            to_store = concat_ds.copy()
+            for var in to_store.variables:
+                to_store[var].encoding.clear()
+
+            to_store.to_zarr(f"{path_to_file}.zarr", mode="w", consolidated=True)
         elif format_file == "netcdf":
             concat_ds.to_netcdf(f"{path_to_file}.nc4")
 
     return concat_ds
 
 
+def calculate_labor_batch_damages(batch, ec, input_path, save_path):
+    print(f"Processing batch={batch} damages in {os.getpid()}")
+    concatenate_labor_damages(
+        input_path=input_path,
+        save_path=save_path,
+        ec_cls=ec,
+        variable="rebased",
+        val_type="wage-levels",
+        format_file="zarr",
+        query=f"exists==True&batch=='batch{batch}'",
+    )
+    print("Saved!")
+
+
+def calculate_labor_damages(
+    path_econ="/shares/gcp/integration/float32/dscim_input_data/econvars/zarrs/integration-econ-bc39.zarr",
+    input_path="/shares/gcp/outputs/labor/impacts-woodwork/mc_correct_rebasing_for_integration",
+    save_path="/shares/gcp/integration/float32/input_data_histclim/labor_data/replication/",
+):
+    ec = EconVars(path_econ)
+    # process in 3 rounds to limit memory usage
+    for i in range(0, 3):
+        partial_func = partial(
+            calculate_labor_batch_damages,
+            input_path=input_path,
+            save_path=save_path,
+            ec=ec,
+        )
+        print("Processing batches:")
+        print(list(range(i * 5, i * 5 + 5)))
+        p_umap(partial_func, list(range(i * 5, i * 5 + 5)))
+
+
 def compute_ag_damages(
     input_path,
     save_path,
@@ -561,13 +598,66 @@ def concatenate_energy_damages(
         logger.info(f"Concatenating and processing {i}")
 
         if format_file == "zarr":
-            concat_ds.to_zarr(f"{path_to_file}.zarr", consolidated=True, mode="w")
+            to_store = concat_ds.copy()
+            for var in to_store.variables:
+                to_store[var].encoding.clear()
+
+            to_store.to_zarr(f"{path_to_file}.zarr", mode="w", consolidated=True)
         elif format_file == "netcdf":
             concat_ds.to_netcdf(f"{path_to_file}.nc4")
 
     return concat_ds
 
 
+def calculate_energy_batch_damages(batch, ec, input_path, save_path):
+    print(f"Processing batch={batch} damages in {os.getpid()}")
+    concatenate_energy_damages(
+        input_path=input_path,
+        file_prefix="TINV_clim_integration_total_energy",
+        save_path=save_path,
+        ec_cls=ec,
+        variable="rebased",
+        format_file="zarr",
+        query=f"exists==True&batch=='batch{batch}'",
+    )
+    print("Saved!")
+
+
+def calculate_energy_damages(
+    re_calculate=True,
+    path_econ="/shares/gcp/integration/float32/dscim_input_data/econvars/zarrs/integration-econ-bc39.zarr",
+    input_path="/shares/gcp/outputs/energy_pixel_interaction/impacts-blueghost/integration_resampled",
+    save_path="/shares/gcp/integration/float32/input_data_histclim/energy_data/replication_2022aug/",
+):
+    ec = EconVars(path_econ)
+
+    if re_calculate:
+        read_energy_files_parallel(
+            input_path=input_path,
+            save_path=save_path,
+            ec_cls=ec,
+            seed="TINV_clim_integration_total_energy_delta",
+        )
+        read_energy_files_parallel(
+            input_path=input_path,
+            save_path=save_path,
+            ec_cls=ec,
+            seed="TINV_clim_integration_total_energy_histclim",
+        )
+
+    # process in 3 rounds to limit memory usage
+    for i in range(0, 3):
+        partial_func = partial(
+            calculate_energy_batch_damages,
+            input_path=input_path,
+            save_path=save_path,
+            ec=ec,
+        )
+        print("Processing batches:")
+        print(list(range(i * 5, i * 5 + 5)))
+        p_umap(partial_func, list(range(i * 5, i * 5 + 5)))
+
+
 def prep_mortality_damages(
     gcms,
     paths,

dscim/utils/menu_runs.py

@@ -35,9 +35,6 @@ def run_ssps(
     AR,
     masks=[None],
     fair_dims_list=[["simulation"]],
-    global_cons=False,
-    factors=False,
-    marginal_damages=False,
     order="damage_function",
 ):
 
@@ -109,67 +106,6 @@ def run_ssps(
         menu_item = MENU_OPTIONS[menu_option](**kwargs)
         menu_item.order_plate(order)
 
-        if global_cons:
-            menu_item.global_consumption_no_pulse.to_netcdf(
-                f"{save_path}/{menu_option}_{discount_type}_eta{menu_item.eta}_rho{menu_item.rho}_global_consumption_no_pulse.nc4"
-            )
-            menu_item.global_consumption.to_netcdf(
-                f"{save_path}/{menu_option}_{discount_type}_eta{menu_item.eta}_rho{menu_item.rho}_global_consumption.nc4"
-            )
-
-        if marginal_damages:
-            md = (
-                menu_item.global_consumption_no_pulse
-                - menu_item.global_consumption_pulse
-            ) * menu_item.climate.conversion
-            md = md.rename("marginal_damages").to_dataset()
-            for var in md.variables:
-                md[var].encoding.clear()
-            md.chunk(
-                {
-                    "discount_type": 1,
-                    "weitzman_parameter": 1,
-                    "ssp": 1,
-                    "model": 1,
-                    "gas": 1,
-                    "year": 10,
-                }
-            ).to_zarr(
-                f"{save_path}/{menu_option}_{discount_type}_eta{menu_item.eta}_rho{menu_item.rho}_uncollapsed_marginal_damages.zarr",
-                consolidated=True,
-                mode="w",
-            )
-
-        if factors:
-
-            # holding population constant
-            # from 2100 to 2300 with 2099 values
-            pop = menu_item.collapsed_pop.sum("region")
-            pop = pop.reindex(
-                year=range(pop.year.min().values, menu_item.ext_end_year + 1),
-                method="ffill",
-            )
-
-            df = menu_item.calculate_discount_factors(
-                menu_item.global_consumption_no_pulse / pop
-            ).to_dataset(name="discount_factor")
-            for var in df.variables:
-                df[var].encoding.clear()
-            df.chunk(
-                {
-                    "discount_type": 1,
-                    "weitzman_parameter": 1,
-                    "ssp": 1,
-                    "model": 1,
-                    # "gas":1,
-                    "year": 10,
-                }
-            ).to_zarr(
-                f"{save_path}/{menu_option}_{discount_type}_eta{menu_item.eta}_rho{menu_item.rho}_uncollapsed_discount_factors.zarr",
-                consolidated=True,
-                mode="w",
-            )
-
 
 def run_rff(
     sectors,
@@ -179,8 +115,6 @@ def run_rff(
     config,
     USA,
     global_cons=True,
-    factors=True,
-    marginal_damages=True,
     order="scc",
 ):
 
@@ -228,64 +162,3 @@ def run_rff(
 
         menu_item = MENU_OPTIONS[menu_option](**kwargs)
         menu_item.order_plate(order)
-
-        if global_cons:
-            menu_item.global_consumption_no_pulse.to_netcdf(
-                f"{save_path}/{menu_option}_{discount_type}_eta{menu_item.eta}_rho{menu_item.rho}_global_consumption_no_pulse.nc4"
-            )
-            menu_item.global_consumption.to_netcdf(
-                f"{save_path}/{menu_option}_{discount_type}_eta{menu_item.eta}_rho{menu_item.rho}_global_consumption.nc4"
-            )
-
-        if marginal_damages:
-            md = (
-                (
-                    (
-                        menu_item.global_consumption_no_pulse
-                        - menu_item.global_consumption_pulse
-                    )
-                    * menu_item.climate.conversion
-                )
-                .rename("marginal_damages")
-                .to_dataset()
-            )
-
-            for var in md.variables:
-                md[var].encoding.clear()
-
-            md.chunk(
-                {
-                    "discount_type": 1,
-                    "weitzman_parameter": 14,
-                    "runid": 10000,
-                    "gas": 1,
-                    "year": 10,
-                }
-            ).to_zarr(
-                f"{save_path}/{menu_option}_{discount_type}_eta{menu_item.eta}_rho{menu_item.rho}_uncollapsed_marginal_damages.zarr",
-                consolidated=True,
-                mode="w",
-            )
-        if factors:
-
-            f = menu_item.calculate_discount_factors(
-                menu_item.global_consumption_no_pulse / menu_item.pop
-            ).to_dataset(name="discount_factor")
-
-            for var in f.variables:
-                f[var].encoding.clear()
-
-            f.chunk(
-                {
-                    "discount_type": 1,
-                    "weitzman_parameter": 14,
-                    "runid": 10000,
-                    "gas": 1,
-                    "region": 1,
-                    "year": 10,
-                }
-            ).to_zarr(
-                f"{save_path}/{menu_option}_{discount_type}_eta{menu_item.eta}_rho{menu_item.rho}_uncollapsed_discount_factors.zarr",
-                consolidated=True,
-                mode="w",
-            )