NumPy zonal stats: return a data array of calculated stats (#685)

* zonal_stats returns a DataArray * flake8

Author: thuydotm

xrspatial/tests/test_zonal.py

@@ -11,7 +11,7 @@
 from xrspatial import zonal_stats as stats
 from xrspatial.zonal import regions
 
-from .general_checks import create_test_raster, has_cuda_and_cupy
+from .general_checks import create_test_raster, general_output_checks, has_cuda_and_cupy
 
 
 @pytest.fixture
@@ -60,6 +60,40 @@ def result_default_stats():
  return expected_result
 
 
+@pytest.fixture
+def result_default_stats_dataarray():
+ expected_result = np.array(
+ [[[0., 0., 1., 1., 2., 2., 2.4, 2.4],
+ [0., 0., 1., 1., 2., 2., 2.4, 2.4],
+ [0., 0., 1., 1., 2., np.nan, 2.4, 2.4]],
+
+ [[0., 0., 1., 1., 2., 2., 3., 3.],
+ [0., 0., 1., 1., 2., 2., 3., 3.],
+ [0., 0., 1., 1., 2., np.nan, 3., 3.]],
+
+ [[0., 0., 1., 1., 2., 2., 0., 0.],
+ [0., 0., 1., 1., 2., 2., 0., 0.],
+ [0., 0., 1., 1., 2., np.nan, 0., 0.]],
+
+ [[0., 0., 6., 6., 8., 8., 12., 12.],
+ [0., 0., 6., 6., 8., 8., 12., 12.],
+ [0., 0., 6., 6., 8., np.nan, 12., 12.]],
+
+ [[0., 0., 0., 0., 0., 0., 1.2, 1.2],
+ [0., 0., 0., 0., 0., 0., 1.2, 1.2],
+ [0., 0., 0., 0., 0., np.nan, 1.2, 1.2]],
+
+ [[0., 0., 0., 0., 0., 0., 1.44, 1.44],
+ [0., 0., 0., 0., 0., 0., 1.44, 1.44],
+ [0., 0., 0., 0., 0., np.nan, 1.44, 1.44]],
+
+ [[5., 5., 6., 6., 4., 4., 5., 5.],
+ [5., 5., 6., 6., 4., 4., 5., 5.],
+ [5., 5., 6., 6., 4., np.nan, 5., 5.]]]
+ )
+ return expected_result
+
+
 @pytest.fixture
 def result_zone_ids_stats():
  zone_ids = [0, 3]
@@ -76,6 +110,41 @@ def result_zone_ids_stats():
  return zone_ids, expected_result
 
 
+@pytest.fixture
+def result_zone_ids_stats_dataarray():
+ zone_ids = [0, 3]
+ expected_result = np.array(
+ [[[0., 0., np.nan, np.nan, np.nan, np.nan, 2.4, 2.4],
+ [0., 0., np.nan, np.nan, np.nan, np.nan, 2.4, 2.4],
+ [0., 0., np.nan, np.nan, np.nan, np.nan, 2.4, 2.4]],
+
+ [[0., 0., np.nan, np.nan, np.nan, np.nan, 3., 3.],
+ [0., 0., np.nan, np.nan, np.nan, np.nan, 3., 3.],
+ [0., 0., np.nan, np.nan, np.nan, np.nan, 3., 3.]],
+
+ [[0., 0., np.nan, np.nan, np.nan, np.nan, 0., 0.],
+ [0., 0., np.nan, np.nan, np.nan, np.nan, 0., 0.],
+ [0., 0., np.nan, np.nan, np.nan, np.nan, 0., 0.]],
+
+ [[0., 0., np.nan, np.nan, np.nan, np.nan, 12., 12.],
+ [0., 0., np.nan, np.nan, np.nan, np.nan, 12., 12.],
+ [0., 0., np.nan, np.nan, np.nan, np.nan, 12., 12.]],
+
+ [[0., 0., np.nan, np.nan, np.nan, np.nan, 1.2, 1.2],
+ [0., 0., np.nan, np.nan, np.nan, np.nan, 1.2, 1.2],
+ [0., 0., np.nan, np.nan, np.nan, np.nan, 1.2, 1.2]],
+
+ [[0., 0., np.nan, np.nan, np.nan, np.nan, 1.44, 1.44],
+ [0., 0., np.nan, np.nan, np.nan, np.nan, 1.44, 1.44],
+ [0., 0., np.nan, np.nan, np.nan, np.nan, 1.44, 1.44]],
+
+ [[5., 5., np.nan, np.nan, np.nan, np.nan, 5., 5.],
+ [5., 5., np.nan, np.nan, np.nan, np.nan, 5., 5.],
+ [5., 5., np.nan, np.nan, np.nan, np.nan, 5., 5.]]])
+
+ return zone_ids, expected_result
+
+
 def _double_sum(values):
  return values.sum() * 2
 
@@ -96,6 +165,22 @@ def result_custom_stats():
  return nodata_values, zone_ids, expected_result
 
 
+@pytest.fixture
+def result_custom_stats_dataarray():
+ zone_ids = [1, 2]
+ nodata_values = 0
+ expected_result = np.array(
+ [[[np.nan, np.nan, 12., 12., 16., 16., np.nan, np.nan],
+ [np.nan, np.nan, 12., 12., 16., 16., np.nan, np.nan],
+ [np.nan, np.nan, 12., 12., 16., np.nan, np.nan, np.nan]],
+
+ [[np.nan, np.nan, 0., 0., 0., 0., np.nan, np.nan],
+ [np.nan, np.nan, 0., 0., 0., 0., np.nan, np.nan],
+ [np.nan, np.nan, 0., 0., 0., np.nan, np.nan, np.nan]]]
+ )
+ return nodata_values, zone_ids, expected_result
+
+
 @pytest.fixture
 def result_count_crosstab_2d():
  zone_ids = [1, 2, 3]
@@ -174,6 +259,22 @@ def test_default_stats(backend, data_zones, data_values_2d, result_default_stats
  check_results(backend, df_result, result_default_stats)
 
 
+@pytest.mark.parametrize("backend", ['numpy'])
+def test_default_stats_dataarray(
+ backend, data_zones, data_values_2d, result_default_stats_dataarray
+):
+ dataarray_result = stats(
+ zones=data_zones, values=data_values_2d, return_type='xarray.DataArray'
+ )
+ general_output_checks(
+ data_values_2d,
+ dataarray_result,
+ result_default_stats_dataarray,
+ verify_dtype=False,
+ verify_attrs=False,
+ )
+
+
 @pytest.mark.parametrize("backend", ['numpy', 'dask+numpy', 'cupy'])
 def test_zone_ids_stats(backend, data_zones, data_values_2d, result_zone_ids_stats):
  if backend == 'cupy' and not has_cuda_and_cupy():
@@ -184,6 +285,19 @@ def test_zone_ids_stats(backend, data_zones, data_values_2d, result_zone_ids_sta
  check_results(backend, df_result, expected_result)
 
 
+@pytest.mark.parametrize("backend", ['numpy'])
+def test_zone_ids_stats_dataarray(
+ backend, data_zones, data_values_2d, result_zone_ids_stats_dataarray
+):
+ zone_ids, expected_result = result_zone_ids_stats_dataarray
+ dataarray_result = stats(
+ zones=data_zones, values=data_values_2d, zone_ids=zone_ids, return_type='xarray.DataArray'
+ )
+ general_output_checks(
+ data_values_2d, dataarray_result, expected_result, verify_dtype=False, verify_attrs=False
+ )
+
+
 @pytest.mark.parametrize("backend", ['numpy', 'cupy'])
 def test_custom_stats(backend, data_zones, data_values_2d, result_custom_stats):
  # ---- custom stats (NumPy and CuPy only) ----
@@ -203,6 +317,23 @@ def test_custom_stats(backend, data_zones, data_values_2d, result_custom_stats):
  check_results(backend, df_result, expected_result)
 
 
+@pytest.mark.parametrize("backend", ['numpy'])
+def test_custom_stats_dataarray(backend, data_zones, data_values_2d, result_custom_stats_dataarray):
+ # ---- custom stats returns a xr.DataArray (NumPy only) ----
+ custom_stats = {
+ 'double_sum': _double_sum,
+ 'range': _range,
+ }
+ nodata_values, zone_ids, expected_result = result_custom_stats_dataarray
+ dataarray_result = stats(
+ zones=data_zones, values=data_values_2d, stats_funcs=custom_stats,
+ zone_ids=zone_ids, nodata_values=nodata_values, return_type='xarray.DataArray'
+ )
+ general_output_checks(
+ data_values_2d, dataarray_result, expected_result, verify_dtype=False, verify_attrs=False
+ )
+
+
 @pytest.mark.parametrize("backend", ['numpy', 'dask+numpy'])
 def test_count_crosstab_2d(backend, data_zones, data_values_2d, result_count_crosstab_2d):
  zone_ids, cat_ids, expected_result = result_count_crosstab_2d

xrspatial/zonal.py

@@ -105,7 +105,7 @@ def _sort_and_stride(zones, values, unique_zones):
  sorted_zones = sorted_zones[np.isfinite(sorted_zones)]
  zone_breaks = _strides(sorted_zones, unique_zones)
 
- return values_by_zones, zone_breaks
+ return sorted_indices, values_by_zones, zone_breaks
 
 
 def _calc_stats(
@@ -123,8 +123,7 @@ def _calc_stats(
  if unique_zones[i] in zone_ids:
  zone_values = values_by_zones[start:end]
  # filter out non-finite and nodata_values
- zone_values = zone_values[
- np.isfinite(zone_values) & (zone_values != nodata_values)]
+ zone_values = zone_values[np.isfinite(zone_values) & (zone_values != nodata_values)]
  if len(zone_values) > 0:
  results[i] = func(zone_values)
  start = end
@@ -141,13 +140,8 @@ def _single_stats_func(
  nodata_values: Union[int, float] = None,
 ) -> pd.DataFrame:
 
- values_by_zones, zone_breaks = _sort_and_stride(
- zones_block, values_block, unique_zones
- )
- results = _calc_stats(
- values_by_zones, zone_breaks,
- unique_zones, zone_ids, func, nodata_values
- )
+ _, values_by_zones, zone_breaks = _sort_and_stride(zones_block, values_block, unique_zones)
+ results = _calc_stats(values_by_zones, zone_breaks, unique_zones, zone_ids, func, nodata_values)
  return results
 
 
@@ -224,19 +218,15 @@ def _stats_dask_numpy(
  stats_dict['mean'] = _dask_mean(stats_dict['sum'], stats_dict['count'])
  if 'std' in stats_funcs:
  stats_dict['std'] = _dask_std(
- stats_dict['sum_squares'], stats_dict['sum'] ** 2,
- stats_dict['count']
+ stats_dict['sum_squares'], stats_dict['sum'] ** 2, stats_dict['count']
  )
  if 'var' in stats_funcs:
  stats_dict['var'] = _dask_var(
- stats_dict['sum_squares'], stats_dict['sum'] ** 2,
- stats_dict['count']
+ stats_dict['sum_squares'], stats_dict['sum'] ** 2, stats_dict['count']
  )
 
  # generate dask dataframe
- stats_df = dd.concat(
- [dd.from_dask_array(s) for s in stats_dict.values()], axis=1
- )
+ stats_df = dd.concat([dd.from_dask_array(s) for s in stats_dict.values()], axis=1)
  # name columns
  stats_df.columns = stats_dict.keys()
  # select columns
@@ -259,7 +249,8 @@ def _stats_numpy(
  zone_ids: List[Union[int, float]],
  stats_funcs: Dict,
  nodata_values: Union[int, float],
-) -> pd.DataFrame:
+ return_type: str,
+) -> Union[pd.DataFrame, np.ndarray]:
 
  # find ids for all zones
  unique_zones = np.unique(zones[np.isfinite(zones)])
@@ -271,23 +262,40 @@ def _stats_numpy(
  # remove zones that do not exist in `zones` raster
  zone_ids = [z for z in zone_ids if z in unique_zones]
 
- selected_indexes = [i for i, z in enumerate(unique_zones) if z in zone_ids]
- values_by_zones, zone_breaks = _sort_and_stride(
- zones, values, unique_zones
- )
-
- stats_dict = {}
- stats_dict["zone"] = zone_ids
- for stats in stats_funcs:
- func = stats_funcs.get(stats)
- stats_dict[stats] = _calc_stats(
- values_by_zones, zone_breaks,
- unique_zones, zone_ids, func, nodata_values
- )
- stats_dict[stats] = stats_dict[stats][selected_indexes]
+ sorted_indices, values_by_zones, zone_breaks = _sort_and_stride(zones, values, unique_zones)
+ if return_type == 'pandas.DataFrame':
+ stats_dict = {}
+ stats_dict["zone"] = zone_ids
+ selected_indexes = [i for i, z in enumerate(unique_zones) if z in zone_ids]
+ for stats in stats_funcs:
+ func = stats_funcs.get(stats)
+ stats_dict[stats] = _calc_stats(
+ values_by_zones, zone_breaks,
+ unique_zones, zone_ids, func, nodata_values
+ )
+ stats_dict[stats] = stats_dict[stats][selected_indexes]
+ result = pd.DataFrame(stats_dict)
 
- stats_df = pd.DataFrame(stats_dict)
- return stats_df
+ else:
+ result = np.full((len(stats_funcs), values.size), np.nan)
+ zone_ids_map = {z: i for i, z in enumerate(unique_zones) if z in zone_ids}
+ stats_id = 0
+ for stats in stats_funcs:
+ func = stats_funcs.get(stats)
+ stats_results = _calc_stats(
+ values_by_zones, zone_breaks,
+ unique_zones, zone_ids, func, nodata_values
+ )
+ for zone in zone_ids:
+ iz = zone_ids_map[zone] # position of zone in unique_zones
+ if iz == 0:
+ zs = sorted_indices[: zone_breaks[iz]]
+ else:
+ zs = sorted_indices[zone_breaks[iz-1]: zone_breaks[iz]]
+ result[stats_id][zs] = stats_results[iz]
+ stats_id += 1
+ result = result.reshape(len(stats_funcs), *values.shape)
+ return result
 
 
 def _stats_cupy(
@@ -391,7 +399,8 @@ def stats(
  "count",
  ],
  nodata_values: Union[int, float] = None,
-) -> Union[pd.DataFrame, dd.DataFrame]:
+ return_type: str = 'pandas.DataFrame',
+) -> Union[pd.DataFrame, dd.DataFrame, xr.DataArray]:
  """
  Calculate summary statistics for each zone defined by a `zones`
  dataset, based on `values` aggregate.
@@ -438,6 +447,11 @@ def stats(
  Cells with `nodata_values` do not belong to any zone,
  and thus excluded from calculation.
 
+ return_type: str, default='pandas.DataFrame'
+ Format of returned data. If `zones` and `values` numpy backed xarray DataArray,
+ allowed values are 'pandas.DataFrame', and 'xarray.DataArray'.
+ Otherwise, only 'pandas.DataFrame' is supported.
+
  Returns
  -------
  stats_df : Union[pandas.DataFrame, dask.dataframe.DataFrame]
@@ -568,17 +582,25 @@ def stats(
  stats_funcs_dict = stats_funcs.copy()
 
  mapper = ArrayTypeFunctionMapping(
- numpy_func=_stats_numpy,
+ numpy_func=lambda *args: _stats_numpy(*args, return_type=return_type),
  dask_func=_stats_dask_numpy,
  cupy_func=_stats_cupy,
  dask_cupy_func=lambda *args: not_implemented_func(
  *args, messages='stats() does not support dask with cupy backed DataArray' # noqa
  ),
  )
- stats_df = mapper(values)(
- zones.data, values.data, zone_ids, stats_funcs_dict, nodata_values
+ result = mapper(values)(
+ zones.data, values.data, zone_ids, stats_funcs_dict, nodata_values,
  )
- return stats_df
+
+ if return_type == 'xarray.DataArray':
+ return xr.DataArray(
+ result,
+ coords={'stats': list(stats_funcs_dict.keys()), **values.coords},
+ dims=('stats', *values.dims),
+ attrs=values.attrs
+ )
+ return result
 
 
 def _find_cats(values, cat_ids, nodata_values):
@@ -680,7 +702,7 @@ def _crosstab_numpy(
  for cat in cat_ids:
  crosstab_dict[cat] = []
 
- values_by_zones, zone_breaks = _sort_and_stride(
+ _, values_by_zones, zone_breaks = _sort_and_stride(
  zones, values, unique_zones
  )
 
@@ -731,7 +753,7 @@ def _single_chunk_crosstab(
  for cat in cat_ids:
  results[cat] = []
 
- values_by_zones, zone_breaks = _sort_and_stride(
+ _, values_by_zones, zone_breaks = _sort_and_stride(
  zones_block, values_block, unique_zones
  )