Skip to content

Commit baf204a

Browse files
dcheriandcherian
committed
Support multiple axes in reductions.
1 parent 66b55a4 commit baf204a

File tree

3 files changed

+132
-45
lines changed

3 files changed

+132
-45
lines changed

‎cf_xarray/accessor.py‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import functools
22
import inspect
3+
import itertools
34
from collections import ChainMap
45
from typing import Callable, List, Mapping, MutableMapping, Optional, Set, Tuple, Union
56

@@ -230,7 +231,7 @@ def _get_measure(
230231

231232
#: Default mappers for common keys.
232233
_DEFAULT_KEY_MAPPERS: Mapping[str, Mapper] = {
233-
"dim": _get_axis_coord_single,
234+
"dim": _get_axis_coord,
234235
"coord": _get_axis_coord_single,
235236
"group": _get_axis_coord_single,
236237
"weights": _get_measure_variable, # type: ignore
@@ -395,9 +396,16 @@ def _rewrite_values(self, kwargs, key_mappers: dict, var_kws):
395396

396397
else:
397398
# things like sum which have dim
398-
updates[key] = [mapper(self._obj, v, False, v) for v in value]
399-
if len(updates[key]) == 1:
400-
updates[key] = updates[key][0]
399+
newvalue = [mapper(self._obj, v, False, v) for v in value]
400+
if len(newvalue) == 1:
401+
# works for groupby("time")
402+
newvalue = newvalue[0]
403+
else:
404+
# Mappers return list by default
405+
# for input dim=["lat", "X"], newvalue=[["lat"], ["lon"]],
406+
# so we deal with that here.
407+
newvalue = list(itertools.chain(*newvalue))
408+
updates[key] = newvalue
401409

402410
kwargs.update(updates)
403411

‎cf_xarray/tests/test_accessor.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,13 @@ def test_wrapped_classes(obj, attr, xrkwargs, cfkwargs):
7373
assert_identical(expected, actual)
7474

7575

76+
@pytest.mark.parametrize("obj", objects)
77+
def test_groupby_reduce_multiple_dims(obj):
78+
actual = obj.cf.groupby("time.month").mean(["lat", "X"])
79+
expected = obj.groupby("time.month").mean(["lat", "lon"])
80+
assert_identical(actual, expected)
81+
82+
7683
@pytest.mark.parametrize("obj", dataarrays)
7784
def test_weighted(obj):
7885
with raise_if_dask_computes(max_computes=2):
@@ -105,6 +112,10 @@ def test_kwargs_expand_key_to_multiple_keys():
105112
expected = ds.isel(x1=5, y1=3, x2=5, y2=3)
106113
assert_identical(actual, expected)
107114

115+
actual = ds.cf.mean("X")
116+
expected = ds.mean(["x1", "x2"])
117+
assert_identical(actual, expected)
118+
108119
actual = ds.cf.coarsen(X=10, Y=5)
109120
expected = ds.coarsen(x1=10, y1=5, x2=10, y2=5)
110121
assert_identical(actual.mean(), expected.mean())

‎doc/examples/introduction.ipynb‎

Lines changed: 109 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,15 @@
11
{
22
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {
6+
"toc": true
7+
},
8+
"source": [
9+
"<h1>Table of Contents<span class=\"tocSkip\"></span></h1>\n",
10+
"<div class=\"toc\"><ul class=\"toc-item\"><li><span><a href=\"#Introduction-to-cf_xarray\" data-toc-modified-id=\"Introduction-to-cf_xarray-1\"><span class=\"toc-item-num\">1&nbsp;&nbsp;</span>Introduction to <code>cf_xarray</code></a></span><ul class=\"toc-item\"><li><span><a href=\"#What-attributes-have-been-discovered?\" data-toc-modified-id=\"What-attributes-have-been-discovered?-1.1\"><span class=\"toc-item-num\">1.1&nbsp;&nbsp;</span>What attributes have been discovered?</a></span></li><li><span><a href=\"#Feature:-Accessing-coordinate-variables\" data-toc-modified-id=\"Feature:-Accessing-coordinate-variables-1.2\"><span class=\"toc-item-num\">1.2&nbsp;&nbsp;</span>Feature: Accessing coordinate variables</a></span></li><li><span><a href=\"#Feature:-Rewriting-arguments\" data-toc-modified-id=\"Feature:-Rewriting-arguments-1.3\"><span class=\"toc-item-num\">1.3&nbsp;&nbsp;</span>Feature: Rewriting arguments</a></span><ul class=\"toc-item\"><li><span><a href=\"#Slicing\" data-toc-modified-id=\"Slicing-1.3.1\"><span class=\"toc-item-num\">1.3.1&nbsp;&nbsp;</span>Slicing</a></span></li><li><span><a href=\"#Reductions\" data-toc-modified-id=\"Reductions-1.3.2\"><span class=\"toc-item-num\">1.3.2&nbsp;&nbsp;</span>Reductions</a></span></li><li><span><a href=\"#Plotting\" data-toc-modified-id=\"Plotting-1.3.3\"><span class=\"toc-item-num\">1.3.3&nbsp;&nbsp;</span>Plotting</a></span></li><li><span><a href=\"#Resample-&amp;-groupby\" data-toc-modified-id=\"Resample-&amp;-groupby-1.3.4\"><span class=\"toc-item-num\">1.3.4&nbsp;&nbsp;</span>Resample &amp; groupby</a></span></li><li><span><a href=\"#Rolling-&amp;-coarsen\" data-toc-modified-id=\"Rolling-&amp;-coarsen-1.3.5\"><span class=\"toc-item-num\">1.3.5&nbsp;&nbsp;</span>Rolling &amp; coarsen</a></span><ul class=\"toc-item\"><li><span><a href=\"#miscellaneous-features\" data-toc-modified-id=\"miscellaneous-features-1.3.5.1\"><span class=\"toc-item-num\">1.3.5.1&nbsp;&nbsp;</span>miscellaneous features</a></span></li></ul></li></ul></li><li><span><a href=\"#Feature:-Weight-by-Cell-Measures\" data-toc-modified-id=\"Feature:-Weight-by-Cell-Measures-1.4\"><span class=\"toc-item-num\">1.4&nbsp;&nbsp;</span>Feature: Weight by Cell Measures</a></span></li></ul></li></ul></div>"
11+
]
12+
},
313
{
414
"cell_type": "markdown",
515
"metadata": {},
@@ -14,8 +24,8 @@
1424
"execution_count": null,
1525
"metadata": {
1626
"ExecuteTime": {
17-
"end_time": "2020-06-12T20:35:44.967831Z",
18-
"start_time": "2020-06-12T20:35:44.943306Z"
27+
"end_time": "2020-06-19T16:44:52.277430Z",
28+
"start_time": "2020-06-19T16:44:51.354881Z"
1929
}
2030
},
2131
"outputs": [],
@@ -37,8 +47,8 @@
3747
"execution_count": null,
3848
"metadata": {
3949
"ExecuteTime": {
40-
"end_time": "2020-06-12T20:33:37.322927Z",
41-
"start_time": "2020-06-12T20:33:37.262000Z"
50+
"end_time": "2020-06-19T16:44:52.331138Z",
51+
"start_time": "2020-06-19T16:44:52.279190Z"
4252
}
4353
},
4454
"outputs": [],
@@ -59,8 +69,8 @@
5969
"execution_count": null,
6070
"metadata": {
6171
"ExecuteTime": {
62-
"end_time": "2020-06-12T20:40:55.378756Z",
63-
"start_time": "2020-06-12T20:40:55.338717Z"
72+
"end_time": "2020-06-19T16:44:52.399087Z",
73+
"start_time": "2020-06-19T16:44:52.333063Z"
6474
}
6575
},
6676
"outputs": [],
@@ -94,7 +104,12 @@
94104
{
95105
"cell_type": "code",
96106
"execution_count": null,
97-
"metadata": {},
107+
"metadata": {
108+
"ExecuteTime": {
109+
"end_time": "2020-06-19T16:44:52.449912Z",
110+
"start_time": "2020-06-19T16:44:52.400460Z"
111+
}
112+
},
98113
"outputs": [],
99114
"source": [
100115
"multiple = xr.Dataset()\n",
@@ -120,8 +135,8 @@
120135
"execution_count": null,
121136
"metadata": {
122137
"ExecuteTime": {
123-
"end_time": "2020-06-12T20:25:12.597651Z",
124-
"start_time": "2020-06-12T20:25:12.554432Z"
138+
"end_time": "2020-06-19T16:44:52.480112Z",
139+
"start_time": "2020-06-19T16:44:52.451337Z"
125140
}
126141
},
127142
"outputs": [],
@@ -143,8 +158,8 @@
143158
"execution_count": null,
144159
"metadata": {
145160
"ExecuteTime": {
146-
"end_time": "2020-06-12T20:27:09.803560Z",
147-
"start_time": "2020-06-12T20:27:09.770516Z"
161+
"end_time": "2020-06-19T16:44:52.498434Z",
162+
"start_time": "2020-06-19T16:44:52.481531Z"
148163
}
149164
},
150165
"outputs": [],
@@ -164,8 +179,8 @@
164179
"execution_count": null,
165180
"metadata": {
166181
"ExecuteTime": {
167-
"end_time": "2020-06-12T20:41:54.117735Z",
168-
"start_time": "2020-06-12T20:41:54.087549Z"
182+
"end_time": "2020-06-19T16:44:52.522360Z",
183+
"start_time": "2020-06-19T16:44:52.500354Z"
169184
}
170185
},
171186
"outputs": [],
@@ -183,7 +198,12 @@
183198
{
184199
"cell_type": "code",
185200
"execution_count": null,
186-
"metadata": {},
201+
"metadata": {
202+
"ExecuteTime": {
203+
"end_time": "2020-06-19T16:44:52.553921Z",
204+
"start_time": "2020-06-19T16:44:52.526299Z"
205+
}
206+
},
187207
"outputs": [],
188208
"source": [
189209
"multiple.cf.describe()"
@@ -203,8 +223,8 @@
203223
"execution_count": null,
204224
"metadata": {
205225
"ExecuteTime": {
206-
"end_time": "2020-06-12T20:41:57.554031Z",
207-
"start_time": "2020-06-12T20:41:57.504506Z"
226+
"end_time": "2020-06-19T16:44:52.575193Z",
227+
"start_time": "2020-06-19T16:44:52.555069Z"
208228
}
209229
},
210230
"outputs": [],
@@ -224,8 +244,8 @@
224244
"execution_count": null,
225245
"metadata": {
226246
"ExecuteTime": {
227-
"end_time": "2020-06-12T20:42:02.439458Z",
228-
"start_time": "2020-06-12T20:42:02.364630Z"
247+
"end_time": "2020-06-19T16:44:52.608225Z",
248+
"start_time": "2020-06-19T16:44:52.576642Z"
229249
}
230250
},
231251
"outputs": [],
@@ -245,8 +265,8 @@
245265
"execution_count": null,
246266
"metadata": {
247267
"ExecuteTime": {
248-
"end_time": "2020-06-12T20:43:09.230608Z",
249-
"start_time": "2020-06-12T20:43:09.189175Z"
268+
"end_time": "2020-06-19T16:44:52.640868Z",
269+
"start_time": "2020-06-19T16:44:52.609404Z"
250270
}
251271
},
252272
"outputs": [],
@@ -277,8 +297,8 @@
277297
"execution_count": null,
278298
"metadata": {
279299
"ExecuteTime": {
280-
"end_time": "2020-06-12T20:28:53.616306Z",
281-
"start_time": "2020-06-12T20:28:53.562495Z"
300+
"end_time": "2020-06-19T16:44:52.677400Z",
301+
"start_time": "2020-06-19T16:44:52.641979Z"
282302
}
283303
},
284304
"outputs": [],
@@ -296,12 +316,60 @@
296316
{
297317
"cell_type": "code",
298318
"execution_count": null,
299-
"metadata": {},
319+
"metadata": {
320+
"ExecuteTime": {
321+
"end_time": "2020-06-19T16:44:52.700101Z",
322+
"start_time": "2020-06-19T16:44:52.679004Z"
323+
}
324+
},
300325
"outputs": [],
301326
"source": [
302327
"multiple.cf.isel(X=1, Y=1)"
303328
]
304329
},
330+
{
331+
"cell_type": "markdown",
332+
"metadata": {},
333+
"source": [
334+
"### Reductions"
335+
]
336+
},
337+
{
338+
"cell_type": "code",
339+
"execution_count": null,
340+
"metadata": {
341+
"ExecuteTime": {
342+
"end_time": "2020-06-19T16:44:52.749269Z",
343+
"start_time": "2020-06-19T16:44:52.701198Z"
344+
}
345+
},
346+
"outputs": [],
347+
"source": [
348+
"ds.air.cf.mean(\"X\")"
349+
]
350+
},
351+
{
352+
"cell_type": "markdown",
353+
"metadata": {},
354+
"source": [
355+
"Expanding to multiple dimensions is also supported"
356+
]
357+
},
358+
{
359+
"cell_type": "code",
360+
"execution_count": null,
361+
"metadata": {
362+
"ExecuteTime": {
363+
"end_time": "2020-06-19T16:44:52.777454Z",
364+
"start_time": "2020-06-19T16:44:52.750613Z"
365+
}
366+
},
367+
"outputs": [],
368+
"source": [
369+
"# takes the mean along [\"x1\", \"x2\"]\n",
370+
"multiple.cf.mean(\"X\")"
371+
]
372+
},
305373
{
306374
"cell_type": "markdown",
307375
"metadata": {},
@@ -314,8 +382,8 @@
314382
"execution_count": null,
315383
"metadata": {
316384
"ExecuteTime": {
317-
"end_time": "2020-06-12T20:28:59.835481Z",
318-
"start_time": "2020-06-12T20:28:59.603378Z"
385+
"end_time": "2020-06-19T16:44:53.214185Z",
386+
"start_time": "2020-06-19T16:44:52.779161Z"
319387
}
320388
},
321389
"outputs": [],
@@ -328,8 +396,8 @@
328396
"execution_count": null,
329397
"metadata": {
330398
"ExecuteTime": {
331-
"end_time": "2020-06-12T20:30:26.160180Z",
332-
"start_time": "2020-06-12T20:30:25.980074Z"
399+
"end_time": "2020-06-19T16:44:53.396332Z",
400+
"start_time": "2020-06-19T16:44:53.216108Z"
333401
}
334402
},
335403
"outputs": [],
@@ -349,8 +417,8 @@
349417
"execution_count": null,
350418
"metadata": {
351419
"ExecuteTime": {
352-
"end_time": "2020-06-12T20:49:14.502662Z",
353-
"start_time": "2020-06-12T20:49:14.002534Z"
420+
"end_time": "2020-06-19T16:44:53.937352Z",
421+
"start_time": "2020-06-19T16:44:53.397781Z"
354422
}
355423
},
356424
"outputs": [],
@@ -370,8 +438,8 @@
370438
"execution_count": null,
371439
"metadata": {
372440
"ExecuteTime": {
373-
"end_time": "2020-06-12T20:29:50.966391Z",
374-
"start_time": "2020-06-12T20:29:50.317513Z"
441+
"end_time": "2020-06-19T16:44:54.577478Z",
442+
"start_time": "2020-06-19T16:44:53.938825Z"
375443
}
376444
},
377445
"outputs": [],
@@ -384,8 +452,8 @@
384452
"execution_count": null,
385453
"metadata": {
386454
"ExecuteTime": {
387-
"end_time": "2020-06-12T20:30:55.508515Z",
388-
"start_time": "2020-06-12T20:30:55.433502Z"
455+
"end_time": "2020-06-19T16:44:54.644922Z",
456+
"start_time": "2020-06-19T16:44:54.578741Z"
389457
}
390458
},
391459
"outputs": [],
@@ -405,8 +473,8 @@
405473
"execution_count": null,
406474
"metadata": {
407475
"ExecuteTime": {
408-
"end_time": "2020-06-12T20:31:11.369963Z",
409-
"start_time": "2020-06-12T20:31:11.240738Z"
476+
"end_time": "2020-06-19T16:44:54.690873Z",
477+
"start_time": "2020-06-19T16:44:54.646190Z"
410478
}
411479
},
412480
"outputs": [],
@@ -442,8 +510,8 @@
442510
"execution_count": null,
443511
"metadata": {
444512
"ExecuteTime": {
445-
"end_time": "2020-06-12T20:34:38.685926Z",
446-
"start_time": "2020-06-12T20:34:38.627758Z"
513+
"end_time": "2020-06-19T16:44:54.750667Z",
514+
"start_time": "2020-06-19T16:44:54.693023Z"
447515
}
448516
},
449517
"outputs": [],
@@ -465,8 +533,8 @@
465533
"execution_count": null,
466534
"metadata": {
467535
"ExecuteTime": {
468-
"end_time": "2020-06-12T20:47:31.384370Z",
469-
"start_time": "2020-06-12T20:47:31.347731Z"
536+
"end_time": "2020-06-19T16:44:54.771517Z",
537+
"start_time": "2020-06-19T16:44:54.752078Z"
470538
}
471539
},
472540
"outputs": [],
@@ -487,8 +555,8 @@
487555
"execution_count": null,
488556
"metadata": {
489557
"ExecuteTime": {
490-
"end_time": "2020-06-12T20:47:08.559480Z",
491-
"start_time": "2020-06-12T20:47:08.371283Z"
558+
"end_time": "2020-06-19T16:44:54.961814Z",
559+
"start_time": "2020-06-19T16:44:54.772773Z"
492560
}
493561
},
494562
"outputs": [],

0 commit comments

Comments
 (0)