MODEL01_ViT_48h_tune_with_CRPS.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0de92920-99ef-48bb-b6db-54df00ae5796",
   "metadata": {},
   "source": [
    "# ViT fine-tuning with data aug and CRPS valid loss\n",
    "\n",
    "This notebook continues the ViT training with CRPS-based validation loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "16a9dfa9-e564-444d-84d4-d0c3212bb61f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import time\n",
    "import h5py\n",
    "import numpy as np\n",
    "from glob import glob\n",
    "\n",
    "# ------------------------------------------------------- #\n",
    "# Turn-off warnings\n",
    "import logging\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' \n",
    "logging.getLogger(\"tensorflow\").setLevel(logging.ERROR)\n",
    "\n",
    "# ------------------------------------------------------- #\n",
    "# Turn-off tensoflow-specific warnings\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "\n",
    "tf.autograph.set_verbosity(0)\n",
    "tf.get_logger().setLevel('ERROR')\n",
    "\n",
    "# ------------------------------------------------------- #\n",
    "# Import customized modules and settings\n",
    "sys.path.insert(0, '/glade/u/home/ksha/GAN_proj/')\n",
    "sys.path.insert(0, '/glade/u/home/ksha/GAN_proj/libs/')\n",
    "\n",
    "from namelist import *\n",
    "import data_utils as du\n",
    "import model_utils as mu\n",
    "import verif_utils as vu\n",
    "\n",
    "mu.set_seeds(888)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c4c003be-d01f-4518-9b93-df3de5509a81",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3cc6258c-cda5-4236-aec0-1759e6f7331f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_precip(x):\n",
    "    x[x<0] = 0\n",
    "    return 10*(np.exp(x) - 1)\n",
    "\n",
    "def verif_func(model, valid_GEFS, valid_CCPA_full, land_mask_CCPA):\n",
    "    \n",
    "    L_valid = valid_GEFS.shape[0]\n",
    "    N_leads = EN = valid_GEFS.shape[1]\n",
    "    EN = valid_GEFS.shape[2]\n",
    "\n",
    "    shape_full = (L_valid, N_leads, EN, 224, 464)\n",
    "    GEFS_full = np.empty(shape_full)\n",
    "    pred_ = np.empty(valid_GEFS.shape[:-1]+(4,))\n",
    "\n",
    "    for iens in range(EN):\n",
    "        pred_[:, :, iens, ...] = model.predict(valid_GEFS[:, :, iens, ...], verbose=0)\n",
    "        \n",
    "    for ilead in range(N_leads):\n",
    "        for iens in range(EN):\n",
    "            GEFS_full[:, ilead, iens, ...] = decoder.predict(pred_[:, ilead, iens, ...], \n",
    "                                                             verbose=0)[..., 0]\n",
    "    GEFS_full = to_precip(GEFS_full)\n",
    "    CRPS = np.empty((L_valid, N_leads, 224, 464)); CRPS[...] = np.nan\n",
    "\n",
    "    for ilead in range(N_leads):\n",
    "        crps_ilead, _, _ = vu.CRPS_2d(valid_CCPA_full[:, ilead, ...], \n",
    "                                      GEFS_full[:, ilead, ...], land_mask=land_mask_CCPA)\n",
    "        CRPS[:, ilead, ...] = crps_ilead\n",
    "    return np.nanmean(CRPS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "51e90049-aa12-4e01-8e28-8d93598c8500",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters\n",
    "filter_nums = [64, 128] # number of convolution kernels per down-/upsampling layer \n",
    "latent_dim = 4 # number of latent feature channels\n",
    "activation = 'gelu' # activation function\n",
    "num_embeddings = 128 #128 # number of the VQ codes\n",
    "\n",
    "input_size = (224, 464, 1) # size of MRMS input\n",
    "latent_size = (14, 29, latent_dim) # size of compressed latent features\n",
    "\n",
    "drop_encode = False\n",
    "drop_decode = False\n",
    "\n",
    "model_name_encoder_load = model_dir+'models/VQ_VAE_encoder_stack1_tune0'\n",
    "model_name_decoder_load = model_dir+'models/VQ_VAE_decoder_stack1_tune0'\n",
    "\n",
    "encoder = mu.VQ_VAE_encoder(input_size, filter_nums, latent_dim, num_embeddings, activation, drop_encode)\n",
    "\n",
    "W_old = mu.dummy_loader(model_name_encoder_load)\n",
    "encoder.set_weights(W_old)\n",
    "\n",
    "decoder = mu.VQ_VAE_decoder(latent_size, filter_nums, activation, drop_decode)\n",
    "\n",
    "W_old = mu.dummy_loader(model_name_decoder_load)\n",
    "decoder.set_weights(W_old)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58c66595-cd34-4abe-a459-57c5eb1caff4",
   "metadata": {},
   "source": [
    "## Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9cb6a0bf-d943-4a52-b6c2-fba3b78a9605",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ------------------------------------------------------- #\n",
    "# Hyperparameters\n",
    "# input size for the 48h models\n",
    "latent_size = (14, 29, 4)\n",
    "input_size = (8, 14, 29, 4)\n",
    "output_size = (8, 14, 29, 4)\n",
    "\n",
    "N_ens = 31\n",
    "\n",
    "# 0-48 hr settings\n",
    "lead_name = '0_48'\n",
    "ilead_start = 0\n",
    "ilead_end = 8\n",
    "N_leads = ilead_end - ilead_start\n",
    "pad_timelag = 2\n",
    "\n",
    "# # 54-96 hr settings\n",
    "# lead_name = '54_96'\n",
    "# ilead_start = 8\n",
    "# ilead_end = 16\n",
    "# N_leads = ilead_end - ilead_start\n",
    "# pad_timelag = 2\n",
    "\n",
    "# # 102-144 hr settings\n",
    "# lead_name = '102_144'\n",
    "# ilead_start = 16\n",
    "# ilead_end = 24\n",
    "# N_leads = ilead_end - ilead_start\n",
    "# pad_timelag = 2\n",
    "\n",
    "# ============================= #\n",
    "# Tuned hyperparameters\n",
    "patch_size = (1, 1, 1) # (time, space, space)\n",
    "N_heads = 4\n",
    "N_layers = 8\n",
    "project_dim = 128\n",
    "# ============================= #\n",
    "\n",
    "load_weights = True\n",
    "\n",
    "# location of the previous weights\n",
    "model_name_load = model_dir+'baseline/ViT3d_{}_depth{}_patch{}{}{}_dim{}_heads{}_tune2'.format(\n",
    "    lead_name, N_layers, patch_size[0], patch_size[1], patch_size[2], project_dim, N_heads)\n",
    "# location for saving new weights\n",
    "model_name_save = model_dir+'baseline/ViT3d_{}_depth{}_patch{}{}{}_dim{}_heads{}_tune2'.format(\n",
    "    lead_name, N_layers, patch_size[0], patch_size[1], patch_size[2], project_dim, N_heads)\n",
    "\n",
    "# Training setups\n",
    "epochs = 9999\n",
    "batch_size = 64 #64\n",
    "N_batch = 32\n",
    "lrs = mu.cosine_schedule(N_batch, l_min=1e-6, l_max=5e-5)\n",
    "\n",
    "aug_timelag = True\n",
    "aug_revert = True\n",
    "\n",
    "if aug_timelag:\n",
    "    N_pad = pad_timelag # + pad_timelag0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b8aa7e2c-a06d-403a-92b5-5797a320934a",
   "metadata": {},
   "outputs": [],
   "source": [
    "with h5py.File(save_dir+'CCPA_domain.hdf', 'r') as h5io:\n",
    "    land_mask_CCPA = h5io['land_mask_CCPA'][...]\n",
    "    \n",
    "land_mask_CCPA = land_mask_CCPA == 1.0\n",
    "ccpa_shape = land_mask_CCPA.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33b880aa-c679-4d41-aaa8-844514be8b3f",
   "metadata": {},
   "source": [
    "## Validation set prep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a5e77dcd-79d5-45c5-b374-76b3831a3c12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ------------------------------------------------------- #\n",
    "# Validation set\n",
    "BATCH_dir = camp_dir+'BATCH_ViT_members_opt/'\n",
    "filenames = sorted(glob(BATCH_dir+'*npy'))\n",
    "\n",
    "L_valid = 50\n",
    "filenames_valid = filenames[::2][:L_valid] #[::10]\n",
    "\n",
    "valid_GEFS = np.empty((L_valid, N_leads, N_ens)+input_size[1:])\n",
    "valid_CCPA = np.empty((L_valid, N_leads,)+output_size[1:])\n",
    "\n",
    "valid_GEFS_raw = np.empty((L_valid, N_leads, N_ens)+(224, 464))\n",
    "valid_CCPA_true = np.empty((L_valid, N_leads)+(224, 464))\n",
    "\n",
    "for i, name_ in enumerate(filenames_valid):\n",
    "    temp_data = np.load(name_, allow_pickle=True)[()]\n",
    "    valid_GEFS[i, ...] = temp_data['GEFS_embed'][ilead_start:ilead_end, ..., 0:4]\n",
    "    valid_CCPA[i, ...] = temp_data['CCPA_embed'][ilead_start:ilead_end, ...]\n",
    "    valid_GEFS_raw[i, ...] = temp_data['GEFS_raw'][ilead_start:ilead_end, ...]\n",
    "    valid_CCPA_true[i, ...] = temp_data['CCPA_true'][ilead_start:ilead_end, ...]\n",
    "\n",
    "valid_CCPA_true = to_precip(valid_CCPA_true)\n",
    "\n",
    "valid_GEFS = valid_GEFS[:, :, ::4, ...]\n",
    "valid_GEFS_raw = valid_GEFS_raw[:, :, ::4, ...]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bde420f0-e4fc-429a-b92d-0c346875fc48",
   "metadata": {},
   "source": [
    "## Model training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dca6d179-7373-493b-8e88-641235b9ce1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_dir = camp_dir+'BATCH_ViT/'\n",
    "filename_train = sorted(glob(BATCH_dir+'*npy'))\n",
    "filename_train = list(set(filename_train) - set(filenames_valid))\n",
    "L_train = len(filename_train)\n",
    "\n",
    "min_del = 0.0\n",
    "max_tol = 3 # early stopping with 2-epoch patience\n",
    "tol = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0bae1181-f376-4a10-a870-af272b168af5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ------------------------------------------------------- #\n",
    "# Training loop\n",
    "strategy = tf.distribute.MirroredStrategy()\n",
    "print('Number of devices: {}'.format(strategy.num_replicas_in_sync))\n",
    "\n",
    "with strategy.scope():\n",
    "    \n",
    "    model = mu.ViT3d_corrector(input_size, output_size, patch_size, project_dim, N_layers, N_heads)\n",
    "    \n",
    "    model.compile(loss=keras.losses.mean_absolute_error, \n",
    "                  optimizer=keras.optimizers.Adam(learning_rate=5e-5))\n",
    "    \n",
    "    # load weights\n",
    "    if load_weights:\n",
    "        W_old = mu.dummy_loader(model_name_load)\n",
    "        model.set_weights(W_old)\n",
    "        \n",
    "    # ----------------------------------------------- #\n",
    "    # Major training loop + training batch generation\n",
    "    \n",
    "    batch_GEFS = np.empty((batch_size, N_leads+N_pad,)+latent_size)\n",
    "    batch_GEFS[...] = np.nan\n",
    "    batch_CCPA = np.empty((batch_size, N_leads+N_pad,)+latent_size)\n",
    "    batch_CCPA[...] = np.nan\n",
    "\n",
    "    batch_GEFS_aug = np.empty((batch_size, N_leads,)+latent_size)\n",
    "    batch_GEFS_aug[...] = np.nan\n",
    "    batch_CCPA_aug = np.empty((batch_size, N_leads,)+latent_size)\n",
    "    batch_CCPA_aug[...] = np.nan\n",
    "    \n",
    "    for i in range(epochs):\n",
    "        \n",
    "        print('epoch = {}'.format(i))\n",
    "        \n",
    "        if i == 0:\n",
    "            record = verif_func(model, valid_GEFS, valid_CCPA_true, land_mask_CCPA)\n",
    "            print('Initial validation loss: {}'.format(record))\n",
    "        \n",
    "        start_time = time.time()\n",
    "        for j in range(N_batch):\n",
    "\n",
    "            tf.keras.backend.set_value(model.optimizer.learning_rate, lrs[j])\n",
    "            \n",
    "            inds_rnd = du.shuffle_ind(L_train)\n",
    "            inds_ = inds_rnd[:batch_size]\n",
    "\n",
    "            ## Note: always train with first few days \n",
    "            for k, ind in enumerate(inds_):\n",
    "                # import batch data\n",
    "                name_ = filename_train[ind]\n",
    "                temp_data = np.load(name_, allow_pickle=True)[()]\n",
    "                batch_GEFS[k, ...] = temp_data['GEFS_embed'][0:8+pad_timelag1, ...]\n",
    "                batch_CCPA[k, ...] = temp_data['CCPA_embed'][0:8+pad_timelag1, ...]\n",
    "\n",
    "            if aug_timelag:\n",
    "                for k in range(batch_size):\n",
    "                    i_start = np.random.randint(0, N_pad)\n",
    "                    batch_GEFS_aug[k, ...] = batch_GEFS[k, i_start:i_start+N_leads, ...]\n",
    "                    batch_CCPA_aug[k, ...] = batch_CCPA[k, i_start:i_start+N_leads, ...]\n",
    "\n",
    "            if aug_revert:\n",
    "                for k in range(batch_size):\n",
    "                    i_revert = np.random.randint(0, 4)\n",
    "                    if i_revert == 4:\n",
    "                        batch_GEFS_aug[k, ...] = batch_GEFS_aug[k, ::-1, ...]\n",
    "                        batch_CCPA_aug[k, ...] = batch_CCPA_aug[k, ::-1, ...]\n",
    "\n",
    "            if (aug_timelag is False) and (aug_revert is False):\n",
    "                batch_GEFS_aug = batch_GEFS[:, pad_timelag0:-pad_timelag1, ...]\n",
    "                batch_CCPA_aug = batch_CCPA[:, pad_timelag0:-pad_timelag1, ...]\n",
    "                \n",
    "            if np.sum(np.isnan(batch_GEFS_aug)) > 0:\n",
    "                raise\n",
    "                \n",
    "            model.train_on_batch(batch_GEFS_aug, batch_CCPA_aug)\n",
    "            \n",
    "        # on epoch-end\n",
    "        record_temp = verif_func(model, valid_GEFS, valid_CCPA_true, land_mask_CCPA)\n",
    "    \n",
    "        if record - record_temp > min_del:\n",
    "            print('Validation loss improved from {} to {}'.format(record, record_temp))\n",
    "            record = record_temp\n",
    "            print(\"Save to {}\".format(model_name_save))\n",
    "            model.save(model_name_save)\n",
    "            \n",
    "        else:\n",
    "            print('Validation loss {} NOT improved'.format(record_temp))\n",
    "        \n",
    "        print(\"--- %s seconds ---\" % (time.time() - start_time))\n",
    "        # mannual callbacks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6f5e993-c5aa-4b17-b016-04bc84d94529",
   "metadata": {},
   "source": [
    "## Monitoring the progress"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83503ea7-0d33-4df0-8940-5b70f58d62f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pred_func(model, valid_GEFS, valid_CCPA_full, land_mask_CCPA):\n",
    "    \n",
    "    L_valid = valid_GEFS.shape[0]\n",
    "    N_leads = EN = valid_GEFS.shape[1]\n",
    "    EN = valid_GEFS.shape[2]\n",
    "\n",
    "    shape_full = (L_valid, N_leads, EN, 224, 464)\n",
    "    GEFS_full = np.empty(shape_full)\n",
    "    pred_ = np.empty(valid_GEFS.shape[:-1]+(4,))\n",
    "\n",
    "    for iens in range(EN):\n",
    "        pred_[:, :, iens, ...] = model.predict(valid_GEFS[:, :, iens, ...], verbose=0)\n",
    "        \n",
    "    for ilead in range(N_leads):\n",
    "        for iens in range(EN):\n",
    "            GEFS_full[:, ilead, iens, ...] = decoder.predict(pred_[:, ilead, iens, ...], \n",
    "                                                             verbose=0)[..., 0]\n",
    "    GEFS_full = to_precip(GEFS_full)\n",
    "    return GEFS_full"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85cd3701-4eaa-4e7a-a8f2-9b474dd9d360",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = mu.ViT3d_corrector(input_size, output_size, patch_size, project_dim, N_layers, N_heads)\n",
    "\n",
    "model.compile(loss=keras.losses.mean_absolute_error, \n",
    "              optimizer=keras.optimizers.Adam(learning_rate=1e-4))\n",
    "\n",
    "W_old = mu.dummy_loader(model_name_save)\n",
    "model.set_weights(W_old)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09f59208-913c-4b27-9ef2-fbe88061b688",
   "metadata": {},
   "outputs": [],
   "source": [
    "verif_func(model, valid_GEFS, valid_CCPA_true, land_mask_CCPA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "cc16e930-4a6c-4238-8f3c-53a542bc7edc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pred_ = pred_func(model, valid_GEFS, valid_CCPA_true, land_mask_CCPA)\n",
    "# plt.pcolormesh(pred_[0, 0, 0, ...], cmap=plt.cm.nipy_spectral_r)\n",
    "# plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c726ffa-82e7-4431-86d1-5f389f6dab98",
   "metadata": {},
   "source": [
    "**RAW GEFS CRPS**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39ae2ac8-e118-44c5-9370-2a2e3624720c",
   "metadata": {},
   "outputs": [],
   "source": [
    "CRPS = np.empty((L_valid, N_leads, 224, 464)); CRPS[...] = np.nan\n",
    "\n",
    "for ilead in range(N_leads):\n",
    "    crps_ilead, _, _ = vu.CRPS_2d(valid_CCPA_true[:, ilead, ...], \n",
    "                                  valid_GEFS_raw[:, ilead, ...], land_mask=land_mask_CCPA)\n",
    "    CRPS[:, ilead, ...] = crps_ilead\n",
    "print(np.nanmean(CRPS))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c64513fb-1d7e-4473-b120-3b31c55b56b2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}