-
Notifications
You must be signed in to change notification settings - Fork 0
/
CMIP6_xgrid_utils.py
167 lines (133 loc) · 4.73 KB
/
CMIP6_xgrid_utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
# -*- coding: utf-8 -*-
"""Operations on cartesian geographical grid.
GIST from https://gist.github.com/dennissergeev/60bf7b03443f1b2c8eb96ce0b1880150
"""
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])