4_Input_processing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c8bbdddb-b9dd-44fa-913a-6667dcde5b7d",
   "metadata": {},
   "source": [
    "# Input processing\n",
    "\n",
    "Hydroclimatic variables typically exhibit certain temporal correlation. To enable the CNN to explore temporal correlation between each input variable and its antecedent conditions, we stack the input image at time t on top of its antecedent conditions to form a 3-D volume. We set the number of lags to 2 (i.e., t − 1, t − 2) after preliminary experiments; thus, each input volume has dimensions 128 × 128 × 3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4b698293-b432-4f4f-b5c2-209eb3692d8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install rioxarray xarray numpy\n",
    "\n",
    "import rioxarray\n",
    "import xarray as xr\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be14aeae-13d0-4a70-a2d8-328276309ce7",
   "metadata": {},
   "source": [
    "#### 1. Stack TWSA<sub>GLDAS</sub>, precipitation and temperature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4c9654cf-636c-4435-b125-f8fc844d518d",
   "metadata": {},
   "outputs": [],
   "source": [
    "gldas_ds = xr.open_mfdataset('./data/GLDAS/processed/TWSA_M_2002_2016_India_03125.nc4').transpose(\"lat\", \"lon\", \"time\")\n",
    "\n",
    "twsa_gldas_da = gldas_ds['twsa_gldas']\n",
    "temp_da = gldas_ds['tair_f_inst']\n",
    "precip_da = gldas_ds['rainf_f_tavg']\n",
    "\n",
    "# normalise input\n",
    "twsa_gldas_da = (twsa_gldas_da - twsa_gldas_da.mean()) / twsa_gldas_da.std()\n",
    "temp_da = (temp_da - temp_da.mean()) / temp_da.std()\n",
    "precip_da = (precip_da - precip_da.mean()) / precip_da.std()\n",
    "#twsa_gldas_da = (twsa_gldas_da - twsa_gldas_da.min()) / (twsa_gldas_da.max() - twsa_gldas_da.min())\n",
    "#temp_da = (temp_da - temp_da.min()) / (temp_da.max() - temp_da.min())\n",
    "#precip_da = (precip_da - precip_da.min()) / (precip_da.max() - precip_da.min())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dfc6ec53-039b-4b6c-9176-296cd05987b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# replace nan with 1e-7 for masking in Keras\n",
    "twsa_gldas_da = twsa_gldas_da.fillna(1e-7)\n",
    "temp_da = temp_da.fillna(1e-7)\n",
    "precip_da = precip_da.fillna(1e-7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "be68a8fd-0c5b-4830-a278-64ce4bf61190",
   "metadata": {},
   "outputs": [],
   "source": [
    "T = twsa_gldas_da.shape[2]\n",
    "\n",
    "for t in range(2,T):\n",
    "    twsa_image = np.stack((twsa_gldas_da[:,:,t], twsa_gldas_da[:,:,t-1], twsa_gldas_da[:,:,t-2]), axis=-1)\n",
    "    np.save(f'./data/processed_input/twsa/image_{t}.npy', twsa_image)\n",
    "\n",
    "    precip_image = np.stack((precip_da[:,:,t], precip_da[:,:,t-1], precip_da[:,:,t-2]), axis=-1)\n",
    "    np.save(f'./data/processed_input/precip/image_{t}.npy', precip_image)\n",
    "    \n",
    "    temp_image = np.stack((temp_da[:,:,t], temp_da[:,:,t-1], temp_da[:,:,t-2]), axis=-1)\n",
    "    np.save(f'./data/processed_input/temp/image_{t}.npy', temp_image)"
   ]
  }
 ],
 "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}