create_hypertuning_directories.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports & Paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import numpy as np\n",
    "import os\n",
    "import pathlib\n",
    "import re\n",
    "import shutil\n",
    "from typing import Any, Mapping, Sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Path to working directory.\n",
    "working_dir = pathlib.Path('.')\n",
    "\n",
    "# Paths to run directories.\n",
    "run_dir = working_dir / 'runs/'\n",
    "hypertuning_dir = run_dir / 'pub' / 'hyper_tuning'\n",
    "\n",
    "if os.path.isdir(hypertuning_dir):\n",
    "    shutil.rmtree(hypertuning_dir)\n",
    "os.mkdir(hypertuning_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Experiment Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is the hypersweep. This must be modified by hand.\n",
    "assimilation_windows = [1, 3, 5, 20]\n",
    "epochs = [5, 10, 100, 1000]\n",
    "histories = [1, 5, 20]\n",
    "learning_rates = [1e-4, 1e-3, 1e-2, 5e-2, 1e-1]\n",
    "learning_rate_drop_factors = [0.1, 0.5, 0.9]\n",
    "asssimilation_targets_lists = [\n",
    "    ['c_n'], \n",
    "    # ['h_n'], \n",
    "    # ['x_d'], \n",
    "    # ['x_s'], \n",
    "    # ['c_n', 'h_n'], \n",
    "    # ['c_n', 'x_d'], \n",
    "    # ['x_d', 'x_s'], \n",
    "    # ['c_n', 'h_n', 'x_d', 'x_s']\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Default DA hyperparamters. Mostly not necessary, but some are used.\n",
    "assimilation_config = {\n",
    "    'assimilation_lead_time': 0,\n",
    "    'assimilation_targets': ['c_n'],\n",
    "    'assimilation_window': 5,\n",
    "    'epochs': 1000,\n",
    "    'history': 20,\n",
    "    'learning_rate': 0.05,\n",
    "    'learning_rate_drop_factor': 0.1,\n",
    "    'learning_rate_epoch_drop': 1001,\n",
    "    'loss': 'MSE',\n",
    "    'model_dropout': 0.,\n",
    "    'optimizer': 'Adam',\n",
    "    'predict_last_n': 1,\n",
    "    'regularization': [],\n",
    "    'seq_length': 365,\n",
    "    'target_variables': ['QObs(mm/d)'],05\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Assimilation Test Directories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_assim_dir(\n",
    "    sim_dir: pathlib.Path, \n",
    "    assim_dir: pathlib.Path, \n",
    "    assimilation_config: Mapping[str, Any],\n",
    "    epoch: int,\n",
    "    history: int,\n",
    "    drop_factor: float,\n",
    "    learning_rate: float,\n",
    "    window: int,\n",
    "    targets: Sequence[str],\n",
    "):\n",
    "\n",
    "    # Copy if the directory does not already exist.\n",
    "    if os.path.isdir(assim_dir):\n",
    "        shutil.rmtree(assim_dir)\n",
    "    shutil.copytree(sim_dir, assim_dir)\n",
    "    print(f'Finished copying {assim_dir}')\n",
    "    \n",
    "    # Read the config file to modify it.\n",
    "    config_file = f'{assim_dir}/config.yml'\n",
    "    with open(config_file, \"r\") as file:  \n",
    "        yaml_file_data = file.read()  \n",
    "\n",
    "    yaml_file_data = yaml_file_data + f'assimilation_config:\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  assimilation_lead_time: 1\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  assimilation_targets:\\n'\n",
    "    for target in targets:\n",
    "        yaml_file_data = yaml_file_data + f'  - {target}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  assimilation_window: {window}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  epochs: {epoch}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  history: {history}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  learning_rate: {learning_rate}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  learning_rate_drop_factor: {drop_factor}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  learning_rate_epoch_drop: {assimilation_config[\"learning_rate_epoch_drop\"]}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  loss: {assimilation_config[\"loss\"]}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  model_dropout: {assimilation_config[\"model_dropout\"]}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  optimizer: {assimilation_config[\"optimizer\"]}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  predict_last_n: {assimilation_config[\"predict_last_n\"]}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  regularization: {assimilation_config[\"regularization\"]}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  seq_length: {assimilation_config[\"seq_length\"]}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  target_variables:\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  - QObs(mm/d)\\n'\n",
    "\n",
    "    # Save the modified config file.\n",
    "    with open(config_file, \"w\") as file:  \n",
    "        file.write(yaml_file_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are 304 runs in the hyper sweep.\n",
      "This will cost 38.0 GPU cyles and take approximately 12.666666666666666 hours.\n"
     ]
    }
   ],
   "source": [
    "# Copy trained simulation directory into assimilation directories.\n",
    "specific_sim_dir = './runs/pub/simulation/pub_simulation_kfold_0_seed_0_1304_022120'\n",
    "\n",
    "count = 0\n",
    "for epoch in epochs:\n",
    "    for history in histories:\n",
    "        for drop in learning_rate_drop_factors:\n",
    "            for window in assimilation_windows:\n",
    "                for rate in learning_rates:\n",
    "                    for t, targets in enumerate(asssimilation_targets_lists):\n",
    "                        if window * history > 300:\n",
    "                            continue\n",
    "                        count += 1\n",
    "                        specific_assim_dir = hypertuning_dir  / f'window_{window}_epoch_{epoch}_history_{history}_drop_{drop}_rate_{rate}_targets_{t}'\n",
    "                        create_assim_dir(\n",
    "                            sim_dir=specific_sim_dir, \n",
    "                            assim_dir=specific_assim_dir,\n",
    "                            epoch = epoch,\n",
    "                            history=history,\n",
    "                            drop_factor=drop,\n",
    "                            learning_rate=rate,\n",
    "                            targets=targets,\n",
    "                            window=window,\n",
    "                            assimilation_config=assimilation_config,\n",
    "                        )\n",
    "\n",
    "print(f'There are {count} runs in the hyper sweep.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Quick & very conservative estimate of runtime.\n",
    "# On my machine (Tesla V100), it takes about 20 minutes to run one experiment with 1000 epochs.\n",
    "# Fewer epochs will make experiments faster. \n",
    "n_gpus = 8\n",
    "hours_per_gpu = 1/3\n",
    "print(f'This will cost {count/n_gpus} GPU cyles and take approximately {count*hours_per_gpu/n_gpus} hours.')"
   ]
  }
 ],
 "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.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}