@@ -1371,36 +1371,51 @@ def test_groupby_multidim_map(self):
1371
1371
)
1372
1372
assert_identical (expected , actual )
1373
1373
1374
- def test_groupby_bins (self ):
1375
- array = DataArray (np .arange (4 ), dims = "dim_0" )
1374
+ @pytest .mark .parametrize ("use_flox" , [True , False ])
1375
+ @pytest .mark .parametrize ("coords" , [np .arange (4 ), np .arange (4 )[::- 1 ], [2 , 0 , 3 , 1 ]])
1376
+ def test_groupby_bins (self , coords : np .typing .ArrayLike , use_flox : bool ) -> None :
1377
+ array = DataArray (
1378
+ np .arange (4 ), dims = "dim_0" , coords = {"dim_0" : coords }, name = "a"
1379
+ )
1376
1380
# the first value should not be part of any group ("right" binning)
1377
1381
array [0 ] = 99
1378
1382
# bins follow conventions for pandas.cut
1379
1383
# http://pandas.pydata.org/pandas-docs/stable/generated/pandas.cut.html
1380
1384
bins = [0 , 1.5 , 5 ]
1381
- bin_coords = pd .cut (array ["dim_0" ], bins ).categories
1382
- expected = DataArray (
1383
- [1 , 5 ], dims = "dim_0_bins" , coords = {"dim_0_bins" : bin_coords }
1385
+
1386
+ df = array .to_dataframe ()
1387
+ df ["dim_0_bins" ] = pd .cut (array ["dim_0" ], bins )
1388
+
1389
+ expected_df = df .groupby ("dim_0_bins" ).sum ()
1390
+ # TODO: can't convert df with IntervalIndex to Xarray
1391
+
1392
+ expected = (
1393
+ expected_df .reset_index (drop = True )
1394
+ .to_xarray ()
1395
+ .assign_coords (index = np .array (expected_df .index ))
1396
+ .rename ({"index" : "dim_0_bins" })["a" ]
1384
1397
)
1385
- actual = array .groupby_bins ("dim_0" , bins = bins ).sum ()
1386
- assert_identical (expected , actual )
1387
1398
1388
- actual = array .groupby_bins ("dim_0" , bins = bins , labels = [1.2 , 3.5 ]).sum ()
1389
- assert_identical (expected .assign_coords (dim_0_bins = [1.2 , 3.5 ]), actual )
1399
+ with xr .set_options (use_flox = use_flox ):
1400
+ actual = array .groupby_bins ("dim_0" , bins = bins ).sum ()
1401
+ assert_identical (expected , actual )
1390
1402
1391
- actual = array .groupby_bins ("dim_0" , bins = bins ). map ( lambda x : x . sum () )
1392
- assert_identical (expected , actual )
1403
+ actual = array .groupby_bins ("dim_0" , bins = bins , labels = [ 1.2 , 3.5 ]). sum ()
1404
+ assert_identical (expected . assign_coords ( dim_0_bins = [ 1.2 , 3.5 ]) , actual )
1393
1405
1394
- # make sure original array dims are unchanged
1395
- assert len ( array . dim_0 ) == 4
1406
+ actual = array . groupby_bins ( "dim_0" , bins = bins ). map ( lambda x : x . sum ())
1407
+ assert_identical ( expected , actual )
1396
1408
1397
- da = xr .DataArray (np .ones ((2 , 3 , 4 )))
1398
- bins = [- 1 , 0 , 1 , 2 ]
1399
- with xr .set_options (use_flox = False ):
1400
- actual = da .groupby_bins ("dim_0" , bins ).mean (...)
1401
- with xr .set_options (use_flox = True ):
1402
- expected = da .groupby_bins ("dim_0" , bins ).mean (...)
1403
- assert_allclose (actual , expected )
1409
+ # make sure original array dims are unchanged
1410
+ assert len (array .dim_0 ) == 4
1411
+
1412
+ da = xr .DataArray (np .ones ((2 , 3 , 4 )))
1413
+ bins = [- 1 , 0 , 1 , 2 ]
1414
+ with xr .set_options (use_flox = False ):
1415
+ actual = da .groupby_bins ("dim_0" , bins ).mean (...)
1416
+ with xr .set_options (use_flox = True ):
1417
+ expected = da .groupby_bins ("dim_0" , bins ).mean (...)
1418
+ assert_allclose (actual , expected )
1404
1419
1405
1420
def test_groupby_bins_empty (self ):
1406
1421
array = DataArray (np .arange (4 ), [("x" , range (4 ))])
0 commit comments