{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a006d73d-f70a-4434-86d7-a61eda5e13c9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "prms_channel_flow_graph jit compiling with numba\n"
     ]
    }
   ],
   "source": [
    "import pathlib as pl\n",
    "\n",
    "import numpy as np\n",
    "import pywatershed as pws"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fa48496f-2db3-4642-9a1b-65633e81a864",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain_dir = pl.Path(\"../atm_forcing_domain\")\n",
    "nb_output_dir = pl.Path(\"./output\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6735552b-8713-4e61-b875-ebd8438ad17b",
   "metadata": {},
   "outputs": [],
   "source": [
    "control = pws.Control.load_prms(\n",
    "    domain_dir / \"nhm.control\", warn_unused_options=False\n",
    ")\n",
    "control.options[\"input_dir\"] = domain_dir\n",
    "# not sure why the init time is off by a year compared to the netcdf files\n",
    "control.edit_end_time(np.datetime64(\"2021-09-25T00\"))\n",
    "control.edit_init_start_times(new_init_time=np.datetime64(\"2010-12-31T00\"))\n",
    "control.options[\"netcdf_output_var_names\"] += [\"gwres_flow_vol\", \"sroff_vol\", \"ssres_flow_vol\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f30d87a5-d6f1-4949-8e13-066cc41b5c1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "params = pws.parameters.PrmsParameters.load(domain_dir / \"myparam.param\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "72f1e112-5af5-4c6d-b26c-c3f3eeecfbc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "nhm_processes = [\n",
    "    pws.PRMSSolarGeometry,\n",
    "    pws.PRMSAtmosphere,\n",
    "    pws.PRMSCanopy,\n",
    "    pws.PRMSSnow,\n",
    "    pws.PRMSRunoff,\n",
    "    pws.PRMSSoilzone,\n",
    "    pws.PRMSGroundwater,\n",
    "    pws.PRMSChannel,\n",
    "]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c2a3676e-ce5b-4ce0-b745-0711f82399c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PRMSCanopy jit compiling with numba \n",
      "PRMSSnow jit compiling with numba \n",
      "PRMSRunoff jit compiling with numba \n",
      "PRMSSoilzone jit compiling with numba \n",
      "PRMSGroundwater jit compiling with numba \n",
      "PRMSChannel jit compiling with numba \n"
     ]
    }
   ],
   "source": [
    "nhm = pws.Model(\n",
    "    nhm_processes,\n",
    "    control=control,\n",
    "    parameters=params,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "144bc315-c240-48f5-b377-c91d443c7af7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model initializing NetCDF output\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9e36fdefc7104e0abe3d9f8e1eb61d2f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3921 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wrote file: output/tmaxf.nc\n",
      "Wrote file: output/tminf.nc\n",
      "Wrote file: output/prmx.nc\n",
      "Wrote file: output/hru_ppt.nc\n",
      "Wrote file: output/hru_rain.nc\n",
      "Wrote file: output/hru_snow.nc\n",
      "Wrote file: output/swrad.nc\n",
      "Wrote file: output/potet.nc\n",
      "Wrote file: output/transp_on.nc\n",
      "Wrote file: output/tmaxc.nc\n",
      "Wrote file: output/tavgc.nc\n",
      "Wrote file: output/tminc.nc\n",
      "Wrote file: output/pptmix.nc\n",
      "Wrote file: output/orad_hru.nc\n",
      "model.run(): finalizing\n"
     ]
    }
   ],
   "source": [
    "nhm.run(finalize=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffc0473d-0ca1-484a-811a-571319ca31dd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pws]",
   "language": "python",
   "name": "conda-env-pws-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}