{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demonstration of generating domain and state files for MetSim\n",
    "\n",
    "The first cell is a simple setup column that gets the libraries that we need loaded. Then we grab all of the data that we're going to use and compute what will be the spatial extent of the data. We also set some paths that we will write our files out to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pylab inline\n",
    "import os\n",
    "import sys\n",
    "import glob\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import pandas as pd\n",
    "import metsim\n",
    "\n",
    "# Path to your ascii files\n",
    "data_dir = 'your/path/here'\n",
    "\n",
    "# Where to save the files\n",
    "OUT_DOMAIN = 'your/path/here'\n",
    "OUT_STATE = 'your/path/here'\n",
    "\n",
    "# List of all of the ascii files\n",
    "data_files = os.listdir(data_dir)\n",
    "\n",
    "# Unique listing of latitudes and longitudes to include\n",
    "# This is dependent on your files having naming scheme FILENAME_LAT_LON\n",
    "# You may have to adjust the part inside of the `float` call as necessary\n",
    "lats = np.unique([float(f.split('_')[1]) for f in data_files])\n",
    "lons = np.unique([float(f.split('_')[2]) for f in data_files])\n",
    "\n",
    "# Grid cells that you actually have data for\n",
    "# Again, you may have to adjust this if your file naming scheme is different\n",
    "mask_locs = list(map(lambda l: [float(ll) for ll in l], [x.split('_')[1:] for x in data_files]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generation of the domain dataset\n",
    "\n",
    "The domain dataset has three variables that represent the spatial extent of the run.  The variables we will generate are:\n",
    "\n",
    "* Mask: A simple way to tell if a grid-cell should be processed.  Values greater than 0 will be run.\n",
    "* Elevation: The elevation of the grid cell in meters.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coords = {'lat': lats, 'lon': lons}\n",
    "shape = (len(lats), len(lons))\n",
    "dims = ('lat', 'lon')\n",
    "domain = xr.Dataset(coords=coords)\n",
    "domain['elev'] = xr.DataArray(data=np.full(shape, np.nan), coords=coords, dims=dims)\n",
    "domain['mask'] = xr.DataArray(data=np.full(shape, np.nan), coords=coords, dims=dims)\n",
    "\n",
    "# Add the data\n",
    "for loc in mask_locs:\n",
    "    loc_dict = {'lat': loc[0], 'lon': loc[1]}\n",
    "    domain['elev'].loc[loc_dict] = 1 # your elevation here\n",
    "    domain['mask'].loc[loc_dict] = 1\n",
    "\n",
    "print(domain)\n",
    "domain.to_netcdf(OUT_DOMAIN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generation of the state dataset\n",
    "\n",
    "The state dataset requires the three variables that represent the spatial extent and a 90 day history.\n",
    "\n",
    "There are three required variables.  Each of them should be given as a 90 day record prior to the start date. They are:\n",
    "\n",
    " * Daily minimum temperature\n",
    " * Daily maximum temperature\n",
    " * Daily precipitation\n",
    "\n",
    "In newer revisions of MetSim the SWE variable can be omitted from the state, but we will add it regardless.\n",
    "\n",
    "This just fills the state with random information.  You should **not** do this, but will have to choose how to get this information yourself.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate the dates required\n",
    "metsim_start = pd.datetime(1949, 1, 1)\n",
    "state_start = metsim_start - pd.Timedelta('90 days')\n",
    "state_stop = metsim_start - pd.Timedelta('1 days')\n",
    "dates = pd.date_range(state_start, state_stop)\n",
    "\n",
    "# Required variables that have time\n",
    "varnames = ['prec', 't_min', 't_max']\n",
    "coords = {'time': dates, 'lon': lons, 'lat': lats}\n",
    "shape = (len(dates), len(domain['lat']), len(domain['lon']))\n",
    "dims = ('time', 'lat', 'lon')\n",
    "\n",
    "# Initialize the state dataset\n",
    "state = xr.Dataset(coords=coords)\n",
    "for var in varnames: \n",
    "    state[var] = xr.DataArray(data=np.full(shape, np.nan), coords=coords, dims=dims, name=var)\n",
    "\n",
    "# Now add SWE\n",
    "shape = (len(domain['lat']), len(domain['lon']))\n",
    "coords = {'lon': lons, 'lat': lats}\n",
    "dims = ('lat', 'lon')\n",
    "state['swe'] = xr.DataArray(data=np.full(shape, np.nan), coords=coords, dims=dims, name='swe')\n",
    "\n",
    "# Fill in the data\n",
    "for loc in mask_locs:\n",
    "    loc_dict = {'lat': loc[0], 'lon': loc[1]}\n",
    "    state['t_min'].loc[loc_dict] = np.random.random(90)\n",
    "    state['t_max'].loc[loc_dict] = np.random.random(90)\n",
    "    state['prec'].loc[loc_dict] = np.random.random(90)\n",
    "    state['swe'].loc[loc_dict] = 0.0\n",
    "\n",
    "print(state)\n",
    "state.to_netcdf(OUT_STATE)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.6.8"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autocomplete": false,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}