4 - Predicting with a DLWP-CS model - NCI.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predicting with a DLWP-CS model\n",
    "\n",
    "- Finally we will explore using the advanced functionality in DLWP-CS to make a time-series global weather prediction with our trained DLWP-CS model. \n",
    "- We will save the prediction to a netCDF file and apply inverse scaling to get physical variables back. (It is recommend having this model run on a GPU with at least 4 GB of video memory.)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters\n",
    "\n",
    "Some user-specified parameters. The `scale_file` contains the mean and standard of the data (which was dropped in the cubed sphere remapping). The `map_files` were produced by the cubed sphere remapping. We can re-use them here so we don't have to generate them again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root_directory: /scratch/vp91/mah900/NCI-DLWP-CS/Data/\n",
      "predictor_file: /g/data/dk92/data/NCI-DLWP-CS/Data/NCI_tutorial/NCI_tutorial_z500_t2m_CS.nc\n",
      "scale_file:  /g/data/dk92/data/NCI-DLWP-CS/Data/NCI_tutorial/NCI_tutorial_z500_t2m.nc\n",
      "model:       /scratch/vp91/mah900/NCI-DLWP-CS/Data/dlwp-cs_NCI_tutorial\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' \n",
    "os.environ['TF_GPU_ALLOCATOR']='cuda_malloc_async' \n",
    "os.chdir (f\"/scratch/vp91/{os.environ['USER']}/NCI-DLWP-CS\")\n",
    "\n",
    "root_directory = f\"/scratch/vp91/{os.environ['USER']}/NCI-DLWP-CS/Data/\"\n",
    "# From your selected dir\n",
    "# predictor_file = os.path.join(root_directory, 'NCI_tutorial', 'NCI_tutorial_z500_t2m_CS.nc')\n",
    "# scale_file = os.path.join(root_directory, 'NCI_tutorial', 'NCI_tutorial_z500_t2m.nc')\n",
    "# model = os.path.join(root_directory, 'dlwp-cs_NCI_tutorial')\n",
    "\n",
    "# From dk92\n",
    "predictor_file = '/g/data/dk92/data/NCI-DLWP-CS/Data/NCI_tutorial/NCI_tutorial_z500_t2m_CS.nc'\n",
    "scale_file = '/g/data/dk92/data/NCI-DLWP-CS/Data/NCI_tutorial/NCI_tutorial_z500_t2m.nc' \n",
    "model = '/g/data/dk92/data/NCI-DLWP-CS/Data/dlwp-cs_NCI_tutorial'\n",
    "# map_files = ('map_LL91x180_CS48.nc', 'map_CS48_LL91x180.nc')\n",
    "# map_files  = ('map_LL32x64_CS48.nc', 'map_CS48_LL32x64.nc')\n",
    "map_files  = ('/g/data/dk92/data/NCI-DLWP-CS/map_LL32x64_CS48.nc', '/g/data/dk92/data/NCI-DLWP-CS/map_CS48_LL32x64.nc')\n",
    "\n",
    "print('root_directory:', root_directory)\n",
    "print('predictor_file:', predictor_file)\n",
    "print('scale_file: ', scale_file)\n",
    "print('model:      ', model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- We'll resurrect some parameters from the training tutorial. See that notebook for definitions. \n",
    "- Note that we omit `data_interval` because we simply select only every 6 hours from the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "io_selection = {'varlev': ['z/500', 't2m/0']}\n",
    "add_solar = True\n",
    "io_time_steps = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Specify the data used for prediction. We'll make weekly forecasts in the test set, initialized at 0 UTC. We need to specify a subset of the data that contains all these initializations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import xarray as xr\n",
    "\n",
    "validation_set = pd.date_range('2016-12-31', '2018-12-31', freq='6H')\n",
    "validation_set = np.array(validation_set, dtype='datetime64[ns]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the initialization dates, the numer of foreward forecast hours, and the time step (could be automated...)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "dates = pd.date_range('2017-01-01', '2018-12-31', freq='7D')\n",
    "initialization_dates = xr.DataArray(dates)\n",
    "num_forecast_hours = 5 * 24\n",
    "dt = 6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the DLWP model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from DLWP.util import load_model, remove_chars, is_channels_last\n",
    "\n",
    "dlwp = load_model(model)\n",
    "\n",
    "# File to save the forecast\n",
    "forecast_file = os.path.join(root_directory, 'forecast_%s.nc' % remove_chars(model.split(os.sep)[-1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Open the data and create the data generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_ds = xr.open_dataset(predictor_file)\n",
    "predictor_ds = all_ds.sel(sample=validation_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from DLWP.model import SeriesDataGenerator\n",
    "\n",
    "sequence = dlwp._n_steps if hasattr(dlwp, '_n_steps') and dlwp._n_steps > 1 else None\n",
    "val_generator = SeriesDataGenerator(dlwp, predictor_ds, rank=3, add_insolation=add_solar,\n",
    "                                    input_sel=io_selection, output_sel=io_selection,\n",
    "                                    input_time_steps=io_time_steps, output_time_steps=io_time_steps,\n",
    "                                    shuffle=False, sequence=sequence, batch_size=32,\n",
    "                                    load=False, channels_last=is_channels_last(dlwp))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the estimator and make a prediction\n",
    "\n",
    "We use the handy TimeSeriesEstimator class to intelligently produce a time series forecast. This class depends on the model and the data generator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from DLWP.model import TimeSeriesEstimator\n",
    "\n",
    "estimator = TimeSeriesEstimator(dlwp, val_generator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicting with the model /scratch/vp91/mah900/NCI-DLWP-CS/Data/dlwp-cs_NCI_tutorial ...\n",
      "Time step 1/5\n",
      "4/4 [==============================] - 11s 101ms/step\n",
      "Time step 2/5\n",
      "4/4 [==============================] - 0s 63ms/step\n",
      "Time step 3/5\n",
      "4/4 [==============================] - 0s 64ms/step\n",
      "Time step 4/5\n",
      "4/4 [==============================] - 0s 62ms/step\n",
      "Time step 5/5\n",
      "4/4 [==============================] - 0s 64ms/step\n"
     ]
    }
   ],
   "source": [
    "print('Predicting with the model %s ...' % model)\n",
    "\n",
    "# Select the samples from the initialization dates. The first \"time\" input to the model is actually one time step earlier\n",
    "samples = np.array([int(np.where(val_generator.ds['sample'] == s)[0]) for s in initialization_dates]) \\\n",
    "    - io_time_steps + 1\n",
    "time_series = estimator.predict(num_forecast_hours // dt, samples=samples, verbose=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Transpose if channels_last was used for the model\n",
    "if is_channels_last(dlwp):\n",
    "    time_series = time_series.transpose('f_hour', 'time', 'varlev', 'x0', 'x1', 'x2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Scale the variables back to real data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "if scale_file is None:\n",
    "    scale_ds = predictor_ds\n",
    "else:\n",
    "    scale_ds = xr.open_dataset(scale_file)\n",
    "sel_mean = scale_ds['mean'].sel(io_selection)\n",
    "sel_std = scale_ds['std'].sel(io_selection)\n",
    "time_series = time_series * sel_std + sel_mean"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sometimes the time series output, when saved to netCDF, is not compatible with TempestRemap (for reasons unknown). But there is a function in the DLWP `verify` module that re-formats a time series and produces output that TempestRemap is happy with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from DLWP.verify import add_metadata_to_forecast_cs\n",
    "\n",
    "fh = np.arange(dt, time_series.shape[0] * dt + 1., dt)\n",
    "time_series = add_metadata_to_forecast_cs(\n",
    "    time_series.values,\n",
    "    fh,\n",
    "    predictor_ds.sel(**io_selection).sel(sample=initialization_dates)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save the prediction in cubed sphere format. Drop the string \"varlev\" coordinate for TempestRemap."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "time_series.drop('varlev').to_netcdf(forecast_file + '.cs')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Remap the forecast to a latitude-longitude grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CubeSphereRemap.convert_from_faces: loading data to memory...\n",
      "CubeSphereRemap: assigning new coordinates to dataset\n",
      "CubeSphereRemap.convert_from_faces: exporting data to file /scratch/vp91/mah900/NCI-DLWP-CS/Data/forecast_dlwp-cs_NCI_tutorial.nc.tmp...\n",
      "CubeSphereRemap.convert_from_faces: successfully exported reformatted data\n",
      "CubeSphereRemap: applying inverse map...\n",
      "/opt/conda/envs/nci-dlwp-cs/bin/ApplyOfflineMap --in_data /scratch/vp91/mah900/NCI-DLWP-CS/Data/forecast_dlwp-cs_NCI_tutorial.nc.tmp --out_data /scratch/vp91/mah900/NCI-DLWP-CS/Data/forecast_dlwp-cs_NCI_tutorial.nc --map /g/data/dk92/data/NCI-DLWP-CS/map_CS48_LL32x64.nc --var forecast\n",
      "CubeSphereRemap: successfully inverse remapped data into /scratch/vp91/mah900/NCI-DLWP-CS/Data/forecast_dlwp-cs_NCI_tutorial.nc\n"
     ]
    }
   ],
   "source": [
    "from DLWP.remap import CubeSphereRemap\n",
    "csr = CubeSphereRemap(to_netcdf4=True)\n",
    "csr.assign_maps(*map_files)\n",
    "csr.convert_from_faces(forecast_file + '.cs', forecast_file + '.tmp')\n",
    "csr.inverse_remap(forecast_file + '.tmp', forecast_file, '--var', 'forecast')\n",
    "os.remove(forecast_file + '.tmp')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The End"
   ]
  },
  {
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}