create_da_and_ar_test_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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Path to working directory.\n",
    "working_dir = pathlib.Path('/home/gsnearing_google_com/lstm-data-assimilation')\n",
    "\n",
    "# Paths to run directories.\n",
    "run_dir = working_dir / 'runs/'\n",
    "\n",
    "time_split_run_dir = run_dir / 'time_split'\n",
    "simulation_time_split_run_dir = time_split_run_dir / 'simulation'\n",
    "autoregression_time_split_run_dir = time_split_run_dir / 'autoregression'\n",
    "assimilation_time_split_run_dir = time_split_run_dir / 'assimilation'\n",
    "\n",
    "pub_run_dir = run_dir / 'pub'\n",
    "simulation_pub_run_dir = pub_run_dir / 'simulation'\n",
    "autoregression_pub_run_dir = pub_run_dir / 'autoregression'\n",
    "assimilation_pub_run_dir = pub_run_dir / 'assimilation'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Experiment Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "holdout fractions:  [0.   0.25 0.5  0.75 1.  ]\n"
     ]
    }
   ],
   "source": [
    "# Number of ensemble members. Each ensemble is initialized with an explicit seed.\n",
    "n_ensemble = 1\n",
    "seeds = list(range(n_ensemble))\n",
    "\n",
    "# Lead times.\n",
    "lead_times = [1, 2, 4, 8, 10]\n",
    "\n",
    "# Fractions of data to be withheld during inference.\n",
    "holdout_fractions = np.round(np.linspace(0, 1, 5) * 100) / 100\n",
    "print('holdout fractions: ', holdout_fractions)\n",
    "\n",
    "# Number of PUB k-fold splits.\n",
    "n_pub_kfold = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Default data assimilation hyperparameters.\n",
    "# These are chosen from the hypertuning results.\n",
    "assimilation_config = {\n",
    "    'assimilation_lead_time': 0,\n",
    "    'assimilation_targets': ['c_n'],\n",
    "    'assimilation_window': 5,\n",
    "    'epochs': 100,\n",
    "    'history': 20,\n",
    "    'learning_rate': 0.1,\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)'],\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Assimilation Test Directories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_assim_dir(\n",
    "    sim_dir: pathlib.Path, \n",
    "    assim_dir: pathlib.Path, \n",
    "    holdout: float,\n",
    "    lead: int,\n",
    "    assimilation_config: Mapping[str, Any],\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",
    "    \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",
    "    # Add a random holdout fraction of the shifted input data. \n",
    "    yaml_file_data = yaml_file_data + f'random_holdout_from_dynamic_features: \\n'\n",
    "    yaml_file_data = yaml_file_data + f'  QObs(mm/d): \\n'\n",
    "    yaml_file_data = yaml_file_data + f'    missing_fraction: {holdout} \\n'\n",
    "    yaml_file_data = yaml_file_data + f'    mean_missing_length: 5 \\n'\n",
    "\n",
    "    yaml_file_data = yaml_file_data + f'assimilation_config:\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  assimilation_lead_time: {lead}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  assimilation_targets:\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  - c_n\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  assimilation_window: {assimilation_config[\"assimilation_window\"]}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  epochs: {assimilation_config[\"epochs\"]}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  history: {assimilation_config[\"history\"]}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  learning_rate: {assimilation_config[\"learning_rate\"]}\\n'\n",
    "    yaml_file_data = yaml_file_data + f'  learning_rate_drop_factor: {assimilation_config[\"learning_rate_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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy trained simulation directory into assimilation directories.\n",
    "for seed in seeds:\n",
    "    specific_sim_dir = glob.glob(str(simulation_time_split_run_dir / f'simulation_seed_{seed}_*'))[0]\n",
    "    for holdout in holdout_fractions:\n",
    "        for lead in lead_times:\n",
    "            specific_assim_dir = assimilation_time_split_run_dir / f'assimilation_seed_{seed}_holdout_{holdout}_lead_{lead}'\n",
    "            create_assim_dir(\n",
    "                sim_dir=specific_sim_dir, \n",
    "                assim_dir=specific_assim_dir,\n",
    "                holdout=holdout,\n",
    "                lead=lead,\n",
    "                assimilation_config=assimilation_config,\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy trained simulation directory into assimilation directories.\n",
    "for seed in seeds:\n",
    "    for kfold in range(n_pub_kfold):\n",
    "        specific_sim_dir = glob.glob(str(simulation_pub_run_dir / f'pub_simulation_kfold_{kfold}_seed_{seed}_*'))[0]\n",
    "        for holdout in [0.5]:\n",
    "            for lead in [1]:\n",
    "                specific_assim_dir = assimilation_pub_run_dir / f'assimilation_kfold_{kfold}_seed_{seed}_holdout_{holdout}_lead_{lead}'\n",
    "                create_assim_dir(\n",
    "                    sim_dir=specific_sim_dir, \n",
    "                    assim_dir=specific_assim_dir,\n",
    "                    holdout=holdout,\n",
    "                    lead=lead,\n",
    "                    assimilation_config=assimilation_config,\n",
    "                )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Autoregression Test Directories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_autoregression_dir(\n",
    "    ar_dir: pathlib.Path, \n",
    "    test_dir: pathlib.Path, \n",
    "    train_holdout: float,\n",
    "    test_holdout: float,\n",
    "    lead: int,\n",
    "    seed: int,\n",
    "):\n",
    "    # Copy if the directory does not already exist.\n",
    "    try:\n",
    "        shutil.copytree(ar_dir, test_dir)\n",
    "    except:\n",
    "        pass\n",
    "                        \n",
    "    # Read the config file to modify it.\n",
    "    config_file = f'{test_dir}/config.yml'\n",
    "    with open(config_file, \"r\") as file:  \n",
    "        yaml_file_data = file.read()  \n",
    "\n",
    "    # Add random holdout for inference (autoregression).\n",
    "    yaml_file_data = yaml_file_data.replace(f'missing_fraction: {train_holdout}',\n",
    "                                            f'missing_fraction: {test_holdout}')\n",
    "\n",
    "    # Change all directory paths.\n",
    "    pattern = f'autoregression_lead_{lead}_holdout_{train_holdout}_seed_{seed}_\\d+_\\d+'\n",
    "    replace = f'autoregression_lead_{lead}_train_holdout_{train_holdout}_test_holdout_{test_holdout}_seed_{seed}'\n",
    "    yaml_file_data = re.sub(pattern, replace, yaml_file_data)\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy trained simulation directory into autoregression directories.\n",
    "for seed in seeds:\n",
    "    for lead in lead_times:\n",
    "        for train_holdout in holdout_fractions:\n",
    "            for test_holdout in holdout_fractions:\n",
    "                print(seed, lead, train_holdout, test_holdout)\n",
    "                specific_ar_dir = glob.glob(str(autoregression_time_split_run_dir / f'autoregression_lead_{lead}_holdout_{train_holdout}_seed_{seed}_*'))[0]\n",
    "                test_ar_dir = autoregression_time_split_run_dir / f'autoregression_lead_{lead}_train_holdout_{train_holdout}_test_holdout_{test_holdout}_seed_{seed}'\n",
    "                create_autoregression_dir(\n",
    "                    ar_dir=specific_ar_dir, \n",
    "                    test_dir=test_ar_dir, \n",
    "                    train_holdout=train_holdout,\n",
    "                    test_holdout=test_holdout,\n",
    "                    lead=lead,\n",
    "                    seed=seed,\n",
    "                )"
   ]
  }
 ],
 "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
}