Skip to content

Instantly share code, notes, and snippets.

@dennissergeev
Last active May 10, 2024 10:00
Show Gist options
  • Save dennissergeev/60bf7b03443f1b2c8eb96ce0b1880150 to your computer and use it in GitHub Desktop.
Save dennissergeev/60bf7b03443f1b2c8eb96ce0b1880150 to your computer and use it in GitHub Desktop.
Calculate a spatial mean of xarray's DataArray
# -*- coding: utf-8 -*-
"""Operations on cartesian geographical grid."""
import numpy as np
EARTH_RADIUS = 6371000.0 # m
def _guess_bounds(points, bound_position=0.5):
"""
Guess bounds of grid cells.
Simplified function from iris.coord.Coord.
Parameters
----------
points: numpy.array
Array of grid points of shape (N,).
bound_position: float, optional
Bounds offset relative to the grid cell centre.
Returns
-------
Array of shape (N, 2).
"""
diffs = np.diff(points)
diffs = np.insert(diffs, 0, diffs[0])
diffs = np.append(diffs, diffs[-1])
min_bounds = points - diffs[:-1] * bound_position
max_bounds = points + diffs[1:] * (1 - bound_position)
return np.array([min_bounds, max_bounds]).transpose()
def _quadrant_area(radian_lat_bounds, radian_lon_bounds, radius_of_earth):
"""
Calculate spherical segment areas.
Taken from SciTools iris library.
Area weights are calculated for each lat/lon cell as:
.. math::
r^2 (lon_1 - lon_0) ( sin(lat_1) - sin(lat_0))
The resulting array will have a shape of
*(radian_lat_bounds.shape[0], radian_lon_bounds.shape[0])*
The calculations are done at 64 bit precision and the returned array
will be of type numpy.float64.
Parameters
----------
radian_lat_bounds: numpy.array
Array of latitude bounds (radians) of shape (M, 2)
radian_lon_bounds: numpy.array
Array of longitude bounds (radians) of shape (N, 2)
radius_of_earth: float
Radius of the Earth (currently assumed spherical)
Returns
-------
Array of grid cell areas of shape (M, N).
"""
# ensure pairs of bounds
if (
radian_lat_bounds.shape[-1] != 2
or radian_lon_bounds.shape[-1] != 2
or radian_lat_bounds.ndim != 2
or radian_lon_bounds.ndim != 2
):
raise ValueError("Bounds must be [n,2] array")
# fill in a new array of areas
radius_sqr = radius_of_earth ** 2
radian_lat_64 = radian_lat_bounds.astype(np.float64)
radian_lon_64 = radian_lon_bounds.astype(np.float64)
ylen = np.sin(radian_lat_64[:, 1]) - np.sin(radian_lat_64[:, 0])
xlen = radian_lon_64[:, 1] - radian_lon_64[:, 0]
areas = radius_sqr * np.outer(ylen, xlen)
# we use abs because backwards bounds (min > max) give negative areas.
return np.abs(areas)
def grid_cell_areas(lon1d, lat1d, radius=EARTH_RADIUS):
"""
Calculate grid cell areas given 1D arrays of longitudes and latitudes
for a planet with the given radius.
Parameters
----------
lon1d: numpy.array
Array of longitude points [degrees] of shape (M,)
lat1d: numpy.array
Array of latitude points [degrees] of shape (M,)
radius: float, optional
Radius of the planet [metres] (currently assumed spherical)
Returns
-------
Array of grid cell areas [metres**2] of shape (M, N).
"""
lon_bounds_radian = np.deg2rad(_guess_bounds(lon1d))
lat_bounds_radian = np.deg2rad(_guess_bounds(lat1d))
area = _quadrant_area(lat_bounds_radian, lon_bounds_radian, radius)
return area
def calc_spatial_mean(
xr_da, lon_name="longitude", lat_name="latitude", radius=EARTH_RADIUS
):
"""
Calculate spatial mean of xarray.DataArray with grid cell weighting.
Parameters
----------
xr_da: xarray.DataArray
Data to average
lon_name: str, optional
Name of x-coordinate
lat_name: str, optional
Name of y-coordinate
radius: float
Radius of the planet [metres], currently assumed spherical (not important anyway)
Returns
-------
Spatially averaged xarray.DataArray.
"""
lon = xr_da[lon_name].values
lat = xr_da[lat_name].values
area_weights = grid_cell_areas(lon, lat, radius=radius)
aw_factor = area_weights / area_weights.max()
return (xr_da * aw_factor).mean(dim=[lon_name, lat_name])
def calc_spatial_integral(
xr_da, lon_name="longitude", lat_name="latitude", radius=EARTH_RADIUS
):
"""
Calculate spatial integral of xarray.DataArray with grid cell weighting.
Parameters
----------
xr_da: xarray.DataArray
Data to average
lon_name: str, optional
Name of x-coordinate
lat_name: str, optional
Name of y-coordinate
radius: float
Radius of the planet [metres], currently assumed spherical (not important anyway)
Returns
-------
Spatially averaged xarray.DataArray.
"""
lon = xr_da[lon_name].values
lat = xr_da[lat_name].values
area_weights = grid_cell_areas(lon, lat, radius=radius)
return (xr_da * area_weights).sum(dim=[lon_name, lat_name])
@dennissergeev
Copy link
Author

dennissergeev commented Jun 6, 2019

A much simpler version, without using the area of grid cells:

def calc_spatial_mean(xr_da, lon_name="longitude", lat_name="latitude"):
    """
    Calculate spatial mean of xarray.DataArray with latitude weighting.
    
    Parameters
    ----------
    xr_da: xarray.DataArray
        Data to average
    lon_name: str, optional
        Name of x-coordinate
    lat_name: str, optional
        Name of y-coordinate

    Returns
    -------
    Spatially averaged xarray.DataArray.
    """
    coslat = np.cos(np.deg2rad(xr_da[lat_name]))
    return (xr_da * coslat).sum(dim=[lon_name, lat_name]) / (
        coslat.sum(lat_name) * len(xr_da[lon_name])
    )

Example:

import xarray as xr

ds = xr.open_dataset("/path/to/data")

mean_precip = calc_spatial_mean(ds.precipitation_flux)

or, if the dataset has different names for longitude and latitude:

mean_precip = calc_spatial_mean(ds.precipitation_flux, lat_name="lat", lon_name="lon")

@dennissergeev
Copy link
Author

A function to calculate meridional averages of iris cubes:

def calc_meridional_mean(cube, lat_name="latitude"):
    coslat = np.cos(np.deg2rad(cube.coord(lat_name).points))
    coslat2d = iris.util.broadcast_to_shape(coslat, cube.shape, (0,))
    cube_mean = (cube * coslat2d).collapsed(lat_name, iris.analysis.SUM) / np.sum(coslat)
    return cube_mean

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment