Added ROMS ocean model example notebook

azure-pipelines.yml

@@ -64,7 +64,7 @@ jobs:
   steps:
   - template: ci/azure/install.yml
   - bash: |
-      source activate test_env
+      source activate xarray-tests
       mypy . || exit 0
     displayName: mypy type checks
 
@@ -76,7 +76,7 @@ jobs:
     parameters:
       env_file: doc/environment.yml
   - bash: |
-      source activate test_env
+      source activate xarray-tests
       cd doc
       sphinx-build -n -j auto -b html -d _build/doctrees . _build/html
     displayName: Build HTML docs
@@ -89,6 +89,6 @@ jobs:
   steps:
   - template: ci/azure/install.yml
   - bash: |
-      source activate test_env
+      source activate xarray-tests
       pytest properties
     displayName: Property based tests

ci/azure/install.yml

@@ -6,11 +6,11 @@ steps:
 - template: add-conda-to-path.yml
 
 - bash: |
-    conda env create -n test_env --file ${{ parameters.env_file }}
+    conda env create -n xarray-tests --file ${{ parameters.env_file }}
   displayName: Install conda dependencies
 
 - bash: |
-    source activate test_env
+    source activate xarray-tests
     pip install -f https://7933911d6844c6c53a7d-47bd50c35cd79bd838daf386af554a83.ssl.cf2.rackcdn.com \
         --no-deps \
         --pre \
@@ -29,12 +29,12 @@ steps:
   displayName: Install upstream dev dependencies
 
 - bash: |
-    source activate test_env
+    source activate xarray-tests
     pip install --no-deps -e .
   displayName: Install xarray
 
 - bash: |
-    source activate test_env
+    source activate xarray-tests
     conda info -a
     conda list
     python xarray/util/print_versions.py

ci/azure/unit-tests.yml

@@ -3,14 +3,14 @@ steps:
 - template: install.yml
 
 - bash: |
-    source activate test_env
+    source activate xarray-tests
     python -OO -c "import xarray"
   displayName: Import xarray
 
 # Work around for allowed test failures:
 # https://github.com/microsoft/azure-pipelines-tasks/issues/9302
 - bash: |
-    source activate test_env
+    source activate xarray-tests
     pytest xarray \
     --junitxml=junit/test-results.xml \
     --cov=xarray \

ci/requirements/py35-min.yml

@@ -1,4 +1,4 @@
-name: test_env
+name: xarray-tests
 dependencies:
   - python=3.5.0
   - pytest

ci/requirements/py36-hypothesis.yml

@@ -1,4 +1,4 @@
-name: test_env
+name: xarray-tests
 channels:
   - conda-forge
 dependencies:

ci/requirements/py36.yml

@@ -1,4 +1,4 @@
-name: test_env
+name: xarray-tests
 channels:
   - conda-forge
 dependencies:

ci/requirements/py37-windows.yml

@@ -1,4 +1,4 @@
-name: test_env
+name: xarray-tests
 channels:
   - conda-forge
 dependencies:

ci/requirements/py37.yml

@@ -1,4 +1,4 @@
-name: test_env
+name: xarray-tests
 channels:
   - conda-forge
 dependencies:

doc/contributing.rst

@@ -152,10 +152,10 @@ We'll now kick off a two-step process:
 
    # Create and activate the build environment
    conda env create -f ci/requirements/py36.yml
-   conda activate test_env
+   conda activate xarray-tests
 
    # or with older versions of Anaconda:
-   source activate test_env
+   source activate xarray-tests
 
    # Build and install xarray
    pip install -e .

doc/examples.rst

@@ -8,3 +8,4 @@ Examples
     examples/monthly-means
     examples/multidimensional-coords
     examples/visualization_gallery
+    examples/ROMS_ocean_model

doc/examples/ROMS_ocean_model.ipynb

@@ -0,0 +1,226 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# ROMS Ocean Model Example"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The Regional Ocean Modeling System ([ROMS](http://myroms.org)) is an open source hydrodynamic model that is used for simulating currents and water properties in coastal and estuarine regions. ROMS is one of a few standard ocean models, and it has an active user community.\n",
+    "\n",
+    "ROMS uses a regular C-Grid in the horizontal, similar to other structured grid ocean and atmospheric models, and a stretched vertical coordinate (see [the ROMS documentation](https://www.myroms.org/wiki/Vertical_S-coordinate) for more details). Both of these require special treatment when using `xarray` to analyze ROMS ocean model output. This example notebook shows how to create a lazily evaluated vertical coordinate, and make some basic plots. The `xgcm` package is required to do analysis that is aware of the horizontal C-Grid."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import cartopy.crs as ccrs\n",
+    "import cartopy.feature as cfeature\n",
+    "import matplotlib.pyplot as plt\n",
+    "%matplotlib inline\n",
+    "\n",
+    "import xarray as xr"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Load a sample ROMS file. This is a subset of a full model available at \n",
+    "\n",
+    "    http://barataria.tamu.edu/thredds/catalog.html?dataset=txla_hindcast_agg\n",
+    "    \n",
+    "The subsetting was done using the following command on one of the output files:\n",
+    "\n",
+    "    #open dataset\n",
+    "    ds = xr.open_dataset('/d2/shared/TXLA_ROMS/output_20yr_obc/2001/ocean_his_0015.nc')\n",
+    "    \n",
+    "    # Turn on chunking to activate dask and parallelize read/write.\n",
+    "    ds = ds.chunk({'ocean_time': 1})\n",
+    "    \n",
+    "    # Pick out some of the variables that will be included as coordinates\n",
+    "    ds = ds.set_coords(['Cs_r', 'Cs_w', 'hc', 'h', 'Vtransform'])\n",
+    "    \n",
+    "    # Select a a subset of variables. Salt will be visualized, zeta is used to \n",
+    "    # calculate the vertical coordinate\n",
+    "    variables = ['salt', 'zeta']\n",
+    "    ds[variables].isel(ocean_time=slice(47, None, 7*24), \n",
+    "                       xi_rho=slice(300, None)).to_netcdf('ROMS_example.nc', mode='w')\n",
+    "\n",
+    "So, the `ROMS_example.nc` file contains a subset of the grid, one 3D variable, and two time steps."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Load in ROMS dataset as an xarray object"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# load in the file\n",
+    "ds = xr.tutorial.open_dataset('ROMS_example.nc', chunks={'ocean_time': 1})\n",
+    "\n",
+    "# This is a way to turn on chunking and lazy evaluation. Opening with mfdataset, or \n",
+    "# setting the chunking in the open_dataset would also achive this.\n",
+    "ds"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Add a lazilly calculated vertical coordinates\n",
+    "\n",
+    "Write equations to calculate the vertical coordinate. These will be only evaluated when data is requested. Information about the ROMS vertical coordinate can be found (here)[https://www.myroms.org/wiki/Vertical_S-coordinate]\n",
+    "\n",
+    "In short, for `Vtransform==2` as used in this example, \n",
+    "\n",
+    "$Z_0 = (h_c \\, S + h \\,C) / (h_c + h)$\n",
+    "\n",
+    "$z = Z_0 (\\zeta + h) + \\zeta$\n",
+    "\n",
+    "where the variables are defined as in the link above."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if ds.Vtransform == 1:\n",
+    "    Zo_rho = ds.hc * (ds.s_rho - ds.Cs_r) + ds.Cs_r * ds.h\n",
+    "    z_rho = Zo_rho + ds.zeta * (1 + Zo_rho/ds.h)\n",
+    "elif ds.Vtransform == 2:\n",
+    "    Zo_rho = (ds.hc * ds.s_rho + ds.Cs_r * ds.h) / (ds.hc + ds.h)\n",
+    "    z_rho = ds.zeta + (ds.zeta + ds.h) * Zo_rho\n",
+    "\n",
+    "ds.coords['z_rho'] = z_rho.transpose()   # needing transpose seems to be an xarray bug\n",
+    "ds.salt"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### A naive vertical slice\n",
+    "\n",
+    "Create a slice using the s-coordinate as the vertical dimension is typically not very informative."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "ds.salt.isel(xi_rho=50, ocean_time=0).plot()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can feed coordinate information to the plot method to give a more informative cross-section that uses the depths. Note that we did not need to slice the depth or longitude information separately, this was done automatically as the variable was sliced."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "section = ds.salt.isel(xi_rho=50, eta_rho=slice(0, 167), ocean_time=0)\n",
+    "section.plot(x='lon_rho', y='z_rho', figsize=(15, 6), clim=(25, 35))\n",
+    "plt.ylim([-100, 1]);"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### A plan view\n",
+    "\n",
+    "Now make a naive plan view, without any projection information, just using lon/lat as x/y. This looks OK, but will appear compressed because lon and lat do not have an aspect constrained by the projection."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds.salt.isel(s_rho=-1, ocean_time=0).plot(x='lon_rho', y='lat_rho')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "And let's use a projection to make it nicer, and add a coast."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "proj = ccrs.LambertConformal(central_longitude=-92, central_latitude=29)\n",
+    "fig = plt.figure(figsize=(15, 5))\n",
+    "ax = plt.axes(projection=proj)\n",
+    "ds.salt.isel(s_rho=-1, ocean_time=0).plot(x='lon_rho', y='lat_rho', \n",
+    "                                          transform=ccrs.PlateCarree())\n",
+    "\n",
+    "coast_10m = cfeature.NaturalEarthFeature('physical', 'land', '10m',\n",
+    "                                        edgecolor='k', facecolor='0.8')\n",
+    "ax.add_feature(coast_10m)"
+   ]
+  },
+  {
+   "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.6.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

doc/whats-new.rst

@@ -24,6 +24,11 @@ New functions/methods
 Enhancements
 ~~~~~~~~~~~~
 
+- Added example notebook demonstrating use of xarray with Regional Ocean
+  Modeling System (ROMS) ocean hydrodynamic model output.
+  (:issue:`3116`).
+  By `Robert Hetland <https://github.com/hetland>`
+
 Bug fixes
 ~~~~~~~~~
 

xarray/tutorial.py

@@ -33,8 +33,9 @@ def open_dataset(name, cache=True, cache_dir=_default_cache_dir,
     Parameters
     ----------
     name : str
-        Name of the netcdf file containing the dataset
-        ie. 'air_temperature'
+        Name of the file containing the dataset. If no suffix is given, assumed
+        to be netCDF ('.nc' is appended)
+        e.g. 'air_temperature'
     cache_dir : string, optional
         The directory in which to search for and write cached data.
     cache : boolean, optional
@@ -51,10 +52,13 @@ def open_dataset(name, cache=True, cache_dir=_default_cache_dir,
     xarray.open_dataset
 
     """
+    root, ext = _os.path.splitext(name)
+    if not ext:
+        ext = '.nc'
+    fullname = root + ext
     longdir = _os.path.expanduser(cache_dir)
-    fullname = name + '.nc'
     localfile = _os.sep.join((longdir, fullname))
-    md5name = name + '.md5'
+    md5name = fullname + '.md5'
     md5file = _os.sep.join((longdir, md5name))
 
     if not _os.path.exists(localfile):