Skip to content
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.

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

Optimized Geometry Functions using Numba #1072

Merged
merged 35 commits into from
Nov 20, 2024
Merged
Show file tree
Hide file tree
Changes from 30 commits
Commits
Show all changes
35 commits
Select commit Hold shift + click to select a range
0a84d62
re-write bounds to use Numba
philipc2 Nov 10, 2024
8e6c5a6
more numba
philipc2 Nov 10, 2024
a9a33d1
fix all tests but two
philipc2 Nov 10, 2024
35cece6
remove unused helper
philipc2 Nov 10, 2024
d833107
fix all tests but one
philipc2 Nov 10, 2024
99da011
recursion limit
philipc2 Nov 11, 2024
9b439c3
recursion limit
philipc2 Nov 11, 2024
976cb57
cache and parallel
philipc2 Nov 11, 2024
248ebd1
Merge branch 'main' into philipc2/bounds-numba
philipc2 Nov 11, 2024
51e71fc
Merge branch 'main' into philipc2/bounds-numba
philipc2 Nov 11, 2024
dcf4ff5
fix failing test
philipc2 Nov 12, 2024
19bf853
Merge branch 'philipc2/bounds-numba' of https://github.com/UXARRAY/ux…
philipc2 Nov 12, 2024
a919681
fix all tests but one
philipc2 Nov 12, 2024
0d714e1
fix final test
philipc2 Nov 12, 2024
7e86e5e
Clean up commented out code
philipc2 Nov 12, 2024
4371562
cleanup comments
philipc2 Nov 12, 2024
5c08d53
more comment cleanup
philipc2 Nov 12, 2024
8646f1f
update use of face bounds
philipc2 Nov 14, 2024
884ed87
fix failing test
philipc2 Nov 14, 2024
7d5860e
update docstrings (1)
philipc2 Nov 15, 2024
3cd53bd
update API
philipc2 Nov 15, 2024
efe52c4
update docstrings (2)
philipc2 Nov 15, 2024
2c4edb4
Merge branch 'main' into philipc2/bounds-numba
philipc2 Nov 16, 2024
9d66dcc
only use bounding box and raise warning for first computation
philipc2 Nov 16, 2024
f26fdc1
add exception for non-face-centered data variables
philipc2 Nov 16, 2024
ff8fa33
update docstring
philipc2 Nov 18, 2024
3262dcb
Merge branch 'main' into philipc2/bounds-numba
philipc2 Nov 18, 2024
e2dd04c
Merge branch 'main' into philipc2/bounds-numba
philipc2 Nov 18, 2024
286383b
pre-commit
philipc2 Nov 18, 2024
5655878
update docstring and parameter
philipc2 Nov 19, 2024
ce97728
add warning when translating numba functions for the first time
philipc2 Nov 20, 2024
18a00d3
Merge branch 'main' into philipc2/bounds-numba
philipc2 Nov 20, 2024
604ad13
update numba check
philipc2 Nov 20, 2024
9ff3cba
Merge branch 'philipc2/bounds-numba' of https://github.com/UXARRAY/ux…
philipc2 Nov 20, 2024
5a465b8
remove print statement
philipc2 Nov 20, 2024
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
2 changes: 2 additions & 0 deletions docs/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,8 @@ Descriptors
Grid.descriptors
Grid.face_areas
Grid.bounds
Grid.face_bounds_lon
Grid.face_bounds_lat
Grid.edge_node_distances
Grid.edge_face_distances
Grid.antimeridian_face_indices
Expand Down
34 changes: 17 additions & 17 deletions test/test_arcs.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
import uxarray as ux

from uxarray.grid.coordinates import _lonlat_rad_to_xyz
from uxarray.grid.arcs import point_within_gca
from uxarray.grid.arcs import point_within_gca, _point_within_gca_cartesian

try:
import constants
Expand Down Expand Up @@ -38,27 +38,27 @@ def test_pt_within_gcr(self):

pt_same_lon_in = _lonlat_rad_to_xyz(0.0, 0.0)
with self.assertRaises(ValueError):
point_within_gca(pt_same_lon_in, gcr_180degree_cart)
_point_within_gca_cartesian(pt_same_lon_in, gcr_180degree_cart)

# Test when the point and the GCA all have the same longitude
gcr_same_lon_cart = [
_lonlat_rad_to_xyz(0.0, 1.5),
_lonlat_rad_to_xyz(0.0, -1.5)
]
pt_same_lon_in = _lonlat_rad_to_xyz(0.0, 0.0)
self.assertTrue(point_within_gca(pt_same_lon_in, gcr_same_lon_cart))
self.assertTrue(_point_within_gca_cartesian(pt_same_lon_in, gcr_same_lon_cart))

pt_same_lon_out = _lonlat_rad_to_xyz(0.0, 1.500000000000001)
res = point_within_gca(pt_same_lon_out, gcr_same_lon_cart)
res = _point_within_gca_cartesian(pt_same_lon_out, gcr_same_lon_cart)
self.assertFalse(res)

pt_same_lon_out_2 = _lonlat_rad_to_xyz(0.1, 1.0)
res = point_within_gca(pt_same_lon_out_2, gcr_same_lon_cart)
res = _point_within_gca_cartesian(pt_same_lon_out_2, gcr_same_lon_cart)
self.assertFalse(res)

# And if we increase the digital place by one, it should be true again
pt_same_lon_out_add_one_place = _lonlat_rad_to_xyz(0.0, 1.5000000000000001)
res = point_within_gca(pt_same_lon_out_add_one_place, gcr_same_lon_cart)
res = _point_within_gca_cartesian(pt_same_lon_out_add_one_place, gcr_same_lon_cart)
self.assertTrue(res)

# Normal case
Expand All @@ -69,7 +69,7 @@ def test_pt_within_gcr(self):
-0.997]])
pt_cart_within = np.array(
[0.25616109352676675, 0.9246590335292105, -0.010021496695000144])
self.assertTrue(point_within_gca(pt_cart_within, gcr_cart_2, True))
self.assertTrue(_point_within_gca_cartesian(pt_cart_within, gcr_cart_2, True))

# Test other more complicate cases : The anti-meridian case

Expand All @@ -80,16 +80,16 @@ def test_pt_within_gcr_antimeridian(self):
gcr_cart = np.array([[0.351, -0.724, 0.593], [0.617, 0.672, 0.410]])
pt_cart = np.array(
[0.9438777657502077, 0.1193199333436068, 0.922714737029319])
self.assertTrue(point_within_gca(pt_cart, gcr_cart, is_directed=True))
self.assertTrue(_point_within_gca_cartesian(pt_cart, gcr_cart, is_directed=True))
# If we swap the gcr, it should throw a value error since it's larger than 180 degree
gcr_cart_flip = np.array([[0.617, 0.672, 0.410], [0.351, -0.724,
0.593]])
with self.assertRaises(ValueError):
point_within_gca(pt_cart, gcr_cart_flip, is_directed=True)
_point_within_gca_cartesian(pt_cart, gcr_cart_flip, is_directed=True)

# If we flip the gcr in the undirected mode, it should still work
self.assertTrue(
point_within_gca(pt_cart, gcr_cart_flip, is_directed=False))
_point_within_gca_cartesian(pt_cart, gcr_cart_flip, is_directed=False))

# 2nd anti-meridian case
# GCR vertex0 in radian : [4.104711496596806, 0.5352983676533828],
Expand All @@ -100,9 +100,9 @@ def test_pt_within_gcr_antimeridian(self):
pt_cart_within = np.array(
[0.6136726305712109, 0.28442243941920053, -0.365605190899831])
self.assertFalse(
point_within_gca(pt_cart_within, gcr_cart_1, is_directed=True))
_point_within_gca_cartesian(pt_cart_within, gcr_cart_1, is_directed=True))
self.assertFalse(
point_within_gca(pt_cart_within, gcr_cart_1, is_directed=False))
_point_within_gca_cartesian(pt_cart_within, gcr_cart_1, is_directed=False))

# The first case should not work and the second should work
v1_rad = [0.1, 0.0]
Expand All @@ -112,10 +112,10 @@ def test_pt_within_gcr_antimeridian(self):
gcr_cart = np.array([v1_cart, v2_cart])
pt_cart = _lonlat_rad_to_xyz(0.01, 0.0)
with self.assertRaises(ValueError):
point_within_gca(pt_cart, gcr_cart, is_directed=True)
_point_within_gca_cartesian(pt_cart, gcr_cart, is_directed=True)
gcr_car_flipped = np.array([v2_cart, v1_cart])
self.assertTrue(
point_within_gca(pt_cart, gcr_car_flipped, is_directed=True))
_point_within_gca_cartesian(pt_cart, gcr_car_flipped, is_directed=True))

def test_pt_within_gcr_cross_pole(self):
gcr_cart = np.array([[0.351, 0.0, 0.3], [-0.351, 0.0, 0.3]])
Expand All @@ -125,7 +125,7 @@ def test_pt_within_gcr_cross_pole(self):
# Normalize the point abd the GCA
pt_cart = pt_cart / np.linalg.norm(pt_cart)
gcr_cart = np.array([x / np.linalg.norm(x) for x in gcr_cart])
self.assertTrue(point_within_gca(pt_cart, gcr_cart, is_directed=False))
self.assertTrue(_point_within_gca_cartesian(pt_cart, gcr_cart, is_directed=False))

gcr_cart = np.array([[0.351, 0.0, 0.3], [-0.351, 0.0, -0.6]])
pt_cart = np.array(
Expand All @@ -134,6 +134,6 @@ def test_pt_within_gcr_cross_pole(self):
# When the point is not within the GCA
pt_cart = pt_cart / np.linalg.norm(pt_cart)
gcr_cart = np.array([x / np.linalg.norm(x) for x in gcr_cart])
self.assertFalse(point_within_gca(pt_cart, gcr_cart, is_directed=False))
self.assertFalse(_point_within_gca_cartesian(pt_cart, gcr_cart, is_directed=False))
with self.assertRaises(ValueError):
point_within_gca(pt_cart, gcr_cart, is_directed=True)
_point_within_gca_cartesian(pt_cart, gcr_cart, is_directed=True)
59 changes: 26 additions & 33 deletions test/test_cross_sections.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,103 +35,98 @@ class TestQuadHex:
All four faces intersect a constant latitude of 0.0
"""

@pytest.mark.parametrize("use_spherical_bounding_box", [True, False])
def test_constant_lat_cross_section_grid(self, use_spherical_bounding_box):


def test_constant_lat_cross_section_grid(self):

uxgrid = ux.open_grid(quad_hex_grid_path)

grid_top_two = uxgrid.cross_section.constant_latitude(lat=0.1, use_spherical_bounding_box=use_spherical_bounding_box)
grid_top_two = uxgrid.cross_section.constant_latitude(lat=0.1, )

assert grid_top_two.n_face == 2

grid_bottom_two = uxgrid.cross_section.constant_latitude(lat=-0.1, use_spherical_bounding_box=use_spherical_bounding_box)
grid_bottom_two = uxgrid.cross_section.constant_latitude(lat=-0.1, )

assert grid_bottom_two.n_face == 2

grid_all_four = uxgrid.cross_section.constant_latitude(lat=0.0, use_spherical_bounding_box=use_spherical_bounding_box)
grid_all_four = uxgrid.cross_section.constant_latitude(lat=0.0, )

assert grid_all_four.n_face == 4

with pytest.raises(ValueError):
# no intersections found at this line
uxgrid.cross_section.constant_latitude(lat=10.0, use_spherical_bounding_box=use_spherical_bounding_box)
uxgrid.cross_section.constant_latitude(lat=10.0, )

@pytest.mark.parametrize("use_spherical_bounding_box", [False])
def test_constant_lon_cross_section_grid(self, use_spherical_bounding_box):
def test_constant_lon_cross_section_grid(self):
uxgrid = ux.open_grid(quad_hex_grid_path)

grid_left_two = uxgrid.cross_section.constant_longitude(lon=-0.1, use_spherical_bounding_box=use_spherical_bounding_box)
grid_left_two = uxgrid.cross_section.constant_longitude(lon=-0.1, )

assert grid_left_two.n_face == 2

grid_right_two = uxgrid.cross_section.constant_longitude(lon=0.2, use_spherical_bounding_box=use_spherical_bounding_box)
grid_right_two = uxgrid.cross_section.constant_longitude(lon=0.2, )

assert grid_right_two.n_face == 2

with pytest.raises(ValueError):
# no intersections found at this line
uxgrid.cross_section.constant_longitude(lon=10.0)

@pytest.mark.parametrize("use_spherical_bounding_box", [False])
def test_constant_lat_cross_section_uxds(self, use_spherical_bounding_box):
def test_constant_lat_cross_section_uxds(self):
uxds = ux.open_dataset(quad_hex_grid_path, quad_hex_data_path)
uxds.uxgrid.normalize_cartesian_coordinates()

da_top_two = uxds['t2m'].cross_section.constant_latitude(lat=0.1, use_spherical_bounding_box=use_spherical_bounding_box)
da_top_two = uxds['t2m'].cross_section.constant_latitude(lat=0.1, )

nt.assert_array_equal(da_top_two.data, uxds['t2m'].isel(n_face=[1, 2]).data)

da_bottom_two = uxds['t2m'].cross_section.constant_latitude(lat=-0.1, use_spherical_bounding_box=use_spherical_bounding_box)
da_bottom_two = uxds['t2m'].cross_section.constant_latitude(lat=-0.1, )

nt.assert_array_equal(da_bottom_two.data, uxds['t2m'].isel(n_face=[0, 3]).data)

da_all_four = uxds['t2m'].cross_section.constant_latitude(lat=0.0, use_spherical_bounding_box=use_spherical_bounding_box)
da_all_four = uxds['t2m'].cross_section.constant_latitude(lat=0.0, )

nt.assert_array_equal(da_all_four.data , uxds['t2m'].data)

with pytest.raises(ValueError):
# no intersections found at this line
uxds['t2m'].cross_section.constant_latitude(lat=10.0, use_spherical_bounding_box=use_spherical_bounding_box)
uxds['t2m'].cross_section.constant_latitude(lat=10.0, )


@pytest.mark.parametrize("use_spherical_bounding_box", [False])
def test_constant_lon_cross_section_uxds(self, use_spherical_bounding_box):
def test_constant_lon_cross_section_uxds(self):
uxds = ux.open_dataset(quad_hex_grid_path, quad_hex_data_path)
uxds.uxgrid.normalize_cartesian_coordinates()

da_left_two = uxds['t2m'].cross_section.constant_longitude(lon=-0.1, use_spherical_bounding_box=use_spherical_bounding_box)
da_left_two = uxds['t2m'].cross_section.constant_longitude(lon=-0.1, )

nt.assert_array_equal(da_left_two.data, uxds['t2m'].isel(n_face=[0, 2]).data)

da_right_two = uxds['t2m'].cross_section.constant_longitude(lon=0.2, use_spherical_bounding_box=use_spherical_bounding_box)
da_right_two = uxds['t2m'].cross_section.constant_longitude(lon=0.2, )

nt.assert_array_equal(da_right_two.data, uxds['t2m'].isel(n_face=[1, 3]).data)

with pytest.raises(ValueError):
# no intersections found at this line
uxds['t2m'].cross_section.constant_longitude(lon=10.0, use_spherical_bounding_box=use_spherical_bounding_box)
uxds['t2m'].cross_section.constant_longitude(lon=10.0, )


class TestCubeSphere:
@pytest.mark.parametrize("use_spherical_bounding_box", [True, False])
def test_north_pole(self, use_spherical_bounding_box):

def test_north_pole(self):
uxgrid = ux.open_grid(cube_sphere_grid)

lats = [89.85, 89.9, 89.95, 89.99]

for lat in lats:
cross_grid = uxgrid.cross_section.constant_latitude(lat=lat, use_spherical_bounding_box=use_spherical_bounding_box)
cross_grid = uxgrid.cross_section.constant_latitude(lat=lat, )
# Cube sphere grid should have 4 faces centered around the pole
assert cross_grid.n_face == 4
@pytest.mark.parametrize("use_spherical_bounding_box", [True, False])
def test_south_pole(self, use_spherical_bounding_box):

def test_south_pole(self):
uxgrid = ux.open_grid(cube_sphere_grid)

lats = [-89.85, -89.9, -89.95, -89.99]

for lat in lats:
cross_grid = uxgrid.cross_section.constant_latitude(lat=lat, use_spherical_bounding_box=use_spherical_bounding_box)
cross_grid = uxgrid.cross_section.constant_latitude(lat=lat, )
# Cube sphere grid should have 4 faces centered around the pole
assert cross_grid.n_face == 4

Expand All @@ -152,8 +147,7 @@ def test_constant_lat(self):

candidate_faces = constant_lat_intersections_face_bounds(
lat=const_lat,
face_min_lat_rad=bounds_rad[:, 0, 0],
face_max_lat_rad=bounds_rad[:, 0, 1],
face_bounds_lat=bounds_rad[:, 0],
)

# Expected output
Expand All @@ -176,8 +170,7 @@ def test_constant_lat_out_of_bounds(self):

candidate_faces = constant_lat_intersections_face_bounds(
lat=const_lat,
face_min_lat_rad=bounds_rad[:, 0, 0],
face_max_lat_rad=bounds_rad[:, 0, 1],
face_bounds_lat=bounds_rad[:, 0],
)

assert len(candidate_faces) == 0
Loading
Loading