notebook showing subsetting and resampling of ERA5

era5_resample.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "016999eb-6c63-48cc-b637-8012de56f30c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U flox xarray obstore dask distributed zarr\n",
    "import flox \n",
    "import xarray as xr\n",
    "from distributed import Client \n",
    "from zarr.storage import ObjectStore\n",
    "import obstore as obs\n",
    "from obstore.store import from_url\n",
    "import zarr \n",
    "import xarray as xr\n",
    "import icechunk\n",
    "from icechunk.xarray import to_icechunk\n",
    "from dask.diagnostics import ProgressBar\n",
    "zarr.config.set({\"async.concurrency\": 128})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02d1d856-55b4-4d94-b9b8-ea2835479945",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(n_workers=32)\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "002bda4d-7850-4e79-b805-e16c5923c510",
   "metadata": {},
   "outputs": [],
   "source": [
    "url = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3'\n",
    "\n",
    "store = from_url(url)\n",
    "zarr_store = ObjectStore(store, read_only=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00990da1-8d74-42a2-8496-8825be61b00a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "ds = xr.open_dataset(\n",
    "    zarr_store,\n",
    "    engine=\"zarr\",\n",
    "    chunks={}\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22bce642-329f-4242-ad25-2109bebd32cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c437ec3d-5f19-4af8-93c0-cd6839803858",
   "metadata": {},
   "outputs": [],
   "source": [
    "# My data preprocessing\n",
    "# Select variables from ERA 5\n",
    "list_vars = [\"10m_u_component_of_wind\", \"10m_v_component_of_wind\"]\n",
    "ds_sel = ds[list_vars]\n",
    "# Select timelapse\n",
    "date_start, date_end = \"1983-01-01\", \"2024-12-31\"\n",
    "# date_start, date_end = \"1983-01-01\", \"1984-12-31\"\n",
    "\n",
    "\n",
    "ds_time = ds_sel.sel(time=slice(date_start, date_end))\n",
    "# Resample from hourly to daily\n",
    "latmin, latmax = -15, -45\n",
    "lonmin, lonmax = 115, 145\n",
    "ds_sel = ds_time.sel( \n",
    "        latitude = slice(latmin, latmax),\n",
    "        longitude = slice(lonmin, lonmax)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9b5e829-0c36-4b46-a54b-6903ab5e4164",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_res = ds_sel.resample(time=\"1D\").mean()\n",
    "ds_res\n",
    "\n",
    "# chunking is: (1, 721, 1440) is that OK for analysis?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2805cfda-b367-407d-8653-ae8d85bfe426",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_res = ds_res.chunk({'time':730,'longitude':61,'latitude':61})\n",
    "ds_res\n",
    "# ~ 10 MB chunks, ~2 years, 4 spatial chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "611e6b47-2ebf-467e-bd59-cbc2c40c1432",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14b645ef-2cc3-4a2a-b822-83aa05dcc97a",
   "metadata": {},
   "outputs": [],
   "source": [
    "storage = icechunk.gcs_storage(\n",
    "    bucket=\"leap-persistent\",\n",
    "    prefix=\"data-library/era5_resampled/10m-u-v-wind_1983-2024_spatial_subset.icechunk\",\n",
    "    from_env=True\n",
    ")\n",
    "repo = icechunk.Repository.open_or_create(storage)\n",
    "\n",
    "session = repo.writable_session(\"main\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b23637f-756e-486d-9a62-3b2477cb438a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# ~ 2 hrs and 10 mins with 4 workers\n",
    "with ProgressBar():\n",
    "    to_icechunk(ds_res.drop_encoding(), session)\n",
    "    session.commit(\"resample_era5\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c1997a4-3bfa-4968-bcfe-ca4db0a0666b",
   "metadata": {},
   "outputs": [],
   "source": [
    "session_ro = repo.readonly_session(\"main\")\n",
    "rtds = xr.open_zarr(session_ro.store)\n",
    "rtds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0086348a-b745-4b6a-ac34-56dfe5f23190",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot map\n",
    "rtds['10m_u_component_of_wind'].isel(time=0).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a6aff6a-f847-4ca3-beab-fa181b49b127",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot time series\n",
    "rtds['10m_u_component_of_wind'].sel(latitude=-25, longitude=130).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e057fd5a-a134-44db-b746-2e78e8e88970",
   "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.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}