From 1a68daa24a4840334027fd2621fc1882abf19743 Mon Sep 17 00:00:00 2001
From: Philip Chmielowiec <67855069+philipc2@users.noreply.github.com>
Date: Fri, 9 Aug 2024 00:59:06 -0500
Subject: [PATCH] Optimize Face Bounds Computation  (#877)

* add benchmark for face bounds helpers

* add benchmark for face bounds helpers

* optimize point_within_gca

* disable fma for benchmark

* numba implementations of isclose and allclose

* dot and cross numba

* dot and cross numba

* update dot

* update dot

* try to fix dot (1)

* gca_gca_extreme

* numba xyz to latlon rad with no norm

* disable FMA

* enable FMA

* add more dot calls

* remove leftover code

* fix tests

* fix failing tests

* update isclose call

* disable FMA for benchmark

* fix broken isclose call

* add docstrings, option to enable and disable fma

* add some vectorization

* numba all

* point_within_gca numba

* fix error in tests

* better function name

* revert asv.conf.json
---
 benchmarks/face_bounds.py     |   2 -
 test/test_arcs.py             |  28 +++---
 test/test_grid.py             |   9 +-
 uxarray/__init__.py           |  19 ++++
 uxarray/constants.py          |   2 +
 uxarray/grid/arcs.py          | 170 +++++++++++++++++++---------------
 uxarray/grid/coordinates.py   |  52 +++++++++++
 uxarray/grid/geometry.py      |  71 +++++++-------
 uxarray/grid/grid.py          |   1 +
 uxarray/grid/intersections.py |  39 ++++----
 uxarray/utils/computing.py    |  59 ++++++++++++
 11 files changed, 308 insertions(+), 144 deletions(-)

diff --git a/benchmarks/face_bounds.py b/benchmarks/face_bounds.py
index c41602b6b..b249e7b99 100644
--- a/benchmarks/face_bounds.py
+++ b/benchmarks/face_bounds.py
@@ -10,8 +10,6 @@
 grid_scrip = current_path / "test" / "meshfiles" / "scrip" / "outCSne8" / "outCSne8.nc"
 grid_mpas= current_path / "test" / "meshfiles" / "mpas" / "QU" / "oQU480.231010.nc"
 
-
-
 class FaceBounds:
 
     params = [grid_quad_hex, grid_geoflow, grid_scrip, grid_mpas]
diff --git a/test/test_arcs.py b/test/test_arcs.py
index c9859e724..17d04fe57 100644
--- a/test/test_arcs.py
+++ b/test/test_arcs.py
@@ -36,7 +36,7 @@ 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(np.asarray(pt_same_lon_in), np.asarray(gcr_180degree_cart))
 
         gcr_180degree_cart = [
             _lonlat_rad_to_xyz(0.0, np.pi / 2.0),
@@ -45,7 +45,7 @@ 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(np.asarray(pt_same_lon_in), np.asarray(gcr_180degree_cart))
 
         # Test when the point and the GCA all have the same longitude
         gcr_same_lon_cart = [
@@ -53,19 +53,19 @@ def test_pt_within_gcr(self):
             _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(np.asarray(pt_same_lon_in), np.asarray(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(np.asarray(pt_same_lon_out), np.asarray(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(np.asarray(pt_same_lon_out_2), np.asarray(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(np.asarray(pt_same_lon_out_add_one_place), np.asarray(gcr_same_lon_cart))
         self.assertTrue(res)
 
         # Normal case
@@ -76,7 +76,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(np.asarray(pt_cart_within), np.asarray(gcr_cart_2), True))
 
         # Test other more complicate cases : The anti-meridian case
 
@@ -87,16 +87,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(np.asarray(pt_cart), np.asarray(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(np.asarray(pt_cart), np.asarray(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(np.asarray(pt_cart), np.asarray(gcr_cart_flip), is_directed=False))
 
         # 2nd anti-meridian case
         # GCR vertex0 in radian : [4.104711496596806, 0.5352983676533828],
@@ -107,9 +107,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(np.asarray(pt_cart_within), np.asarray(gcr_cart_1), is_directed=True))
         self.assertFalse(
-            point_within_gca(pt_cart_within, gcr_cart_1, is_directed=False))
+            point_within_gca(np.asarray(pt_cart_within), np.asarray(gcr_cart_1), is_directed=False))
 
         # The first case should not work and the second should work
         v1_rad = [0.1, 0.0]
@@ -119,10 +119,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(np.asarray(pt_cart), np.asarray(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(np.asarray(pt_cart), np.asarray(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]])
diff --git a/test/test_grid.py b/test/test_grid.py
index 10faba696..57eb3f3b8 100644
--- a/test/test_grid.py
+++ b/test/test_grid.py
@@ -952,8 +952,7 @@ def test_populate_bounds_GCA_mix(self):
         face_bounds = bounds_xarray.values
         nt.assert_allclose(grid.bounds.values, expected_bounds, atol=ERROR_TOLERANCE)
 
-    def test_populate_bounds_MPAS(self):
-        xrds = xr.open_dataset(self.gridfile_mpas)
-        uxgrid = ux.Grid.from_dataset(xrds, use_dual=True)
-        bounds_xarray = uxgrid.bounds
-        pass
+    def test_opti_bounds(self):
+        import uxarray
+        uxgrid = ux.open_grid(gridfile_CSne8)
+        bounds = uxgrid.bounds
diff --git a/uxarray/__init__.py b/uxarray/__init__.py
index ed5c31cad..d15b38075 100644
--- a/uxarray/__init__.py
+++ b/uxarray/__init__.py
@@ -1,6 +1,8 @@
 # Sets the version of uxarray currently installeds
 # Attempt to import the needed modules
 
+import uxarray.constants
+
 from .core.api import open_grid, open_dataset, open_mfdataset
 
 from .core.dataset import UxDataset
@@ -21,6 +23,21 @@
     # Placeholder version incase an error occurs, such as the library isn't installed
     __version__ = "999"
 
+# Flag for enabling FMA instructions across the package
+
+
+def enable_fma():
+    """Enables Fused-Multiply-Add (FMA) instructions using the ``pyfma``
+    package."""
+    uxarray.constants.ENABLE_FMA = True
+
+
+def disable_fma():
+    """Disable Fused-Multiply-Add (FMA) instructions using the ``pyfma``
+    package."""
+    uxarray.constants.ENABLE_FMA = False
+
+
 __all__ = (
     "open_grid",
     "open_dataset",
@@ -30,4 +47,6 @@
     "INT_DTYPE",
     "INT_FILL_VALUE",
     "Grid",
+    "enable_fma",
+    "disable_fma",
 )
diff --git a/uxarray/constants.py b/uxarray/constants.py
index d42befac4..a8baa9c1c 100644
--- a/uxarray/constants.py
+++ b/uxarray/constants.py
@@ -15,4 +15,6 @@
 ENABLE_JIT_CACHE = True
 ENABLE_JIT = True
 
+ENABLE_FMA = False
+
 GRID_DIMS = ["n_node", "n_edge", "n_face"]
diff --git a/uxarray/grid/arcs.py b/uxarray/grid/arcs.py
index 381d02e41..63344473e 100644
--- a/uxarray/grid/arcs.py
+++ b/uxarray/grid/arcs.py
@@ -2,9 +2,13 @@
 
 # from uxarray.grid.coordinates import node_xyz_to_lonlat_rad, normalize_in_place
 
-from uxarray.grid.coordinates import _xyz_to_lonlat_rad, _normalize_xyz
+from uxarray.grid.coordinates import _xyz_to_lonlat_rad_no_norm, _normalize_xyz_scalar
 from uxarray.constants import ERROR_TOLERANCE
 
+from uxarray.utils.computing import isclose, cross, dot
+
+from numba import njit
+
 
 def _to_list(obj):
     if not isinstance(obj, list):
@@ -17,78 +21,29 @@ def _to_list(obj):
     return obj
 
 
-def point_within_gca(pt, gca_cart, is_directed=False):
-    """Check if a point lies on a given Great Circle Arc (GCA). The anti-
-    meridian case is also considered.
-
-    Parameters
-    ----------
-    pt : numpy.ndarray (float)
-        Cartesian coordinates of the point.
-    gca_cart : numpy.ndarray of shape (2, 3), (np.float or gmpy2.mpfr)
-        Cartesian coordinates of the Great Circle Arc (GCR).
-    is_directed : bool, optional, default = False
-        If True, the GCA is considered to be directed, which means it can only from v0-->v1. If False, the GCA is undirected,
-        and we will always assume the small circle (The one less than 180 degree) side is the GCA. The default is False.
-        For the case of the anti-podal case, the direction is v_0--> the pole point that on the same hemisphere as v_0-->v_1
-
-    Returns
-    -------
-    bool
-        True if the point lies between the two endpoints of the GCR, False otherwise.
-
-    Raises
-    ------
-    ValueError
-        If the input GCR spans exactly 180 degrees (π radians), as this GCR can have multiple planes.
-        In such cases, consider breaking the GCR into two separate GCRs.
-
-    ValueError
-        If the input GCR spans more than 180 degrees (π radians).
-        In such cases, consider breaking the GCR into two separate GCRs.
-
-    Notes
-    -----
-    The function checks if the given point is on the Great Circle Arc by considering its cartesian coordinates and
-    accounting for the anti-meridian case.
-
-    The anti-meridian case occurs when the GCR crosses the anti-meridian (0 longitude).
-    In this case, the function handles scenarios where the GCA spans across more than 180 degrees, requiring specific operation.
-
-    The function relies on the `_angle_of_2_vectors` and `is_between` functions to perform the necessary calculations.
-
-    Please ensure that the input coordinates are in radians and adhere to the ERROR_TOLERANCE value for floating-point comparisons.
-    """
-    # Convert the cartesian coordinates to lonlat coordinates
-    pt_lonlat = np.array(_xyz_to_lonlat_rad(pt[0], pt[1], pt[2]))
-    GCRv0_lonlat = np.array(
-        _xyz_to_lonlat_rad(gca_cart[0][0], gca_cart[0][1], gca_cart[0][2])
-    )
-    GCRv1_lonlat = np.array(
-        _xyz_to_lonlat_rad(gca_cart[1][0], gca_cart[1][1], gca_cart[1][2])
-    )
-
-    # Convert the list to np.float64
-    gca_cart[0] = np.array(gca_cart[0], dtype=np.float64)
-    gca_cart[1] = np.array(gca_cart[1], dtype=np.float64)
-
-    # First if the input GCR is exactly 180 degree, we throw an exception, since this GCR can have multiple planes
+@njit
+def _point_within_gca_body(
+    angle, gca_cart, pt, GCRv0_lonlat, GCRv1_lonlat, pt_lonlat, is_directed
+):
     angle = _angle_of_2_vectors(gca_cart[0], gca_cart[1])
-    if np.allclose(angle, np.pi, rtol=0, atol=ERROR_TOLERANCE):
+    if isclose(angle, np.pi, rtol=0.0, atol=ERROR_TOLERANCE):
         raise ValueError(
             "The input Great Circle Arc is exactly 180 degree, this Great Circle Arc can have multiple planes. "
             "Consider breaking the Great Circle Arc"
             "into two Great Circle Arcs"
         )
 
-    if not np.allclose(
-        np.dot(np.cross(gca_cart[0], gca_cart[1]), pt), 0, rtol=0, atol=ERROR_TOLERANCE
+    if not isclose(
+        dot(cross(np.asarray(gca_cart[0]), np.asarray(gca_cart[1])), pt),
+        0,
+        rtol=0.0,
+        atol=ERROR_TOLERANCE,
     ):
         return False
 
-    if np.isclose(GCRv0_lonlat[0], GCRv1_lonlat[0], rtol=0, atol=ERROR_TOLERANCE):
+    if isclose(GCRv0_lonlat[0], GCRv1_lonlat[0], rtol=0.0, atol=ERROR_TOLERANCE):
         # If the pt and the GCA are on the same longitude (the y coordinates are the same)
-        if np.isclose(GCRv0_lonlat[0], pt_lonlat[0], rtol=0, atol=ERROR_TOLERANCE):
+        if isclose(GCRv0_lonlat[0], pt_lonlat[0], rtol=0.0, atol=ERROR_TOLERANCE):
             # Now use the latitude to determine if the pt falls between the interval
             return in_between(GCRv0_lonlat[1], pt_lonlat[1], GCRv1_lonlat[1])
         else:
@@ -96,16 +51,16 @@ def point_within_gca(pt, gca_cart, is_directed=False):
             return False
 
     # If the longnitude span is exactly 180 degree, then the GCA goes through the pole point
-    if np.isclose(
-        abs(GCRv1_lonlat[0] - GCRv0_lonlat[0]), np.pi, rtol=0, atol=ERROR_TOLERANCE
+    if isclose(
+        abs(GCRv1_lonlat[0] - GCRv0_lonlat[0]), np.pi, rtol=0.0, atol=ERROR_TOLERANCE
     ):
         # Special case, if the pt is on the pole point, then set its longitude to the GCRv0_lonlat[0]
-        if np.isclose(abs(pt_lonlat[1]), np.pi / 2, rtol=0, atol=ERROR_TOLERANCE):
+        if isclose(abs(pt_lonlat[1]), np.pi / 2, rtol=0.0, atol=ERROR_TOLERANCE):
             pt_lonlat[0] = GCRv0_lonlat[0]
-        if not np.isclose(
-            GCRv0_lonlat[0], pt_lonlat[0], rtol=0, atol=ERROR_TOLERANCE
-        ) and not np.isclose(
-            GCRv1_lonlat[0], pt_lonlat[0], rtol=0, atol=ERROR_TOLERANCE
+        if not isclose(
+            GCRv0_lonlat[0], pt_lonlat[0], rtol=0.0, atol=ERROR_TOLERANCE
+        ) and not isclose(
+            GCRv1_lonlat[0], pt_lonlat[0], rtol=0.0, atol=ERROR_TOLERANCE
         ):
             return False
         else:
@@ -169,6 +124,69 @@ def point_within_gca(pt, gca_cart, is_directed=False):
             )
 
 
+def point_within_gca(pt, gca_cart, is_directed=False):
+    """Check if a point lies on a given Great Circle Arc (GCA). The anti-
+    meridian case is also considered.
+
+    Parameters
+    ----------
+    pt : numpy.ndarray (float)
+        Cartesian coordinates of the point.
+    gca_cart : numpy.ndarray of shape (2, 3), (np.float or gmpy2.mpfr)
+        Cartesian coordinates of the Great Circle Arc (GCR).
+    is_directed : bool, optional, default = False
+        If True, the GCA is considered to be directed, which means it can only from v0-->v1. If False, the GCA is undirected,
+        and we will always assume the small circle (The one less than 180 degree) side is the GCA. The default is False.
+        For the case of the anti-podal case, the direction is v_0--> the pole point that on the same hemisphere as v_0-->v_1
+
+    Returns
+    -------
+    bool
+        True if the point lies between the two endpoints of the GCR, False otherwise.
+
+    Raises
+    ------
+    ValueError
+        If the input GCR spans exactly 180 degrees (π radians), as this GCR can have multiple planes.
+        In such cases, consider breaking the GCR into two separate GCRs.
+
+    ValueError
+        If the input GCR spans more than 180 degrees (π radians).
+        In such cases, consider breaking the GCR into two separate GCRs.
+
+    Notes
+    -----
+    The function checks if the given point is on the Great Circle Arc by considering its cartesian coordinates and
+    accounting for the anti-meridian case.
+
+    The anti-meridian case occurs when the GCR crosses the anti-meridian (0 longitude).
+    In this case, the function handles scenarios where the GCA spans across more than 180 degrees, requiring specific operation.
+
+    The function relies on the `_angle_of_2_vectors` and `is_between` functions to perform the necessary calculations.
+
+    Please ensure that the input coordinates are in radians and adhere to the ERROR_TOLERANCE value for floating-point comparisons.
+    """
+    # Convert the cartesian coordinates to lonlat coordinates
+    pt_lonlat = np.array(_xyz_to_lonlat_rad_no_norm(pt[0], pt[1], pt[2]))
+    GCRv0_lonlat = np.array(
+        _xyz_to_lonlat_rad_no_norm(gca_cart[0][0], gca_cart[0][1], gca_cart[0][2])
+    )
+    GCRv1_lonlat = np.array(
+        _xyz_to_lonlat_rad_no_norm(gca_cart[1][0], gca_cart[1][1], gca_cart[1][2])
+    )
+    gca_cart = np.asarray(gca_cart)
+
+    # First if the input GCR is exactly 180 degree, we throw an exception, since this GCR can have multiple planes
+    angle = _angle_of_2_vectors(gca_cart[0], gca_cart[1])
+
+    out = _point_within_gca_body(
+        angle, gca_cart, pt, GCRv0_lonlat, GCRv1_lonlat, pt_lonlat, is_directed
+    )
+
+    return out
+
+
+@njit
 def in_between(p, q, r) -> bool:
     """Determines whether the number q is between p and r.
 
@@ -190,6 +208,7 @@ def in_between(p, q, r) -> bool:
     return p <= q <= r or r <= q <= p
 
 
+@njit
 def _decide_pole_latitude(lat1, lat2):
     """Determine the pole latitude based on the latitudes of two points on a
     Great Circle Arc (GCA).
@@ -229,6 +248,7 @@ def _decide_pole_latitude(lat1, lat2):
     return closest_pole
 
 
+@njit
 def _angle_of_2_vectors(u, v):
     """Calculate the angle between two 3D vectors u and v in radians. Can be
     used to calcualte the span of a GCR.
@@ -281,22 +301,24 @@ def extreme_gca_latitude(gca_cart, extreme_type):
         raise ValueError("extreme_type must be either 'max' or 'min'")
 
     n1, n2 = gca_cart
-    dot_n1_n2 = np.dot(n1, n2)
+
+    dot_n1_n2 = dot(np.asarray(n1), np.asarray(n2))
     denom = (n1[2] + n2[2]) * (dot_n1_n2 - 1.0)
     d_a_max = (n1[2] * dot_n1_n2 - n2[2]) / denom
 
     d_a_max = (
         np.clip(d_a_max, 0, 1)
-        if np.isclose(d_a_max, [0, 1], atol=ERROR_TOLERANCE).any()
+        if isclose(d_a_max, 0, atol=ERROR_TOLERANCE)
+        or isclose(d_a_max, 1, atol=ERROR_TOLERANCE)
         else d_a_max
     )
 
-    _, lat_n1 = _xyz_to_lonlat_rad(n1[0], n1[1], n1[2])
-    _, lat_n2 = _xyz_to_lonlat_rad(n2[0], n2[1], n2[2])
+    _, lat_n1 = _xyz_to_lonlat_rad_no_norm(n1[0], n1[1], n1[2])
+    _, lat_n2 = _xyz_to_lonlat_rad_no_norm(n2[0], n2[1], n2[2])
 
     if 0 < d_a_max < 1:
         node3 = (1 - d_a_max) * n1 + d_a_max * n2
-        node3 = np.array(_normalize_xyz(node3[0], node3[1], node3[2]))
+        node3 = np.array(_normalize_xyz_scalar(node3[0], node3[1], node3[2]))
         d_lat_rad = np.arcsin(np.clip(node3[2], -1, 1))
 
         return (
diff --git a/uxarray/grid/coordinates.py b/uxarray/grid/coordinates.py
index ceb8acebc..d34f0fc20 100644
--- a/uxarray/grid/coordinates.py
+++ b/uxarray/grid/coordinates.py
@@ -10,6 +10,8 @@
 
 from numba import njit
 
+import math
+
 
 @njit(cache=True)
 def _lonlat_rad_to_xyz(
@@ -25,6 +27,47 @@ def _lonlat_rad_to_xyz(
     return x, y, z
 
 
+@njit
+def _xyz_to_lonlat_rad_no_norm(
+    x: Union[np.ndarray, float],
+    y: Union[np.ndarray, float],
+    z: Union[np.ndarray, float],
+):
+    """Converts a Cartesian x,y,z coordinates into Spherical latitude and
+    longitude without normalization, decorated with Numba.
+
+    Parameters
+    ----------
+    x : float
+        Cartesian x coordinate
+    y: float
+        Cartesiain y coordinate
+    z: float
+        Cartesian z coordinate
+
+
+    Returns
+    -------
+    lon : float
+        Longitude in radians
+    lat: float
+        Latitude in radians
+    """
+
+    lon = math.atan2(y, x)
+    lat = math.asin(z)
+
+    # set longitude range to [0, pi]
+    lon = np.mod(lon, 2 * np.pi)
+
+    z_mask = np.abs(z) > 1.0 - ERROR_TOLERANCE
+
+    lat = np.where(z_mask, np.sign(z) * np.pi / 2, lat)
+    lon = np.where(z_mask, 0.0, lon)
+
+    return lon, lat
+
+
 def _xyz_to_lonlat_rad(
     x: Union[np.ndarray, float],
     y: Union[np.ndarray, float],
@@ -126,6 +169,15 @@ def _normalize_xyz(
     return x_norm, y_norm, z_norm
 
 
+@njit
+def _normalize_xyz_scalar(x: float, y: float, z: float):
+    denom = np.linalg.norm(np.asarray(np.array([x, y, z]), dtype=np.float64), ord=2)
+    x_norm = x / denom
+    y_norm = y / denom
+    z_norm = z / denom
+    return x_norm, y_norm, z_norm
+
+
 def _populate_node_latlon(grid) -> None:
     """Populates the latitude and longitude coordinates of a Grid (`node_lon`,
     `node_lat`)"""
diff --git a/uxarray/grid/geometry.py b/uxarray/grid/geometry.py
index d2c43e8f2..08229d1a0 100644
--- a/uxarray/grid/geometry.py
+++ b/uxarray/grid/geometry.py
@@ -6,6 +6,8 @@
     _get_cartesian_face_edge_nodes,
     _get_lonlat_rad_face_edge_nodes,
 )
+
+from uxarray.utils.computing import isclose, allclose, all
 import warnings
 import pandas as pd
 import xarray as xr
@@ -467,7 +469,7 @@ def _check_intersection(ref_edge, edges):
         intersection_point = gca_gca_intersection(ref_edge, edge)
 
         if intersection_point.size != 0:
-            if np.allclose(intersection_point, pole_point, atol=ERROR_TOLERANCE):
+            if allclose(intersection_point, pole_point, atol=ERROR_TOLERANCE):
                 return True
             intersection_count += 1
 
@@ -488,9 +490,9 @@ def _classify_polygon_location(face_edge_cart):
         Returns either 'North', 'South' or 'Equator' based on the polygon's location.
     """
     z_coords = face_edge_cart[:, :, 2]
-    if np.all(z_coords > 0):
+    if all(z_coords > 0):
         return "North"
-    elif np.all(z_coords < 0):
+    elif all(z_coords < 0):
         return "South"
     else:
         return "Equator"
@@ -584,7 +586,7 @@ def _insert_pt_in_latlonbox(old_box, new_pt, is_lon_periodic=True):
     >>> _insert_pt_in_latlonbox(np.array([[1.0, 2.0], [3.0, 4.0]]),np.array([1.5, 3.5]))
     array([[1.0, 2.0], [3.0, 4.0]])
     """
-    if np.all(new_pt == INT_FILL_VALUE):
+    if all(new_pt == INT_FILL_VALUE):
         return old_box
 
     latlon_box = np.copy(old_box)  # Create a copy of the old box
@@ -612,18 +614,17 @@ def _insert_pt_in_latlonbox(old_box, new_pt, is_lon_periodic=True):
     # Check for pole points and update latitudes
     is_pole_point = (
         lon_pt == INT_FILL_VALUE
-        and np.isclose(
-            new_pt[0], [0.5 * np.pi, -0.5 * np.pi], atol=ERROR_TOLERANCE
-        ).any()
+        and isclose(new_pt[0], 0.5 * np.pi, atol=ERROR_TOLERANCE)
+        or isclose(new_pt[0], -0.5 * np.pi, atol=ERROR_TOLERANCE)
     )
 
     if is_pole_point:
         # Check if the new point is close to the North Pole
-        if np.isclose(new_pt[0], 0.5 * np.pi, atol=ERROR_TOLERANCE):
+        if isclose(new_pt[0], 0.5 * np.pi, atol=ERROR_TOLERANCE):
             latlon_box[0][1] = 0.5 * np.pi
 
         # Check if the new point is close to the South Pole
-        elif np.isclose(new_pt[0], -0.5 * np.pi, atol=ERROR_TOLERANCE):
+        elif isclose(new_pt[0], -0.5 * np.pi, atol=ERROR_TOLERANCE):
             latlon_box[0][0] = -0.5 * np.pi
 
         return latlon_box
@@ -769,9 +770,7 @@ def _populate_face_latlon_bound(
             )
 
             # Check if the node matches the pole point or if the pole point is within the edge_cart
-            if np.allclose(
-                n1_cart, pole_point, atol=ERROR_TOLERANCE
-            ) or point_within_gca(
+            if allclose(n1_cart, pole_point, atol=ERROR_TOLERANCE) or point_within_gca(
                 pole_point, np.array([n1_cart, n2_cart]), is_directed=False
             ):
                 is_center_pole = False
@@ -850,16 +849,16 @@ def _populate_face_latlon_bound(
             )
 
             # Insert extreme latitude points into the latlonbox if they differ from the node latitudes
-            if not np.isclose(
+            if not isclose(
                 node1_lat_rad, lat_max, atol=ERROR_TOLERANCE
-            ) and not np.isclose(node2_lat_rad, lat_max, atol=ERROR_TOLERANCE):
+            ) and not isclose(node2_lat_rad, lat_max, atol=ERROR_TOLERANCE):
                 # Insert the maximum latitude
                 face_latlon_array = _insert_pt_in_latlonbox(
                     face_latlon_array, np.array([lat_max, node1_lon_rad])
                 )
-            elif not np.isclose(
+            elif not isclose(
                 node1_lat_rad, lat_min, atol=ERROR_TOLERANCE
-            ) and not np.isclose(node2_lat_rad, lat_min, atol=ERROR_TOLERANCE):
+            ) and not isclose(node2_lat_rad, lat_min, atol=ERROR_TOLERANCE):
                 # Insert the minimum latitude
                 face_latlon_array = _insert_pt_in_latlonbox(
                     face_latlon_array, np.array([lat_min, node1_lon_rad])
@@ -939,7 +938,7 @@ def _populate_bounds(
     intervals_tuple_list = []
     intervals_name_list = []
 
-    faces_edges_cartesian = _get_cartesian_face_edge_nodes(
+    face_edges_cartesian = _get_cartesian_face_edge_nodes(
         grid.face_node_connectivity.values,
         grid.n_face,
         grid.n_max_face_edges,
@@ -948,7 +947,7 @@ def _populate_bounds(
         grid.node_z.values,
     )
 
-    faces_edges_lonlat_rad = _get_lonlat_rad_face_edge_nodes(
+    face_edges_lonlat_rad = _get_lonlat_rad_face_edge_nodes(
         grid.face_node_connectivity.values,
         grid.n_face,
         grid.n_max_face_edges,
@@ -956,17 +955,26 @@ def _populate_bounds(
         grid.node_lat.values,
     )
 
-    for face_idx, face_nodes in enumerate(grid.face_node_connectivity):
-        face_edges_cartesian = faces_edges_cartesian[face_idx]
+    face_node_connectivity = grid.face_node_connectivity.values
 
-        # Skip processing if the face is a dummy face
-        if np.any(face_edges_cartesian == INT_FILL_VALUE):
-            continue
+    # # TODO: vectorize dummy face check
+    s = face_edges_cartesian.shape
+    dummy_face_face_edges_cart = np.any(
+        face_edges_cartesian.reshape((s[0], s[1] * s[2] * s[3])) == INT_FILL_VALUE,
+        axis=1,
+    )
+
+    s = face_edges_lonlat_rad.shape
+    dummy_face_face_edges_latlon = np.any(
+        face_edges_lonlat_rad.reshape((s[0], s[1] * s[2] * s[3])) == INT_FILL_VALUE,
+        axis=1,
+    )
 
-        face_edges_lonlat_rad = faces_edges_lonlat_rad[face_idx]
+    dummy_faces = dummy_face_face_edges_cart | dummy_face_face_edges_latlon
 
-        # Skip processing if the face is a dummy face
-        if np.any(face_edges_lonlat_rad == INT_FILL_VALUE):
+    for face_idx, face_nodes in enumerate(face_node_connectivity):
+        if dummy_faces[face_idx]:
+            # skip dummy faces
             continue
 
         is_GCA_list = (
@@ -974,14 +982,15 @@ def _populate_bounds(
         )
 
         temp_latlon_array[face_idx] = _populate_face_latlon_bound(
-            face_edges_cartesian,
-            face_edges_lonlat_rad,
+            face_edges_cartesian[face_idx],
+            face_edges_lonlat_rad[face_idx],
             is_latlonface=is_latlonface,
             is_GCA_list=is_GCA_list,
         )
 
-        assert temp_latlon_array[face_idx][0][0] != temp_latlon_array[face_idx][0][1]
-        assert temp_latlon_array[face_idx][1][0] != temp_latlon_array[face_idx][1][1]
+        # # do we need this ?
+        # assert temp_latlon_array[face_idx][0][0] != temp_latlon_array[face_idx][0][1]
+        # assert temp_latlon_array[face_idx][1][0] != temp_latlon_array[face_idx][1][1]
         lat_array = temp_latlon_array[face_idx][0]
 
         # Now store the latitude intervals in the tuples
@@ -1001,7 +1010,7 @@ def _populate_bounds(
         attrs={
             "cf_role": "face_latlon_bounds",
             "_FillValue": INT_FILL_VALUE,
-            "long_name": "Provides the latitude and longitude bounds for each face in radians.",
+            "long_name": "Latitude and Longitude bounds for each face in radians.",
             "start_index": INT_DTYPE(0),
             "latitude_intervalsIndex": intervalsIndex,
             "latitude_intervals_name_map": df_intervals_map,
diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py
index 5b92c1026..14e0245d8 100644
--- a/uxarray/grid/grid.py
+++ b/uxarray/grid/grid.py
@@ -891,6 +891,7 @@ def bounds(self):
                 "Constructing of `Grid.bounds` has not been optimized, which may lead to a long execution time."
             )
             _populate_bounds(self)
+
         return self._ds["bounds"]
 
     @property
diff --git a/uxarray/grid/intersections.py b/uxarray/grid/intersections.py
index 3d3b03f42..d57969db5 100644
--- a/uxarray/grid/intersections.py
+++ b/uxarray/grid/intersections.py
@@ -4,10 +4,12 @@
 from uxarray.grid.arcs import point_within_gca
 import platform
 import warnings
-from uxarray.utils.computing import cross_fma
+from uxarray.utils.computing import cross_fma, allclose, cross, dot
 
+import uxarray.constants
 
-def gca_gca_intersection(gca1_cart, gca2_cart, fma_disabled=False):
+
+def gca_gca_intersection(gca1_cart, gca2_cart):
     """Calculate the intersection point(s) of two Great Circle Arcs (GCAs) in a
     Cartesian coordinate system.
 
@@ -54,12 +56,13 @@ def gca_gca_intersection(gca1_cart, gca2_cart, fma_disabled=False):
     w0, w1 = gca1_cart
     v0, v1 = gca2_cart
 
-    # Compute normals and orthogonal bases using FMA
-    if fma_disabled:
-        w0w1_norm = np.cross(w0, w1)
-        v0v1_norm = np.cross(v0, v1)
-        cross_norms = np.cross(w0w1_norm, v0v1_norm)
+    if not uxarray.constants.ENABLE_FMA:
+        # Compute normals and orthogonal bases without FMA
+        w0w1_norm = cross(w0, w1)
+        v0v1_norm = cross(v0, v1)
+        cross_norms = cross(w0w1_norm, v0v1_norm)
     else:
+        # Compute normals and orthogonal bases using FMA
         w0w1_norm = cross_fma(w0, w1)
         v0v1_norm = cross_fma(v0, v1)
         cross_norms = cross_fma(w0w1_norm, v0v1_norm)
@@ -73,29 +76,29 @@ def gca_gca_intersection(gca1_cart, gca2_cart, fma_disabled=False):
             )
 
     # Check perpendicularity conditions and floating-point arithmetic limitations
-    if not np.allclose(
-        np.dot(w0w1_norm, w0), 0, atol=ERROR_TOLERANCE
-    ) or not np.allclose(np.dot(w0w1_norm, w1), 0, atol=ERROR_TOLERANCE):
+    if not allclose(dot(w0w1_norm, w0), 0, atol=ERROR_TOLERANCE) or not allclose(
+        dot(w0w1_norm, w1), 0, atol=ERROR_TOLERANCE
+    ):
         warnings.warn(
             "The current input data cannot be computed accurately using floating-point arithmetic. Use with caution."
         )
 
-    if not np.allclose(
-        np.dot(v0v1_norm, v0), 0, atol=ERROR_TOLERANCE
-    ) or not np.allclose(np.dot(v0v1_norm, v1), 0, atol=ERROR_TOLERANCE):
+    if not allclose(dot(v0v1_norm, v0), 0, atol=ERROR_TOLERANCE) or not allclose(
+        dot(v0v1_norm, v1), 0, atol=ERROR_TOLERANCE
+    ):
         warnings.warn(
             "The current input data cannot be computed accurately using floating-point arithmetic.  Use with caution. "
         )
 
-    if not np.allclose(
-        np.dot(cross_norms, v0v1_norm), 0, atol=ERROR_TOLERANCE
-    ) or not np.allclose(np.dot(cross_norms, w0w1_norm), 0, atol=ERROR_TOLERANCE):
+    if not allclose(
+        dot(cross_norms, v0v1_norm), 0, atol=ERROR_TOLERANCE
+    ) or not allclose(dot(cross_norms, w0w1_norm), 0, atol=ERROR_TOLERANCE):
         warnings.warn(
             "The current input data cannot be computed accurately using floating-point arithmetic. Use with caution. "
         )
 
     # If the cross_norms is zero, the two GCAs are parallel
-    if np.allclose(cross_norms, 0, atol=ERROR_TOLERANCE):
+    if allclose(cross_norms, 0, atol=ERROR_TOLERANCE):
         return np.array([])
 
     # Normalize the cross_norms
@@ -156,7 +159,7 @@ def gca_constLat_intersection(
     x1, x2 = gca_cart
 
     if fma_disabled:
-        n = np.cross(x1, x2)
+        n = cross(x1, x2)
 
     else:
         # Raise a warning for Windows users
diff --git a/uxarray/utils/computing.py b/uxarray/utils/computing.py
index 3d801e929..f9929deb4 100644
--- a/uxarray/utils/computing.py
+++ b/uxarray/utils/computing.py
@@ -1,6 +1,65 @@
 import numpy as np
 import sys
 
+from numba import njit
+
+
+@njit
+def all(a):
+    """Numba decorated implementation of ``np.all()``
+
+    See Also
+    --------
+    numpy.all
+    """
+
+    return np.all(a)
+
+
+@njit
+def isclose(a, b, rtol=1e-05, atol=1e-08):
+    """Numba decorated implementation of ``np.isclose()``
+
+    See Also
+    --------
+    numpy.isclose
+    """
+
+    return np.isclose(a, b, rtol=rtol, atol=atol)
+
+
+@njit
+def allclose(a, b, rtol=1e-05, atol=1e-08):
+    """Numba decorated implementation of ``np.allclose()``
+
+    See Also
+    --------
+    numpy.allclose
+    """
+    return np.allclose(a, b, rtol=rtol, atol=atol)
+
+
+@njit
+def cross(a, b):
+    """Numba decorated implementation of ``np.cross()``
+
+    See Also
+    --------
+    numpy.cross
+    """
+    return np.cross(a, b)
+
+
+@njit
+def dot(a, b):
+    """Numba decorated implementation of ``np.dot()``
+
+    See Also
+    --------
+    numpy.dot
+    """
+    return np.dot(a, b)
+
 
 def _fmms(a, b, c, d):
     """