{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c255444e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import azarr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "877f84f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# see https://pangeo-forge.org/dashboard/feedstock/88\n",
    "url = 'https://ncsa.osn.xsede.org/Pangeo/pangeo-forge/pangeo-forge/EOBS-feedstock/eobs-surface-downwelling.zarr'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8d8648f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "z = azarr.open_consolidated(url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "14a398cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['latitude', 'longitude', 'qq', 'time']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "26eb51c1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<azarr.core.AArray '/latitude' (460,) float64 read-only>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z.latitude"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "579c853f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<coroutine object AArray._get_selection at 0x10e36eec0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z.latitude[:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8b444563",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.3319680690765381"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get two three arrays with a single context switch to async\n",
    "# and only one latency wait\n",
    "import asyncio\n",
    "import time\n",
    "\n",
    "t0 = time.time()\n",
    "lat, lot, t = await asyncio.gather(\n",
    "    z.latitude[:], \n",
    "    z.longitude[:],\n",
    "    z.time[:]\n",
    ")\n",
    "t1 = time.time()\n",
    "t1 - t0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b23cc442",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([25.04986065, 25.14986065, 25.24986065, 25.34986065, 25.44986065,\n",
       "       25.54986065, 25.64986065, 25.74986065, 25.84986065, 25.94986065,\n",
       "       26.04986065, 26.14986065, 26.24986065, 26.34986065, 26.44986065,\n",
       "       26.54986064, 26.64986064, 26.74986064, 26.84986064, 26.94986064,\n",
       "       27.04986064, 27.14986064, 27.24986064, 27.34986064, 27.44986064,\n",
       "       27.54986064, 27.64986064, 27.74986064, 27.84986064, 27.94986064,\n",
       "       28.04986064, 28.14986064, 28.24986064, 28.34986064, 28.44986064,\n",
       "       28.54986064, 28.64986064, 28.74986064, 28.84986064, 28.94986064,\n",
       "       29.04986063, 29.14986063, 29.24986063, 29.34986063, 29.44986063,\n",
       "       29.54986063, 29.64986063, 29.74986063, 29.84986063, 29.94986063,\n",
       "       30.04986063, 30.14986063, 30.24986063, 30.34986063, 30.44986063,\n",
       "       30.54986063, 30.64986063, 30.74986063, 30.84986063, 30.94986063,\n",
       "       31.04986063, 31.14986063, 31.24986063, 31.34986063, 31.44986063,\n",
       "       31.54986062, 31.64986062, 31.74986062, 31.84986062, 31.94986062,\n",
       "       32.04986062, 32.14986062, 32.24986062, 32.34986062, 32.44986062,\n",
       "       32.54986062, 32.64986062, 32.74986062, 32.84986062, 32.94986062,\n",
       "       33.04986062, 33.14986062, 33.24986062, 33.34986062, 33.44986062,\n",
       "       33.54986062, 33.64986062, 33.74986062, 33.84986062, 33.94986062,\n",
       "       34.04986061, 34.14986061, 34.24986061, 34.34986061, 34.44986061,\n",
       "       34.54986061, 34.64986061, 34.74986061, 34.84986061, 34.94986061])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lat[:100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "391dd82b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((25933, 460, 700), (40, 460, 700))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z.qq.shape, z.qq.chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8bf24946",
   "metadata": {},
   "outputs": [],
   "source": [
    "# or opposite ends of the same array\n",
    "slice1, slice2 = await asyncio.gather(\n",
    "    z.qq[:2, 200:250, 300:350],\n",
    "    z.qq[-1, 200:250, 300:350]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "cb6fccd2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[78., 77., 75., ..., 49., 46., 46.],\n",
       "       [76., 75., 74., ..., 53., 47., 46.],\n",
       "       [73., 73., 72., ..., 56., 48., 47.],\n",
       "       ...,\n",
       "       [59., 59., 58., ..., 24., 27., 30.],\n",
       "       [52., 52., 51., ..., 23., 26., 28.],\n",
       "       [45., 45., 44., ..., 22., 25., 27.]], dtype=float32)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "slice2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04912aa9",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}