Vectorize Coordinate Conversion Functions (#748)

* initial work on vectorization

* start overhaul of code to support new functions

* update existing code and tests to use new functions

* fix tests (1)

* fix tests (2)

* fix tests (3)

* fix tests (4)

* fix tests (5)

* fix tests (6)

* Update API reference

* generalize implementation

* update _lonlat_rad_to_xyz

* update function signature

* vectorize conversion in tests, when possible

docs/internal_api/index.rst

@@ -147,15 +147,17 @@ Coordinates
 .. autosummary::
    :toctree: generated/
 
-   grid.coordinates._get_lonlat_from_xyz
-   grid.coordinates._get_xyz_from_lonlat
-   grid.coordinates._populate_cartesian_xyz_coord
-   grid.coordinates._populate_lonlat_coord
+   grid.coordinates._lonlat_rad_to_xyz
+   grid.coordinates._xyz_to_lonlat_rad
+   grid.coordinates._xyz_to_lonlat_deg
+   grid.coordinates._normalize_xyz
+   grid.coordinates._populate_node_latlon
+   grid.coordinates._populate_node_xyz
    grid.coordinates._populate_face_centroids
-   grid.coordinates._construct_face_centroids
-   grid.coordinates._set_desired_longitude_range
    grid.coordinates._populate_edge_centroids
+   grid.coordinates._construct_face_centroids
    grid.coordinates._construct_edge_centroids
+   grid.coordinates._set_desired_longitude_range
 
 
 Arcs

docs/user_api/index.rst

@@ -369,16 +369,6 @@ Connectivity
 
    grid.connectivity.close_face_nodes
 
-Coordinates
------------
-.. autosummary::
-   :toctree: generated/
-
-   grid.coordinates.node_lonlat_rad_to_xyz
-   grid.coordinates.node_xyz_to_lonlat_rad
-   grid.coordinates.normalize_in_place
-
-
 Arcs
 ----
 .. autosummary::

test/test_arcs.py

@@ -9,7 +9,7 @@
 
 import uxarray as ux
 
-from uxarray.grid.coordinates import node_lonlat_rad_to_xyz
+from uxarray.grid.coordinates import _lonlat_rad_to_xyz
 from uxarray.grid.arcs import point_within_gca
 
 try:
@@ -31,43 +31,40 @@ class TestIntersectionPoint(TestCase):
     def test_pt_within_gcr(self):
         # The GCR that's eexactly 180 degrees will have Value Error raised
         gcr_180degree_cart = [
-            ux.grid.coordinates.node_lonlat_rad_to_xyz([0.0, 0.0]),
-            ux.grid.coordinates.node_lonlat_rad_to_xyz([np.pi, 0.0])
+            _lonlat_rad_to_xyz(0.0, 0.0),
+            _lonlat_rad_to_xyz(np.pi, 0.0)
         ]
-        pt_same_lon_in = ux.grid.coordinates.node_lonlat_rad_to_xyz([0.0, 0.0])
+        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)
 
         gcr_180degree_cart = [
-            ux.grid.coordinates.node_lonlat_rad_to_xyz([0.0, np.pi / 2.0]),
-            ux.grid.coordinates.node_lonlat_rad_to_xyz([0.0, -np.pi / 2.0])
+            _lonlat_rad_to_xyz(0.0, np.pi / 2.0),
+            _lonlat_rad_to_xyz(0.0, -np.pi / 2.0)
         ]
 
-        pt_same_lon_in = ux.grid.coordinates.node_lonlat_rad_to_xyz([0.0, 0.0])
+        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)
 
         # Test when the point and the GCA all have the same longitude
         gcr_same_lon_cart = [
-            ux.grid.coordinates.node_lonlat_rad_to_xyz([0.0, 1.5]),
-            ux.grid.coordinates.node_lonlat_rad_to_xyz([0.0, -1.5])
+            _lonlat_rad_to_xyz(0.0, 1.5),
+            _lonlat_rad_to_xyz(0.0, -1.5)
         ]
-        pt_same_lon_in = ux.grid.coordinates.node_lonlat_rad_to_xyz([0.0, 0.0])
+        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))
 
-        pt_same_lon_out = ux.grid.coordinates.node_lonlat_rad_to_xyz(
-            [0.0, 1.500000000000001])
+        pt_same_lon_out = _lonlat_rad_to_xyz(0.0, 1.500000000000001)
         res = point_within_gca(pt_same_lon_out, gcr_same_lon_cart)
         self.assertFalse(res)
 
-        pt_same_lon_out_2 = ux.grid.coordinates.node_lonlat_rad_to_xyz(
-            [0.1, 1.0])
+        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)
         self.assertFalse(res)
 
         # And if we increase the digital place by one, it should be true again
-        pt_same_lon_out_add_one_place = ux.grid.coordinates.node_lonlat_rad_to_xyz(
-            [0.0, 1.5000000000000001])
+        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)
         self.assertTrue(res)
 
@@ -117,10 +114,10 @@ def test_pt_within_gcr_antimeridian(self):
         # The first case should not work and the second should work
         v1_rad = [0.1, 0.0]
         v2_rad = [2 * np.pi - 0.1, 0.0]
-        v1_cart = ux.grid.coordinates.node_lonlat_rad_to_xyz(v1_rad)
-        v2_cart = ux.grid.coordinates.node_lonlat_rad_to_xyz(v2_rad)
+        v1_cart = _lonlat_rad_to_xyz(v1_rad[0], v1_rad[1])
+        v2_cart = _lonlat_rad_to_xyz(v2_rad[0], v1_rad[1])
         gcr_cart = np.array([v1_cart, v2_cart])
-        pt_cart = ux.grid.coordinates.node_lonlat_rad_to_xyz([0.01, 0.0])
+        pt_cart = _lonlat_rad_to_xyz(0.01, 0.0)
         with self.assertRaises(ValueError):
             point_within_gca(pt_cart, gcr_cart, is_directed=True)
         gcr_car_flipped = np.array([v2_cart, v1_cart])

test/test_centroids.py

@@ -4,7 +4,7 @@
 import numpy.testing as nt
 import uxarray as ux
 from pathlib import Path
-from uxarray.grid.coordinates import _populate_face_centroids, _populate_edge_centroids, normalize_in_place
+from uxarray.grid.coordinates import _populate_face_centroids, _populate_edge_centroids, _normalize_xyz
 
 current_path = Path(os.path.dirname(os.path.realpath(__file__)))
 
@@ -21,8 +21,8 @@ def test_centroids_from_mean_verts_triangle(self):
 
         # Calculate the expected centroid
         expected_centroid = np.mean(test_triangle, axis=0)
-        [norm_x, norm_y, norm_z] = normalize_in_place(
-            [expected_centroid[0], expected_centroid[1], expected_centroid[2]])
+        norm_x, norm_y, norm_z = _normalize_xyz(
+            expected_centroid[0], expected_centroid[1], expected_centroid[2])
 
         # Open the dataset and find the centroids
         grid = ux.open_grid(test_triangle)
@@ -42,8 +42,8 @@ def test_centroids_from_mean_verts_pentagon(self):
 
         # Calculate the expected centroid
         expected_centroid = np.mean(test_triangle, axis=0)
-        [norm_x, norm_y, norm_z] = normalize_in_place(
-            [expected_centroid[0], expected_centroid[1], expected_centroid[2]])
+        norm_x, norm_y, norm_z = _normalize_xyz(
+            expected_centroid[0], expected_centroid[1], expected_centroid[2])
 
         # Open the dataset and find the centroids
         grid = ux.open_grid(test_triangle)
@@ -65,7 +65,8 @@ def test_centroids_from_mean_verts_scrip(self):
 
         _populate_face_centroids(uxgrid, repopulate=True)
 
-        computed_face_x = (uxgrid.face_lon.values + 180) % 360 - 180
+        # computed_face_x = (uxgrid.face_lon.values + 180) % 360 - 180
+        computed_face_x = uxgrid.face_lon.values
         computed_face_y = uxgrid.face_lat.values
 
         nt.assert_array_almost_equal(expected_face_x, computed_face_x)

test/test_computing.py

@@ -1,7 +1,7 @@
 from unittest import TestCase
 import numpy.testing as nt
 import numpy as np
-from uxarray.grid.coordinates import normalize_in_place
+from uxarray.grid.coordinates import _normalize_xyz
 import uxarray.utils.computing as ac_utils
 from uxarray.constants import ERROR_TOLERANCE
 
@@ -12,8 +12,8 @@ class TestCrossProduct(TestCase):
     one."""
 
     def test_cross_fma(self):
-        v1 = np.array(normalize_in_place([1.0, 2.0, 3.0]))
-        v2 = np.array(normalize_in_place([4.0, 5.0, 6.0]))
+        v1 = np.array(_normalize_xyz(*[1.0, 2.0, 3.0]))
+        v2 = np.array(_normalize_xyz(*[4.0, 5.0, 6.0]))
 
         np_cross = np.cross(v1, v2)
         fma_cross = ac_utils.cross_fma(v1, v2)
@@ -26,8 +26,8 @@ class TestDotProduct(TestCase):
     one."""
 
     def test_dot_fma(self):
-        v1 = np.array(normalize_in_place([1.0, 0.0, 0.0]), dtype=np.float64)
-        v2 = np.array(normalize_in_place([1.0, 0.0, 0.0]), dtype=np.float64)
+        v1 = np.array(_normalize_xyz(*[1.0, 0.0, 0.0]), dtype=np.float64)
+        v2 = np.array(_normalize_xyz(*[1.0, 0.0, 0.0]), dtype=np.float64)
 
         np_dot = np.dot(v1, v2)
         fma_dot = ac_utils.dot_fma(v1, v2)