xarray_demo.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4ea0e923-6158-49f2-ad09-adca7f5df8b1",
   "metadata": {},
   "source": [
    "# Simple Xarray demo\n",
    "This provides simple example of using Xarray. See https://docs.xarray.dev for more examples and detailed documentation.\n",
    "\n",
    "Choose **Help -> Show Contextual Help** to get interactive help."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0b0d24f-ebc6-453e-adb5-4ac79f2b86f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbc64b62-2d31-4b04-bf34-290a9dfd13da",
   "metadata": {},
   "source": [
    "## Manipulations requiring only NetCDF metadata access\n",
    "The commands in this section run quickly, as they use the metadata in the NetCDF files, not the full volume of numerical data (which is a much larger)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afd35038-e2ae-499a-8ddd-1d88e305bc08",
   "metadata": {},
   "source": [
    "### Create a virtual dataset by concatenating many files along time axis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "288d9016-008c-4343-ba8a-065c42f31457",
   "metadata": {},
   "source": [
    "Define the file paths to open. Note the wildcard (`*`) - this refers to 61 individual NetCDF files. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b76bd218-65c0-43d3-8b21-1aa6b957c582",
   "metadata": {},
   "outputs": [],
   "source": [
    "files = '/g/data/ik11/outputs/access-om2/1deg_jra55_iaf_omip2_cycle1/output*/ocean/ocean_month.nc'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cd98b21-ff6f-4b12-a72c-8d74deb678b8",
   "metadata": {},
   "source": [
    "Create an Xarray dataset by concatenating these files on the time axis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "685f1318-01db-447a-8316-47bebcb19019",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xr.open_mfdataset(files, parallel=True)  # concatenate many NetCDF files into one virtual dataset object"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "413f8ea2-2cea-4d8e-8b21-05817180df21",
   "metadata": {},
   "source": [
    "Display a representation of the resulting dataset. Click the little triangles to see details.\n",
    "Note that the dataset contains over 100 data variables. Click the icons to the right of each variable for metadata and information on the data each one contains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5405a68f-b02e-431d-9338-0ebd3216d94f",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37090af0-3d79-4074-b421-8e49f1c4dbcc",
   "metadata": {},
   "source": [
    "### Extract and subset temperature data\n",
    "Any variable can be extracted by name using this dot notation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cdaa455-1dfd-43a1-ae60-e99a098c43a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = ds.pot_temp  # select 4D potential temperature dataarray from dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c93d53f-7275-4750-81f7-ba66d595f975",
   "metadata": {},
   "source": [
    "Display a representation of the variable. Note that it is 4-dimensional, of size 732x50x300x360. This would be over 14GB, but hasn't been loaded into memory yet.\n",
    "\n",
    "From the `time` axis you can see it contains 732 monthly means (61 years, from 1958 to 2018).\n",
    "\n",
    "Note: The spatial dimensions are 50x300x360, as this is from a coarse 1° model. 0.1° model data is 135 times larger (spatially 75x2700x3600), so 61 years would be about 2000GB. This is more than can fit into a node's memory, but it could still be handled with this code because only the metadata is needed at this stage (the data itself doesn't need to be loaded)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a0f4935-0f2c-4aa4-8af1-6ee39c0b86e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "t"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb52b1ce-e176-454d-9c7d-4a67aaa4d298",
   "metadata": {},
   "source": [
    "Now we select a subset of the data, specified by calendar dates rather than indices. (The requested date range extends beyond the data range, so only a subset of these dates is returned.) Note that the memory requirement has dropped to about 2GB. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16dbc433-c4a9-4fee-a6a6-efb162b42dce",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = t.sel(time=slice('2010-01-01', '2022-01-01'))  # subset on time axis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "469aa0db-b314-441c-bcf1-65798d3810e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "t"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6782be2-76f3-4dcb-91a1-f79a5160e8e2",
   "metadata": {},
   "source": [
    "The data is global, but let's also restrict the spatial range to focus on the region near Australia. Again, we can specify this using physical units rather than grid indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffa1b5ec-aa0b-4dfa-ab2c-f645649b4b0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = t.sel(xt_ocean=slice(-270, -150)).sel(yt_ocean=slice(-60, 20))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "827dec64-2b12-4219-9823-57f581b15550",
   "metadata": {},
   "source": [
    "The data volume is further reduced"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9163dfeb-6738-4ca8-97c3-b0c889fa958c",
   "metadata": {},
   "outputs": [],
   "source": [
    "t"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d0717ed-74b4-4507-86d8-b388bcfd6a36",
   "metadata": {},
   "source": [
    "Let's now just select data from 50m depth. Note that the vertical grid doesn't have data at exactly that level (check this by clicking the cylindrical icon for `st_ocean` above), so we interpolate. Use interactive help to see what other interpolation methods are available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82bc7fd9-c0bb-4301-ab90-1e84f5a2f7c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = t.interp(st_ocean=50, method='linear')  # extract 2D interpolated data at 50m depth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0651eb63-e8f7-4538-a395-fb39e1b1fb5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "t"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2254dbe2-d3d3-4a60-b837-948c0f259f7d",
   "metadata": {},
   "source": [
    "We can get a monthly climatology (the avrage of every January, and the average of every February, etc etc) by using `groupby` and then `mean`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4fb86bf-4fe3-4be8-844f-58e55728a2cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = t.groupby('time.month').mean()  # monthly averages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "033f0eeb-8cad-4106-bf80-1ee1cc1855a6",
   "metadata": {},
   "source": [
    "now we have a `month` dimension instead of `time`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edc718a5-0047-4a80-a468-3538169aabec",
   "metadata": {},
   "outputs": [],
   "source": [
    "t"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "605ac542-b14d-4071-935f-f4ee34ee7ea1",
   "metadata": {},
   "source": [
    "We can then select the first month to get a January average, and convert from K to °C:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11890197-f2ce-44d6-9a64-30bf910a15bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = t.sel(month=1) - 273.15   # January mean, converted to °C"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0af01ea3-100a-48dc-a426-1cd9e1a89631",
   "metadata": {},
   "source": [
    "This is now a really small amount of data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2de0c5d-33a1-474d-b296-2df954a093e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "t"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2161d68-e7c4-4b6d-beeb-ab7793e5f497",
   "metadata": {},
   "source": [
    "## Plotting result: requires data access and calculation\n",
    "None of the manipulation and calculation we've specified so far (subsetting, interpolation, grouping, averaging) has actually taken place - this is **lazy evaluation**. But now we want to plot it, so the deferred calculation will need to be undertaken. But it's still quick, because only the required subset of data actually needs to be read.\n",
    "\n",
    "We're just doing a simple plot here, but there are [more sophisticated things you can do](https://docs.xarray.dev/en/latest/user-guide/plotting.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "710e493d-2b8d-4b61-8fea-9c050c00db37",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "t.plot()  # plot – this is when data access and calculation occur"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7681c447-b8ac-4c04-9083-712bc05fd55d",
   "metadata": {},
   "source": [
    "## Exercise\n",
    "Extract the salinity variable and plot the 1980 average in the Atlantic ocean at 100m depth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10d8ca6a-65a4-4710-bc81-7a11cfc6c540",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:analysis3-23.04] *",
   "language": "python",
   "name": "conda-env-analysis3-23.04-py"
  },
  "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.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}