Add weight threshold option for spatial averaging #672

pochedls · 2024-06-28T19:58:43Z

Description

This PR adds an optional argument that requires a minimum fraction of data be available to perform a spatial average. The initial PR is for spatial averaging only (it would need to be expanded to handle temporal averaging).

Closes [Enhancement]: Add weight threshold option for averaging operations #531

Checklist

My code follows the style guidelines of this project
I have performed a self-review of my own code
My changes generate no new warnings
Any dependent changes have been merged and published in downstream modules

If applicable:

I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass with my changes (locally and CI/CD build)
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
I have noted that this is a breaking change for a major release (fix or feature that would cause existing functionality to not work as expected)

codecov · 2024-06-28T20:02:25Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 100.00%. Comparing base (27396e5) to head (9331a02).

Additional details and impacted files

@@            Coverage Diff            @@
##              main      #672   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           15        15           
  Lines         1609      1634   +25     
=========================================
+ Hits          1609      1634   +25

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚨 Try these New Features:

Flaky Tests Detection - Detect and resolve failed and flaky tests

pochedls · 2024-06-28T20:01:29Z

xcdat/spatial.py

+            # zero out cells with missing values in data_var
+            weights = xr.where(~np.isnan(data_var), weights, 0)
+            # sum all weights (including zero for missing values)
+            weight_sum_masked = weights.sum(dim=dim)  # type: ignore


Needed to add type: ignore for unclear reasons...

I think dim is expected to be a string but we pass a Hashable (dict key).

# type: ignore is fine here or set dim=str(dim) to remove that comment

I don't think this should be a string. The loop above creates a list of strings (e.g., ["lat", "lon"]), which is supposedly okay according to xarray docs.

From mypy:

error: Argument "dim" to "sum" of "DataArrayAggregations" has incompatible type "list[str | list[str]]"; expected "str | Collection[Hashable] | ellipsis | None" [arg-type]

I tried tuple(dims), but that didn't work either. I'll leave this alone for now (but if you know how to fix the mypy complaint, please do!).

We can just keep # type: ignore. The docstring for dim says:

Name of dimension[s] along which to apply sum. For e.g. dim="x" or dim=["x", "y"]. If “…” or None, will reduce over all dimensions.

Although the type annotation for dim is Dims = Union[str, Collection[Hashable], "ellipsis", None], the example includes a list (dim=["x", "y"]).

pochedls · 2024-06-28T20:02:15Z

xcdat/spatial.py

@@ -716,6 +725,9 @@ def _averager(self, data_var: xr.DataArray, axis: List[SpatialAxis]):
            Data variable inside a Dataset.
        axis : List[SpatialAxis]
            List of axis dimensions to average over.
+        required_weight : optional, float


Is required_weight a good parameter name? What should the default be (currently zero to make default behavior backwards compatible).

Maybe required_weight_pct? required_weight sounds like a boolean parameter.

I think 0 is the appropriate default value.

Actually I like None as a default value to indicate the arg is not set by default.

None is okay with me (given my comment above). I prefer required_weight (since this isn't a percentage, shorter is better, and this doesn't come across as a boolean to me – I think require_weight sounds more Boolean-like than required_weight).

How about the alternative: minimum_weight?

I can see how required_weight doesn't sound like a boolean. I like minimum_weight, it sounds clearer. How about min_weight?

min_weight is good with me!

xcdat/spatial.py

pochedls · 2024-06-28T20:12:14Z

tests/test_spatial.py

+    def test_spatial_average_with_required_weight_as_None(self):
+        ds = self.ds.copy()
+
+        result = ds.spatial.average(
+            "ts",
+            axis=["X", "Y"],
+            lat_bounds=(-5.0, 5),
+            lon_bounds=(-170, -120.1),
+            required_weight=None,
+        )
+
+        expected = self.ds.copy()
+        expected["ts"] = xr.DataArray(
+            data=np.array([2.25, 1.0, 1.0]),
+            coords={"time": expected.time},
+            dims="time",
+        )
+
+        xr.testing.assert_allclose(result, expected)


This test is related to the code chunk in _averager that deals with the case where required_weight=None. I don't know why it would ever be None (it is supposed to be Optional[float] = 0., so I don't know how it could be None). If you don't have those conditionals you get an issue (maybe from mypy) about comparing None to int | float with the lines that check if required_weight > 0..

So if we could ensure that required_weight is always of type float we could remove this test and other conditionals in _averager. That would be ideal...

Using Optional indicates that the argument can also be None. You can just specify float as the type annotation if None is not expected.

I like None as the default value

I see. From the docs:

Note that this is not the same concept as an optional argument, which is one that has a default.

None is converted to 0. later. I guess None might feel different to the user than 0., so I'll leave this alone.

tomvothecoder · 2024-07-30T21:29:24Z

I'm going to open a separate PR to address this enhancement with the temporal APIs because those require more work and we can merge this earlier when it is ready.

xcdat/spatial.py

tomvothecoder · 2024-07-30T22:42:37Z

I started PR #683 for temporal operations. I used some of your code and split them up into reusable functions. We can think about making these functions generalizable across the spatial and temporal classes.

Check them out here: https://github.com/xCDAT/xcdat/pull/683/files

pochedls

Review review and apply suggestions.

pochedls · 2024-08-23T17:40:02Z

tests/test_spatial.py

+    def test_spatial_average_with_required_weight_as_None(self):
+        ds = self.ds.copy()
+
+        result = ds.spatial.average(
+            "ts",
+            axis=["X", "Y"],
+            lat_bounds=(-5.0, 5),
+            lon_bounds=(-170, -120.1),
+            required_weight=None,
+        )
+
+        expected = self.ds.copy()
+        expected["ts"] = xr.DataArray(
+            data=np.array([2.25, 1.0, 1.0]),
+            coords={"time": expected.time},
+            dims="time",
+        )
+
+        xr.testing.assert_allclose(result, expected)


I see. From the docs:

Note that this is not the same concept as an optional argument, which is one that has a default.

None is converted to 0. later. I guess None might feel different to the user than 0., so I'll leave this alone.

pochedls · 2024-08-23T17:43:30Z

xcdat/spatial.py

@@ -716,6 +725,9 @@ def _averager(self, data_var: xr.DataArray, axis: List[SpatialAxis]):
            Data variable inside a Dataset.
        axis : List[SpatialAxis]
            List of axis dimensions to average over.
+        required_weight : optional, float


None is okay with me (given my comment above). I prefer required_weight (since this isn't a percentage, shorter is better, and this doesn't come across as a boolean to me – I think require_weight sounds more Boolean-like than required_weight).

How about the alternative: minimum_weight?

xcdat/spatial.py

pochedls · 2024-08-23T17:47:07Z

xcdat/spatial.py

+            # zero out cells with missing values in data_var
+            weights = xr.where(~np.isnan(data_var), weights, 0)
+            # sum all weights (including zero for missing values)
+            weight_sum_masked = weights.sum(dim=dim)  # type: ignore


I don't think this should be a string. The loop above creates a list of strings (e.g., ["lat", "lon"]), which is supposedly okay according to xarray docs.

From mypy:

error: Argument "dim" to "sum" of "DataArrayAggregations" has incompatible type "list[str | list[str]]"; expected "str | Collection[Hashable] | ellipsis | None" [arg-type]

I tried tuple(dims), but that didn't work either. I'll leave this alone for now (but if you know how to fix the mypy complaint, please do!).

tomvothecoder · 2024-08-27T22:49:56Z

I've created general functions for weight thresholds here. You can copy them over to utils.py in your branch and then use them in the SpatialAccessor class as needed.

tomvothecoder · 2024-08-28T18:39:23Z

xcdat/spatial.py

+        # need weights to match data_var dimensionality
+        if minimum_weight > 0.0:
+            weights, data_var = xr.broadcast(weights, data_var)


I think we can remove this line.

Per docs for xr.broadcast:

xarray objects automatically broadcast against each other in arithmetic operations, so this function should not be necessary for normal use.

I think the weights are probably of dimensions [lat, lon] and this casts them to [time, lat, lon]. That way when you sum them below you get [time] (instead of a dimensionless scalar) which is then used to figure out which values do not meet the weight threshold and need to be masked.

Xarray should be handling the automatic broadcasting when using the xr.where() API to mask the data using weights. Otherwise, it would probably break with an error that the dimensions don't align.

According to the docs for xr.where():

Performs xarray-like broadcasting across input arguments.

The unit tests pass with this line commented out. I also didn't implement xr.broadcast() explicitly in PR #683.

tomvothecoder · 2024-09-05T17:48:01Z

Rebased this branch on latest main and pushed commit 34b570d (#672) with code review updates (refer to commit message for list of changes).

Co-authored-by: Tom Vo <tomvothecoder@gmail.com>

- Rename arg `minimum_weight` to `min_weight` - Add `_get_masked_weights()` and `_validate_min_weight()` to `utils.py` - Update `SpatialAccessor` to use `_get_masked_weights()` and `_validate_min_weight()` - Replace type annotation `Optional` with `|` - Extract `_mask_var_with_with_threshold()` from `_averager()` for readability

tomvothecoder

@pochedls Your PR is ready for review again.

I actually think it is okay to keep this PR and #683 as separate PRs. I've ported the utilities over from #683 and updated spatial.py to use them.

tomvothecoder · 2024-11-21T20:34:25Z

xcdat/spatial.py

+        # TODO: This conditional might not be needed because Xarray will
+        # automatically broadcast the weights to the data variable for
+        # operations such as .mean() and .where().
+        if min_weight > 0.0:
+            weights, data_var = xr.broadcast(weights, data_var)


Suggested change

# TODO: This conditional might not be needed because Xarray will

# automatically broadcast the weights to the data variable for

# operations such as .mean() and .where().

if min_weight > 0.0:

weights, data_var = xr.broadcast(weights, data_var)

Xarray will automatically broadcast to align the shapes of weights and data_var.

github-actions bot added the type: enhancement New enhancement request label Jun 28, 2024

pochedls commented Jun 28, 2024

View reviewed changes

tomvothecoder changed the title ~~Set (optional) weight threshold for averaging operations~~ Add weight threshold option for spatial averaging Jul 30, 2024

tomvothecoder mentioned this pull request Jul 30, 2024

Add weight threshold option for temporal operations #683

Open

9 tasks

tomvothecoder reviewed Jul 30, 2024

View reviewed changes

xcdat/spatial.py Outdated Show resolved Hide resolved

pochedls commented Aug 23, 2024

View reviewed changes

tomvothecoder added this to the FY24Q4 (07/01/24 - 09/30/24) milestone Aug 28, 2024

tomvothecoder reviewed Aug 28, 2024

View reviewed changes

tomvothecoder force-pushed the feature/531_min_weight_for_average branch from b56befe to 34b570d Compare September 5, 2024 17:47

tomvothecoder modified the milestones: FY24Q4 (07/01/24 - 09/30/24), FY25Q1 (10/01/24 - 12/31/24) Oct 7, 2024

tomvothecoder added the project: seats-fy24 label Oct 10, 2024

pochedls and others added 7 commits November 21, 2024 12:30

initial attempt at #531 (for spatial averaging)

ea08915

cleanup print statement and complete code coverage

127e711

Apply review suggestion.

a2fee04

Co-authored-by: Tom Vo <tomvothecoder@gmail.com>

update required_weight argument (to minimum_weight)

0602ea6

update tests for minimum_weight parameter

1c5c3c1

Fix TypeError for optional region arg in TypedDict

82cbfe2

tomvothecoder force-pushed the feature/531_min_weight_for_average branch from 7c27e6f to 82cbfe2 Compare November 21, 2024 20:32

tomvothecoder reviewed Nov 21, 2024

View reviewed changes

Remove rebase comment

9331a02

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add weight threshold option for spatial averaging #672

Add weight threshold option for spatial averaging #672

pochedls commented Jun 28, 2024 •

edited by tomvothecoder

Loading

codecov bot commented Jun 28, 2024 •

edited

Loading

pochedls Jun 28, 2024

tomvothecoder Jul 30, 2024 •

edited

Loading

pochedls Aug 23, 2024

tomvothecoder Aug 27, 2024 •

edited

Loading

pochedls Jun 28, 2024

tomvothecoder Jul 30, 2024

tomvothecoder Jul 30, 2024

pochedls Aug 23, 2024

tomvothecoder Aug 27, 2024

pochedls Aug 27, 2024

pochedls Jun 28, 2024

tomvothecoder Jul 30, 2024

tomvothecoder Jul 30, 2024

pochedls Aug 23, 2024

tomvothecoder commented Jul 30, 2024

tomvothecoder commented Jul 30, 2024

pochedls left a comment

pochedls Aug 23, 2024

pochedls Aug 23, 2024

pochedls Aug 23, 2024

tomvothecoder commented Aug 27, 2024

tomvothecoder Aug 28, 2024 •

edited

Loading

pochedls Aug 28, 2024

tomvothecoder Sep 5, 2024

tomvothecoder Sep 5, 2024

tomvothecoder commented Sep 5, 2024 •

edited

Loading

tomvothecoder left a comment

tomvothecoder Nov 21, 2024

Add weight threshold option for spatial averaging #672

Are you sure you want to change the base?

Add weight threshold option for spatial averaging #672

Conversation

pochedls commented Jun 28, 2024 • edited by tomvothecoder Loading

Description

Checklist

codecov bot commented Jun 28, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

tomvothecoder Jul 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tomvothecoder Aug 27, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tomvothecoder commented Jul 30, 2024

tomvothecoder commented Jul 30, 2024

pochedls left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tomvothecoder commented Aug 27, 2024

tomvothecoder Aug 28, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tomvothecoder commented Sep 5, 2024 • edited Loading

tomvothecoder left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pochedls commented Jun 28, 2024 •

edited by tomvothecoder

Loading

codecov bot commented Jun 28, 2024 •

edited

Loading

tomvothecoder Jul 30, 2024 •

edited

Loading

tomvothecoder Aug 27, 2024 •

edited

Loading

tomvothecoder Aug 28, 2024 •

edited

Loading

tomvothecoder commented Sep 5, 2024 •

edited

Loading