Enable overriding params file with kwargs in execute_pyciam (#17)

* enable override of params file in execute_pyciam

* ruff formatting

Author: bolliger32

README.md

@@ -56,7 +56,6 @@ pyCIAM contains only a handful of public functions that a user will want to empl
 * `execute_pyciam`: This is the end-to-end function that represents the most likely entrypoint for users. All other public functions are called by this one.
 * `create_surge_lookup`: Creates a lookup table that can be leveraged to build a 2D linear spline function for calculating damages from extreme sea levels. This can drastically reduce computational expense required for simulation on large ensembles of sea level rise trajectories.
 * `load_ciam_inputs`: An I/O function to load SLIIDERS-like input data, storm surge damage lookup table (if specified), model parameters, and process/format these data for inclusion in pyCIAM.
-* `load_diaz_inputs`: A similar function to load a preprocessed SLIIDERS-like input dataset that is generated from the same input data used in Diaz 2016. This is used to generate comparisons to the Diaz 2016 results within Depsky et al. 2023.
 * `calc_costs`: This is the main computation engine in pyCIAM. It computes costs for all cost types, regions, years, socioeconomic and SLR trajectories, and adaptation case. It does *not* compute the optimal adaptation case, which must be computed afterward, for reasons described below.
 * `select_optimal_case`: This function calculates the optimal adaptation choice for a given region and returns the associated costs and NPV.
 

notebooks/data-processing/create-diaz-pyCIAM-inputs.ipynb

@@ -304,8 +304,7 @@
     "    {\n",
     "        \"long_name\": \"Surge Damage Coefficients\",\n",
     "        \"description\": (\n",
-    "            \"Coefficients used in the original Diaz 2016 paper to estimate surge \"\n",
-    "            \"damage\"\n",
+    "            \"Coefficients used in the original Diaz 2016 paper to estimate surge damage\"\n",
     "        ),\n",
     "    }\n",
     ")\n",

notebooks/data-processing/slr/sweet.ipynb

@@ -13,7 +13,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -32,7 +32,7 @@
     "import xarray as xr\n",
     "from shared import open_dataarray, open_dataset, open_zarr, save\n",
     "\n",
-    "from pyCIAM import spherical_nearest_neighbor as snn"
+    "from pyCIAM.utils import spherical_nearest_neighbor as snn"
    ]
   },
   {

notebooks/models/fit-movefactor.ipynb

@@ -1279,7 +1279,7 @@
     "    min_R_noadapt=this_refA.drop(\"case\"),\n",
     "    return_year0_hts=False,\n",
     "    return_RH_heights=True,\n",
-    "    **get_model_kwargs(True)\n",
+    "    **get_model_kwargs(True),\n",
     ")\n",
     "\n",
     "hts = RH_heights.sel(case=SOLVCASES).isel(case=this_refA.case)\n",

notebooks/models/run-pyCIAM-diaz2016.ipynb

@@ -31,8 +31,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from pathlib import Path\n",
-    "\n",
     "import distributed as dd\n",
     "import pandas as pd\n",
     "from pyCIAM.run import execute_pyciam\n",

notebooks/post-processing/pyCIAM-results-figures.ipynb

@@ -4524,18 +4524,18 @@
     "seg_adm_pop_2050 = inputs.pop_2019.isel(elev=pyCIAM_elevs).sum(\n",
     "    \"elev\"\n",
     ") * inputs.pop_scale.sel(country=inputs.seg_country, year=2050)\n",
-    "print(f'SSP Pops in 2050 (bn): {(seg_adm_pop_2050.sum(\"seg_adm\") / 1e9).values}')\n",
+    "print(f\"SSP Pops in 2050 (bn): {(seg_adm_pop_2050.sum('seg_adm') / 1e9).values}\")\n",
     "\n",
     "seg_adm_K_2050 = inputs.K_2019.isel(elev=pyCIAM_elevs).sum(\"elev\") * inputs.K_scale.sel(\n",
     "    country=inputs.seg_country, year=2050\n",
     ")\n",
     "print(\n",
     "    \"(IIASA) SSP Ks in 2050 ($Tn): \"\n",
-    "    f'{(seg_adm_K_2050.sel(iam=\"IIASA\").sum(\"seg_adm\") / 1e12).values}'\n",
+    "    f\"{(seg_adm_K_2050.sel(iam='IIASA').sum('seg_adm') / 1e12).values}\"\n",
     ")\n",
     "print(\n",
     "    \"(OECD) SSP Ks in 2050 ($Tn): \"\n",
-    "    f'{(seg_adm_K_2050.sel(iam=\"OECD\").sum(\"seg_adm\") / 1e12).values}'\n",
+    "    f\"{(seg_adm_K_2050.sel(iam='OECD').sum('seg_adm') / 1e12).values}\"\n",
     ")"
    ]
   },
@@ -4867,7 +4867,7 @@
     "        linewidth=0,\n",
     "        alpha=0.8,\n",
     "        zorder=10,\n",
-    "        **kwargs\n",
+    "        **kwargs,\n",
     "    )\n",
     "\n",
     "    if filename is not None:\n",

pyCIAM/__init__.py

@@ -1,4 +1,11 @@
-from pyCIAM.io import load_ciam_inputs, load_diaz_inputs
+from pyCIAM.io import load_ciam_inputs
 from pyCIAM.run import calc_costs, execute_pyciam, select_optimal_case
 from pyCIAM.surge.lookup import create_surge_lookup
-from pyCIAM.utils import spherical_nearest_neighbor
+
+__all__ = [
+    "load_ciam_inputs",
+    "calc_costs",
+    "execute_pyciam",
+    "select_optimal_case",
+    "create_surge_lookup",
+]

pyCIAM/io.py

@@ -9,7 +9,6 @@
 """
 
 import tempfile
-from collections.abc import Iterable
 from io import BytesIO
 from pathlib import Path
 from zipfile import ZipFile
@@ -548,9 +547,9 @@ def save_to_zarr_region(ds_in, store, already_aligned=False, storage_options={})
         n_valid = len(valid_ixs)
         st = valid_ixs[0]
         end = valid_ixs[-1]
-        assert (
-            end - st == n_valid - 1
-        ), f"Indices are not continuous along dimension {r}"
+        assert end - st == n_valid - 1, (
+            f"Indices are not continuous along dimension {r}"
+        )
         regions[r] = slice(st, end + 1)
 
     # align coords
@@ -692,7 +691,7 @@ def load_ciam_inputs(
         seg_vals,
         # dropping the "refA_scenario_selectors" b/c this doesn't need to be added to
         # the input dataset object
-        constants=params[params.map(type) != dict].to_dict(),
+        constants=params[params.map(type) != dict].to_dict(),  # noqa: E721
         seg_var=seg_var,
         selectors=selectors,
         storage_options=storage_options,
@@ -791,7 +790,7 @@ def load_diaz_inputs(
     inputs = prep_sliiders(
         input_store,
         seg_vals,
-        constants=params[params.map(type) != dict].to_dict(),
+        constants=params[params.map(type) != dict].to_dict(),  # noqa: E721
         seg_var="seg",
         calc_popdens_with_wetland_area=False,
         storage_options=storage_options,
@@ -861,7 +860,7 @@ def download_and_extract_partial_zip(lpath, url, zip_glob, n_retries=5):
         if not fl.is_file():
             retries = 0
             while retries < n_retries:
-                print(f"...Downloading {fl.name} (attempt {retries+1}/{n_retries})")
+                print(f"...Downloading {fl.name} (attempt {retries + 1}/{n_retries})")
                 try:
                     data = z.cat_file(fr)
                     break

pyCIAM/run.py

@@ -15,7 +15,7 @@
 import pandas as pd
 import xarray as xr
 from cloudpathlib import AnyPath, CloudPath
-from distributed import Client, wait
+from distributed import Client
 from rhg_compute_tools.xarray import dataarray_from_delayed
 
 from pyCIAM.constants import CASE_DICT, CASES, COSTTYPES, PLIST, RLIST, SOLVCASES
@@ -951,6 +951,7 @@ def execute_pyciam(
     diaz_config=False,
     dask_client_func=Client,
     storage_options=None,
+    params_override={},
     **model_kwargs,
 ):
     """Execute the full pyCIAM model. The following inputs are assumed:
@@ -1082,6 +1083,8 @@ def execute_pyciam(
         reflected in `storage_options["token"]`. Other cloud storage providers will have
         different authentication methods and have not yet been tested with this
         function.
+    params_override : dict, default {}
+        Used to override params specified in `params_path`
     **model_kwargs
         Passed directly to :py:func:`pyCIAM.calc_costs`
     """
@@ -1115,6 +1118,7 @@ def execute_pyciam(
 
     # read parameters
     params = pd.read_json(params_path)["values"]
+    params.update(params_override)
 
     # determine whether to check for finished jobs
     if output_path is None:
@@ -1276,6 +1280,10 @@ def execute_pyciam(
                 compute=False,
                 mode="w",
                 storage_options=storage_options,
+                encoding={
+                    "costs": {"fill_value": "NaN"},
+                    "optimal_case": {"fill_value": 255},
+                },
             )
 
     ####################################################
@@ -1469,8 +1477,6 @@ def execute_pyciam(
         .costs.notnull()
         .all()
     )
-    client.cluster.close()
-    client.close()
     if remove_tmpfile:
         if isinstance(tmp_output_path, CloudPath):
             tmp_output_path.rmtree()