Create Zarr files on Jetstream object store

create_zarr_jetstream.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!conda install --yes -c conda-forge zarr s3fs xarray netcdf4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Open a dataset from local disk, in this case NetCDF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xr.open_mfdataset('Sandy_ocean_his_nc4.nc',\n",
    "                       concat_dim='ensemble').persist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "368.837296"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds.nbytes / 1e6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test just storing it locally in a Zarr file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.backends.zarr.ZarrStore at 0x7f65cc3862b0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds.to_zarr(store='Sandy_ocean_his_nc42.zarr', mode='w')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "182M\tSandy_ocean_his_nc4.nc\n",
      "215M\tSandy_ocean_his_nc4.zarr\n",
      "215M\tSandy_ocean_his_nc42.zarr\n"
     ]
    }
   ],
   "source": [
    "!du -sh Sandy*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0m\u001b[01;34mAkk_bak\u001b[0m/        \u001b[01;34mgls_c2\u001b[0m/       \u001b[01;34mlon_v\u001b[0m/          \u001b[01;34mnRST\u001b[0m/        \u001b[01;34mtheta_s\u001b[0m/\r\n",
      "\u001b[01;34mAkp_bak\u001b[0m/        \u001b[01;34mgls_c3m\u001b[0m/      \u001b[01;34mLsshCLM\u001b[0m/        \u001b[01;34mntimes\u001b[0m/      \u001b[01;34mtke\u001b[0m/\r\n",
      "\u001b[01;34mAkt_bak\u001b[0m/        \u001b[01;34mgls_c3p\u001b[0m/      \u001b[01;34mLtracerCLM\u001b[0m/     \u001b[01;34mntsAVG\u001b[0m/      \u001b[01;34mTnudg\u001b[0m/\r\n",
      "\u001b[01;34mAKv\u001b[0m/            \u001b[01;34mgls_cmu0\u001b[0m/     \u001b[01;34mLtracerSponge\u001b[0m/  \u001b[01;34mocean_time\u001b[0m/  \u001b[01;34mTnudg_SSS\u001b[0m/\r\n",
      "\u001b[01;34mAkv_bak\u001b[0m/        \u001b[01;34mgls_Kmin\u001b[0m/     \u001b[01;34mLtracerSrc\u001b[0m/     \u001b[01;34momega\u001b[0m/       \u001b[01;34mTobc_in\u001b[0m/\r\n",
      "\u001b[01;34mangle\u001b[0m/          \u001b[01;34mgls_m\u001b[0m/        \u001b[01;34mLuvSponge\u001b[0m/      \u001b[01;34mPair\u001b[0m/        \u001b[01;34mTobc_out\u001b[0m/\r\n",
      "\u001b[01;34mbustr\u001b[0m/          \u001b[01;34mgls_n\u001b[0m/        \u001b[01;34mLuvSrc\u001b[0m/         \u001b[01;34mpm\u001b[0m/          \u001b[01;34mu\u001b[0m/\r\n",
      "\u001b[01;34mbvstr\u001b[0m/          \u001b[01;34mgls_p\u001b[0m/        \u001b[01;34mLwave\u001b[0m/          \u001b[01;34mpn\u001b[0m/          \u001b[01;34mubar\u001b[0m/\r\n",
      "\u001b[01;34mCharnok_alpha\u001b[0m/  \u001b[01;34mgls_Pmin\u001b[0m/     \u001b[01;34mLwavep\u001b[0m/         \u001b[01;34mPwave_bot\u001b[0m/   \u001b[01;34mUwave_rms\u001b[0m/\r\n",
      "\u001b[01;34mCrgBan_cw\u001b[0m/      \u001b[01;34mgls_sigk\u001b[0m/     \u001b[01;34mlwrad\u001b[0m/          \u001b[01;34mPwave_top\u001b[0m/   \u001b[01;34mUwind\u001b[0m/\r\n",
      "\u001b[01;34mCs_r\u001b[0m/           \u001b[01;34mgls_sigp\u001b[0m/     \u001b[01;34mLwSrc\u001b[0m/          \u001b[01;34mrain\u001b[0m/        \u001b[01;34mUwind_eastward\u001b[0m/\r\n",
      "\u001b[01;34mCs_w\u001b[0m/           \u001b[01;34mgrid\u001b[0m/         \u001b[01;34mM2nudg\u001b[0m/         \u001b[01;34mrdrg\u001b[0m/        \u001b[01;34mv\u001b[0m/\r\n",
      "\u001b[01;34mdstart\u001b[0m/         \u001b[01;34mh\u001b[0m/            \u001b[01;34mM2obc_in\u001b[0m/       \u001b[01;34mrdrg2\u001b[0m/       \u001b[01;34mvbar\u001b[0m/\r\n",
      "\u001b[01;34mdt\u001b[0m/             \u001b[01;34mhc\u001b[0m/           \u001b[01;34mM2obc_out\u001b[0m/      \u001b[01;34mrho0\u001b[0m/        \u001b[01;34mVstretching\u001b[0m/\r\n",
      "\u001b[01;34mdtfast\u001b[0m/         \u001b[01;34mHwave\u001b[0m/        \u001b[01;34mM3nudg\u001b[0m/         \u001b[01;34msalt\u001b[0m/        \u001b[01;34mVtransform\u001b[0m/\r\n",
      "\u001b[01;34mDwave\u001b[0m/          \u001b[01;34mlatent\u001b[0m/       \u001b[01;34mM3obc_in\u001b[0m/       \u001b[01;34msensible\u001b[0m/    \u001b[01;34mVwind\u001b[0m/\r\n",
      "\u001b[01;34mel\u001b[0m/             \u001b[01;34mlat_psi\u001b[0m/      \u001b[01;34mM3obc_out\u001b[0m/      \u001b[01;34mshflux\u001b[0m/      \u001b[01;34mVwind_northward\u001b[0m/\r\n",
      "\u001b[01;34mEminusP\u001b[0m/        \u001b[01;34mlat_rho\u001b[0m/      \u001b[01;34mmask_psi\u001b[0m/       \u001b[01;34mspherical\u001b[0m/   \u001b[01;34mw\u001b[0m/\r\n",
      "\u001b[01;34mevaporation\u001b[0m/    \u001b[01;34mlat_u\u001b[0m/        \u001b[01;34mmask_rho\u001b[0m/       \u001b[01;34ms_rho\u001b[0m/       \u001b[01;34mxl\u001b[0m/\r\n",
      "\u001b[01;34mf\u001b[0m/              \u001b[01;34mlat_v\u001b[0m/        \u001b[01;34mmask_u\u001b[0m/         \u001b[01;34mssflux\u001b[0m/      \u001b[01;34mzeta\u001b[0m/\r\n",
      "\u001b[01;34mFalpha\u001b[0m/         \u001b[01;34mLm2CLM\u001b[0m/       \u001b[01;34mmask_v\u001b[0m/         \u001b[01;34msustr\u001b[0m/       \u001b[01;34mZnudg\u001b[0m/\r\n",
      "\u001b[01;34mFbeta\u001b[0m/          \u001b[01;34mLm3CLM\u001b[0m/       \u001b[01;34mnAVG\u001b[0m/           \u001b[01;34msvstr\u001b[0m/       \u001b[01;34mZob\u001b[0m/\r\n",
      "\u001b[01;34mFgamma\u001b[0m/         \u001b[01;34mLnudgeM2CLM\u001b[0m/  \u001b[01;34mndefAVG\u001b[0m/        \u001b[01;34ms_w\u001b[0m/         \u001b[01;34mZos\u001b[0m/\r\n",
      "\u001b[01;34mFSobc_in\u001b[0m/       \u001b[01;34mLnudgeM3CLM\u001b[0m/  \u001b[01;34mndefHIS\u001b[0m/        \u001b[01;34mswrad\u001b[0m/       \u001b[01;34mZos_hsig_alpha\u001b[0m/\r\n",
      "\u001b[01;34mFSobc_out\u001b[0m/      \u001b[01;34mLnudgeTCLM\u001b[0m/   \u001b[01;34mndtfast\u001b[0m/        \u001b[01;34msz_alpha\u001b[0m/\r\n",
      "\u001b[01;34mgamma2\u001b[0m/         \u001b[01;34mlon_psi\u001b[0m/      \u001b[01;34mnHIS\u001b[0m/           \u001b[01;34mTcline\u001b[0m/\r\n",
      "\u001b[01;34mgls\u001b[0m/            \u001b[01;34mlon_rho\u001b[0m/      \u001b[01;34mnl_tnu2\u001b[0m/        \u001b[01;34mtemp\u001b[0m/\r\n",
      "\u001b[01;34mgls_c1\u001b[0m/         \u001b[01;34mlon_u\u001b[0m/        \u001b[01;34mnl_visc2\u001b[0m/       \u001b[01;34mtheta_b\u001b[0m/\r\n"
     ]
    }
   ],
   "source": [
    "%ls Sandy*4.zarr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Write to object store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import s3fs\n",
    "fs = s3fs.S3FileSystem(client_kwargs=dict(endpoint_url=\"JETSTREAM_SWIFT_ENDPOINT\"))\n",
    "d = s3fs.S3Map(\"sandy_zarr2\", s3=fs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.backends.zarr.ZarrStore at 0x7f65ab0e2cc0>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds.to_zarr(store=d, mode='w')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['sandy_zarr2/.zattrs', 'sandy_zarr2/.zgroup', 'sandy_zarr2/AKv', 'sandy_zarr2/Akk_bak', 'sandy_zarr2/Akp_bak', 'sandy_zarr2/Akt_bak', 'sandy_zarr2/Akv_bak', 'sandy_zarr2/Charnok_alpha', 'sandy_zarr2/CrgBan_cw', 'sandy_zarr2/Cs_r', 'sandy_zarr2/Cs_w', 'sandy_zarr2/Dwave', 'sandy_zarr2/EminusP', 'sandy_zarr2/FSobc_in', 'sandy_zarr2/FSobc_out', 'sandy_zarr2/Falpha', 'sandy_zarr2/Fbeta', 'sandy_zarr2/Fgamma', 'sandy_zarr2/Hwave', 'sandy_zarr2/Lm2CLM', 'sandy_zarr2/Lm3CLM', 'sandy_zarr2/LnudgeM2CLM', 'sandy_zarr2/LnudgeM3CLM', 'sandy_zarr2/LnudgeTCLM', 'sandy_zarr2/LsshCLM', 'sandy_zarr2/LtracerCLM', 'sandy_zarr2/LtracerSponge', 'sandy_zarr2/LtracerSrc', 'sandy_zarr2/LuvSponge', 'sandy_zarr2/LuvSrc', 'sandy_zarr2/LwSrc', 'sandy_zarr2/Lwave', 'sandy_zarr2/Lwavep', 'sandy_zarr2/M2nudg', 'sandy_zarr2/M2obc_in', 'sandy_zarr2/M2obc_out', 'sandy_zarr2/M3nudg', 'sandy_zarr2/M3obc_in', 'sandy_zarr2/M3obc_out', 'sandy_zarr2/Pair', 'sandy_zarr2/Pwave_bot', 'sandy_zarr2/Pwave_top', 'sandy_zarr2/Tcline', 'sandy_zarr2/Tnudg', 'sandy_zarr2/Tnudg_SSS', 'sandy_zarr2/Tobc_in', 'sandy_zarr2/Tobc_out', 'sandy_zarr2/Uwave_rms', 'sandy_zarr2/Uwind', 'sandy_zarr2/Uwind_eastward', 'sandy_zarr2/Vstretching', 'sandy_zarr2/Vtransform', 'sandy_zarr2/Vwind', 'sandy_zarr2/Vwind_northward', 'sandy_zarr2/Znudg', 'sandy_zarr2/Zob', 'sandy_zarr2/Zos', 'sandy_zarr2/Zos_hsig_alpha', 'sandy_zarr2/angle', 'sandy_zarr2/bustr', 'sandy_zarr2/bvstr', 'sandy_zarr2/dstart', 'sandy_zarr2/dt', 'sandy_zarr2/dtfast', 'sandy_zarr2/el', 'sandy_zarr2/evaporation', 'sandy_zarr2/f', 'sandy_zarr2/gamma2', 'sandy_zarr2/gls', 'sandy_zarr2/gls_Kmin', 'sandy_zarr2/gls_Pmin', 'sandy_zarr2/gls_c1', 'sandy_zarr2/gls_c2', 'sandy_zarr2/gls_c3m', 'sandy_zarr2/gls_c3p', 'sandy_zarr2/gls_cmu0', 'sandy_zarr2/gls_m', 'sandy_zarr2/gls_n', 'sandy_zarr2/gls_p', 'sandy_zarr2/gls_sigk', 'sandy_zarr2/gls_sigp', 'sandy_zarr2/grid', 'sandy_zarr2/h', 'sandy_zarr2/hc', 'sandy_zarr2/lat_psi', 'sandy_zarr2/lat_rho', 'sandy_zarr2/lat_u', 'sandy_zarr2/lat_v', 'sandy_zarr2/latent', 'sandy_zarr2/lon_psi', 'sandy_zarr2/lon_rho', 'sandy_zarr2/lon_u', 'sandy_zarr2/lon_v', 'sandy_zarr2/lwrad', 'sandy_zarr2/mask_psi', 'sandy_zarr2/mask_rho', 'sandy_zarr2/mask_u', 'sandy_zarr2/mask_v', 'sandy_zarr2/nAVG', 'sandy_zarr2/nHIS', 'sandy_zarr2/nRST', 'sandy_zarr2/ndefAVG', 'sandy_zarr2/ndefHIS', 'sandy_zarr2/ndtfast', 'sandy_zarr2/nl_tnu2', 'sandy_zarr2/nl_visc2', 'sandy_zarr2/ntimes', 'sandy_zarr2/ntsAVG', 'sandy_zarr2/ocean_time', 'sandy_zarr2/omega', 'sandy_zarr2/pm', 'sandy_zarr2/pn', 'sandy_zarr2/rain', 'sandy_zarr2/rdrg', 'sandy_zarr2/rdrg2', 'sandy_zarr2/rho0', 'sandy_zarr2/s_rho', 'sandy_zarr2/s_w', 'sandy_zarr2/salt', 'sandy_zarr2/sensible', 'sandy_zarr2/shflux', 'sandy_zarr2/spherical', 'sandy_zarr2/ssflux', 'sandy_zarr2/sustr', 'sandy_zarr2/svstr', 'sandy_zarr2/swrad', 'sandy_zarr2/sz_alpha', 'sandy_zarr2/temp', 'sandy_zarr2/theta_b', 'sandy_zarr2/theta_s', 'sandy_zarr2/tke', 'sandy_zarr2/u', 'sandy_zarr2/ubar', 'sandy_zarr2/v', 'sandy_zarr2/vbar', 'sandy_zarr2/w', 'sandy_zarr2/xl', 'sandy_zarr2/zeta']\n"
     ]
    }
   ],
   "source": [
    "print(fs.ls(\"sandy_zarr2\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [default]",
   "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}