Feature/weighted

doc/api.rst

@@ -165,6 +165,7 @@ Computation
    Dataset.groupby_bins
    Dataset.rolling
    Dataset.rolling_exp
+   Dataset.weighted
    Dataset.coarsen
    Dataset.resample
    Dataset.diff
@@ -340,6 +341,7 @@ Computation
    DataArray.groupby_bins
    DataArray.rolling
    DataArray.rolling_exp
+   DataArray.weighted
    DataArray.coarsen
    DataArray.dt
    DataArray.resample

doc/computation.rst

@@ -243,6 +243,69 @@ You can also use ``construct`` to compute a weighted rolling sum:
 
 .. _comput.coarsen:
 
+Weighted array reductions
+=========================
+
+``DataArray`` and ``Dataset`` objects include :py:meth:`~xarray.DataArray.weighted`
+and :py:meth:`~xarray.Dataset.weighted` array reduction methods. They currently
+support weighted ``sum`` and weighted ``mean``.
+
+.. ipython:: python
+
+  coords = dict(month=('month', [1, 2, 3]))
+
+  prec = xr.DataArray([1.1, 1.0, 0.9], dims=('month', ), coords=coords)
+  weights = xr.DataArray([31, 28, 31], dims=('month', ), coords=coords)
+
+Create a weighted object:
+
+.. ipython:: python
+
+  weighted_prec = prec.weighted(weights)
+  weighted_prec
+
+Calculate the weighted sum:
+
+.. ipython:: python
+
+  weighted_prec.sum()
+
+Calculate the weighted mean:
+
+.. ipython:: python
+
+        weighted_prec.mean(dim="month")
+
+The weighted sum corresponds to:
+
+.. ipython:: python
+
+  weighted_sum = (prec * weights).sum()
+  weighted_sum
+
+and the weighted mean to:
+
+.. ipython:: python
+
+  weighted_mean = weighted_sum / weights.sum()
+  weighted_mean
+
+However, the functions also take missing values in the data into account:
+
+.. ipython:: python
+
+  data = xr.DataArray([np.NaN, 2, 4])
+  weights = xr.DataArray([8, 1, 1])
+
+  data.weighted(weights).mean()
+
+Using ``(data * weights).sum() / weights.sum()`` would (incorrectly) result
+in 0.6.
+
+.. note::
+  ``weights`` must be a ``DataArray`` and cannot contain missing values.
+  Missing values can be replaced manually by `weights.fillna(0)`.
+
 Coarsen large arrays
 ====================
 

doc/examples.rst

@@ -6,6 +6,7 @@ Examples
 
     examples/weather-data
     examples/monthly-means
+    examples/area_weighted_temperature
     examples/multidimensional-coords
     examples/visualization_gallery
     examples/ROMS_ocean_model

doc/examples/area_weighted_temperature.ipynb

@@ -0,0 +1,163 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Compare weighted and unweighted mean temperature\n",
+    "\n",
+    "\n",
+    "Author: [Mathias Hauser](https://github.com/mathause/)\n",
+    "\n",
+    "The data used for this example can be found in the [xarray-data](https://github.com/pydata/xarray-data) repository. You may need to change the path to `air_temperature` below.\n",
+    "\n",
+    "We use the air_temperature example dataset to calculate the area-weighted temperature over its domain. This dataset has a regular latitude/ longitude grid, thus the gridcell area decreases towards the pole. For this grid we can use the cosine of the latitude as proxy for the grid cell area.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib inline\n",
+    "\n",
+    "import cartopy.crs as ccrs\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "\n",
+    "import xarray as xr"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Data\n",
+    "\n",
+    "Load the data, convert to celsius, and resample to daily values"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds = xr.tutorial.load_dataset(\"air_temperature\")\n",
+    "\n",
+    "# to celsius\n",
+    "air = ds.air - 273.15\n",
+    "\n",
+    "# resample from 6-hourly to daily values\n",
+    "air = air.resample(time=\"D\").mean()\n",
+    "\n",
+    "air"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Plot the first timestep:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "projection = ccrs.LambertConformal(central_longitude=-95, central_latitude=45)\n",
+    "\n",
+    "f, ax = plt.subplots(subplot_kw=dict(projection=projection))\n",
+    "\n",
+    "air.isel(time=0).plot(transform=ccrs.PlateCarree(), cbar_kwargs=dict(shrink=0.7))\n",
+    "ax.coastlines()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Creating weights\n",
+    "\n",
+    "For a for a rectangular grid the cosine of the latitude is proportional to the grid cell area."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "weights = np.cos(np.deg2rad(air.lat))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Weighted mean"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "air_weighted = air.weighted(weights).mean((\"lon\", \"lat\"))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Plot: comparison with unweighted mean\n",
+    "\n",
+    "Note how the weighted mean temperature is higher than the unweighted."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "air_weighted.plot(label=\"weighted\")\n",
+    "air.mean((\"lon\", \"lat\")).plot(label=\"unweighted\")\n",
+    "\n",
+    "plt.legend()"
+   ]
+  },
+  {
+   "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.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

doc/whats-new.rst

@@ -25,6 +25,9 @@ Breaking changes
 New Features
 ~~~~~~~~~~~~
 
+- Weighted array reductions are now supported via the new :py:meth:`DataArray.weighted`
+  and :py:meth:`Dataset.weighted` methods. By `Mathias Hauser <https://github.com/mathause>`_
+  (:issue:`422`).
 - Added support for :py:class:`pandas.DatetimeIndex`-style rounding of
   ``cftime.datetime`` objects directly via a :py:class:`CFTimeIndex` or via the
   :py:class:`~core.accessor_dt.DatetimeAccessor`.

xarray/core/common.py

@@ -748,6 +748,25 @@ def groupby_bins(
             },
         )
 
+    def weighted(self, weights):
+        """
+        Weighted operations.
+
+        Parameters
+        ----------
+        weights : DataArray
+            An array of weights associated with the values in this Dataset.
+            Each value in the data contributes to the reduction operation
+            according to its associated weight.
+
+        Note
+        ----
+        ``weights`` must be a ``DataArray`` and cannot contain missing values.
+        Missing values can be replaced by ``weights.fillna(0)``.
+        """
+
+        return self._weighted_cls(self, weights)
+
     def rolling(
         self,
         dim: Mapping[Hashable, int] = None,

xarray/core/dataarray.py

@@ -33,6 +33,7 @@
     resample,
     rolling,
     utils,
+    weighted,
 )
 from .accessor_dt import CombinedDatetimelikeAccessor
 from .accessor_str import StringAccessor
@@ -258,6 +259,7 @@ class DataArray(AbstractArray, DataWithCoords):
     _rolling_cls = rolling.DataArrayRolling
     _coarsen_cls = rolling.DataArrayCoarsen
     _resample_cls = resample.DataArrayResample
+    _weighted_cls = weighted.DataArrayWeighted
 
     dt = property(CombinedDatetimelikeAccessor)
 

xarray/core/dataset.py

@@ -46,6 +46,7 @@
     resample,
     rolling,
     utils,
+    weighted,
 )
 from .alignment import _broadcast_helper, _get_broadcast_dims_map_common_coords, align
 from .common import (
@@ -457,6 +458,7 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
     _rolling_cls = rolling.DatasetRolling
     _coarsen_cls = rolling.DatasetCoarsen
     _resample_cls = resample.DatasetResample
+    _weighted_cls = weighted.DatasetWeighted
 
     def __init__(
         self,