add zonal_anomaly with shared face-band weight kernel#1508
Open
rajeeja wants to merge 2 commits into
Open
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds a new UxDataArray.zonal_anomaly() API to compute per-face zonal anomalies (subtracting each face’s latitude-band mean) and refactors conservative zonal-mean geometry into a shared face–band overlap kernel.
Changes:
- Added
UxDataArray.zonal_anomaly()public method to return an anomaly field on the original unstructured grid. - Introduced
_compute_face_band_weights()to share conservative face–latitude-band overlap calculations. - Refactored conservative zonal-mean implementation to use the shared weight kernel.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 7 comments.
| File | Description |
|---|---|
uxarray/core/zonal.py |
Adds shared face–band weight computation and implements zonal anomaly backend logic. |
uxarray/core/dataarray.py |
Exposes zonal_anomaly() as a new UxDataArray method. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
259
to
262
| bands = np.asarray(bands, dtype=float) | ||
| if bands.ndim != 1 or bands.size < 2: | ||
| raise ValueError("bands must be 1D with at least two edges") | ||
|
|
||
| nb = bands.size - 1 | ||
|
|
||
| # Initialize result array | ||
| shape = list(uxda.shape) | ||
| shape[face_axis] = nb | ||
| if isinstance(uxda.data, da.Array): | ||
| result = da.zeros(shape, dtype=uxda.dtype) | ||
| else: | ||
| result = np.zeros(shape, dtype=uxda.dtype) | ||
| W = np.zeros((uxgrid.n_face, nb), dtype=float) | ||
|
|
Comment on lines
+376
to
+386
| band_means = np.full(nb, np.nan) | ||
| for bi in range(nb): | ||
| overlapping = np.nonzero(W[:, bi] > 0)[0] | ||
| if overlapping.size == 0: | ||
| continue | ||
| w = W[overlapping, bi] | ||
| total = w.sum() | ||
| if total > 0: | ||
| vals = uxda.isel(n_face=overlapping, ignore_grid=True).values | ||
| band_means[bi] = (w * vals).sum() / total | ||
|
|
Comment on lines
+407
to
+415
| band_means = np.full(nb, np.nan) | ||
| for bi in range(nb): | ||
| mask = band_indices == bi | ||
| if mask.any(): | ||
| band_means[bi] = float( | ||
| uxda.isel( | ||
| n_face=np.nonzero(mask)[0], ignore_grid=True | ||
| ).values.mean() | ||
| ) |
Comment on lines
+404
to
+406
| face_lats = uxda.uxgrid.face_lat.values | ||
| band_indices = np.clip(np.digitize(face_lats, bands) - 1, 0, nb - 1) | ||
|
|
| # Broadcast face_means to match uxda shape (face axis may not be last) | ||
| shape = [1] * uxda.ndim | ||
| shape[face_axis] = n_face | ||
| return uxda.values - face_means.reshape(shape) |
Comment on lines
+771
to
+833
| def zonal_anomaly(self, lat=(-90, 90, 10), conservative: bool = False): | ||
| """Compute the zonal anomaly: each face value minus the mean of its latitude band. | ||
|
|
||
| Returns a new ``UxDataArray`` with the same dimensions as the input, | ||
| where each face holds its original value minus the zonal mean of the | ||
| latitude band it belongs to. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| lat : tuple or array-like, default=(-90, 90, 10) | ||
| Latitude band specification: | ||
| - tuple (start, end, step): band edges via np.linspace(start, end, n) | ||
| - array-like: explicit band edges in degrees | ||
| conservative : bool, default=False | ||
| If True, uses area-weighted band means and blends across bands for | ||
| faces that straddle a band boundary, reusing the face-band weight | ||
| matrix computed for zonal_mean so no geometry is duplicated. | ||
| If False, assigns each face to a band by its centroid latitude. | ||
|
|
||
| Returns | ||
| ------- | ||
| UxDataArray | ||
| Same dimensions as input with per-face band mean subtracted. | ||
|
|
||
| Examples | ||
| -------- | ||
| >>> uxds["var"].zonal_anomaly() | ||
| >>> uxds["var"].zonal_anomaly(lat=(-60, 60, 5), conservative=True) | ||
| """ | ||
| if not self._face_centered(): | ||
| raise ValueError( | ||
| "Zonal anomaly is only supported for face-centered data variables." | ||
| ) | ||
|
|
||
| if isinstance(lat, tuple): | ||
| start, end, step = lat | ||
| if step <= 0: | ||
| raise ValueError("Step size must be positive.") | ||
| num_points = int(round((end - start) / step)) + 1 | ||
| edges = np.linspace(start, end, num_points) | ||
| edges = np.clip(edges, -90, 90) | ||
| elif isinstance(lat, (list, np.ndarray)): | ||
| edges = np.asarray(lat, dtype=float) | ||
| else: | ||
| raise ValueError( | ||
| "Invalid value for 'lat'. Must be a tuple (start, end, step) or array-like band edges." | ||
| ) | ||
|
|
||
| if edges.ndim != 1 or edges.size < 2: | ||
| raise ValueError("Band edges must be 1D with at least two values.") | ||
|
|
||
| res = _compute_zonal_anomaly(self, edges, conservative=conservative) | ||
|
|
||
| return UxDataArray( | ||
| res, | ||
| dims=self.dims, | ||
| coords=self.coords, | ||
| name=self.name + "_zonal_anomaly" | ||
| if self.name is not None | ||
| else "zonal_anomaly", | ||
| attrs={"zonal_anomaly": True, "conservative": conservative}, | ||
| uxgrid=self.uxgrid, | ||
| ) |
Comment on lines
+320
to
+323
| if isinstance(uxda.data, da.Array): | ||
| result = da.full(shape, np.nan, dtype=float) | ||
| else: | ||
| result = np.full(shape, np.nan, dtype=float) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
UxDataArray.zonal_anomaly()returning a same-dimension array where each face value has its latitude-band mean subtracted.To avoid duplicating the expensive geometric intersection work already done by
zonal_mean, the face-band overlap areas are extracted into a shared kernel_compute_face_band_weights(uxgrid, bands)that bothzonal_meanandzonal_anomalycall. The conservative path blends band means across straddling faces using the same weight matrix; the default (non-conservative) path assigns each face to a band by centroid latitude.Fixes #1247