New example: plotting in spherical coordinates (#417)

* add example of plots in spherical coordinates

* fix a bug in spherical.py

* move generalised spherical-to-cartesian functions to utilities/features.py

* add tests for grid_from_sph_coords() and transform_vectors_sph_to_cart()

* remove redundant import statement from spherical.py

examples/02-plot/spherical.py

@@ -0,0 +1,160 @@
+"""
+Plot data in spherical coordinates
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Generate and visualize meshes from data in longitude-latitude coordinates.
+"""
+
+import pyvista as pv
+import numpy as np
+
+
+def _cell_bounds(points, bound_position=0.5):
+    """
+    Calculate coordinate cell boundaries.
+
+    Parameters
+    ----------
+    points: numpy.array
+        One-dimensional array of uniformy spaced values of shape (M,)
+    bound_position: bool, optional
+        The desired position of the bounds relative to the position
+        of the points.
+
+    Returns
+    -------
+    bounds: numpy.array
+        Array of shape (M+1,)
+
+    Examples
+    --------
+    >>> a = np.arange(-1, 2.5, 0.5)
+    >>> a
+    array([-1. , -0.5,  0. ,  0.5,  1. ,  1.5,  2. ])
+    >>> cell_bounds(a)
+    array([-1.25, -0.75, -0.25,  0.25,  0.75,  1.25,  1.75,  2.25])
+    """
+    assert points.ndim == 1, "Only 1D points are allowed"
+    diffs = np.diff(points)
+    delta = diffs[0] * bound_position
+    bounds = np.concatenate([[points[0] - delta], points + delta])
+    return bounds
+
+
+# First, create some dummy data
+
+# Approximate radius of the Earth
+RADIUS = 6371.0
+
+# Longitudes and latitudes
+x = np.arange(0, 360, 5)
+y = np.arange(-90, 91, 10)
+y_polar = 90.0 - y  # grid_from_sph_coords() expects polar angle
+
+xx, yy = np.meshgrid(x, y)
+
+
+# x- and y-components of the wind vector
+u_vec = np.cos(np.radians(xx))  # zonal
+v_vec = np.sin(np.radians(yy))  # meridional
+
+# Scalar data
+scalar = u_vec ** 2 + v_vec ** 2
+
+# Create arrays of grid cell boundaries, which have shape of (x.shape[0] + 1)
+xx_bounds = _cell_bounds(x)
+yy_bounds = _cell_bounds(y_polar)
+# Vertical levels
+# in this case a single level slighly above the surface of a sphere
+levels = [RADIUS * 1.01]
+
+###############################################################################
+# Create a structured grid
+grid_scalar = pv.grid_from_sph_coords(xx_bounds, yy_bounds, levels)
+
+# And fill its cell arrays with the scalar data
+grid_scalar.cell_arrays["example"] = np.array(scalar).swapaxes(-2, -1).ravel("C")
+
+# Make a plot
+p = pv.Plotter()
+p.add_mesh(pv.Sphere(radius=RADIUS))
+p.add_mesh(grid_scalar, clim=[0.1, 2.0], opacity=0.5, cmap="plasma")
+p.show()
+
+
+###############################################################################
+# Visualize vectors in spherical coordinates
+# Vertical wind
+w_vec = np.random.rand(*u_vec.shape)
+
+wind_level = [RADIUS * 1.2]
+
+# Sequence of axis indices for transpose()
+# (1, 0) for 2D arrays
+# (2, 1, 0) for 3D arrays
+inv_axes = [*range(u_vec.ndim)[::-1]]
+
+# Transform vectors to cartesian coordinates
+vectors = np.stack(
+    [
+        i.transpose(inv_axes).swapaxes(-2, -1).ravel("C")
+        for i in pv.transform_vectors_sph_to_cart(
+            x,
+            y_polar,
+            wind_level,
+            u_vec.transpose(inv_axes),
+            -v_vec.transpose(inv_axes),  # Minus sign because y-vector in polar coords is required
+            w_vec.transpose(inv_axes),
+        )
+    ],
+    axis=1,
+)
+
+# Scale vectors to make them visible
+vectors *= RADIUS * 0.1
+
+# Create a grid for the vectors
+grid_winds = pv.grid_from_sph_coords(x, y_polar, wind_level)
+
+# Add vectors to the grid
+grid_winds.point_arrays["example"] = vectors
+
+# Show the result
+p = pv.Plotter()
+p.add_mesh(pv.Sphere(radius=RADIUS))
+p.add_mesh(grid_winds.glyph(orient="example", scale="example", tolerance=0.005))
+p.show()
+
+
+###############################################################################
+# Isurfaces of 3D data in spherical coordinates
+
+# Number of vertical levels
+nlev = 10
+
+# Dummy 3D scalar data
+scalar_3d = (
+    scalar.repeat(nlev).reshape((*scalar.shape, nlev)) * np.arange(nlev)[np.newaxis, np.newaxis, :]
+).transpose(2, 0, 1)
+
+
+z_scale = 10
+z_offset = RADIUS * 1.1
+
+# Now it's not a single level but an array of levels
+levels = z_scale * (np.arange(scalar_3d.shape[0] + 1)) ** 2 + z_offset
+
+# Create a structured grid by transforming coordinates
+grid_scalar_3d = pv.grid_from_sph_coords(xx_bounds, yy_bounds, levels)
+
+# Add data to the grid
+grid_scalar_3d.cell_arrays["example"] = np.array(scalar_3d).swapaxes(-2, -1).ravel("C")
+
+# Create a set of isosurfaces
+surfaces = grid_scalar_3d.cell_data_to_point_data().contour(isosurfaces=[1, 5, 10, 15])
+
+# Show the result
+p = pv.Plotter()
+p.add_mesh(pv.Sphere(radius=RADIUS))
+p.add_mesh(surfaces)
+p.show()

pyvista/utilities/features.py

@@ -30,7 +30,7 @@ def create_grid(dataset, dimensions=(101, 101, 101)):
     """
     bounds = np.array(dataset.bounds)
     if dimensions is None:
-        # TODO: we should implement an algorithm to automaitcally deterime an
+        # TODO: we should implement an algorithm to automatically determine an
         # "optimal" grid size by looking at the sparsity of the points in the
         # input dataset - I actaully think VTK might have this implemented
         # somewhere
@@ -52,3 +52,67 @@ def single_triangle():
     points[2] = [0.5, 0.707, 0]
     cells = np.array([[3, 0, 1, 2]], ctypes.c_long)
     return pyvista.PolyData(points, cells)
+
+
+def grid_from_sph_coords(theta, phi, r):
+    """
+    Create a structured grid from arrays of spherical coordinates.
+
+    Parameters
+    ----------
+    theta: array-like
+        Azimuthal angle in degrees [0, 360)
+    phi: array-like
+        Polar (zenith) angle in degrees [0, 180]
+    r: array-like
+        Distance (radius) from the point of origin
+
+    Returns
+    -------
+    pyvista.StructuredGrid
+    """
+    x, y, z = np.meshgrid(np.radians(theta), np.radians(phi), r)
+    # Transform grid to cartesian coordinates
+    x_cart = z * np.sin(y) * np.cos(x)
+    y_cart = z * np.sin(y) * np.sin(x)
+    z_cart = z * np.cos(y)
+    # Make a grid object
+    return pyvista.StructuredGrid(x_cart, y_cart, z_cart)
+
+
+def transform_vectors_sph_to_cart(theta, phi, r, u, v, w):
+    """
+    Transform vectors from spherical (r, phi, theta) to cartesian coordinates (z, y, x).
+
+    Note the "reverse" order of arrays's axes, commonly used in geosciences.
+
+    Parameters
+    ----------
+    theta: array-like
+        Azimuthal angle in degrees [0, 360) of shape (M,)
+    phi: array-like
+        Polar (zenith) angle in degrees [0, 180] of shape (N,)
+    r: array-like
+        Distance (radius) from the point of origin of shape (P,)
+    u: array-like
+        X-component of the vector of shape (P, N, M)
+    v: array-like
+        Y-component of the vector of shape (P, N, M)
+    w: array-like
+        Z-component of the vector of shape (P, N, M)
+
+    Returns
+    -------
+    u_t, v_t, w_t: array-like
+        Arrays of transformed x-, y-, z-components, respectively.
+    """
+    xx, yy, _ = np.meshgrid(np.radians(theta), np.radians(phi), r, indexing="ij")
+    th, ph = xx.squeeze(), yy.squeeze()
+
+    # Transform wind components from spherical to cartesian coordinates
+    # https://en.wikipedia.org/wiki/Vector_fields_in_cylindrical_and_spherical_coordinates
+    u_t = np.sin(ph) * np.cos(th) * w + np.cos(ph) * np.cos(th) * v - np.sin(th) * u
+    v_t = np.sin(ph) * np.sin(th) * w + np.cos(ph) * np.sin(th) * v + np.cos(th) * u
+    w_t = np.cos(ph) * w - np.sin(ph) * v
+
+    return u_t, v_t, w_t

tests/test_utilities.py

@@ -127,3 +127,42 @@ def test_lines_from_points():
     cells = poly.lines
     assert np.allclose(cells[0], [2, 0,1])
     assert np.allclose(cells[1], [2, 1,2])
+
+
+def test_grid_from_sph_coords():
+    x = np.arange(0.0, 360.0, 40.0)  # longitude
+    y = np.arange(0.0, 181.0, 60.0)  # colatitude
+    z = [1]  # elevation (radius)
+    g = pyvista.grid_from_sph_coords(x, y, z)
+    assert g.n_cells == 24
+    assert g.n_points == 36
+    assert np.allclose(
+        g.bounds,
+        [
+            -0.8137976813493738,
+            0.8660254037844387,
+            -0.8528685319524434,
+            0.8528685319524433,
+            -1.0,
+            1.0,
+        ],
+    )
+    assert np.allclose(g.points[1], [0.8660254037844386, 0.0, 0.5])
+    z = np.linspace(10, 30, 3)
+    g = pyvista.grid_from_sph_coords(x, y, z)
+    assert g.n_cells == 48
+    assert g.n_points == 108
+    assert np.allclose(g.points[0], [0.0, 0.0, 10.0])
+
+
+def test_transform_vectors_sph_to_cart():
+    lon = np.arange(0.0, 360.0, 40.0)  # longitude
+    lat = np.arange(0.0, 181.0, 60.0)  # colatitude
+    lev = [1]  # elevation (radius)
+    u, v = np.meshgrid(lon, lat, indexing="ij")
+    w = u ** 2 - v ** 2
+    uu, vv, ww = pyvista.transform_vectors_sph_to_cart(lon, lat, lev, u, v, w)
+    assert np.allclose(
+        [uu[-1, -1], vv[-1, -1], ww[-1, -1]],
+        [67.80403533828323, 360.8359915416445, -70000.0],
+    )