[PR]: Add Z axis support for spatial averaging #606
Conversation
|
I'd appreciate testing with real-world data (and comparisons with CDAT). |
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #606 +/- ##
=========================================
Coverage 100.00% 100.00%
=========================================
Files 15 15
Lines 1588 1595 +7
=========================================
+ Hits 1588 1595 +7 ☔ View full report in Codecov by Sentry. |
| xr.testing.assert_allclose(result["so"], expected) | ||
|
|
||
| # check that vertical averaging returns the correct weights | ||
| expected = xr.DataArray( | ||
| data=np.array([2000, 2000, 1000, 0.0]), | ||
| coords={"time": ds.time, "lev": ds.lev, "lat": ds.lat, "lon": ds.lon}, | ||
| dims=["lev"], | ||
| attrs={"xcdat_bounds": True}, | ||
| ) | ||
|
|
||
| xr.testing.assert_allclose(result["lev_wts"], expected) |
There was a problem hiding this comment.
Are two-in-one tests permitted?
There was a problem hiding this comment.
Yeah there's no problem with that here. As long as the tests are relatively easy to maintain.
| @@ -73,10 +73,12 @@ def average( | |||
| keep_weights: bool = False, | |||
| lat_bounds: Optional[RegionAxisBounds] = None, | |||
| lon_bounds: Optional[RegionAxisBounds] = None, | |||
| lev_bounds: Optional[RegionAxisBounds] = None, | |||
There was a problem hiding this comment.
I think lev_bounds is a good generic name, but I'm open to other possibilities.
| def _get_vertical_weights( | ||
| self, domain_bounds: xr.DataArray, region_bounds: Optional[np.ndarray] | ||
| ) -> xr.DataArray: | ||
| """Gets weights for the vertical axis. | ||
|
|
||
| This method scales the domain to a region (if selected) and returns weights | ||
| proportional to the difference between each pair of level bounds. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| domain_bounds : xr.DataArray | ||
| The array of bounds for the vertical domain. | ||
| region_bounds : Optional[np.ndarray] | ||
| The array of bounds for vertical selection. | ||
|
|
||
| Returns | ||
| ------- | ||
| xr.DataArray | ||
| The vertical axis weights. | ||
| """ | ||
| if region_bounds is not None: | ||
| domain_bounds = self._scale_domain_to_region(domain_bounds, region_bounds) | ||
|
|
||
| weights = self._calculate_weights(domain_bounds) | ||
| return weights | ||
|
|
There was a problem hiding this comment.
I think this is sufficiently different from _get_longitude_weights (which deals with the prime meridian) and _get_latitude_weights (sine of bounds) to justify its own function.
tomvothecoder
changed the title
|
Tom:
|
tomvothecoder
modified the milestones:
FY24Q2 (01/01/24 - 03/31/24),
FY24Q3 (4/01/24 - 06/30/24)
|
I'm going to pick this back up for review. We might be able to get this in for v0.7.0. |
tomvothecoder
force-pushed
the
feature/596-vertical-averaging
branch
from
aa9677d to
98ea3d7
Compare
| so = ds["so"] | ||
| so[:] = np.array([1, 2, 3, 4]) | ||
| ds["so"] = so |
There was a problem hiding this comment.
| so = ds["so"] | |
| so[:] = np.array([1, 2, 3, 4]) | |
| ds["so"] = so | |
| ds["so"].values = np.array([1, 2, 3, 4]) |
Assigning the numpy array directly to the DataArray will work too
Can you provide me some datasets to test with? I wrote a script and tested a single E3SM dataset, which produced nearly identical results (1e-16 max rel diff). |
Do you have xsearch running? I think you could look for monthly, historical data, for the |
I'll check on |
tomvothecoder
modified the milestones:
FY24Q3 (4/01/24 - 06/30/24),
FY24Q4 (07/01/24 - 09/30/24)
tomvothecoder
modified the milestones:
FY24Q4 (07/01/24 - 09/30/24),
FY25Q1 (10/01/24 - 12/31/24)
Description
This PR adds functionality to compute vertical averages.
Checklist
If applicable: