catch notebooks up with changes

Author: jmccreight

examples/01_multi-process_models.ipynb

@@ -241,7 +241,7 @@
     "    time_step=np.timedelta64(24, \"h\"),\n",
     "    options={\n",
     "        \"input_dir\": domain_dir,\n",
-    "        \"budget_type\": None,\n",
+    "        \"budget_type\": \"error\",\n",
     "        \"netcdf_output_dir\": nb_output_dir / \"nhm_memory\",\n",
     "    },\n",
     ")\n",
@@ -349,8 +349,10 @@
    "source": [
     "run_dir = pl.Path(nb_output_dir / \"nhm_yaml\")\n",
     "run_dir.mkdir(exist_ok=True)\n",
-    "control_yaml_file = run_dir / \"control.yml\"\n",
-    "control.to_yaml(control_yaml_file)"
+    "control_yaml_file = run_dir / \"control.yaml\"\n",
+    "control_yaml = deepcopy(control)\n",
+    "control_yaml.options[\"netcdf_output_dir\"] = nb_output_dir / \"nhm_yaml\"\n",
+    "control_yaml.to_yaml(control_yaml_file)"
    ]
   },
   {
@@ -366,6 +368,16 @@
     "We add the option `netcdf_output_dir` to the control since we assume we wont be able to do so at run time. Note that this option and the `input_dir` option are `pathlib.Path` objects. These are not what we want to write to file. We want their string version. We could do `str()` on each one by hand, but it will be more handy to write a small, recursive function to do this on a supplied dictionary since this will be a recurring task with the model dictionary we will create after the control YAML file."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8044016c-7728-4149-9be5-f52b8f43975d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "control"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -471,7 +483,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model_dict_yaml_file = run_dir / \"model_dict.yml\"\n",
+    "model_dict_yaml_file = run_dir / \"model_dict.yaml\"\n",
     "with open(model_dict_yaml_file, \"w\") as file:\n",
     "    _ = yaml.dump(model_dict, file)"
    ]
@@ -491,7 +503,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "! cat 01_multi-process_models/nhm_yaml/control.yml"
+    "! cat 01_multi-process_models/nhm_yaml/control.yaml"
    ]
   },
   {
@@ -501,7 +513,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "! cat 01_multi-process_models/nhm_yaml/model_dict.yml"
+    "! cat 01_multi-process_models/nhm_yaml/model_dict.yaml"
    ]
   },
   {
@@ -519,8 +531,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model_yml = pws.Model.from_yaml(model_dict_yaml_file)\n",
-    "model_yml"
+    "model_yaml = pws.Model.from_yaml(model_dict_yaml_file)\n",
+    "model_yaml"
    ]
   },
   {
@@ -533,7 +545,7 @@
     "show_params = not (platform == \"darwin\" and processor() == \"arm\")\n",
     "try:\n",
     "    pws.analysis.ModelGraph(\n",
-    "        model_yml,\n",
+    "        model_yaml,\n",
     "        hide_variables=False,\n",
     "        process_colors=palette,\n",
     "        show_params=show_params,\n",
@@ -558,8 +570,8 @@
    "outputs": [],
    "source": [
     "%%time\n",
-    "model_yml.run()\n",
-    "model_yml.finalize()"
+    "model_yaml.run()\n",
+    "model_yaml.finalize()"
    ]
   },
   {
@@ -584,11 +596,23 @@
    "outputs": [],
    "source": [
     "mem_out_dir = nb_output_dir / \"nhm_memory\"\n",
-    "yml_out_dir = nb_output_dir / \"nhm_yaml\"\n",
+    "yaml_out_dir = nb_output_dir / \"nhm_yaml\"\n",
     "mem_files = sorted(mem_out_dir.glob(\"*.nc\"))\n",
-    "yml_files = sorted(yml_out_dir.glob(\"*.nc\"))\n",
+    "yaml_files = sorted(yaml_out_dir.glob(\"*.nc\"))\n",
     "# We get all the same output files\n",
-    "assert set([ff.name for ff in mem_files]) == set([ff.name for ff in yml_files])"
+    "assert set([ff.name for ff in mem_files]) == set(\n",
+    "    [ff.name for ff in yaml_files]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "310a8d35-5ddb-4ef1-b4a1-65a823d4aa7c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "set([ff.name for ff in yaml_files]) - set([ff.name for ff in mem_files])"
    ]
   },
   {
@@ -604,16 +628,29 @@
     "Now compare the values of all variables:"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "523633e0-bb6e-4fa7-8a7a-bd2de3d602b0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "nb_output_dir.resolve()"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "db0d920e-0227-45ba-92c6-8ffacab9ca3c",
    "metadata": {},
    "outputs": [],
    "source": [
-    "for mf, yf in zip(mem_files, yml_files):\n",
+    "for mf, yf in zip(mem_files, yaml_files):\n",
     "    var = mf.with_suffix(\"\").name\n",
-    "    # print(var)\n",
+    "\n",
+    "    if \"budget\" in var.lower():\n",
+    "        continue\n",
+    "\n",
     "    mda = xr.open_dataset(mf)[var]\n",
     "    yda = xr.open_dataset(yf)[var]\n",
     "    xr.testing.assert_equal(mda, yda)\n",
@@ -681,7 +718,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "run_dir = pl.Path(nb_output_dir / \"yml_less_output\").resolve()\n",
+    "run_dir = pl.Path(nb_output_dir / \"yaml_less_output\").resolve()\n",
     "run_dir.mkdir(exist_ok=True)"
    ]
   },
@@ -704,7 +741,7 @@
     "]\n",
     "print(control_cp)  # .to_dict(), sort_dicts=False)\n",
     "\n",
-    "control_yaml_file = run_dir / \"control.yml\"\n",
+    "control_yaml_file = run_dir / \"control.yaml\"\n",
     "control_cp.to_yaml(control_yaml_file)"
    ]
   },
@@ -725,7 +762,7 @@
    "source": [
     "model_dict_copy = deepcopy(model_dict)\n",
     "model_dict_copy[\"control\"] = str(control_yaml_file)\n",
-    "model_dict_yaml_file = run_dir / \"model_dict.yml\""
+    "model_dict_yaml_file = run_dir / \"model_dict.yaml\""
    ]
   },
   {
@@ -928,7 +965,7 @@
     "control_cp = deepcopy(control)\n",
     "control_cp.options[\"input_dir\"] = yaml_output_dir.resolve()\n",
     "control_cp.options[\"netcdf_output_dir\"] = run_dir.resolve()\n",
-    "control_yaml_file = run_dir / \"control.yml\"\n",
+    "control_yaml_file = run_dir / \"control.yaml\"\n",
     "control_cp.to_yaml(control_yaml_file)\n",
     "pprint(control.to_dict(), sort_dicts=False)"
    ]
@@ -949,7 +986,7 @@
    "outputs": [],
    "source": [
     "model_dict[\"control\"] = str(control_yaml_file)\n",
-    "model_dict_yaml_file = run_dir / \"model_dict.yml\"\n",
+    "model_dict_yaml_file = run_dir / \"model_dict.yaml\"\n",
     "keep_procs = [\"prms_soilzone\", \"prms_groundwater\", \"prms_channel\"]\n",
     "model_dict[\"model_order\"] = keep_procs\n",
     "for kk in list(model_dict.keys()):\n",
@@ -1205,14 +1242,6 @@
     "* Regan, R. S., Markstrom, S. L., Hay, L. E., Viger, R. J., Norton, P. A., Driscoll, J. M., & LaFontaine, J. H. (2018). Description of the national hydrologic model for use with the precipitation-runoff modeling system (prms) (No. 6-B9). US Geological Survey.\n",
     "* Regan, R.S., Markstrom, S.L., LaFontaine, J.H., 2022, PRMS version 5.2.1: Precipitation-Runoff Modeling System (PRMS): U.S. Geological Survey Software Release, 02/10/2022."
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "ff4dbac1-046f-4354-94e6-2610269c2de4",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {

examples/02_prms_legacy_models.ipynb

@@ -69,11 +69,7 @@
   {
    "cell_type": "markdown",
    "id": "39abd69f-13a3-4f23-9678-a57c0b1f848d",
-   "metadata": {
-    "jupyter": {
-     "source_hidden": true
-    }
-   },
+   "metadata": {},
    "source": [
     "The domain directory is where we have all the required inputs to run this model (among others) and `nb_output_dir` is where this notebook will write its output. "
    ]
@@ -493,14 +489,6 @@
     "%%time\n",
     "submodel.run(finalize=True)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "cfcc3aa8-cc8f-40b9-9a72-124567c2c8bf",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {

examples/03_compare_pws_prms.ipynb

@@ -96,7 +96,7 @@
     "calc_method: str = \"numba\"\n",
     "budget_type: str = None\n",
     "\n",
-    "run_prms: bool = True  ## always forced/overwrite\n",
+    "run_prms: bool = False  ## always forced/overwrite\n",
     "\n",
     "run_pws: bool = True  # run if the output does not exist on disk\n",
     "force_pws_run: bool = True  # if it exists on disk, re-run it and overwrite?"
@@ -194,7 +194,22 @@
     "        input_dir = domain_dir / \"output\"\n",
     "        run_dir = nb_output_dir / f\"{domain_name}_subset_nhm\"\n",
     "\n",
-    "    control = pws.Control.load(domain_dir / \"control.test\")\n",
+    "    control = pws.Control.load_prms(\n",
+    "        domain_dir / \"control.test\", warn_unused_options=False\n",
+    "    )\n",
+    "    output_var_names = control.options[\"netcdf_output_var_names\"] + [\n",
+    "        \"soltab_horad_potsw\",\n",
+    "        \"soltab_potsw\",\n",
+    "        \"soltab_sunhrs\",\n",
+    "    ]\n",
+    "    control.options = control.options | {\n",
+    "        \"input_dir\": input_dir,\n",
+    "        \"budget_type\": budget_type,\n",
+    "        \"calc_method\": calc_method,\n",
+    "        \"netcdf_output_dir\": run_dir,\n",
+    "        \"netcdf_output_var_names\": output_var_names,\n",
+    "    }\n",
+    "\n",
     "    params = pws.parameters.PrmsParameters.load(domain_dir / \"myparam.param\")\n",
     "\n",
     "    if run_dir.exists():\n",
@@ -205,13 +220,6 @@
     "\n",
     "    print(f\"PWS writing output to {run_dir}\")\n",
     "\n",
-    "    control.options = control.options | {\n",
-    "        \"input_dir\": input_dir,\n",
-    "        \"budget_type\": budget_type,\n",
-    "        \"calc_method\": calc_method,\n",
-    "        \"netcdf_output_dir\": run_dir,\n",
-    "    }\n",
-    "\n",
     "    nhm = pws.Model(\n",
     "        nhm_processes,\n",
     "        control=control,\n",
@@ -248,11 +256,16 @@
     "    var_meta = pws.meta.find_variables(var_name)[var_name]\n",
     "    ylabel = f\"{fill(var_meta['desc'], 40)}\\n({var_meta['units']})\"\n",
     "\n",
+    "    pws_file = run_dir / f\"{var_name}.nc\"\n",
     "    prms_file = domain_dir / f\"output/{var_name}.nc\"\n",
     "    if not prms_file.exists():\n",
     "        return None\n",
+    "    if not pws_file.exists():\n",
+    "        print(f\"PWS file {pws_file} DNE, skipping.\")\n",
+    "        return None\n",
+    "\n",
     "    prms_var = xr.open_dataarray(prms_file)\n",
-    "    pws_var = xr.open_dataarray(run_dir / f\"{var_name}.nc\")\n",
+    "    pws_var = xr.open_dataarray(pws_file)\n",
     "\n",
     "    if rmse_min is not None:\n",
     "        if \"time\" in prms_var.dims:\n",
@@ -313,7 +326,9 @@
     "        return None\n",
     "    prms = xr.open_dataarray(prms_file, decode_timedelta=False)\n",
     "    pws_file = run_dir / f\"{var_name}.nc\"\n",
-    "    assert pws_file.exists()\n",
+    "    if not pws_file.exists():\n",
+    "        print(f\"PWS file '{prms_file}' DNE, skipping.\")\n",
+    "        return None\n",
     "    nhm_after = xr.open_dataarray(pws_file, decode_timedelta=False)\n",
     "    if \"time\" in prms.dims:\n",
     "        time_dim = \"time\"\n",
@@ -549,7 +564,7 @@
    "outputs": [],
    "source": [
     "if pws.PRMSGroundwater in nhm_processes:\n",
-    "    compare_var_timeseries(\"gwres_flow_vol\")"
+    "    compare_var_timeseries(\"gwres_flow\")"
    ]
   },
   {

pywatershed/atmosphere/prms_atmosphere.py

@@ -138,12 +138,11 @@ def __init__(
         self._calculated = False
 
         self._netcdf_initialized = False
-        self.netcdf_output_dir = netcdf_output_dir
-        if self.netcdf_output_dir:
+        if self._netcdf_output_dir:
             self.initialize_netcdf(
-                output_dir=pl.Path(netcdf_output_dir),
-                separate_variables=netcdf_separate_files,
-                output_vars=netcdf_output_vars,
+                output_dir=pl.Path(self._netcdf_output_dir),
+                separate_variables=self._netcdf_separate_files,
+                output_vars=self._netcdf_output_vars,
             )
 
         return
@@ -835,7 +834,7 @@ def initialize_netcdf(
         output_vars: list = None,
         **kwargs,
     ):
-        if self._netcdf_initialized:
+        if self._netcdf_initialized and self.control.options["verbosity"] > 5:
             msg = (
                 f"{self.name} class previously initialized netcdf output "
                 f"in {self._netcdf_output_dir}"

pywatershed/atmosphere/prms_solar_geometry.py

@@ -97,12 +97,11 @@ def __init__(
         self._calculated = False
 
         self._netcdf_initialized = False
-        if self.netcdf_output_dir:
-            self._calculate_all_time()
+        if self._netcdf_output_dir:
             self.initialize_netcdf(
-                output_dir=pl.Path(netcdf_output_dir),
-                separate_variables=netcdf_separate_files,
-                output_vars=netcdf_output_vars,
+                output_dir=pl.Path(self._netcdf_output_dir),
+                separate_variables=self._netcdf_separate_files,
+                output_vars=self._netcdf_output_vars,
             )
 
         return
@@ -481,7 +480,7 @@ def initialize_netcdf(
         output_vars: list = None,
         **kwargs,
     ):
-        if self._netcdf_initialized:
+        if self._netcdf_initialized and self.control.options["verbosity"] > 5:
             msg = (
                 f"{self.name} class previously initialized netcdf output "
                 f"in {self._netcdf_output_dir}"