Enhance compute_MVBS feasability (#1470)

* add range var max and reindex to compute_MVBS parameters and drop flox maximum version

* add tests for compute MVBS changes

* add todo for compute NASC

* add np.nan fill value

* Update echopype/commongrid/api.py

Co-authored-by: Wu-Jung Lee <[email protected]>

---------

Co-authored-by: Wu-Jung Lee <[email protected]>

Author: ctuguinay

echopype/commongrid/api.py

@@ -32,9 +32,12 @@ def compute_MVBS(
     range_var: Literal["echo_range", "depth"] = "echo_range",
     range_bin: str = "20m",
     ping_time_bin: str = "20s",
-    method="map-reduce",
-    skipna=True,
+    method: str = "map-reduce",
+    reindex: bool = False,
+    skipna: bool = True,
+    fill_value: float = np.nan,
     closed: Literal["left", "right"] = "left",
+    range_var_max: str = None,
     **flox_kwargs,
 ):
     """
@@ -63,11 +66,23 @@ def compute_MVBS(
         The flox strategy for reduction of dask arrays only.
         See flox `documentation <https://flox.readthedocs.io/en/latest/implementation.html>`_
         for more details.
+    reindex: bool, default False
+        If False, reindex after the blockwise stage. If True, reindex at the blockwise stage.
+        Generally, `reindex=False` results in less memory at the cost of computation speed.
+        Can only be used when method='map-reduce'.
+        See flox `documentation <https://flox.readthedocs.io/en/latest/implementation.html>`_
+        for more details.
     skipna: bool, default True
         If true, the mean operation skips NaN values.
         Else, the mean operation includes NaN values.
+    fill_value: float, default np.nan
+        Fill value when no group data exists to aggregate.
     closed: {'left', 'right'}, default 'left'
         Which side of bin interval is closed.
+    range_var_max: str, default None
+        Range variable maximum. Can be true range variable maximum or the maximum depth the
+        user wishes to regrid to. If known, users can pass in range variable maximum to
+        ensure that `compute_MVBS` can lazily run without any computation.
     **flox_kwargs
         Additional keyword arguments to be passed
         to flox reduction function.
@@ -76,6 +91,8 @@ def compute_MVBS(
     -------
     A dataset containing bin-averaged Sv
     """
+    if method != "map-reduce" and reindex is not None:
+        raise ValueError(f"Passing in reindex={reindex} is only allowed when method='map_reduce'.")
 
     # Setup and validate
     # * Sv dataset must contain specified range_var
@@ -86,10 +103,15 @@ def compute_MVBS(
     if not isinstance(ping_time_bin, str):
         raise TypeError("ping_time_bin must be a string")
 
-    # create bin information for echo_range
-    # this computes the echo range max since there might NaNs in the data
-    echo_range_max = ds_Sv[range_var].max()
-    range_interval = np.arange(0, echo_range_max + range_bin, range_bin)
+    # Create bin information for the range variable
+    if range_var_max is None:
+        # This computes the range variable max since there might be NaNs in the data
+        range_var_max = ds_Sv[range_var].max(skipna=True)
+    else:
+        # Parse string and small increase to ensure that we get the bin
+        # corresponding to range_var_max
+        range_var_max = _parse_x_bin(range_var_max) + 1e-8
+    range_interval = np.arange(0, range_var_max + range_bin, range_bin)
 
     # create bin information needed for ping_time
     d_index = (
@@ -109,7 +131,9 @@ def compute_MVBS(
         ping_interval,
         range_var=range_var,
         method=method,
+        reindex=reindex,
         skipna=skipna,
+        fill_value=fill_value,
         **flox_kwargs,
     )
 
@@ -275,6 +299,7 @@ def compute_NASC(
 ) -> xr.Dataset:
     """
     Compute Nautical Areal Scattering Coefficient (NASC) from an Sv dataset.
+    TODO: Add range_var_max and reindex parameters to match `compute_MVBS`.
 
     Parameters
     ----------

echopype/commongrid/utils.py

@@ -20,7 +20,9 @@ def compute_raw_MVBS(
     ping_interval: Union[pd.IntervalIndex, np.ndarray],
     range_var: Literal["echo_range", "depth"] = "echo_range",
     method="map-reduce",
+    reindex=False,
     skipna=True,
+    fill_value=np.nan,
     **flox_kwargs,
 ):
     """
@@ -46,9 +48,17 @@ def compute_raw_MVBS(
         The flox strategy for reduction of dask arrays only.
         See flox `documentation <https://flox.readthedocs.io/en/latest/implementation.html>`_
         for more details.
+    reindex: bool, default False
+        If False, reindex after the blockwise stage. If True, reindex at the blockwise stage.
+        Generally, `reindex=False` results in less memory at the cost of computation speed.
+        Can only be used when method='map-reduce'.
+        See flox `documentation <https://flox.readthedocs.io/en/latest/implementation.html>`_
+        for more details.
     skipna: bool, default True
         If true, the mean operation skips NaN values.
         Else, the mean operation includes NaN values.
+    fill_value: float, default np.nan
+        Fill value when no group data exists to aggregate.
     **flox_kwargs
         Additional keyword arguments to be passed
         to flox reduction function.
@@ -69,8 +79,10 @@ def compute_raw_MVBS(
         x_var=x_var,
         range_var=range_var,
         method=method,
+        reindex=reindex,
         func="nanmean" if skipna else "mean",
         skipna=skipna,
+        fill_value=fill_value,
         **flox_kwargs,
     )
 
@@ -495,8 +507,10 @@ def _groupby_x_along_channels(
     x_var: Literal["ping_time", "distance_nmi"] = "ping_time",
     range_var: Literal["echo_range", "depth"] = "echo_range",
     method: str = "map-reduce",
+    reindex: bool = False,
     func: str = "nanmean",
     skipna: bool = True,
+    fill_value: float = np.nan,
     **flox_kwargs,
 ) -> xr.Dataset:
     """
@@ -534,6 +548,12 @@ def _groupby_x_along_channels(
         The flox strategy for reduction of dask arrays only.
         See flox `documentation <https://flox.readthedocs.io/en/latest/implementation.html>`_
         for more details.
+    reindex: bool, default False
+        If False, reindex after the blockwise stage. If True, reindex at the blockwise stage.
+        Generally, `reindex=False` results in less memory at the cost of computation speed.
+        Can only be used when method='map-reduce'.
+        See flox `documentation <https://flox.readthedocs.io/en/latest/implementation.html>`_
+        for more details.
     func: str, default 'nanmean'
         The aggregation function used for reducing the data array.
         By default, 'nanmean' is used. Other options can be found in the flox `documentation
@@ -543,6 +563,8 @@ def _groupby_x_along_channels(
         Else, aggregation function includes NaN values.
         Note that if ``func`` is set to 'mean' and ``skipna`` is set to True, then aggregation
         will have the same behavior as if func is set to 'nanmean'.
+    fill_value: float, default np.nan
+        Fill value when no group data exists to aggregate.
     **flox_kwargs
         Additional keyword arguments to be passed
         to flox reduction function.
@@ -593,8 +615,10 @@ def _groupby_x_along_channels(
         expected_groups=(None, x_interval, range_interval),
         isbin=[False, True, True],
         method=method,
+        reindex=reindex,
         func=func,
         skipna=skipna,
+        fill_value=fill_value,
         **flox_kwargs,
     )
     return sv_mean

echopype/tests/commongrid/test_commongrid_api.py

@@ -555,3 +555,33 @@ def test_assign_actual_range(request):
             round(float(ds_MVBS["Sv"].max().values), 2),
         ]
     )
+
+
+@pytest.mark.integration
+def test_compute_MVBS_range_var_max(request):
+    """
+    Tests compute_MVBS when user specified range_var_max is passed in.
+    """
+    # Grab mock Sv dataset
+    ds_Sv = request.getfixturevalue("mock_Sv_dataset_regular")
+
+    # Compute MVBS
+    ds_MVBS = ep.commongrid.compute_MVBS(ds_Sv, range_bin="1m", range_var_max="8m")
+
+    # Ensure that last echo range value matches range variable maximum
+    assert ds_MVBS["echo_range"].max().compute() == 8
+
+
+@pytest.mark.integration
+def test_compute_reindex_non_NaN_not_map_reduce(request):
+    """
+    Tests compute_MVBS when user passes in reindex non-NaN without method as map-reduce.
+    """
+    # Grab mock Sv dataset
+    ds_Sv = request.getfixturevalue("mock_Sv_dataset_regular")
+
+    # Compute MVBS with invalid parameters
+    for method in ["blockwise", "cohorts"]:
+        for reindex in [True, False]:
+            with pytest.raises(ValueError, match=f"Passing in reindex={reindex} is only allowed when method='map_reduce'."):
+                ep.commongrid.compute_MVBS(ds_Sv, method=method, reindex=reindex)

requirements.txt

@@ -2,7 +2,7 @@ aiohttp
 bottleneck
 dask[array,distributed]
 dask-image
-flox>=0.7.2,<1.0.0
+flox>=0.7.2
 fsspec
 geopy
 jinja2