Repository: tobac-project/tobac

Issue:
Issue with calculate_area
I get the following error while using "tobac.calculate_area". At first, it looked like the issue was resolved with the latest release (1.5.4) but I still encounter it when using the function on files that were created using an older tobac release. 

After installing the latest release - the function runs when first processing the input file (dbz... .nc) using tobac but the area variable only includes "NaN" for all timesteps (see Track... .h5). When trying to re-run the function on the Track file, the same error pops up again. 

Input file and output files created after installing the latest release are uploaded here: https://drive.google.com/drive/u/0/folders/1ARx-GB59t9PXzRUJoD80jc07OIKp9OyS








---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[82], line 1
----> 1 tobac.calculate_area(Features,Mask_dbz)

File ~/.local/lib/python3.11/site-packages/tobac/utils/decorators.py:279, in xarray_to_iris.<locals>.xarray_to_iris_i.<locals>.wrapper(*args, **kwargs)
    276             output = xarray.DataArray.from_iris(output)
    278 else:
--> 279     output = func(*args, **kwargs)
    280 # print(output)
    281 return output

File ~/.local/lib/python3.11/site-packages/tobac/analysis/spatial.py:440, in calculate_area(features, mask, method_area, vertical_coord)
    437 # Area needs to be a dataarray for get_statistics from mask, but otherwise dims/coords don't actually matter
    438 area = xr.DataArray(area, dims=("a", "b"))
--> 440 features = get_statistics_from_mask(
    441     features,
    442     mask,
    443     area,
    444     statistic={"area": np.sum},
    445     default=np.nan,
    446     collapse_dim=collapse_dim,
    447 )
    449 return features

File ~/.local/lib/python3.11/site-packages/tobac/utils/decorators.py:180, in iris_to_xarray.<locals>.iris_to_xarray_i.<locals>.wrapper(*args, **kwargs)
    177 kwargs_new = _conv_kwargs_iris_to_xarray(kwargs)
    178 # print(args)
    179 # print(kwargs)
--> 180 output = func(*args, **kwargs_new)
    181 if type(output) == tuple:
    182     output = tuple(
    183         [
    184             (
   (...)
    190         ]
    191     )

File ~/.local/lib/python3.11/site-packages/tobac/utils/bulk_statistics.py:326, in get_statistics_from_mask(features, segmentation_mask, statistic, index, default, id_column, collapse_dim, *fields)
    322         continue
    323     else:
    324         # make sure that features are not double-defined
    325         step_statistics.append(
--> 326             get_statistics(
    327                 features_t,
    328                 segmentation_mask_t,
    329                 *fields_t,
    330                 statistic=statistic,
    331                 default=default,
    332                 index=index,
    333                 id_column=id_column,
    334                 collapse_axis=collapse_axis,
    335             )
    336         )
    338 features = pd.concat(step_statistics)
    340 return features

File ~/.local/lib/python3.11/site-packages/tobac/utils/bulk_statistics.py:134, in get_statistics(features, labels, statistic, index, default, id_column, collapse_axis, *fields)
    131 index_in_features = np.isin(index, features[id_column])
    133 # set negative markers to 0 as they are unsegmented
--> 134 bins = np.cumsum(np.bincount(np.maximum(labels.ravel(), 0)))
    135 argsorted = np.argsort(labels.ravel())
    137 # Create lambdas to get (ravelled) label locations using argsorted and bins

File <__array_function__ internals>:200, in cumsum(*args, **kwargs)

File /opt/conda/lib/python3.11/site-packages/dask/array/core.py:1760, in Array.__array_function__(self, func, types, args, kwargs)
   1757 if has_keyword(da_func, "like"):
   1758     kwargs["like"] = self
-> 1760 return da_func(*args, **kwargs)

File /opt/conda/lib/python3.11/site-packages/dask/array/reductions.py:1617, in cumsum(x, axis, dtype, out, method)
   1604 @derived_from(np)
   1605 def cumsum(x, axis=None, dtype=None, out=None, method="sequential"):
   1606     """Dask added an additional keyword-only argument ``method``.
   1607 
   1608     method : {'sequential', 'blelloch'}, optional
   (...)
   1615           scheduler, and hardware.  More benchmarking is necessary.
   1616     """
-> 1617     return cumreduction(
   1618         np.cumsum,
   1619         _cumsum_merge,
   1620         0,
   1621         x,
   1622         axis,
   1623         dtype,
   1624         out=out,
   1625         method=method,
   1626         preop=np.sum,
   1627     )

File /opt/conda/lib/python3.11/site-packages/dask/array/reductions.py:1541, in cumreduction(func, binop, ident, x, axis, dtype, out, method, preop)
   1536     raise ValueError(
   1537         f'Invalid method for cumreduction.  Expected "sequential" or "blelloch".  Got: {method!r}'
   1538     )
   1540 if axis is None:
-> 1541     x = x.flatten().rechunk(chunks=x.npartitions)
   1542     axis = 0
   1543 if dtype is None:

File /opt/conda/lib/python3.11/site-packages/dask/array/core.py:2767, in Array.rechunk(self, chunks, threshold, block_size_limit, balance, method)
   2757 """Convert blocks in dask array x for new chunks.
   2758 
   2759 Refer to :func:`dask.array.rechunk` for full documentation.
   (...)
   2763 dask.array.rechunk : equivalent function
   2764 """
   2765 from dask.array.rechunk import rechunk  # avoid circular import
-> 2767 return rechunk(self, chunks, threshold, block_size_limit, balance, method)

File /opt/conda/lib/python3.11/site-packages/dask/array/rechunk.py:347, in rechunk(x, chunks, threshold, block_size_limit, balance, method)
    345 if isinstance(chunks, (tuple, list)):
    346     chunks = tuple(lc if lc is not None else rc for lc, rc in zip(chunks, x.chunks))
--> 347 chunks = normalize_chunks(
    348     chunks, x.shape, limit=block_size_limit, dtype=x.dtype, previous_chunks=x.chunks
    349 )
    351 # Now chunks are tuple of tuples
    352 ndim = x.ndim

File /opt/conda/lib/python3.11/site-packages/dask/array/core.py:3101, in normalize_chunks(chunks, shape, limit, dtype, previous_chunks)
   3098     chunks = tuple(c if c not in {None, -1} else s for c, s in zip(chunks, shape))
   3100 if chunks and shape is not None:
-> 3101     chunks = sum(
   3102         (
   3103             blockdims_from_blockshape((s,), (c,))
   3104             if not isinstance(c, (tuple, list))
   3105             else (c,)
   3106             for s, c in zip(shape, chunks)
   3107         ),
   3108         (),
   3109     )
   3110 for c in chunks:
   3111     if not c:

File /opt/conda/lib/python3.11/site-packages/dask/array/core.py:3103, in <genexpr>(.0)
   3098     chunks = tuple(c if c not in {None, -1} else s for c, s in zip(chunks, shape))
   3100 if chunks and shape is not None:
   3101     chunks = sum(
   3102         (
-> 3103             blockdims_from_blockshape((s,), (c,))
   3104             if not isinstance(c, (tuple, list))
   3105             else (c,)
   3106             for s, c in zip(shape, chunks)
   3107         ),
   3108         (),
   3109     )
   3110 for c in chunks:
   3111     if not c:

File /opt/conda/lib/python3.11/site-packages/dask/array/core.py:1260, in blockdims_from_blockshape(shape, chunks)
   1258     raise TypeError("Must supply shape= keyword argument")
   1259 if np.isnan(sum(shape)) or np.isnan(sum(chunks)):
-> 1260     raise ValueError(
   1261         "Array chunk sizes are unknown. shape: %s, chunks: %s%s"
   1262         % (shape, chunks, unknown_chunk_message)
   1263     )
   1264 if not all(map(is_integer, chunks)):
   1265     raise ValueError("chunks can only contain integers.")

ValueError: Array chunk sizes are unknown. shape: (nan,), chunks: (1,)

A possible solution: https://docs.dask.org/en/latest/array-chunks.html#unknown-chunks
Summary: to compute chunks sizes, use

   x.compute_chunk_sizes()  # for Dask Array `x`
   ddf.to_dask_array(lengths=True)  # for Dask DataFrame `ddf`

Discussion:
@w-k-jones: Hi @sgupta92 , thanks for reporting this! I can recreate the second error when using `tobac.calculate_area` when the segmentation mask is a dask array, but I can't recreate the first issue you reported when all the results for calculate area are `NaN`. Could you post the script you used to create the tracks file that you provided?

For the second issue, it is caused by passing a dask dataarray to the `tobac.calculate_area`, and can be resolved in the mean time by calling the `.compute()` method on the datarray before passing it to calculate_area. This should be a simple issue to fix, as it is caused by using `.data` instead of `.values` in bulk_statistics https://github.com/tobac-project/tobac/blob/706b3ee559bcc01c208b8a4406fd8d0a5b2883c1/tobac/utils/bulk_statistics.py#L304 @sgupta92: Hi @w-k-jones Thanks for the response! For some reason, the solutions you provided did not work. But you're right the issue is caused by passing a dask dataarray to the function. This was confusing at first because I was simply loading the segmentation mask created by tobac as an iris cube. Looking at the iris documentation, it seems iris_load_cube loads lazy data in the form of a dask dataarray (https://scitools-iris.readthedocs.io/en/latest/userguide/real_and_lazy_data.html). This led me to the following workaround. Adding the first line loads the data into the memory and the calculate_area function is passed without error. However, the function only outputs a very small subset of the track file with actual values for area (or no data at all for different dates).

> Mask_dbz.data
> tobac.calculate_area(Track,Mask_dbz)

I tried loading the segmentation mask as an iris cube, loading real data, and then converting to xarray dataset and get the same small subset of the track file with actual values for area:

> Mask_dbz.data
> x=xr.DataArray.from_iris(Mask_dbz)
> tobac.calculate_area(Track,x)

I tried loading the segmentation mask as an xarray dataset (to avoid iris completely) and I get the following error:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[45], line 5
      1 # Mask_dbz.data
      2 # x=xr.DataArray.from_iris(Mask_dbz)
      3 # tobac.calculate_area(Track,x)
----> 5 tobac.calculate_area(Track,Mask_dbz)
      7 # for k in range(0,Mask_dbz.shape[0]):
      8 #     print(np.nanmax(Mask_dbz[k].data))

File ~/.local/lib/python3.11/site-packages/tobac/utils/decorators.py:279, in xarray_to_iris.<locals>.xarray_to_iris_i.<locals>.wrapper(*args, **kwargs)
    276             output = xarray.DataArray.from_iris(output)
    278 else:
--> 279     output = func(*args, **kwargs)
    280 # print(output)
    281 return output

File ~/.local/lib/python3.11/site-packages/tobac/analysis/spatial.py:381, in calculate_area(features, mask, method_area, vertical_coord)
    378 features["area"] = np.nan
    380 # Get the first time step of mask to remove time dimension of calculated areas
--> 381 mask_slice = next(mask.slices_over("time"))
    382 is_3d = len(mask_slice.core_data().shape) == 3
    383 if is_3d:

File /opt/conda/lib/python3.11/site-packages/xarray/core/common.py:286, in AttrAccessMixin.__getattr__(self, name)
    284         with suppress(KeyError):
    285             return source[name]
--> 286 raise AttributeError(
    287     f"{type(self).__name__!r} object has no attribute {name!r}"
    288 )

AttributeError: 'Dataset' object has no attribute 'slices_over'@sgupta92: I came up with a not-efficient workaround: 
1. For each row of the Track file, find the grid points within the segmentation mask that represent the feature at that time step. 
2. Multiply number of grid points and area of one grid point (in this example = 0.25 sq. km.) and add this area value to the Track file.

> for i in Track.index:
>     Track.at[i,'area']=np.shape(np.where(Mask_dbz[Track.frame[i]].data==Track.feature[i]))[1]*0.25 

CAUTION: This iterates over all rows of the Track file with the potential to get very slow. So, it is just a workaround when dealing with a small subset of Tracks for particular applications.@w-k-jones: Thanks for the update. I previously tested the data loaded using `xarray`, but can confirm the error is reproducible when loading mask_dbz as an iris cube. I expect we didn't run into this previously as by default `xarray.DataArray.from_iris` converts the data to a masked array (in the same way that `iris.Cube.data` does), but that can cause errors for integer data in older versions of xarray, so we switched to copying the data as is in `v1.5.4` which preserves the data as a dask array causing this error. I'll have a look into fixing this more generally

Regarding the workaround, [`scipy.ndimage.labeled_comprehension`](https://docs.scipy.org/doc/scipy-1.14.1/reference/generated/scipy.ndimage.labeled_comprehension.html) can be used to perform functions on labelled regions of an array much more efficiently, e.g. for this case:
```
Track["area"] = labeled_comprehension(
    Mask_dbz.data.ravel(), Mask_dbz.data.ravel(), Track.feature, lambda x: np.size(x)*0.25, np.float32, 0
)
```@sgupta92: Thanks for the better workaround, I'll try it! 

I should note I ran into the original error with some of the previous releases of tobac, as well. I just got back to working on this task recently and then posted the issue here.@w-k-jones: The original error should be fixed with the imminent `v1.5.5` release

Solution patch:
diff --git a/.github/workflows/matrix_ci.yml b/.github/workflows/matrix_ci.yml
index 953acc27..9dfdca9d 100644
--- a/.github/workflows/matrix_ci.yml
+++ b/.github/workflows/matrix_ci.yml
@@ -35,7 +35,7 @@ jobs:
           cache-downloads: true
           channels: conda-forge
           channel-priority: strict
-          python-version: ${{ matrix.python-version }}
+          create-args: python=${{ matrix.python-version }}
 
       - name: Fetch all history for all tags and branches
         run: |
diff --git a/tobac/merge_split.py b/tobac/merge_split.py
index 012856b9..1f9d8a50 100644
--- a/tobac/merge_split.py
+++ b/tobac/merge_split.py
@@ -7,6 +7,7 @@
     
 """
 
+from __future__ import annotations
 import logging
 from typing import Optional
 from typing_extensions import Literal
diff --git a/tobac/plotting.py b/tobac/plotting.py
index 63cc6c09..81930f51 100644
--- a/tobac/plotting.py
+++ b/tobac/plotting.py
@@ -14,6 +14,7 @@
    12(11), 4551-4570.
 """
 
+from __future__ import annotations
 import warnings
 import logging
 
diff --git a/tobac/segmentation.py b/tobac/segmentation.py
index 4697a25d..8509f184 100644
--- a/tobac/segmentation.py
+++ b/tobac/segmentation.py
@@ -30,6 +30,7 @@
    12(11), 4551-4570.
 """
 
+from __future__ import annotations
 import copy
 import logging
 
diff --git a/tobac/tracking.py b/tobac/tracking.py
index 7b21d1db..b95234ec 100644
--- a/tobac/tracking.py
+++ b/tobac/tracking.py
@@ -19,6 +19,7 @@
    12(11), 4551-4570.
 """
 
+from __future__ import annotations
 import logging
 from operator import is_
 import numpy as np
diff --git a/tobac/utils/bulk_statistics.py b/tobac/utils/bulk_statistics.py
index daa430eb..ba192d45 100644
--- a/tobac/utils/bulk_statistics.py
+++ b/tobac/utils/bulk_statistics.py
@@ -4,6 +4,7 @@
 
 """
 
+from __future__ import annotations
 import logging
 import warnings
 from functools import partial
@@ -301,7 +302,7 @@ def get_statistics_from_mask(
 
     for tt in pd.to_datetime(segmentation_mask.time):
         # select specific timestep
-        segmentation_mask_t = segmentation_mask.sel(time=tt, method="nearest").data
+        segmentation_mask_t = segmentation_mask.sel(time=tt, method="nearest").values
         fields_t = (
             (
                 field.sel(
diff --git a/tobac/utils/decorators.py b/tobac/utils/decorators.py
index 8d304e71..6486c26a 100644
--- a/tobac/utils/decorators.py
+++ b/tobac/utils/decorators.py
@@ -1,6 +1,7 @@
 """Decorators for use with other tobac functions
 """
 
+from __future__ import annotations
 import functools
 import warnings
 
diff --git a/tobac/utils/general.py b/tobac/utils/general.py
index 89ab6aa1..5ddbbb53 100644
--- a/tobac/utils/general.py
+++ b/tobac/utils/general.py
@@ -2,6 +2,7 @@
 
 """
 
+from __future__ import annotations
 import copy
 import logging
 from typing_extensions import Literal


Test patch:
diff --git a/tobac/tests/test_tracking.py b/tobac/tests/test_tracking.py
index 52fab643..70be714c 100644
--- a/tobac/tests/test_tracking.py
+++ b/tobac/tests/test_tracking.py
@@ -2,6 +2,7 @@
 Test for the trackpy tracking functions
 """
 
+from __future__ import annotations
 import datetime
 import copy
 
diff --git a/tobac/tests/test_utils_bulk_statistics.py b/tobac/tests/test_utils_bulk_statistics.py
index d50bf539..cc4ae463 100644
--- a/tobac/tests/test_utils_bulk_statistics.py
+++ b/tobac/tests/test_utils_bulk_statistics.py
@@ -1,8 +1,11 @@
 from datetime import datetime
+
 import numpy as np
+import dask.array as da
 import pandas as pd
-import pytest
 import xarray as xr
+import pytest
+
 import tobac
 import tobac.utils as tb_utils
 import tobac.testing as tb_test
@@ -702,3 +705,85 @@ def test_get_statistics_from_mask_collapse_dim():
             statistic=statistics_sum,
             collapse_dim="not_a_dim",
         )
+
+
+def test_bulk_statistics_dask():
+    """
+    Test dask input for labels and fields is handled correctly
+    """
+
+    test_labels = da.array(
+        [
+            [
+                [0, 0, 0, 0, 0],
+                [0, 1, 0, 2, 0],
+                [0, 1, 0, 2, 0],
+                [0, 1, 0, 0, 0],
+                [0, 0, 0, 0, 0],
+            ],
+            [
+                [0, 0, 0, 0, 0],
+                [0, 3, 0, 0, 0],
+                [0, 3, 0, 4, 0],
+                [0, 3, 0, 4, 0],
+                [0, 0, 0, 0, 0],
+            ],
+        ],
+        dtype=int,
+    )
+
+    test_labels = xr.DataArray(
+        test_labels,
+        dims=("time", "y", "x"),
+        coords={
+            "time": [datetime(2000, 1, 1), datetime(2000, 1, 1, 0, 5)],
+            "y": np.arange(5),
+            "x": np.arange(5),
+        },
+    )
+
+    test_values = da.array(
+        [
+            [
+                [0, 0, 0, 0, 0],
+                [0, 1, 0, 2, 0],
+                [0, 2, 0, 2, 0],
+                [0, 3, 0, 0, 0],
+                [0, 0, 0, 0, 0],
+            ],
+            [
+                [0, 0, 0, 0, 0],
+                [0, 2, 0, 0, 0],
+                [0, 3, 0, 3, 0],
+                [0, 4, 0, 2, 0],
+                [0, 0, 0, 0, 0],
+            ],
+        ]
+    )
+
+    test_values = xr.DataArray(
+        test_values, dims=test_labels.dims, coords=test_labels.coords
+    )
+
+    test_features = pd.DataFrame(
+        {
+            "feature": [1, 2, 3, 4],
+            "frame": [0, 0, 1, 1],
+            "time": [
+                datetime(2000, 1, 1),
+                datetime(2000, 1, 1),
+                datetime(2000, 1, 1, 0, 5),
+                datetime(2000, 1, 1, 0, 5),
+            ],
+        }
+    )
+
+    statistics_size = {"size": np.size}
+
+    expected_size_result = np.array([3, 2, 3, 2])
+
+    bulk_statistics_output = tb_utils.get_statistics_from_mask(
+        test_features, test_labels, test_values, statistic=statistics_size
+    )
+
+    assert np.all(bulk_statistics_output["size"] == expected_size_result)