DATA01_preprocess_GEFS_reforecast_mean.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0d99e926-b62b-44e7-a37c-76c7acd0e898",
   "metadata": {},
   "source": [
    "# GEFS reforecast preprocessing\n",
    "\n",
    "This notebook converts GEFS reforecast grib2 files into h5 format\n",
    "\n",
    "* `$camp_dir/GFS_reforecast/GEFS_AVE_y{}.hdf`\n",
    "* `$camp_dir/GFS_reforecast/GEFS_STD_y{}.hdf`\n",
    "* `$camp_dir/GFS_reforecast/GEFS_c00_y{}.hdf`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "76f99011-74fa-4b3c-bc05-9e3e9f945477",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import time\n",
    "import h5py\n",
    "import pygrib\n",
    "\n",
    "import numpy as np\n",
    "from glob import glob\n",
    "\n",
    "from scipy.interpolate import RegularGridInterpolator\n",
    "from datetime import datetime, timedelta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "33ddc8be-bce9-4530-8afb-5713dd70046a",
   "metadata": {},
   "outputs": [],
   "source": [
    "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c0548e7e-f7a3-46bb-93e8-f6b7b3e74146",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9972901e-3bd5-4eb0-8686-765cb64b82a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "filenames = glob(camp_dir+'wget_GEFS_reforecast/*grib2')\n",
    "ens_name = ['c00', 'p01', 'p02', 'p03', 'p04']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4556b7c4-d7a9-48c6-a6ee-97985a70fdaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "with h5py.File(save_dir+'CCPA_domain.hdf', 'r') as h5io:\n",
    "    lon_CCPA = h5io['lon_CCPA'][...]\n",
    "    lat_CCPA = h5io['lat_CCPA'][...]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c472256f-1b7a-4d88-b603-f74237457ea4",
   "metadata": {},
   "outputs": [],
   "source": [
    "with pygrib.open(filenames[0]) as grbio:\n",
    "    lat_GFS, lon_GFS = grbio[1].latlons()\n",
    "lat_GFS = lat_GFS[:360, 720:]\n",
    "lon_GFS = lon_GFS[:360:, 720:]\n",
    "\n",
    "lat_GFS = np.flipud(lat_GFS)\n",
    "lon_GFS -= 360"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "571b8b58-86a0-419c-8009-9826fdc85b0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "years = np.arange(2002, 2020, 1)\n",
    "grb_inds = np.arange(2, 58, 2)\n",
    "N_leads = len(grb_inds)\n",
    "grid_shape = lon_CCPA.shape\n",
    "grid_shape_gfs = lat_GFS.shape\n",
    "gfs_name = camp_dir+'wget_GEFS_reforecast/apcp_sfc_{}_{}.grib2'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c740661d-1d3d-4cac-a605-2b1ab12fbfc3",
   "metadata": {},
   "source": [
    "## Preprocess ensemble mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dc97796d-9f49-4912-bc14-c667686cb9de",
   "metadata": {},
   "outputs": [],
   "source": [
    "year = 2002\n",
    "\n",
    "if year % 4 == 0:\n",
    "    N_days = 366\n",
    "else:\n",
    "    N_days = 365\n",
    "\n",
    "GFS_save = np.zeros((N_days, N_leads,)+grid_shape)\n",
    "\n",
    "base = datetime(year, 1, 1)\n",
    "date_list = [base + timedelta(days=d) for d in range(N_days)]\n",
    "\n",
    "for i, dt in enumerate(date_list):\n",
    "    dt_str = datetime.strftime(dt, '%Y%m%d00')\n",
    "    \n",
    "    ens_mean_temp = np.zeros((N_leads,)+grid_shape_gfs)\n",
    "    \n",
    "    count = 0\n",
    "    for ens in ens_name:\n",
    "        # Identify the *.grib2 file name\n",
    "        name_ = gfs_name.format(dt_str, ens)\n",
    "\n",
    "        # g\n",
    "        if os.path.isfile(name_):\n",
    "            count += 1\n",
    "            with pygrib.open(name_) as grbio:\n",
    "                for ilead, ind in enumerate(grb_inds):\n",
    "                    apcp_temp = grbio[int(ind)].values\n",
    "                    apcp_temp = apcp_temp[:360, 720:]\n",
    "                    apcp_temp = np.flipud(apcp_temp)\n",
    "                    # <----- accumulate ensebmle mean\n",
    "                    ens_mean_temp[ilead, ...] = ens_mean_temp[ilead, ...] + apcp_temp\n",
    "        else:\n",
    "            print('Missing {}'.format(name_))\n",
    "    if count == 0:\n",
    "        print('Missing all members on {}'.format(dt_str))\n",
    "        GFS_save[i, ...] = np.nan\n",
    "    else:    \n",
    "        ens_mean_temp = ens_mean_temp/count # <------- divide to get the mean value\n",
    "        \n",
    "        for ilead in range(N_leads):\n",
    "            lr_to_hr = RegularGridInterpolator((lat_GFS[:, 0], lon_GFS[0, :]), ens_mean_temp[ilead, ...], \n",
    "                                               bounds_error=False, fill_value=None)\n",
    "            ens_mean_interp = lr_to_hr((lat_CCPA, lon_CCPA))\n",
    "                    \n",
    "            GFS_save[i, ilead, ...] = ens_mean_interp\n",
    "\n",
    "    # tuple_save = (GFS_save,)\n",
    "    # label_save = ['GEFS_APCP',]\n",
    "    # du.save_hdf5(tuple_save, label_save, \n",
    "    #              camp_dir+'GFS_reforecast/', \n",
    "    #              'GEFS_AVE_y{}.hdf'.format(year))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "186c36a9-ed43-44d7-b996-9a2be9074061",
   "metadata": {},
   "source": [
    "## Pre-process individual members (x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "193c4b76-e380-4d58-bb91-08ab6a8b2857",
   "metadata": {},
   "outputs": [],
   "source": [
    "year = 2006\n",
    "\n",
    "if year % 4 == 0:\n",
    "    N_days = 366\n",
    "else:\n",
    "    N_days = 365\n",
    "\n",
    "GFS_save = np.zeros((N_days, N_leads,)+grid_shape)\n",
    "\n",
    "base = datetime(year, 1, 1)\n",
    "date_list = [base + timedelta(days=d) for d in range(N_days)]\n",
    "\n",
    "for ens in ens_name:\n",
    "    # fill NaN to the empty array\n",
    "    GFS_save[...] = np.nan\n",
    "    \n",
    "    for i, dt in enumerate(date_list):\n",
    "        dt_str = datetime.strftime(dt, '%Y%m%d00')\n",
    "        \n",
    "        ens_temp = np.zeros((N_leads,)+grid_shape_gfs)\n",
    "        ens_temp[...] = np.nan\n",
    "        \n",
    "        # Identify the *.grib2 file name\n",
    "        name_ = gfs_name.format(dt_str, ens)\n",
    "\n",
    "        try:\n",
    "            # start reading grib files\n",
    "            if os.path.isfile(name_):\n",
    "                with pygrib.open(name_) as grbio:\n",
    "                    for ilead, ind in enumerate(grb_inds):\n",
    "                        apcp_temp = grbio[int(ind)].values\n",
    "                        apcp_temp = apcp_temp[:360, 720:]\n",
    "                        apcp_temp = np.flipud(apcp_temp)\n",
    "                        ens_temp[ilead, ...] = apcp_temp\n",
    "            else:\n",
    "                print('Missing {}'.format(name_))\n",
    "                continue;\n",
    "        except:\n",
    "            ens_temp[ilead, ...] = np.nan\n",
    "            \n",
    "        if np.sum(np.isnan(ens_temp)) > 0:\n",
    "            print('{} has NaNs'.format(name_))\n",
    "            continue;\n",
    "        else:\n",
    "            for ilead in range(N_leads):\n",
    "                lr_to_hr = RegularGridInterpolator((lat_GFS[:, 0], lon_GFS[0, :]), ens_temp[ilead, ...], \n",
    "                                                   bounds_error=False, fill_value=None)\n",
    "                ens_interp = lr_to_hr((lat_CCPA, lon_CCPA))\n",
    "                        \n",
    "                GFS_save[i, ilead, ...] = ens_interp\n",
    "\n",
    "    # tuple_save = (GFS_save,)\n",
    "    # label_save = ['GEFS_APCP',]\n",
    "    # du.save_hdf5(tuple_save, label_save, \n",
    "    #              camp_dir+'GFS_reforecast/', \n",
    "    #              'GEFS_{}_y{}.hdf'.format(ens, year))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "103c261a-4f72-4557-b379-340c935b35ce",
   "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
}