@@ -1371,36 +1371,51 @@ def test_groupby_multidim_map(self):
13711371 )
13721372 assert_identical (expected , actual )
13731373
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+ )
13761380 # the first value should not be part of any group ("right" binning)
13771381 array [0 ] = 99
13781382 # bins follow conventions for pandas.cut
13791383 # http://pandas.pydata.org/pandas-docs/stable/generated/pandas.cut.html
13801384 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" ]
13841397 )
1385- actual = array .groupby_bins ("dim_0" , bins = bins ).sum ()
1386- assert_identical (expected , actual )
13871398
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 )
13901402
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 )
13931405
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 )
13961408
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 )
14041419
14051420 def test_groupby_bins_empty (self ):
14061421 array = DataArray (np .arange (4 ), [("x" , range (4 ))])
0 commit comments