Skip to content

Commit e0098bc

Browse files
philipc2hongyuchen1030
philipc2
and
hongyuchen1030
authored
Non-Conservative Zonal Mean (#1004)
* Zonal mean with fast cross sections * docstring * work on zonal * implement get_weights, integrate into calculation * add weights module * revert benchmark changes * work on handling north and south pole case * pole point indices * run pre-commit * remove cached pole indices for now * add user guide notebook * add tests * docstrings * add image and work on spherical example * add case for extreme lat * update API and restore files * delete figure that is not needed * notebook * update notebook * fix failing test * remove non-bounds test case * update docstrings * fix parameter * clean up docstrings * add tests * clean notebook * clean up test comments * remove edge weights for now and update api ref * update extreme docstring * fix docstring * fix docstring * remove unused plot util function * clean up accessor.py and include docstrings * update test * update test case * use correct weights * use bounds directly without converting to degrees * debugging * debugging * remove unused code in Grid * update tests * remove debug print and copy * clean up grid * use fma to prevent singular matrix * use numba for jacobian * update jacobian and work on weight optimization * update bench * revisions * update notebook * update docstring and notebook * use polars instead of pandas * remove comments * update weight assertion * remove commented out code * correct zonal mean at pole test * initial refacor of weights, need to clean up * clean up docstrings * add zonal average alais, update docstrings * update usage * update parameter name * add test case * update user guide * remove unused code * clean up leftover code * remove commented out code * add tests for plot --------- Co-authored-by: Hongyu Chen <[email protected]>
1 parent 0fd0c7e commit e0098bc

File tree

11 files changed

+493
-59
lines changed

11 files changed

+493
-59
lines changed

‎benchmarks/mpas_ocean.py‎

Lines changed: 25 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,10 @@
22
import urllib.request
33
from pathlib import Path
44

5-
import uxarray as ux
6-
75
import numpy as np
86

7+
import uxarray as ux
8+
99
current_path = Path(os.path.dirname(os.path.realpath(__file__)))
1010

1111
data_var = 'bottomDepth'
@@ -33,15 +33,16 @@
3333
class DatasetBenchmark:
3434
"""Class used as a template for benchmarks requiring a ``UxDataset`` in
3535
this module across both resolutions."""
36-
param_names = ['resolution',]
37-
params = [['480km', '120km'],]
36+
param_names = ['resolution', ]
37+
params = [['480km', '120km'], ]
3838

3939
def setup(self, resolution, *args, **kwargs):
4040
self.uxds = ux.open_dataset(file_path_dict[resolution][0], file_path_dict[resolution][1])
4141

4242
def teardown(self, resolution, *args, **kwargs):
4343
del self.uxds
4444

45+
4546
class GridBenchmark:
4647
"""Class used as a template for benchmarks requiring a ``Grid`` in this
4748
module across both resolutions."""
@@ -62,6 +63,7 @@ def time_gradient(self, resolution):
6263
def peakmem_gradient(self, resolution):
6364
grad = self.uxds[data_var].gradient()
6465

66+
6567
class Integrate(DatasetBenchmark):
6668
def time_integrate(self, resolution):
6769
self.uxds[data_var].integrate()
@@ -74,12 +76,10 @@ class GeoDataFrame(DatasetBenchmark):
7476
param_names = DatasetBenchmark.param_names + ['exclude_antimeridian']
7577
params = DatasetBenchmark.params + [[True, False]]
7678

77-
7879
def time_to_geodataframe(self, resolution, exclude_antimeridian):
7980
self.uxds[data_var].to_geodataframe(exclude_antimeridian=exclude_antimeridian)
8081

8182

82-
8383
class ConnectivityConstruction(DatasetBenchmark):
8484
def time_n_nodes_per_face(self, resolution):
8585
self.uxds.uxgrid.n_nodes_per_face
@@ -136,6 +136,7 @@ def time_nearest_neighbor_remapping(self):
136136
def time_inverse_distance_weighted_remapping(self):
137137
self.uxds_480["bottomDepth"].remap.inverse_distance_weighted(self.uxds_120.uxgrid)
138138

139+
139140
class HoleEdgeIndices(DatasetBenchmark):
140141
def time_construct_hole_edge_indices(self, resolution):
141142
ux.grid.geometry._construct_boundary_edge_indices(self.uxds.uxgrid.edge_face_connectivity)
@@ -145,6 +146,7 @@ class DualMesh(DatasetBenchmark):
145146
def time_dual_mesh_construction(self, resolution):
146147
self.uxds.uxgrid.get_dual()
147148

149+
148150
class ConstructFaceLatLon(GridBenchmark):
149151
def time_welzl(self, resolution):
150152
self.uxgrid.construct_face_centers(method='welzl')
@@ -177,9 +179,26 @@ def setup(self, resolution, lat_step):
177179
self.lats = np.arange(-45, 45, lat_step)
178180
_ = self.uxgrid.bounds
179181

182+
class CrossSections(DatasetBenchmark):
183+
param_names = DatasetBenchmark.param_names + ['n_lat']
184+
params = DatasetBenchmark.params + [[1, 2, 4, 8]]
185+
186+
def time_constant_lat_fast(self, resolution, n_lat):
187+
for lat in np.linspace(-89, 89, n_lat):
188+
self.uxds.uxgrid.constant_latitude_cross_section(lat, method='fast')
180189
def teardown(self, resolution, lat_step):
181190
del self.uxgrid
182191

183192
def time_const_lat(self, resolution, lat_step):
184193
for lat in self.lats:
185194
self.uxgrid.cross_section.constant_latitude(lat)
195+
196+
197+
class ZonalAverage(DatasetBenchmark):
198+
def setup(self, resolution, *args, **kwargs):
199+
self.uxds = ux.open_dataset(file_path_dict[resolution][0], file_path_dict[resolution][1])
200+
bounds = self.uxds.uxgrid.bounds
201+
202+
def time_zonal_average(self, resolution):
203+
lat_step = 10
204+
self.uxds['bottomDepth'].zonal_mean(lat=(-45, 45, lat_step))

‎docs/api.rst‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,7 @@ I/O & Conversion
197197
UxDataArray.to_dataset
198198
UxDataArray.from_xarray
199199

200+
200201
UxDataset
201202
-----------
202203

@@ -265,6 +266,8 @@ UxDataArray
265266
UxDataArray.plot
266267
UxDataArray.plot.polygons
267268
UxDataArray.plot.points
269+
UxDataArray.plot.line
270+
UxDataArray.plot.scatter
268271

269272
UxDataset
270273
~~~~~~~~~
@@ -422,6 +425,15 @@ on each face.
422425
UxDataArray.topological_all
423426
UxDataArray.topological_any
424427

428+
Zonal Average
429+
~~~~~~~~~~~~~
430+
.. autosummary::
431+
:toctree: generated/
432+
433+
UxDataArray.zonal_mean
434+
435+
436+
425437
Weighted
426438
~~~~~~~~
427439
.. autosummary::

‎docs/user-guide/zonal-average.ipynb‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"id": "dbade3aaedad9ae7",
6+
"metadata": {},
7+
"source": [
8+
"# Zonal Averaging\n",
9+
"\n",
10+
"This section demonstrates how to perform Zonal Averaging using UXarray.\n"
11+
]
12+
},
13+
{
14+
"cell_type": "code",
15+
"id": "185e2061bc4c75b9",
16+
"metadata": {},
17+
"source": [
18+
"import uxarray as ux\n",
19+
"import numpy as np\n",
20+
"\n",
21+
"uxds = ux.open_dataset(\n",
22+
" \"../../test/meshfiles/ugrid/outCSne30/outCSne30.ug\",\n",
23+
" \"../../test/meshfiles/ugrid/outCSne30/outCSne30_vortex.nc\",\n",
24+
")\n",
25+
"uxds[\"psi\"].plot(cmap=\"inferno\", periodic_elements=\"split\")"
26+
],
27+
"outputs": [],
28+
"execution_count": null
29+
},
30+
{
31+
"cell_type": "markdown",
32+
"id": "d938d659b89bc9cb",
33+
"metadata": {},
34+
"source": [
35+
"## What is a Zonal Average/Mean?\n",
36+
"\n",
37+
"A zonal average (or zonal mean) is a statistical measure that represents the average of a variable along one or more lines of constant latitude. In other words, it's the mean value calculated around the sphere at constant latitudes. \n",
38+
"\n",
39+
"UXarray currently implements a non-conservative Zonal Mean, which weights candidate faces by the length of intersection of a line of constant latitude.\n",
40+
"\n",
41+
"\n",
42+
"```{seealso}\n",
43+
"[NCL Zonal Average](https://www.ncl.ucar.edu/Applications/zonal.shtml)\n",
44+
"```"
45+
]
46+
},
47+
{
48+
"cell_type": "code",
49+
"id": "d342f8f449543994",
50+
"metadata": {},
51+
"source": [
52+
"zonal_mean_psi = uxds[\"psi\"].zonal_mean()\n",
53+
"zonal_mean_psi"
54+
],
55+
"outputs": [],
56+
"execution_count": null
57+
},
58+
{
59+
"cell_type": "markdown",
60+
"id": "65194a38c76e8a62",
61+
"metadata": {},
62+
"source": [
63+
"The default latitude range is between -90 and 90 degrees with a step size of 10 degrees. "
64+
]
65+
},
66+
{
67+
"cell_type": "code",
68+
"id": "b5933beaf2f598ab",
69+
"metadata": {},
70+
"source": [
71+
"(zonal_mean_psi.plot.line() * zonal_mean_psi.plot.scatter(color=\"red\")).opts(\n",
72+
" title=\"Zonal Average Plot (Default)\", xticks=np.arange(-90, 100, 20), xlim=(-95, 95)\n",
73+
")"
74+
],
75+
"outputs": [],
76+
"execution_count": null
77+
},
78+
{
79+
"cell_type": "markdown",
80+
"id": "ba2d641a3076c692",
81+
"metadata": {},
82+
"source": [
83+
"The range of latitudes can be modified by using the `lat` parameter. It accepts:\n",
84+
"\n",
85+
"* **Single scalar**: e.g., `lat=45`\n",
86+
"* **List/array**: e.g., `lat=[10, 20]` or `lat=np.array([10, 20])`\n",
87+
"* **Tuple**: e.g., `(min_lat, max_lat, step)`"
88+
]
89+
},
90+
{
91+
"cell_type": "code",
92+
"id": "4f665827daac1c46",
93+
"metadata": {},
94+
"source": [
95+
"zonal_mean_psi_large = uxds[\"psi\"].zonal_mean(lat=(-90, 90, 1))"
96+
],
97+
"outputs": [],
98+
"execution_count": null
99+
},
100+
{
101+
"cell_type": "code",
102+
"id": "1998f1a55b67100c",
103+
"metadata": {},
104+
"source": [
105+
"(\n",
106+
" zonal_mean_psi_large.plot.line()\n",
107+
" * zonal_mean_psi_large.plot.scatter(color=\"red\", s=1)\n",
108+
").opts(\n",
109+
" title=\"Zonal Average Plot (Larger Sample)\",\n",
110+
" xticks=np.arange(-90, 100, 20),\n",
111+
" xlim=(-95, 95),\n",
112+
")"
113+
],
114+
"outputs": [],
115+
"execution_count": null
116+
},
117+
{
118+
"cell_type": "markdown",
119+
"id": "4b91ebb8f733a318",
120+
"metadata": {},
121+
"source": [
122+
"## Combined Plots\n",
123+
"\n",
124+
"It is often desired to plot the zonal average along side other plots, such as color or contour plots. "
125+
]
126+
},
127+
{
128+
"cell_type": "code",
129+
"id": "cb2255761173d53e",
130+
"metadata": {},
131+
"source": [
132+
"(\n",
133+
" uxds[\"psi\"].plot(\n",
134+
" cmap=\"inferno\",\n",
135+
" periodic_elements=\"split\",\n",
136+
" height=250,\n",
137+
" width=500,\n",
138+
" colorbar=False,\n",
139+
" ylim=(-90, 90),\n",
140+
" )\n",
141+
" + zonal_mean_psi.plot.line(\n",
142+
" x=\"psi_zonal_mean\",\n",
143+
" y=\"latitudes\",\n",
144+
" height=250,\n",
145+
" width=150,\n",
146+
" ylabel=\"\",\n",
147+
" ylim=(-90, 90),\n",
148+
" xlim=(0.8, 1.2),\n",
149+
" xticks=[0.8, 0.9, 1.0, 1.1, 1.2],\n",
150+
" yticks=[-90, -45, 0, 45, 90],\n",
151+
" grid=True,\n",
152+
" )\n",
153+
").opts(title=\"Combined Zonal Average & Raster Plot\")"
154+
],
155+
"outputs": [],
156+
"execution_count": null
157+
}
158+
],
159+
"metadata": {
160+
"kernelspec": {
161+
"display_name": "Python 3 (ipykernel)",
162+
"language": "python",
163+
"name": "python3"
164+
},
165+
"language_info": {
166+
"codemirror_mode": {
167+
"name": "ipython",
168+
"version": 3
169+
},
170+
"file_extension": ".py",
171+
"mimetype": "text/x-python",
172+
"name": "python",
173+
"nbconvert_exporter": "python",
174+
"pygments_lexer": "ipython3",
175+
"version": "3.12.7"
176+
}
177+
},
178+
"nbformat": 4,
179+
"nbformat_minor": 5
180+
}

‎docs/userguide.rst‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,9 @@ These user guides provide detailed explanations of the core functionality in UXa
5252
`Cross-Sections <user-guide/cross-sections.ipynb>`_
5353
Select cross-sections of a grid
5454

55+
`Zonal Means <user-guide/zonal-average.ipynb>`_
56+
Compute the zonal averages across lines of constant latitude
57+
5558
`Remapping <user-guide/remapping.ipynb>`_
5659
Remap (a.k.a Regrid) between unstructured grids
5760

@@ -104,6 +107,7 @@ These user guides provide additional details about specific features in UXarray.
104107
user-guide/advanced-plotting.ipynb
105108
user-guide/subset.ipynb
106109
user-guide/cross-sections.ipynb
110+
user-guide/zonal-average.ipynb
107111
user-guide/remapping.ipynb
108112
user-guide/topological-aggregations.ipynb
109113
user-guide/weighted_mean.ipynb

0 commit comments

Comments
 (0)