Skip to content

DRAFT: Optimized Non-Conservative Zonal Average #1180

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
hongyuchen1030 wants to merge 11 commits into from
Closed

DRAFT: Optimized Non-Conservative Zonal Average #1180

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
a031ea1
initial commit
hongyuchen1030 Mar 12, 2025
3fc6358
Finish all easyckean without touching the latlon bounds
hongyuchen1030 Mar 12, 2025
d2a3aa7
Try adding `face_edges_spherical` and `face_edges_cartesian` as prope…
hongyuchen1030 Mar 13, 2025
bbf52c3
Apply the initial changes
hongyuchen1030 Mar 14, 2025
f994a84
remove the duplicate call for `face_edges_cartesian` calculation
hongyuchen1030 Mar 14, 2025
b392046
Finish all necessary changes
hongyuchen1030 Mar 14, 2025
8621fa5
Avoid unnecessary co-planer check and created `_intersection_within_i…
hongyuchen1030 Apr 7, 2025
0b87406
Run precommit
hongyuchen1030 Apr 7, 2025
c6e5802
fix bugs
hongyuchen1030 Apr 7, 2025
dabeb68
Merge branch 'main' into optimized_nonconservative_avg
hongyuchen1030 Apr 8, 2025
8e7424f
Merge branch 'main' into optimized_nonconservative_avg
philipc2 Apr 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 25 additions & 65 deletions test/test_geometry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from uxarray.grid.coordinates import _populate_node_latlon, _lonlat_rad_to_xyz, _normalize_xyz, _xyz_to_lonlat_rad, \
_xyz_to_lonlat_deg, _xyz_to_lonlat_rad_scalar
from uxarray.grid.arcs import extreme_gca_latitude, extreme_gca_z
from uxarray.grid.utils import _get_cartesian_face_edge_nodes, _get_lonlat_rad_face_edge_nodes
from uxarray.grid.utils import _get_cartesian_faces_edge_nodes, _get_lonlat_rad_faces_edge_nodes

from uxarray.grid.geometry import _populate_face_latlon_bound, _populate_bounds, _pole_point_inside_polygon_cartesian, \
stereographic_projection, inverse_stereographic_projection, point_in_face, haversine_distance
Expand Down Expand Up @@ -1184,14 +1184,9 @@ def test_point_inside():
grid = ux.open_grid(grid_mpas_2)

# Get the face edges of all faces in the grid
faces_edges_cartesian = _get_cartesian_face_edge_nodes(
grid.face_node_connectivity.values,
grid.n_face,
grid.n_max_face_edges,
grid.node_x.values,
grid.node_y.values,
grid.node_z.values,
)
faces_edges_cartesian = _get_cartesian_faces_edge_nodes(grid.face_node_connectivity.values, grid.n_face,
grid.n_max_face_edges, grid.node_x.values,
grid.node_y.values, grid.node_z.values)

# Loop through each face
for i in range(grid.n_face):
Expand All @@ -1209,14 +1204,9 @@ def test_point_outside():
grid = ux.open_grid(grid_mpas_2)

# Get the face edges of all faces in the grid
faces_edges_cartesian = _get_cartesian_face_edge_nodes(
grid.face_node_connectivity.values,
grid.n_face,
grid.n_max_face_edges,
grid.node_x.values,
grid.node_y.values,
grid.node_z.values,
)
faces_edges_cartesian = _get_cartesian_faces_edge_nodes(grid.face_node_connectivity.values, grid.n_face,
grid.n_max_face_edges, grid.node_x.values,
grid.node_y.values, grid.node_z.values)

# Set the point as the face center of a different face than the face tested
point_xyz = np.array([grid.face_x[1].values, grid.face_y[1].values, grid.face_z[1].values])
Expand All @@ -1232,14 +1222,9 @@ def test_point_on_node():
grid = ux.open_grid(grid_mpas_2)

# Get the face edges of all faces in the grid
faces_edges_cartesian = _get_cartesian_face_edge_nodes(
grid.face_node_connectivity.values,
grid.n_face,
grid.n_max_face_edges,
grid.node_x.values,
grid.node_y.values,
grid.node_z.values,
)
faces_edges_cartesian = _get_cartesian_faces_edge_nodes(grid.face_node_connectivity.values, grid.n_face,
grid.n_max_face_edges, grid.node_x.values,
grid.node_y.values, grid.node_z.values)

# Set the point as a node
point_xyz = np.array([*faces_edges_cartesian[0][0][0]])
Expand All @@ -1263,14 +1248,9 @@ def test_point_inside_close():

# Create the grid and face edges
grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True)
faces_edges_cartesian = _get_cartesian_face_edge_nodes(
grid.face_node_connectivity.values,
grid.n_face,
grid.n_max_face_edges,
grid.node_x.values,
grid.node_y.values,
grid.node_z.values,
)
faces_edges_cartesian = _get_cartesian_faces_edge_nodes(grid.face_node_connectivity.values, grid.n_face,
grid.n_max_face_edges, grid.node_x.values,
grid.node_y.values, grid.node_z.values)

# Use point in face to determine if the point is inside or out of the face
assert point_in_face(faces_edges_cartesian[0], point_xyz=point, inclusive=False)
Expand All @@ -1288,14 +1268,9 @@ def test_point_outside_close():

# Create the grid and face edges
grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True)
faces_edges_cartesian = _get_cartesian_face_edge_nodes(
grid.face_node_connectivity.values,
grid.n_face,
grid.n_max_face_edges,
grid.node_x.values,
grid.node_y.values,
grid.node_z.values,
)
faces_edges_cartesian = _get_cartesian_faces_edge_nodes(grid.face_node_connectivity.values, grid.n_face,
grid.n_max_face_edges, grid.node_x.values,
grid.node_y.values, grid.node_z.values)

# Use point in face to determine if the point is inside or out of the face
assert not point_in_face(faces_edges_cartesian[0], point_xyz=point, inclusive=False)
Expand All @@ -1312,14 +1287,9 @@ def test_face_at_pole():

# Create the grid and face edges
grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True)
faces_edges_cartesian = _get_cartesian_face_edge_nodes(
grid.face_node_connectivity.values,
grid.n_face,
grid.n_max_face_edges,
grid.node_x.values,
grid.node_y.values,
grid.node_z.values,
)
faces_edges_cartesian = _get_cartesian_faces_edge_nodes(grid.face_node_connectivity.values, grid.n_face,
grid.n_max_face_edges, grid.node_x.values,
grid.node_y.values, grid.node_z.values)

assert point_in_face(faces_edges_cartesian[0], point_xyz=point, inclusive=True)

Expand All @@ -1334,14 +1304,9 @@ def test_face_at_antimeridian():

# Create the grid and face edges
grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True)
faces_edges_cartesian = _get_cartesian_face_edge_nodes(
grid.face_node_connectivity.values,
grid.n_face,
grid.n_max_face_edges,
grid.node_x.values,
grid.node_y.values,
grid.node_z.values,
)
faces_edges_cartesian = _get_cartesian_faces_edge_nodes(grid.face_node_connectivity.values, grid.n_face,
grid.n_max_face_edges, grid.node_x.values,
grid.node_y.values, grid.node_z.values)

assert point_in_face(faces_edges_cartesian[0], point_xyz=point, inclusive=True)

Expand All @@ -1357,14 +1322,9 @@ def test_face_normal_face():

# Create the grid and face edges
grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True)
faces_edges_cartesian = _get_cartesian_face_edge_nodes(
grid.face_node_connectivity.values,
grid.n_face,
grid.n_max_face_edges,
grid.node_x.values,
grid.node_y.values,
grid.node_z.values,
)
faces_edges_cartesian = _get_cartesian_faces_edge_nodes(grid.face_node_connectivity.values, grid.n_face,
grid.n_max_face_edges, grid.node_x.values,
grid.node_y.values, grid.node_z.values)

assert point_in_face(faces_edges_cartesian[0], point_xyz=point, inclusive=True)

Expand Down
32 changes: 13 additions & 19 deletions test/test_helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
from uxarray.constants import INT_DTYPE, INT_FILL_VALUE
from uxarray.grid.coordinates import _lonlat_rad_to_xyz, _normalize_xyz, _xyz_to_lonlat_rad
from uxarray.grid.arcs import point_within_gca, _angle_of_2_vectors, in_between
from uxarray.grid.utils import _get_cartesian_face_edge_nodes, _get_lonlat_rad_face_edge_nodes
from uxarray.grid.utils import _get_cartesian_faces_edge_nodes, _get_lonlat_rad_faces_edge_nodes
from uxarray.grid.geometry import pole_point_inside_polygon, _pole_point_inside_polygon_cartesian

try:
Expand Down Expand Up @@ -273,9 +273,8 @@ def test_get_cartesian_face_edge_nodes_pipeline():
node_y = grid.node_y.values
node_z = grid.node_z.values

face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes(
face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z
)
face_edges_connectivity_cartesian = _get_cartesian_faces_edge_nodes(face_node_conn, n_face, n_max_face_edges,
node_x, node_y, node_z)

result = _pole_point_inside_polygon_cartesian(
'North', face_edges_connectivity_cartesian[0]
Expand All @@ -298,9 +297,8 @@ def test_get_cartesian_face_edge_nodes_filled_value():
node_y = grid.node_y.values
node_z = grid.node_z.values

face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes(
face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z
)
face_edges_connectivity_cartesian = _get_cartesian_faces_edge_nodes(face_node_conn, n_face, n_max_face_edges,
node_x, node_y, node_z)

result = _pole_point_inside_polygon_cartesian(
'North', face_edges_connectivity_cartesian[0]
Expand Down Expand Up @@ -335,9 +333,8 @@ def test_get_cartesian_face_edge_nodes_filled_value2():
node_y = np.array([v0_cart[1],v1_cart[1],v2_cart[1],v3_cart[1],v4_cart[1]])
node_z = np.array([v0_cart[2],v1_cart[2],v2_cart[2],v3_cart[2],v4_cart[2]])

face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes(
face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z
)
face_edges_connectivity_cartesian = _get_cartesian_faces_edge_nodes(face_node_conn, n_face, n_max_face_edges,
node_x, node_y, node_z)

correct_result = np.array([
[
Expand Down Expand Up @@ -368,9 +365,8 @@ def test_get_lonlat_face_edge_nodes_pipeline():
node_lon = grid.node_lon.values
node_lat = grid.node_lat.values

face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes(
face_node_conn, n_face, n_max_face_edges, node_lon, node_lat
)
face_edges_connectivity_lonlat = _get_lonlat_rad_faces_edge_nodes(face_node_conn, n_face, n_max_face_edges,
node_lon, node_lat)

face_edges_connectivity_lonlat = face_edges_connectivity_lonlat[0]
face_edges_connectivity_cartesian = []
Expand Down Expand Up @@ -398,9 +394,8 @@ def test_get_lonlat_face_edge_nodes_filled_value():
node_lon = grid.node_lon.values
node_lat = grid.node_lat.values

face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes(
face_node_conn, n_face, n_max_face_edges, node_lon, node_lat
)
face_edges_connectivity_lonlat = _get_lonlat_rad_faces_edge_nodes(face_node_conn, n_face, n_max_face_edges,
node_lon, node_lat)

face_edges_connectivity_lonlat = face_edges_connectivity_lonlat[0]
face_edges_connectivity_cartesian = []
Expand Down Expand Up @@ -434,9 +429,8 @@ def test_get_lonlat_face_edge_nodes_filled_value2():
node_lon = np.array([v0_rad[0],v1_rad[0],v2_rad[0],v3_rad[0],v4_rad[0]])
node_lat = np.array([v0_rad[1],v1_rad[1],v2_rad[1],v3_rad[1],v4_rad[1]])

face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes(
face_node_conn, n_face, n_max_face_edges, node_lon, node_lat
)
face_edges_connectivity_lonlat = _get_lonlat_rad_faces_edge_nodes(face_node_conn, n_face, n_max_face_edges,
node_lon, node_lat)

correct_result = np.array([
[
Expand Down
22 changes: 11 additions & 11 deletions test/test_integrate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
_get_faces_constLat_intersection_info, _zonal_face_weights, \
_zonal_face_weights_robust

from uxarray.grid.utils import _get_cartesian_face_edge_nodes
from uxarray.grid.utils import _get_cartesian_faces_edge_nodes

current_path = Path(os.path.dirname(os.path.realpath(__file__)))

Expand Down Expand Up @@ -67,9 +67,9 @@ def test_get_faces_constLat_intersection_info_one_intersection():
latitude_cart = -0.8660254037844386
is_latlonface = False
is_GCA_list = None
unique_intersections, pt_lon_min, pt_lon_max = _get_faces_constLat_intersection_info(face_edges_cart, latitude_cart,
unique_intersections_cart,unique_intersections_lonlat, pt_lon_min, pt_lon_max = _get_faces_constLat_intersection_info(face_edges_cart, latitude_cart,
is_GCA_list, is_latlonface)
assert len(unique_intersections) == 1
assert len(unique_intersections_cart) == 1


def test_get_faces_constLat_intersection_info_encompass_pole():
Expand Down Expand Up @@ -97,9 +97,9 @@ def test_get_faces_constLat_intersection_info_encompass_pole():

is_latlonface = False
is_GCA_list = None
unique_intersections, pt_lon_min, pt_lon_max = _get_faces_constLat_intersection_info(face_edges_cart, latitude_cart,
is_GCA_list, is_latlonface)
assert len(unique_intersections) <= 2 * len(face_edges_cart)
unique_intersections_cart,unique_intersections_lonlat, pt_lon_min, pt_lon_max = _get_faces_constLat_intersection_info(face_edges_cart, latitude_cart,
is_GCA_list, is_latlonface)
assert len(unique_intersections_cart) <= 2 * len(face_edges_cart)


def test_get_faces_constLat_intersection_info_on_pole():
Expand All @@ -119,9 +119,9 @@ def test_get_faces_constLat_intersection_info_on_pole():
latitude_cart = -0.9998476951563913
is_latlonface = False
is_GCA_list = None
unique_intersections, pt_lon_min, pt_lon_max = _get_faces_constLat_intersection_info(face_edges_cart, latitude_cart,
unique_intersections_cart,unique_intersections_lonlat, pt_lon_min, pt_lon_max = _get_faces_constLat_intersection_info(face_edges_cart, latitude_cart,
is_GCA_list, is_latlonface)
assert len(unique_intersections) == 2
assert len(unique_intersections_cart) == 2


def test_get_faces_constLat_intersection_info_near_pole():
Expand All @@ -141,9 +141,9 @@ def test_get_faces_constLat_intersection_info_near_pole():
latitude_deg = np.rad2deg(latitude_rad)
is_latlonface = False
is_GCA_list = None
unique_intersections, pt_lon_min, pt_lon_max = _get_faces_constLat_intersection_info(face_edges_cart, latitude_cart,
unique_intersections_cart,unique_intersections_lonlat, pt_lon_min, pt_lon_max = _get_faces_constLat_intersection_info(face_edges_cart, latitude_cart,
is_GCA_list, is_latlonface)
assert len(unique_intersections) == 1
assert len(unique_intersections_cart) == 1


def test_get_zonal_face_interval():
Expand Down Expand Up @@ -1037,7 +1037,7 @@ def test_compare_zonal_weights():
for gridfile in gridfiles:
uxgrid = ux.open_grid(gridfile)
n_nodes_per_face = uxgrid.n_nodes_per_face.values
face_edge_nodes_xyz = _get_cartesian_face_edge_nodes(
face_edge_nodes_xyz = _get_cartesian_faces_edge_nodes(
uxgrid.face_node_connectivity.values,
uxgrid.n_face,
uxgrid.n_max_face_edges,
Expand Down
12 changes: 4 additions & 8 deletions test/test_intersections.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,20 +161,16 @@ def test_GCA_GCA_north_pole_angled():

def test_GCA_edge_intersection_count():

from uxarray.grid.utils import _get_cartesian_face_edge_nodes
from uxarray.grid.utils import _get_cartesian_faces_edge_nodes

# Generate a normal face that is not crossing the antimeridian or the poles
vertices_lonlat = [[29.5, 11.0], [29.5, 10.0], [30.5, 10.0], [30.5, 11.0]]
vertices_lonlat = np.array(vertices_lonlat)

grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True)
face_edge_nodes_cartesian = _get_cartesian_face_edge_nodes(
grid.face_node_connectivity.values,
grid.n_face,
grid.n_max_face_edges,
grid.node_x.values,
grid.node_y.values,
grid.node_z.values)
face_edge_nodes_cartesian = _get_cartesian_faces_edge_nodes(grid.face_node_connectivity.values, grid.n_face,
grid.n_max_face_edges, grid.node_x.values,
grid.node_y.values, grid.node_z.values)

face_center_xyz = np.array([grid.face_x.values[0], grid.face_y.values[0], grid.face_z.values[0]], dtype=np.float64)
north_pole_xyz = np.array([0.0, 0.0, 1.0], dtype=np.float64)
Expand Down
13 changes: 2 additions & 11 deletions uxarray/core/zonal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@


from uxarray.grid.integrate import _zonal_face_weights, _zonal_face_weights_robust
from uxarray.grid.utils import _get_cartesian_face_edge_nodes


def _compute_non_conservative_zonal_mean(uxda, latitudes, use_robust_weights=False):
Expand All @@ -18,23 +17,15 @@ def _compute_non_conservative_zonal_mean(uxda, latitudes, use_robust_weights=Fal
# Create a NumPy array for storing results
result = np.zeros(shape, dtype=uxda.dtype)

faces_edge_nodes_xyz = _get_cartesian_face_edge_nodes(
uxgrid.face_node_connectivity.values,
uxgrid.n_face,
uxgrid.n_max_face_nodes,
uxgrid.node_x.values,
uxgrid.node_y.values,
uxgrid.node_z.values,
)

bounds = uxgrid.bounds.values
faces_edges_cartesian = uxgrid.faces_edges_cartesian.values

for i, lat in enumerate(latitudes):
face_indices = uxda.uxgrid.get_faces_at_constant_latitude(lat)

z = np.sin(np.deg2rad(lat))

faces_edge_nodes_xyz_candidate = faces_edge_nodes_xyz[face_indices, :, :, :]
faces_edge_nodes_xyz_candidate = faces_edges_cartesian[face_indices, :, :, :]

n_nodes_per_face_candidate = n_nodes_per_face[face_indices]

Expand Down
Loading
Loading