lox: tolerate array / scalar together better

Author: scivision

src/pymap3d/latitude.py

@@ -4,22 +4,24 @@
 import typing
 
 from .ellipsoid import Ellipsoid
-from .utils import sanitize
+from .utils import sanitize, sign
 from . import rcurve
 
 try:
-    from numpy import radians, degrees, tan, sin, exp, pi, sqrt, inf
+    from numpy import radians, degrees, tan, sin, cos, exp, pi, sqrt, inf
     from numpy import arctan as atan, arcsinh as asinh, arctanh as atanh  # noqa: A001
 
     use_numpy = True
 except ImportError:
-    from math import atan, radians, degrees, tan, sin, asinh, atanh, exp, pi, sqrt, inf  # type: ignore
+    from math import atan, radians, degrees, tan, sin, cos, asinh, atanh, exp, pi, sqrt, inf  # type: ignore
 
     use_numpy = False
 
 if typing.TYPE_CHECKING:
     from numpy import ndarray
 
+COS_EPS = 1e-9  # tolerance for angles near abs([90, 270])
+
 __all__ = [
     "geodetic2isometric",
     "isometric2geodetic",
@@ -208,16 +210,22 @@ def geodetic2isometric(
     # isometric_lat = log(tan(a2) * (y ** (e / 2)))
     # isometric_lat = log(tan(a2)) + e/2 * log((1-e*sin(geodetic_lat)) / (1+e*sin(geodetic_lat)))
 
+    coslat = cos(geodetic_lat)
+    i = abs(coslat) <= COS_EPS
+
     try:
-        isometric_lat[abs(geodetic_lat - pi / 2) <= 1e-9] = inf  # type: ignore
-        isometric_lat[abs(-geodetic_lat - pi / 2) <= 1e-9] = -inf   # type: ignore
+        isometric_lat[i] = sign(geodetic_lat[i]) * inf  # type: ignore
     except TypeError:
-        if abs(geodetic_lat - pi / 2) <= 1e-9:  # type: ignore
-            isometric_lat = inf
-        elif abs(-geodetic_lat - pi / 2) <= 1e-9:  # type: ignore
-            isometric_lat = -inf
+        if i:
+            isometric_lat = sign(geodetic_lat) * inf
 
-    return degrees(isometric_lat) if deg else isometric_lat
+    if deg:
+        isometric_lat = degrees(isometric_lat)
+
+    try:
+        return isometric_lat.squeeze()[()]
+    except AttributeError:
+        return isometric_lat
 
 
 def isometric2geodetic(isometric_lat: ndarray, ell: Ellipsoid = None, deg: bool = True) -> ndarray:

src/pymap3d/lox.py

@@ -4,7 +4,17 @@
 import typing
 
 try:
-    from numpy import radians, degrees, cos, arctan2 as atan2, tan, pi, array, atleast_1d
+    from numpy import (
+        radians,
+        degrees,
+        cos,
+        arctan2 as atan2,
+        tan,
+        pi,
+        array,
+        atleast_1d,
+        broadcast_to,
+    )
 except ImportError:
     from math import radians, degrees, cos, atan2, tan, pi  # type: ignore
 
@@ -97,8 +107,6 @@ def loxodrome_inverse(
 
     like Matlab distance('rh',...) and azimuth('rh',...)
 
-    If any of inputs lat1,lon1,lat2,lon2 are arrays, all must be arrays of same shape
-
     Parameters
     ----------
 
@@ -141,6 +149,33 @@ def loxodrome_inverse(
     if deg:
         lat1, lon1, lat2, lon2 = radians(lat1), radians(lon1), radians(lat2), radians(lon2)
 
+    try:
+        lat1 = atleast_1d(lat1)
+        lon1 = atleast_1d(lon1)
+        lat2 = atleast_1d(lat2)
+        lon2 = atleast_1d(lon2)
+
+        if lat1.shape != lon1.shape:
+            if lat1.size == 1:
+                lat1 = broadcast_to(lat1, lon1.shape)
+            elif lon1.size == 1:
+                lon1 = broadcast_to(lon1, lat1.shape)
+            else:
+                raise ValueError("lat1, lon1 must have same shape")
+        if lat2.shape != lon2.shape:
+            if lat2.size == 1:
+                lat2 = broadcast_to(lat2, lon2.shape)
+            elif lon2.size == 1:
+                lon2 = broadcast_to(lon2, lat2.shape)
+            else:
+                raise ValueError("lat2, lon2 must have same shape")
+        if lat2.shape != lat1.shape:
+            lat2 = broadcast_to(lat2, lat1.shape)
+            lon2 = broadcast_to(lon2, lat1.shape)
+
+    except NameError:
+        pass
+
     # compute changes in isometric latitude and longitude between points
     disolat = geodetic2isometric(lat2, deg=False, ell=ell) - geodetic2isometric(
         lat1, deg=False, ell=ell
@@ -165,7 +200,10 @@ def loxodrome_inverse(
     if deg:
         az12 = degrees(az12) % 360.0
 
-    return dist, az12
+    try:
+        return dist.squeeze()[()], az12.squeeze()[()]
+    except AttributeError:
+        return dist, az12
 
 
 def loxodrome_direct(
@@ -182,8 +220,6 @@ def loxodrome_direct(
     like Matlab reckon('rh', ...)
     except that "rng" in meters instead of "arclen" degrees of arc
 
-    If any of inputs lat,lon1,rng are arrays, all must be arrays of same shape
-
     Parameters
     ----------
     lat1 : float
@@ -213,12 +249,20 @@ def loxodrome_direct(
     try:
         lat1 = atleast_1d(lat1)
         rng = atleast_1d(rng)
-        if (abs(lat1) > pi / 2).any():
+
+        if lat1.shape != rng.shape:
+            if rng.size == 1:
+                rng = broadcast_to(rng, lat1.shape)
+            elif lat1.size == 1:
+                lat1 = broadcast_to(lat1, rng.shape)
+            else:
+                raise ValueError("lat1 and rng must each be scalars or same shape")
+        if (abs(lat1) > pi / 2).any():  # type: ignore
             raise ValueError("-90 <= latitude <= 90")
-        if (rng < 0).any():
+        if (rng < 0).any():  # type: ignore
             raise ValueError("ground distance must be >= 0")
     except NameError:
-        if abs(lat1) > pi / 2:
+        if abs(lat1) > pi / 2:  # type: ignore
             raise ValueError("-90 <= latitude <= 90")
         if rng < 0:
             raise ValueError("ground distance must be >= 0")

src/pymap3d/tests/test_latitude.py

@@ -49,7 +49,14 @@ def test_numpy_geodetic_geocentric():
 
 @pytest.mark.parametrize(
     "geodetic_lat, isometric_lat",
-    [(0, 0), (90, inf), (-90, -inf), (45, 50.227466), (-45, -50.227466), (89, 271.275)],
+    [
+        (0, 0),
+        (90, inf),
+        (-90, -inf),
+        (45, 50.227466),
+        (-45, -50.227466),
+        (89, 271.275),
+    ],
 )
 def test_geodetic_isometric(geodetic_lat, isometric_lat):
     isolat = latitude.geodetic2isometric(geodetic_lat)

src/pymap3d/tests/test_rhumb.py

@@ -69,10 +69,12 @@ def test_loxodrome_inverse(lat1, lon1, lat2, lon2, arclen, az):
 def test_numpy_loxodrome_inverse():
     pytest.importorskip("numpy")
     d, a = lox.loxodrome_inverse([40, 40], [-80, -80], 65, -148)
-
     assert d == approx(5248666.209)
     assert a == approx(302.00567)
 
+    d, a = lox.loxodrome_inverse([40, 40], [-80, -80], [65, 65], -148)
+    d, a = lox.loxodrome_inverse([40, 40], [-80, -80], 65, [-148, -148])
+
 
 @pytest.mark.parametrize(
     "lat0,lon0,rng,az,lat1,lon1",
@@ -98,9 +100,13 @@ def test_loxodrome_direct(lat0, lon0, rng, az, lat1, lon1):
 
 def test_numpy_loxodrome_direct():
     pytest.importorskip("numpy")
-    lat, lon = lox.loxodrome_direct([40, 40], [-80, -80], [10000, 10000], [30, 30])
-    assert lat == approx(40.077995)
-    assert lon == approx(-79.941414)
+    lat, lon = lox.loxodrome_direct([40, 40], [-80, -80], [10, 10], [30, 30])
+    assert lat == approx(40.000078)
+    assert lon == approx(-79.99994145)
+
+    lat, lon = lox.loxodrome_direct([40, 40], [-80, -80], 10, 30)
+    lat, lon = lox.loxodrome_direct([40, 40], [-80, -80], [10, 10], 30)
+    lat, lon = lox.loxodrome_direct([40, 40], [-80, -80], 10, [30, 30])
 
 
 @pytest.mark.parametrize("lat,lon", [([0, 45, 90], [0, 45, 90])])