.ipynb_checkpoints/MUR_zarr-checkpoint.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#code adapted from https://medium.com/pangeo/continuously-extending-zarr-datasets-c54fbad3967d\n",
    "#David Brochart\n",
    "\n",
    "import xarray as xr\n",
    "import os\n",
    "import zarr\n",
    "import os.path\n",
    "from os import path\n",
    "from datetime import datetime, timedelta\n",
    "import shutil\n",
    "import subprocess\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "def get_encoding(name):\n",
    "    '''Get encodings from a Zarr archive.\n",
    "    Arguments:\n",
    "        - name: the name of the archive.\n",
    "    Returns:\n",
    "        - encoding: the encodings of the variables.\n",
    "    '''\n",
    "    ds = xr.open_zarr(name)\n",
    "    encoding = {name: ds[name].encoding for name in list(ds.variables)}\n",
    "    return encoding\n",
    "\n",
    "def create_zarr(ds, name, encoding=None):\n",
    "    '''Create a Zarr archive from an Xarray Dataset.\n",
    "    Arguments:\n",
    "        - ds: the Dataset to store.\n",
    "        - name: the name of the Zarr archive.\n",
    "        - encoding: the encoding to use for each variable.\n",
    "    Returns:\n",
    "        - encoding: the encoding used for each variable.\n",
    "    '''\n",
    "    shutil.rmtree(name, ignore_errors=True)\n",
    "    ds = ds.chunk({name: ds[name].shape for name in list(ds.dims)})\n",
    "    ds.to_zarr(name, encoding=encoding)\n",
    "    if encoding is None:\n",
    "        encoding = get_encoding(name)\n",
    "    return encoding\n",
    "\n",
    "def empty_zarr(name, variable=None):\n",
    "    '''Empty the Zarr archive of its data (but not its metadata).\n",
    "    Arguments:\n",
    "        - name: the name of the archive.\n",
    "        - variable: the name of the variable to empty (if None, empty all\n",
    "          variables)\n",
    "    '''\n",
    "    for dname in [f for f in os.listdir(name) if not f.startswith('.')]:\n",
    "        if variable is not None and dname == variable:\n",
    "            for fname in [f for f in os.listdir(f'{name}/{dname}')\n",
    "                          if not f.startswith('.')]:\n",
    "                os.remove(f'{name}/{dname}/{fname}')\n",
    "\n",
    "def append_zarr(src_name, dst_name):\n",
    "    '''Append a Zarr archive to another one.\n",
    "    Arguments:\n",
    "        - src_name: the name of the archive to append.\n",
    "        - dst_name: the name of the archive to be appended to.\n",
    "    '''\n",
    "    zarr_src = zarr.open(src_name, mode='r')\n",
    "    zarr_dst = zarr.open(dst_name, mode='a')\n",
    "    for key in [k for k in zarr_src.array_keys() if k not in ['lat', 'lon']]:\n",
    "        zarr_dst[key].append(zarr_src[key])\n",
    "    empty_zarr('F:/data/sst/jpl_mur/zarr_bucket/mur_v41', 'time')\n",
    "\n",
    "def download_files(datetime_0, datetime_nb):\n",
    "    '''\n",
    "    Download files from FTP server.\n",
    "    Arguments:\n",
    "        - datetime_0: date from which to download.\n",
    "        - datetime_nb: number of dates (~files) to download.\n",
    "    Returns:\n",
    "        - filenames: list of file names to be downloaded.\n",
    "        - datetimes: list of dates corresponding to the downloaded files.\n",
    "    '''\n",
    "\n",
    "  # only downloads files not already local\n",
    "#downloads from mapped podaac drive\n",
    "    dir_local = 'F:/data/sst/jpl_mur/v4.1/'\n",
    "    dir_podaac = 'Z:/OceanTemperature/ghrsst/data/GDS2/L4/GLOB/JPL/MUR/v4.1/'\n",
    "    mur_end = '090000-JPL-L4_GHRSST-SSTfnd-MUR-GLOB-v02.0-fv04.1.nc'\n",
    "\n",
    "    datetimes = [datetime_0 + timedelta(hours=24*i) for i in range(datetime_nb)]\n",
    "    urls, filenames = [], []\n",
    "    for dt in datetimes:\n",
    "        year = str(dt.year)\n",
    "        month = str(dt.month).zfill(2)\n",
    "        day = str(dt.day).zfill(2)\n",
    "        doyear = str(dt.timetuple().tm_yday).zfill(3)\n",
    "        filename = dir_local+year+'/'+doyear+'/'+year+month+day+ mur_end\n",
    "        if path.exists(filename):\n",
    "            filenames.append(filename)\n",
    "        else:\n",
    "            #print('not local:',filename)\n",
    "            filename = dir_podaac+year+'/'+doyear+'/'+year+month+day+ mur_end\n",
    "            if path.exists(filename):\n",
    "                filenames.append(filename)\n",
    "    return filenames, datetimes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "time=datetime(2019, 6, 1, 12)\n",
    "filenames,datetimes=download_files(time,1)\n",
    "print(filenames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dt0 = dt = datetime(2002, 6, 1, 12) # from this date (included)\n",
    "dt1 = datetime(2002, 6, 3, 12)     # to that date (excluded)\n",
    "resume = False  # if True, resume a previous upload\n",
    "                # and dt0 and dt1 must be later than the previous date range\n",
    "fake_gcs = True # if True, won't upload to Google Cloud Storage\n",
    "                # but fake it in local trmm_bucket directory\n",
    "if fake_gcs:\n",
    "    store = 'F:/data/sst/jpl_mur/mur_bucket/'\n",
    "else:\n",
    "    store = gcsfs.GCSMap('pangeo-data/mur_v41')\n",
    "\n",
    "if resume:\n",
    "    time_prev = xr.open_zarr(store).time.values\n",
    "time_nb = 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "while dt < dt1:\n",
    "    #print(f'Downloading {time_nb} files from {dt}...')\n",
    "    filenames, datetimes = download_files(dt, time_nb)\n",
    "    ds = xr.open_mfdataset(filenames,combine='nested',concat_dim='time') #create_dataset(filenames, datetimes)\n",
    "    if not resume and dt == dt0:\n",
    "        encoding = create_zarr(ds, 'F:/data/sst/jpl_mur/mur_bucket/mur_v41')\n",
    "        empty_zarr('F:/data/sst/jpl_mur/mur_bucket/mur_v41', 'time')\n",
    "    else:\n",
    "        if resume:\n",
    "            encoding = get_encoding('F:/data/sst/jpl_mur/mur_bucket/mur_v41_new')\n",
    "        create_zarr(ds, 'F:/data/sst/jpl_mur/mur_bucket/mur_v41_new', encoding)\n",
    "        empty_zarr('F:/data/sst/jpl_mur/mur_bucket/mur_v41')\n",
    "        append_zarr('F:/data/sst/jpl_mur/mur_bucket/mur_v41_new', 'F:/data/sst/jpl_mur/mur_bucket/mur_v41')\n",
    "    print('Uploading...')\n",
    "    #if fake_gcs:\n",
    "    #    subprocess.check_call('mkdir -p F:/data/sst/jpl_mur/mur_bucket/; cp -r mur_v41/* '\n",
    "    #                          'F:/data/sst/jpl_mur/mur_bucket/; cp -r mur_v41/.[^.]* '\n",
    "    #                          'F:/data/sst/jpl_mur/mur_bucket/', shell=True)\n",
    "    #else:\n",
    "    #    subprocess.check_call('gsutil -m cp -r mur_v41/ gs://pangeo-data/'\n",
    "    #                          .split())\n",
    "    dt += timedelta(hours=24*time_nb)\n",
    "\n",
    "time_new = pd.date_range(dt0, dt1-timedelta(hours=24), freq='1D')\n",
    "if resume:\n",
    "    time_var = np.hstack((time_prev, time_new))\n",
    "else:\n",
    "    time_var = time_new\n",
    "time_ds = xr.DataArray(np.zeros(len(time_var)), coords=[time_var], dims=['time']).to_dataset(name='mur_time')\n",
    "shutil.rmtree('F:\\data\\sst\\jpl_mur\\mur_time', ignore_errors=True)\n",
    "time_ds.to_zarr('F:\\data\\sst\\jpl_mur\\mur_time')\n",
    "if fake_gcs:\n",
    "    subprocess.check_call('rm -rf F:/data/sst/jpl_mur/mur_bucket/time/'.split())\n",
    "    subprocess.check_call('cp -r F:/data/sst/jpl_mur/mur_time/time F:/data/sst/jpl_mur/mur_bucket/'.split())\n",
    "    #subprocess.check_call('rm -rf trmm_bucket/time'.split())\n",
    "    #subprocess.check_call('cp -r trmm_time/time trmm_bucket/'.split())\n",
    "else:\n",
    "    subprocess.check_call('gsutil -m rm -rf '\n",
    "                          'gs://pangeo-data/mur_v41/time'.split())\n",
    "    subprocess.check_call('gsutil -m cp -r mur_time/time/ '\n",
    "                          'gs://pangeo-data/mur_v41/'.split())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#    subprocess.check_call('rm -rf F:\\data\\sst\\jpl_mur\\zarr_bucket\\mur_time\\'.split())\n",
    "#    subprocess.check_call('cp -r F:/data/sst/jpl_mur/zarr_time/mur_time F:/data/sst/jpl_mur/zarr_bucket/'.split())\n",
    "subprocess.check_call('rmdir /Q /S F:/data/sst/jpl_mur/zarr_bucket/mur_time/')\n",
    "subprocess.check_call('cp -r F:/data/sst/jpl_mur/zarr_time/mur_time F:/data/sst/jpl_mur/zarr_bucket/'.split())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if fake_gcs:\n",
    "    subprocess.check_call('rm -rf trmm_bucket/time'.split())\n",
    "    subprocess.check_call('cp -r trmm_time/time trmm_bucket/'.split())\n",
    "else:\n",
    "    subprocess.check_call('gsutil -m rm -rf '\n",
    "                          'gs://pangeo-data/trmm_3b42rt/time'.split())\n",
    "    subprocess.check_call('gsutil -m cp -r trmm_time/time/ '\n",
    "                          'gs://pangeo-data/trmm_3b42rt/'.split())\n",
    "subprocess.check_call('rm -rf F:\\data\\sst\\jpl_mur\\zarr_bucket\\time'.split())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}