3d.cnn-example-PyTorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "Collapsed": "false"
   },
   "source": [
    "# NCI WeatherBench-3d: Train a CNN (PyTroch)\n",
    "\n",
    "In this notebook we will go through all the steps required to train a fully convolutional neural network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import os\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "from score import *\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "from dask.diagnostics import ProgressBar\n",
    "from datetime import datetime\n",
    "from dask.distributed import Client\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import train_nn_pytorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(n_workers=12, threads_per_worker=1)  \n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset: '5.625'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "\n",
    "res     = '5.625'\n",
    "datadir = f'/g/data/wb00/NCI-Weatherbench/{res}deg' \n",
    "print (\"Data loading...\" )\n",
    "\n",
    "years = list(range(1999, 2022+1))\n",
    "print (years)\n",
    "z_files = [ file for year in years for file in glob.glob (fr'{datadir}/geopotential/*{year}*')  ] \n",
    "t_files = [ file for year in years for file in glob.glob (fr'{datadir}/temperature/*{year}*')    ] \n",
    "\n",
    "z = xr.open_mfdataset(z_files, combine='by_coords', parallel=True, chunks={'time': 10}).z.sel(level=[500]).load() \n",
    "t = xr.open_mfdataset(t_files, combine='by_coords', parallel=True, chunks={'time': 10}).t.sel(level=[850]).drop('level').load() \n",
    "datasets = [z, t]\n",
    "print (\"Merging ... \")\n",
    "ds = xr.merge(datasets).compute()  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z['time'][:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z['time'][-10:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def data_generate(ds, lead_time, batch_size, variables, train_years, valid_years, test_years):\n",
    "    ds_train = ds.sel(time=slice(*train_years))\n",
    "    ds_valid = ds.sel(time=slice(*valid_years))\n",
    "    ds_test  = ds.sel(time=slice(*test_years))\n",
    "\n",
    "    print (\"Data generation ... \")\n",
    "    dic = {var: 500 for var in variables} #vars}\n",
    "    dg_train = train_nn_pytorch.DataGenerator(ds_train, dic, lead_time, batch_size=batch_size)\n",
    "    dg_valid = train_nn_pytorch.DataGenerator(ds_valid, dic, lead_time, batch_size=batch_size, mean=dg_train.mean,\n",
    "                         std=dg_train.std, shuffle=False)\n",
    "    dg_test =  train_nn_pytorch.DataGenerator(ds_test, dic, lead_time, batch_size=batch_size, mean=dg_train.mean,\n",
    "                         std=dg_train.std, shuffle=False)\n",
    "\n",
    "    print(f'Mean = {dg_train.mean}; Std = {dg_train.std}')\n",
    "    return dg_train, dg_valid, dg_test\n",
    "\n",
    "def train(dg_train_generator, dg_valid_generator, dg_test_generator, model_save_fn):\n",
    "    print (\"Train model ... \") \n",
    "    \n",
    "    model = train_nn_pytorch.Model_cnn(channels, kernels)\n",
    "    \n",
    "    loss_fn = nn.MSELoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n",
    "\n",
    "    print('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    model.to(device)    \n",
    "    \n",
    "    patience = n_patience\n",
    "    best_loss = float('inf')\n",
    "\n",
    "    train_losses = []\n",
    "    valid_losses = []\n",
    "    avg_train_losses = []\n",
    "    avg_valid_losses = []     \n",
    "\n",
    "    for epoch in range(n_epochs):\n",
    "        for X, y in dg_train_generator:\n",
    "            model.train()\n",
    "            X  = torch.squeeze( X )\n",
    "            X = X.to(device)\n",
    "            y_pred = model (X)  \n",
    "            y  = torch.squeeze( y )\n",
    "            y = y.to(device)\n",
    "            loss = loss_fn(y_pred, y)\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            train_losses.append(loss.item())\n",
    "\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            for X_v, y_v in dg_valid_generator:\n",
    "                X_v  = torch.squeeze( X_v )\n",
    "                y_v  = torch.squeeze( y_v )\n",
    "                X_v = X_v.to(device)\n",
    "                y_v = y_v.to(device)\n",
    "\n",
    "                y_v_pred = model(X_v)\n",
    "                loss_v = loss_fn(y_v_pred, y_v)\n",
    "                valid_losses.append(loss_v.item())\n",
    "\n",
    "        train_loss = np.average(train_losses)\n",
    "        valid_loss = np.average(valid_losses)\n",
    "        avg_train_losses.append(train_loss)\n",
    "        avg_valid_losses.append(valid_loss)\n",
    "\n",
    "        epoch_len = len(str(n_epochs))\n",
    "\n",
    "        print_msg = (f'[{epoch:>{epoch_len}}/{n_epochs:>{epoch_len}}] ' +\n",
    "                     f'train_loss: {train_loss:.5f} ' +\n",
    "                     f'valid_loss: {valid_loss:.5f}')\n",
    "\n",
    "        print(print_msg)\n",
    "\n",
    "        # clear  \n",
    "        train_losses = []\n",
    "        valid_losses = []  \n",
    "\n",
    "        if loss_v < best_loss:  \n",
    "            best_loss = loss_v\n",
    "            patience = n_patience  \n",
    "        else:\n",
    "            patience -= 1\n",
    "            if patience == 0:\n",
    "                print(f'Early stopping')\n",
    "                break               \n",
    "\n",
    "    return model  \n",
    "        \n",
    "def evaluate(pred, iterative, test_years):     \n",
    "    print(\"Evaluating forecast ...\")\n",
    "    valid_years = list( range(int(test_years[0]), int(test_years[1])+1 ))\n",
    "    print ('all test_years:', valid_years)\n",
    "    z500_valid_files = [ file for year in valid_years for file in glob.glob (fr'{datadir}/geopotential/*{year}*') ] \n",
    "    t850_valid_files = [ file for year in valid_years for file in glob.glob (fr'{datadir}/temperature/*{year}*')  ]     \n",
    "\n",
    "    z500_valid = load_test_data(z500_valid_files, 'z', slice(*test_years)) \n",
    "    t850_valid = load_test_data(t850_valid_files, 't', slice(*test_years))     \n",
    "    \n",
    "    valid      = xr.merge([z500_valid, t850_valid], compat='override').compute()\n",
    "    \n",
    "    print(train_nn_pytorch.evaluate_iterative_forecast(pred, valid, compute_weighted_rmse).load() if iterative \\\n",
    "                                                else train_nn_pytorch.compute_weighted_rmse(pred, valid).compute())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_years =('2021', '2022')\n",
    "test_years = list (range(int (test_years[0]), int(test_years[1])+1 ))\n",
    "print ('test_years:', test_years)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 72 hours (3 days)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "print (60*\"-\")\n",
    "train_years=('1999', '2015')\n",
    "valid_years=('2016', '2020')\n",
    "test_years =('2021', '2022')\n",
    "batch_size = 32\n",
    "variables = ('z', 't')\n",
    "lead_time = 72\n",
    "channels = [2, 64, 64, 64, 64, 64, 2]\n",
    "kernels = [5, 5, 5, 5, 5, 5]\n",
    "dr = 0\n",
    "save_prefix = 'PyTorch_NCI_tutorial' \n",
    "print ('save_prefix:', save_prefix)\n",
    "model_save_fn = f'/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/saved_models/{save_prefix}_cnn_3d.h5'\n",
    "pred_save_fn  = f'/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/{save_prefix}_cnn_3d.nc'\n",
    "lr = 1e-4\n",
    "iterative = False\n",
    "iterative_lead_time = None\n",
    "n_epochs = 100   \n",
    "n_patience = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "dg_train, dg_valid, dg_test = data_generate(ds, lead_time, batch_size,\n",
    "                                           variables, train_years, valid_years, test_years)\n",
    "\n",
    "params = {'batch_size': 1,\n",
    "          'shuffle': True,\n",
    "          'num_workers': 6}\n",
    "dg_train_generator = torch.utils.data.DataLoader(dg_train, **params)\n",
    "params['shuffle'] = False\n",
    "dg_valid_generator = torch.utils.data.DataLoader(dg_valid, **params)\n",
    "dg_test_generator = torch.utils.data.DataLoader(dg_test, **params)\n",
    "\n",
    "model = train(dg_train_generator, dg_valid_generator, dg_test_generator, model_save_fn)\n",
    "\n",
    "print(f'Saving model weights: {model_save_fn}')\n",
    "torch.save(model.state_dict(), model_save_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "pred = train_nn_pytorch.create_predictions(model, dg_test_generator, \n",
    "                        mean=dg_train.mean.values, std=dg_train.std.values,\n",
    "                        var_dict=dg_test.var_dict,\n",
    "                        valid_time=dg_test.valid_time, \n",
    "                        lat=dg_test.ds.lat, \n",
    "                        lon=dg_test.ds.lon \n",
    "                       )\n",
    "print(f'Saving predictions: {pred_save_fn}')\n",
    "pred.to_netcdf(pred_save_fn)\n",
    "\n",
    "evaluate(pred, iterative, test_years)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 120 hours (5 days)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "\n",
    "print (60*\"-\")\n",
    "batch_size = 32\n",
    "lead_time = 120\n",
    "print ('save_prefix:', save_prefix)\n",
    "model_save_fn= f'/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/saved_models/{save_prefix}_cnn_5d.h5'\n",
    "pred_save_fn = f'/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/{save_prefix}_cnn_5d.nc'\n",
    "lr = 1e-4\n",
    "iterative = False\n",
    "iterative_lead_time = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "\n",
    "dg_train, dg_valid, dg_test = data_generate(ds, lead_time, batch_size,\n",
    "                                           variables, train_years, valid_years, test_years)\n",
    "\n",
    "params = {'batch_size': 1,\n",
    "          'shuffle': True,\n",
    "          'num_workers': 6}\n",
    "dg_train_generator = torch.utils.data.DataLoader(dg_train, **params)\n",
    "params['shuffle'] = False\n",
    "dg_valid_generator = torch.utils.data.DataLoader(dg_valid, **params)\n",
    "dg_test_generator = torch.utils.data.DataLoader(dg_test, **params)\n",
    "\n",
    "model = train(dg_train_generator, dg_valid_generator, dg_test_generator, model_save_fn)\n",
    "\n",
    "print(f'Saving model weights: {model_save_fn}')\n",
    "torch.save(model.state_dict(), model_save_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "pred = train_nn_pytorch.create_predictions(model, dg_test_generator, \n",
    "                        mean=dg_train.mean.values, std=dg_train.std.values,\n",
    "                        var_dict=dg_test.var_dict,\n",
    "                        valid_time=dg_test.valid_time, \n",
    "                        lat=dg_test.ds.lat, \n",
    "                        lon=dg_test.ds.lon \n",
    "                       )\n",
    "print(f'Saving predictions: {pred_save_fn}')\n",
    "pred.to_netcdf(pred_save_fn)\n",
    "\n",
    "evaluate(pred, iterative, test_years)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# fccnn_6h_iter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "\n",
    "print (60*\"-\")\n",
    "batch_size = 32\n",
    "lead_time = 6\n",
    "iterative_lead_time = 5 * 24 \n",
    "print ('save_prefix:', save_prefix)\n",
    "model_save_fn=f'/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/saved_models/{save_prefix}_fccnn_6h_iter.h5'\n",
    "pred_save_fn =f'/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/{save_prefix}_fccnn_6h_iter.nc'\n",
    "lr = 1e-4\n",
    "iterative = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "dg_train, dg_valid, dg_test = data_generate(ds, lead_time, batch_size,\n",
    "                                           variables, train_years, valid_years, test_years)\n",
    "\n",
    "params = {'batch_size': 1,\n",
    "          'shuffle': True,\n",
    "          'num_workers': 6}\n",
    "dg_train_generator = torch.utils.data.DataLoader(dg_train, **params)\n",
    "params['shuffle'] = False\n",
    "dg_valid_generator = torch.utils.data.DataLoader(dg_valid, **params)\n",
    "dg_test_generator = torch.utils.data.DataLoader(dg_test, **params)\n",
    "\n",
    "model = train(dg_train_generator, dg_valid_generator, dg_test_generator, model_save_fn)\n",
    "\n",
    "print(f'Saving model weights: {model_save_fn}')\n",
    "torch.save(model.state_dict(), model_save_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "    \n",
    "pred = train_nn_pytorch.create_iterative_predictions(model, dg_test_generator, \n",
    "                                        max_lead_time= 5*24,\n",
    "                                        mean=dg_train.mean.values, std=dg_train.std.values,\n",
    "                                        var_dict=dg_test.var_dict,\n",
    "                                        valid_time=dg_test.valid_time, \n",
    "                                        lat=dg_test.ds.lat, \n",
    "                                        lon=dg_test.ds.lon,\n",
    "                                        state=dg_test.data[:dg_test.n_samples], \n",
    "                                        lead_time=dg_test.lead_time, \n",
    "                                        init_time=dg_test.init_time\n",
    "                                        )\n",
    "print(f'Saving predictions: {pred_save_fn}')\n",
    "pred.to_netcdf(pred_save_fn)\n",
    "\n",
    "evaluate(pred, iterative, test_years)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}