add point cloud support to_spherical (#285)

src/geovista/common.py

@@ -19,6 +19,7 @@
 
 __all__ = [
     "COASTLINES_RESOLUTION",
+    "DISTANCE_DECIMALS",
     "GV_CELL_IDS",
     "GV_FIELD_CRS",
     "GV_FIELD_NAME",
@@ -529,6 +530,8 @@ def to_spherical(
     longitudes: npt.ArrayLike,
     latitudes: npt.ArrayLike,
     radius: float | None = None,
+    zlevel: int | npt.ArrayLike | None = None,
+    zscale: float | None = None,
     stacked: bool | None = True,
 ) -> np.ndarray:
     """Convert geographic `longitudes` and `latitudes` to cartesian ``xyz`` points.
@@ -541,6 +544,14 @@ def to_spherical(
         The latitude values (degrees) to be converted.
     radius : float, optional
         The radius of the sphere. Defaults to :data:`RADIUS`.
+    zlevel : int or ArrayLike, default=0
+        The z-axis level. Used in combination with the `zscale` to offset the
+        `radius` by a proportional amount i.e., ``radius * zlevel * zscale``.
+        If `zlevel` is not a scalar, then its shape must match or broadcast
+        with the shape of `longitude` and `latitude`.
+    zscale : float, optional
+        The proportional multiplier for z-axis `zlevel`. Defaults to
+        :data:`ZLEVEL_SCALE`.
     stacked : bool, default=True
         Specify whether the resultant xyz points have shape (N, 3).
         Otherwise, they will have shape (3, N).
@@ -555,21 +566,38 @@ def to_spherical(
     .. versionadded:: 0.1.0
 
     """
-    longitudes = np.ravel(longitudes)
-    latitudes = np.ravel(latitudes)
+    longitudes = np.asanyarray(longitudes)
+    latitudes = np.asanyarray(latitudes)
+
+    if (shape := longitudes.shape) != latitudes.shape:
+        emsg = (
+            f"Require longitudes and latitudes with same shape, got {shape} and "
+            f"{latitudes.shape} respectively."
+        )
+        raise ValueError(emsg)
+
     radius = RADIUS if radius is None else abs(float(radius))
+    zscale = ZLEVEL_SCALE if zscale is None else float(zscale)
+    zlevel = np.array([0]) if zlevel is None else np.atleast_1d(zlevel).astype(int)
+
+    try:
+        _ = np.broadcast_shapes(zshape := zlevel.shape, shape)
+    except ValueError as err:
+        emsg = (
+            f"Cannot broadcast zlevel with shape {zshape} to longitude/latitude"
+            f"shape {shape}."
+        )
+        raise ValueError(emsg) from err
+
+    radius += radius * zlevel * zscale
 
     x_rad = np.radians(longitudes)
     y_rad = np.radians(90.0 - latitudes)
-    x = radius * np.sin(y_rad) * np.cos(x_rad)
-    y = radius * np.sin(y_rad) * np.sin(x_rad)
-    z = radius * np.cos(y_rad)
+    x = np.ravel(radius * np.sin(y_rad) * np.cos(x_rad))
+    y = np.ravel(radius * np.sin(y_rad) * np.sin(x_rad))
+    z = np.ravel(radius * np.cos(y_rad))
     xyz = [x, y, z]
-
-    if stacked:
-        xyz = np.vstack(xyz).T
-    else:
-        xyz = np.array(xyz)
+    xyz = np.vstack(xyz).T if stacked else np.array(xyz)
 
     return xyz
 

tests/common/test_to_spherical.py

@@ -0,0 +1,89 @@
+"""Unit-tests for :func:`geovista.common.to_spherical`."""
+import numpy as np
+import numpy.typing as npt
+import pytest
+
+from geovista.common import DISTANCE_DECIMALS, RADIUS, ZLEVEL_SCALE, to_spherical
+
+
+def _distance(
+    pts: npt.ArrayLike, stacked: bool = True, decimals: int = DISTANCE_DECIMALS
+) -> float:
+    """Calculate the mean distance to the xyz-cartesian `pts` from the origin."""
+    if np.ma.isMaskedArray(pts):
+        pts = pts.data
+    if not stacked:
+        pts = np.transpose(pts)
+    nrow, ncol = pts.shape
+    result = np.round(
+        np.sqrt(np.sum(pts.T @ pts * np.identity(ncol)) / nrow), decimals=decimals
+    )
+    return result
+
+
+def test_shape_fail():
+    """Test trap of longitude and latitude shape mismatch."""
+    lons, lats = np.arange(10), np.arange(10).reshape(5, 2)
+    emsg = "Require longitudes and latitudes with same shape"
+    with pytest.raises(ValueError, match=emsg):
+        _ = to_spherical(lons, lats)
+
+
+def test_zlevel_broadcast_fail():
+    """Test trap of zlevel shape can't broadcast with longitude/latitude."""
+    lons, lats = np.arange(10), np.arange(10)
+    zlevel = np.arange(2)
+    emsg = "Cannot broadcast zlevel"
+    with pytest.raises(ValueError, match=emsg):
+        _ = to_spherical(lons, lats, zlevel=zlevel)
+
+
+@pytest.mark.parametrize("stacked", [True, False])
+def test_defaults(lam_uk_sample, stacked):
+    """Test expected defaults are honoured."""
+    result = to_spherical(*lam_uk_sample, stacked=stacked)
+    assert _distance(result, stacked=stacked) == RADIUS
+
+
+@pytest.mark.parametrize("zlevel", range(-5, 6))
+def test_zlevel__scalar(lam_uk_sample, zlevel):
+    """Test spherical z-control with zlevel."""
+    result = to_spherical(*lam_uk_sample, zlevel=zlevel)
+    actual = _distance(result)
+    expected = RADIUS + RADIUS * zlevel * ZLEVEL_SCALE
+    assert np.isclose(actual, expected)
+
+
+@pytest.mark.parametrize(
+    "xy_reshape, z_reshape", [((-1,), (-1, 1)), ((1, 5, 5), (-1, 1, 1))]
+)
+@pytest.mark.parametrize("n_levels", range(3, 11))
+def test_zlevel__broadcast(lam_uk_sample, xy_reshape, z_reshape, n_levels):
+    """Test spherical z-control with zlevel broadcast."""
+    lons, lats = lam_uk_sample
+    (npts,) = lons.shape
+    vlons = np.vstack([lons[:].reshape(*xy_reshape)] * n_levels)
+    vlats = np.vstack([lats[:].reshape(*xy_reshape)] * n_levels)
+    mid = n_levels // 2
+    zlevel = np.arange(-mid, n_levels - mid)
+    result = to_spherical(vlons, vlats, zlevel=zlevel.reshape(*z_reshape))
+    assert result.ndim == 2
+    assert result.shape[0] == n_levels * npts
+    for i in range(n_levels):
+        actual = _distance(result[i * npts : (i + 1) * npts])
+        expected = RADIUS + RADIUS * zlevel[i] * ZLEVEL_SCALE
+        assert np.isclose(actual, expected)
+
+
+@pytest.mark.parametrize("zlevel", [0, 1])
+@pytest.mark.parametrize("zscale", np.linspace(-1, 1, num=5))
+def test_zscale(lam_uk_sample, zlevel, zscale):
+    """Test spherical z-control with zscale."""
+    lons, lats = lam_uk_sample
+    (npts,) = lons.shape
+    result = to_spherical(lons, lats, zlevel=zlevel, zscale=zscale)
+    assert result.ndim == 2
+    assert result.shape[0] == npts
+    actual = _distance(result)
+    expected = RADIUS + RADIUS * zlevel * zscale
+    assert np.isclose(actual, expected)

tests/conftest.py

@@ -2,6 +2,7 @@
 import numpy as np
 import pytest
 
+from geovista.pantry import lam_uk as pantry_lam_uk
 from geovista.samples import lam_uk as sample_lam_uk
 from geovista.samples import lfric as sample_lfric
 from geovista.samples import lfric_sst as sample_lfric_sst
@@ -16,6 +17,13 @@ def lam_uk():
     return mesh
 
 
+@pytest.fixture(scope="session")
+def lam_uk_sample():
+    """Fixture generates a Local Area Model data sample for the UK."""
+    sample = pantry_lam_uk()
+    return sample.lons, sample.lats
+
+
 @pytest.fixture
 def lfric(request):
     """Fixture to provide a cube-sphere mesh."""

tests/geometry/test_load_coastlines.py

@@ -49,18 +49,11 @@ def test_zlevel(resolution, zlevel):
     assert np.isclose(actual, expected)
 
 
+@pytest.mark.parametrize("zlevel", [0, 1])
 @pytest.mark.parametrize("zscale", np.linspace(-1, 1, num=5))
-def test_zscale(resolution, zscale):
+def test_zscale(resolution, zlevel, zscale):
     """Test coastline z-control with zscale with no zlevel."""
-    result = load(resolution=resolution, zscale=zscale)
+    result = load(resolution=resolution, zlevel=zlevel, zscale=zscale)
     actual = distance(result)
-    assert np.isclose(actual, RADIUS)
-
-
-@pytest.mark.parametrize("zscale", np.linspace(-1, 1, num=5))
-def test_zscale__with_zlevel(resolution, zscale):
-    """Test coastline z-control with zscale and zlevel."""
-    result = load(resolution=resolution, zlevel=1, zscale=zscale)
-    actual = distance(result)
-    expected = RADIUS + RADIUS * zscale
+    expected = RADIUS + RADIUS * zlevel * zscale
     assert np.isclose(actual, expected)