DATA00_wget_GEFSv12_individual_members.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a0d95b56-6124-4f90-8183-c17f90af1aa1",
   "metadata": {},
   "source": [
    "# GEFS operational ensemble mean wget cmd generator\n",
    "\n",
    "This notebook generates qsub scripts to download operational GEFS ensemble mean\n",
    "* Source: https://noaa-gefs-pds.s3.amazonaws.com/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e4e88911-1838-477d-934a-391f92d0d53f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from datetime import datetime, timedelta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb961d2-b167-4679-a08c-7da5d163048f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0, '/glade/u/home/ksha/GAN_proj/')\n",
    "from namelist import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ae3f28b0-53c0-4c8f-8210-37ee183d9e36",
   "metadata": {},
   "outputs": [],
   "source": [
    "keys = ['gec00',]\n",
    "\n",
    "for i in range(1, 31, 1):\n",
    "    keys.append('gep{:02d}'.format(i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "35a46fac-f436-41ea-b542-b915de4e21d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "key_leads = []\n",
    "\n",
    "for i in range(6, 144+6, 6):\n",
    "    key_leads.append('{:03d}'.format(i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "78220459-ed92-45ce-9ab8-c18fa4c3a000",
   "metadata": {},
   "outputs": [],
   "source": [
    "base = datetime(2021, 1, 1)\n",
    "date_list = [base + timedelta(days=day) for day in range(365)]\n",
    "\n",
    "base = datetime(2020, 12, 1)\n",
    "date_list = [base + timedelta(days=day) for day in range(31)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "392eea57-18db-4aa1-8a1a-f53c0d6e03a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#key_leads"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64025dab-5d2f-477e-ad2e-63221f6841f4",
   "metadata": {},
   "source": [
    "## Fill-in-the-gap version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "14921191-abd5-4e15-ae37-ea6780b4db41",
   "metadata": {},
   "outputs": [],
   "source": [
    "# check if files exist\n",
    "date_list_fill = []\n",
    "\n",
    "for dt in date_list:\n",
    "    \n",
    "    dt_str = datetime.strftime(dt, '%Y%m%d')\n",
    "    filename = camp_dir+'wget_GEFSv12_idx/{}_done.txt'.format(dt_str)\n",
    "    \n",
    "    if os.path.isfile(filename):\n",
    "        continue;\n",
    "    else:\n",
    "        date_list_fill.append(dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ca3c5879-bfe7-4a8b-b25a-a313f8075231",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "need to download 0 files\n"
     ]
    }
   ],
   "source": [
    "L_fill = len(date_list_fill)\n",
    "print('need to download {} files'.format(L_fill))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "331ebde5-f8f9-483a-8bd1-1ce6556edd96",
   "metadata": {},
   "source": [
    "## Download idx files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b4a3ada5-c995-488a-8a99-94489ee5da32",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_scripts = 3\n",
    "N = int(L_fill/N_scripts) + 1\n",
    "index = np.arange(0, L_fill, N)\n",
    "\n",
    "for i_, i_start in enumerate(index):\n",
    "\n",
    "    f = open(work_dir+'qsub/wget_gefs_{:03d}.sh'.format(i_), 'w') \n",
    "    \n",
    "    heads = '''\n",
    "    #!/bin/bash -l\n",
    "    \n",
    "    #PBS -N wget_gefs_idx\n",
    "    #PBS -A NAML0001\n",
    "    #PBS -l walltime=23:59:59\n",
    "    #PBS -l select=1:ncpus=4:mem=12GB\n",
    "    #PBS -q casper\n",
    "    #PBS -o wget_gefs.log\n",
    "    #PBS -e wget_gefs.err\n",
    "    \n",
    "    cd {}wget_GEFSv12_idx/\n",
    "    \n",
    "    '''.format(camp_dir)\n",
    "    \n",
    "    print(heads, file=f)\n",
    "    \n",
    "    for i in range(i_start, i_start+N, 1):\n",
    "            if i < L_fill:\n",
    "                dt = date_list_fill[i]\n",
    "                dt_str = datetime.strftime(dt, '%Y%m%d')\n",
    "                \n",
    "                for member in keys:\n",
    "                    for ff in key_leads:\n",
    "                        # download idx file\n",
    "                        download_link = 'https://noaa-gefs-pds.s3.amazonaws.com/gefs.{}/00/atmos/pgrb2sp25/{}.t00z.pgrb2s.0p25.f{}.idx'.format(\n",
    "                            dt_str, member, ff)\n",
    "                        \n",
    "                        save_name = '{}_00_{}.t00z.pgrb2s.0p25.f{}.idx'.format(dt_str, member, ff)                        \n",
    "                        print('wget -L -O {} {}'.format(save_name, download_link), file=f)\n",
    "                        \n",
    "                print('touch {}_done.txt'.format(dt_str), file=f)\n",
    "    \n",
    "    f.close()\n",
    "\n",
    "f = open(work_dir+'qsub/wget_gefs_all.sh', 'w')\n",
    "\n",
    "for i_, i_start in enumerate(index):\n",
    "    print('qsub wget_gefs_{:03d}.sh'.format(i_), file=f)\n",
    "    \n",
    "f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d2b5ad94-f56d-4819-a233-417e675d9da7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dt_str = ['20210120', '20210122', '20210227', '20210404', '20210511', '20210617', '20210724', '20210912', '20211006', '20211114', '20211220']\n",
    "# member = ['gep08', 'gep01', 'gep07', 'gep12', 'gep29', 'gep24', 'gep28', 'gep10', 'gep15', 'gep30', 'gep06']\n",
    "# ff = ['084', '012', '084', '060', '090', '036', '156', '114', '108', '114', '132']\n",
    "\n",
    "# for i in range(len(dt_str)):\n",
    "#     download_link = 'https://noaa-gefs-pds.s3.amazonaws.com/gefs.{}/00/atmos/pgrb2sp25/{}.t00z.pgrb2s.0p25.f{}.idx'.format(dt_str[i], member[i], ff[i])\n",
    "#     save_name = '{}_00_{}.t00z.pgrb2s.0p25.f{}.idx'.format(dt_str[i], member[i], ff[i])\n",
    "#     print('wget -L -O {} {}'.format(save_name, download_link))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "893af4bb-99cc-4005-9c7f-e0033e7769f4",
   "metadata": {},
   "source": [
    "## Download APCP based on the idx byte ranges"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7715fb3-7f1d-4e13-b2a9-f0db6d3afc27",
   "metadata": {},
   "source": [
    "```\n",
    "curl -H \"range: bytes=6598193-6836482\" -o test.grib https://noaa-gefs-pds.s3.amazonaws.com/gefs.20210629/00/atmos/pgrb2sp25/gec00.t00z.pgrb2s.0p25.f045\n",
    "```\n",
    "\n",
    "```python\n",
    "import pygrib\n",
    "\n",
    "with pygrib.open('test.grib') as grbio:\n",
    "    var_name = grbio[1]\n",
    "\n",
    "print(var_name)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "eb583d60-1f82-4bec-afa9-c1db83a5e56b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2e0965d1-1738-46a8-bf38-76a5f8a6402a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# check if files exist\n",
    "date_list_fill = []\n",
    "\n",
    "for dt in date_list:\n",
    "    \n",
    "    dt_str = datetime.strftime(dt, '%Y%m%d')\n",
    "    filename = camp_dir+'wget_GEFSv12_members/{}_apcp_done.txt'.format(dt_str)\n",
    "    \n",
    "    if os.path.isfile(filename):\n",
    "        continue;\n",
    "    else:\n",
    "        date_list_fill.append(dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2d148395-c622-40cf-a9ba-d3cdc6f3a368",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "need to download 31 files\n"
     ]
    }
   ],
   "source": [
    "L_fill = len(date_list_fill)\n",
    "print('need to download {} files'.format(L_fill))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d3b2c21f-c193-4e70-9e0e-f10ca54719c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_scripts = 3\n",
    "N = int(L_fill/N_scripts) + 1\n",
    "index = np.arange(0, L_fill, N)\n",
    "\n",
    "for i_, i_start in enumerate(index):\n",
    "\n",
    "    f = open(work_dir+'qsub/wget_gefs_apcp_{:03d}.sh'.format(i_), 'w') \n",
    "    \n",
    "    heads = '''\n",
    "    #!/bin/bash -l\n",
    "    \n",
    "    #PBS -N wget_gefs_idx\n",
    "    #PBS -A NAML0001\n",
    "    #PBS -l walltime=23:59:59\n",
    "    #PBS -l select=1:ncpus=4:mem=12GB\n",
    "    #PBS -q casper\n",
    "    #PBS -o wget_gefs.log\n",
    "    #PBS -e wget_gefs.err\n",
    "    \n",
    "    cd {}wget_GEFSv12_members/\n",
    "    \n",
    "    '''.format(camp_dir)\n",
    "    \n",
    "    print(heads, file=f)\n",
    "    \n",
    "    for i in range(i_start, i_start+N, 1):\n",
    "            if i < L_fill:\n",
    "                dt = date_list_fill[i]\n",
    "                dt_str = datetime.strftime(dt, '%Y%m%d')\n",
    "                \n",
    "                for member in keys:\n",
    "                    for ff in key_leads:\n",
    "\n",
    "                        # get the location of the GEFS idx file\n",
    "                        idx_name = '{}_00_{}.t00z.pgrb2s.0p25.f{}.idx'.format(dt_str, member, ff)  \n",
    "                        filename_idx = camp_dir+'wget_GEFSv12_idx/{}'.format(idx_name)\n",
    "                        \n",
    "                        # read the byte ranges from the idx file\n",
    "                        # no error handlings --> \n",
    "                        #     if the file is missing, the current bash cmd will fail and the next one will start\n",
    "                        APCP_line = 'XXX'\n",
    "                        APCP_line_after = 'XXX'\n",
    "                        try:\n",
    "                            f_lines = open(filename_idx).readlines()\n",
    "                            for i, line in enumerate(f_lines):\n",
    "                                if 'APCP' in line:\n",
    "                                    APCP_line = line\n",
    "                                    APCP_line_after = f_lines[i+1]\n",
    "                            byte_start = re.findall(r'\\d+', APCP_line)[1]\n",
    "                            byte_end = re.findall(r'\\d+', APCP_line_after)[1]\n",
    "                            \n",
    "                            # download idx file\n",
    "                            download_link = 'https://noaa-gefs-pds.s3.amazonaws.com/gefs.{}/00/atmos/pgrb2sp25/{}.t00z.pgrb2s.0p25.f{}'.format(\n",
    "                                dt_str, member, ff)\n",
    "                            \n",
    "                            save_name = '{}_00_{}.t00z.pgrb2s.0p25.f{}'.format(dt_str, member, ff)  \n",
    "                            \n",
    "                            print('curl -H \"range: bytes={}-{}\" -o {} {}'.format(byte_start, byte_end, save_name, download_link), file=f)\n",
    "                        except:\n",
    "                            print('Missing {}'.format(filename_idx))\n",
    "                            \n",
    "                print('touch {}_apcp_done.txt'.format(dt_str), file=f)\n",
    "    \n",
    "    f.close()\n",
    "\n",
    "f = open(work_dir+'qsub/wget_gefs_apcp_all.sh', 'w')\n",
    "\n",
    "for i_, i_start in enumerate(index):\n",
    "    print('qsub wget_gefs_apcp_{:03d}.sh'.format(i_), file=f)\n",
    "    \n",
    "f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "275bcc8c-5f35-4092-97d0-5c247def33c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pygrib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "282a4c16-d2ed-46eb-832d-64e3d18eb241",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1:Total Precipitation:kg m**-2 (accum):regular_ll:surface:level 0:fcst time 0-6 hrs (accum):from 202101010000:hi res cntl fcst\n"
     ]
    }
   ],
   "source": [
    "filename = '20210101_00_gec00.t00z.pgrb2s.0p25.f006'\n",
    "\n",
    "with pygrib.open(filename) as grbio:\n",
    "    #print(dir(grbio))\n",
    "    var_name = grbio[1]\n",
    "\n",
    "print(var_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "355c759f-2490-4c2a-962b-829172ca219d",
   "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
}