Haversine() returns float if end array has two dimensions and output …

…contains only one point (#49) * FIX: HaversineWrapper() can handle different array shapes for same array sizes * HaversineWrapper() returns double or float if only one set of points requested * Add more tests to test_distances.py * Update Haversine() in pyi * Bump to 0.3.6 --------- Co-authored-by: Daniel Stoops <[email protected]>
Stoops-ML · Nov 11, 2024 · 4e2874b · 4e2874b
1 parent ca49a5e
commit 4e2874b
Show file tree

Hide file tree

Showing 4 changed files with 89 additions and 13 deletions.
diff --git a/include/distances.c b/include/distances.c
@@ -81,7 +81,7 @@ HaversineWrapper(PyObject* self, PyObject* args)
         PyErr_SetString(PyExc_ValueError, "Input arrays must be a C contiguous.");
         return NULL;
     }
-    if (!((PyArray_NDIM(rrmStart) == PyArray_NDIM(rrmEnd)) && (PyArray_SIZE(rrmStart) == PyArray_SIZE(rrmEnd)) || ((PyArray_SIZE(rrmStart) == NCOORDSIN3D) && (PyArray_SIZE(rrmStart) < PyArray_SIZE(rrmEnd))))) {
+    if (!((PyArray_NDIM(rrmStart) == PyArray_NDIM(rrmEnd)) && (PyArray_SIZE(rrmStart) == PyArray_SIZE(rrmEnd)) || ((PyArray_SIZE(rrmStart) == NCOORDSIN3D) && (PyArray_SIZE(rrmStart) <= PyArray_SIZE(rrmEnd))))) {
         PyErr_SetString(PyExc_ValueError,
             "Input arrays must have matching size and dimensions or "
             "the start point must be of size three.");
@@ -155,7 +155,12 @@ HaversineWrapper(PyObject* self, PyObject* args)
             "Only 32 and 64 bit float types or all integer are accepted.");
         return NULL;
     }
-    return (PyObject*)result_array;
+    if ((nPoints == 1) && (PyArray_NDIM(rrmEnd) == 2) && (PyArray_TYPE(result_array) == NPY_DOUBLE))
+        return Py_BuildValue("d", *((double*)PyArray_DATA(result_array)));
+    else if ((nPoints == 1) && (PyArray_NDIM(rrmEnd) == 2) && (PyArray_TYPE(result_array) == NPY_FLOAT))
+        return Py_BuildValue("f", *((float*)PyArray_DATA(result_array)));
+    else
+        return (PyObject*)result_array;
 }
 
 // Method definition object for this extension, these argumens mean:

diff --git a/src/transforms84/__init__.py b/src/transforms84/__init__.py
@@ -1 +1 @@
-__version__ = "0.3.5"
+__version__ = "0.3.6"
diff --git a/src/transforms84/distances.pyi b/src/transforms84/distances.pyi
@@ -11,6 +11,7 @@ def Haversine(
         Union[np.float32, np.float64, np.int8, np.int16, np.int32, np.int64]
     ],
     m_radius_sphere: float,
-) -> npt.NDArray[
-    Union[np.float32, np.float64, np.int8, np.int16, np.int32, np.int64]
+) -> Union[
+    float,
+    npt.NDArray[Union[np.float32, np.float64, np.int8, np.int16, np.int32, np.int64]],
 ]: ...
diff --git a/tests/test_distances.py b/tests/test_distances.py
@@ -155,7 +155,8 @@ def test_Haversine_parallel(dtype):
 @pytest.mark.parametrize(
     "dtype0,dtype1", [(np.float64, np.float32), (np.float32, np.float64)]
 )
-def test_Haversine_different_dtypes(dtype0, dtype1):
+def test_Haversine_point_2D_different_dtypes(dtype0, dtype1):
+    """output for 2D arrays end point should be a float"""
     rrm_start_with_height = np.array(
         [[np.deg2rad(33)], [np.deg2rad(34)], [100000]], dtype=dtype0
     )
@@ -165,22 +166,91 @@ def test_Haversine_different_dtypes(dtype0, dtype1):
     out1 = Haversine(rrm_start, rrm_end, WGS84.mean_radius)
     out2 = Haversine(rrm_end, rrm_start_with_height, WGS84.mean_radius)
     out3 = Haversine(rrm_end, rrm_start, WGS84.mean_radius)
+    assert isinstance(out0, float)
+    assert isinstance(out1, float)
+    assert isinstance(out2, float)
+    assert isinstance(out3, float)
+    assert np.isclose(out0, 391225.574516907)
+    assert np.isclose(out1, out0)
+    assert np.isclose(out1, 391225.574516907)
+    assert np.isclose(out2, 391225.574516907)
+    assert np.isclose(out3, 391225.574516907)
+    assert np.isclose(out1, out3)
+
+
+@pytest.mark.parametrize(
+    "dtype0,dtype1", [(np.float64, np.float32), (np.float32, np.float64)]
+)
+def test_Haversine_point_3D_end_point_different_dtypes(dtype0, dtype1):
+    """output for a 3D array end point should be an array"""
+    rrm_start_with_height = np.array(
+        [[np.deg2rad(33)], [np.deg2rad(34)], [100000]], dtype=dtype0
+    )
+    rrm_start = np.array([[np.deg2rad(33)], [np.deg2rad(34)], [0]], dtype=dtype0)
+    rrm_end = np.array([[np.deg2rad(32)], [np.deg2rad(38)], [0]], dtype=dtype1)
+    out0 = Haversine(rrm_start_with_height, rrm_end[None, :], WGS84.mean_radius)
+    out1 = Haversine(rrm_start, rrm_end[None, :], WGS84.mean_radius)
+    out2 = Haversine(rrm_end, rrm_start_with_height[None, :], WGS84.mean_radius)
+    out3 = Haversine(rrm_end, rrm_start[None, :], WGS84.mean_radius)
+    assert isinstance(out0, np.ndarray)
+    assert isinstance(out1, np.ndarray)
+    assert isinstance(out2, np.ndarray)
+    assert isinstance(out3, np.ndarray)
     assert out0.dtype == np.float64
     assert out1.dtype == np.float64
     assert out2.dtype == np.float64
     assert out3.dtype == np.float64
     assert np.isclose(out0, 391225.574516907)
-    assert np.isclose(
-        out1,
-        out0,
-    )
+    assert np.isclose(out1, out0)
     assert np.isclose(out1, 391225.574516907)
     assert np.isclose(out2, 391225.574516907)
     assert np.isclose(out3, 391225.574516907)
-    assert np.isclose(
-        out1,
-        out3,
+    assert np.isclose(out1, out3)
+
+
+@pytest.mark.parametrize(
+    "dtype0,dtype1", [(np.float64, np.float32), (np.float32, np.float64)]
+)
+def test_Haversine_points_different_dtypes(dtype0, dtype1):
+    rrm_start_with_height = np.array(
+        [
+            [[np.deg2rad(33)], [np.deg2rad(34)], [100000]],
+            [[np.deg2rad(33)], [np.deg2rad(34)], [100000]],
+        ],
+        dtype=dtype0,
+    )
+    rrm_start = np.array(
+        [
+            [[np.deg2rad(33)], [np.deg2rad(34)], [0]],
+            [[np.deg2rad(33)], [np.deg2rad(34)], [0]],
+        ],
+        dtype=dtype0,
+    )
+    rrm_end = np.array(
+        [
+            [[np.deg2rad(32)], [np.deg2rad(38)], [0]],
+            [[np.deg2rad(32)], [np.deg2rad(38)], [0]],
+        ],
+        dtype=dtype1,
     )
+    out0 = Haversine(rrm_start_with_height, rrm_end, WGS84.mean_radius)
+    out1 = Haversine(rrm_start, rrm_end, WGS84.mean_radius)
+    out2 = Haversine(rrm_end, rrm_start_with_height, WGS84.mean_radius)
+    out3 = Haversine(rrm_end, rrm_start, WGS84.mean_radius)
+    assert isinstance(out0, np.ndarray)
+    assert isinstance(out1, np.ndarray)
+    assert isinstance(out2, np.ndarray)
+    assert isinstance(out3, np.ndarray)
+    assert out0.dtype == np.float64
+    assert out1.dtype == np.float64
+    assert out2.dtype == np.float64
+    assert out3.dtype == np.float64
+    assert np.all(np.isclose(out0, 391225.574516907))
+    assert np.all(np.isclose(out1, out0))
+    assert np.all(np.isclose(out1, 391225.574516907))
+    assert np.all(np.isclose(out2, 391225.574516907))
+    assert np.all(np.isclose(out3, 391225.574516907))
+    assert np.all(np.isclose(out1, out3))
 
 
 @pytest.mark.parametrize("dtype", [np.float64, np.float32])