Add methods for visualizing tesseroids and tesseroid layers

[mdtanker]

While working on #617 I tried to match all functionality between the prisms and tesseroids. One of the remaining differences between the two is the visualization function for prisms, which doesn't exist for tesseroids.

I think they are even more important for tesseroids, since they are conceptually more complex than prisms, so being able to plot an individual tesseroid, or a tesseroid layer, would be a useful feature.

Below are a few possible starting points for functions:

Plot a single tesseroid:

`plot_tesseroid` function

def plot_tesseroid(tesseroid, plot=True, add_sphere=True, add_edges=True):
    """Plot a tesseroid using PyVista."""
    lon_min, lon_max, lat_min, lat_max, r_inner, r_outer = tesseroid

    # phi (latitude) is defined as 0 at north pole and 180 at south pole in PyVista
    def lat_to_phi(lat):
        return 90 - lat

    min_phi = lat_to_phi(lat_max)
    max_phi = lat_to_phi(lat_min)

    # Create a solid sphere sector (tesseroid)
    tesseroid_mesh = pv.SolidSphere(
        outer_radius=r_outer,
        inner_radius=r_inner,
        start_theta=lon_min,  # longitude
        end_theta=lon_max,  # longitude
        start_phi=min_phi,  # latitude
        end_phi=max_phi,  # latitude
        # phi_resolution=2,
        # theta_resolution=2,
        # radius_resolution=2,
    )

    if plot:
        # Plot the tesseroid
        plotter = pv.Plotter()

        if add_sphere:
            # add wireframe sphere if entire Earth/planet for context
            sphere_mesh = pv.Sphere(
                radius=r_outer,
                # phi_resolution=2, # latitude resolution
                # theta_resolution=2, # longitude resolution
                # radius_resolution=2, # radial resolution
            )
            plotter.add_mesh(
                sphere_mesh,
                color="white",
                style="wireframe",
                opacity=0.1,
            )

        plotter.add_mesh(
            tesseroid_mesh,
            color="lightblue",
            show_edges=False,
            opacity=1,
        )

        if add_edges:
            # add wire mesh to show tesseroid edges
            mesh = pv.SolidSphere(
                outer_radius=r_outer,
                inner_radius=r_inner,
                start_theta=lon_min,  # longitude
                end_theta=lon_max,  # longitude
                start_phi=min_phi,  # latitude
                end_phi=max_phi,  # latitude
                phi_resolution=2,
                theta_resolution=2,
                radius_resolution=2,
            )
            # edges = mesh.extract_all_edges()
            edges = mesh.extract_feature_edges()
            plotter.add_mesh(
                edges,
                line_width=10,
                color="black",
            )

        plotter.add_axes()

        _mean_phi, mean_theta = (min_phi + max_phi) / 2, (lon_min + lon_max) / 2

        # set camera to look at equator from prime meridian
        cam_pos = pv.grid_from_sph_coords(
            # theta = [0], # longitude
            phi=[90],  # latitude
            theta=[mean_theta + 20],  # longitude
            # phi = [mean_phi+20], # latitude
            r=[r_outer * 10],  # distance from the origin
        )

        plotter.set_position(cam_pos.points)
        plotter.set_viewup((0, 0, 1))
        plotter.show()

    return tesseroid_mesh

Which is used like this:

ellipsoid = bl.WGS84
mean_radius = ellipsoid.mean_radius
tesseroid = (-70, -50, -40, -20, mean_radius - 3000e3, mean_radius)

plot_tesseroid(tesseroid, add_sphere=False)

Plot multiple tesseroids:

`plot_tesseroid` function

def plot_tesseroids(tesseroids):
    """Plot tesseroids using PyVista."""

    # Plot the tesseroid
    plotter = pv.Plotter()

    # plot tesseroids
    for tesseroid in tesseroids:
        lon_min, lon_max, lat_min, lat_max, r_inner, r_outer = tesseroid

        # phi (latitude) is defined as 0 at north pole and 180 at south pole in PyVista
        def lat_to_phi(lat):
            return 90 - lat

        min_phi = lat_to_phi(lat_max)
        max_phi = lat_to_phi(lat_min)

        # Create a solid sphere sector (tesseroid)
        tesseroid_mesh = pv.SolidSphere(
            outer_radius=r_outer,
            inner_radius=r_inner,
            start_theta=lon_min,  # longitude
            end_theta=lon_max,  # longitude
            start_phi=min_phi,  # latitude
            end_phi=max_phi,  # latitude
            # phi_resolution=2,
            # theta_resolution=2,
            # radius_resolution=2,
        )

        plotter.add_mesh(
            tesseroid_mesh,
            color="lightblue",
            show_edges=False,
            opacity=1,
        )

        # # add wire mesh to show tesseroid edges
        # mesh = pv.SolidSphere(
        #     outer_radius=r_outer,
        #     inner_radius=r_inner,
        #     start_theta=lon_min,  # longitude
        #     end_theta=lon_max,    # longitude
        #     start_phi=min_phi,    # latitude
        #     end_phi=max_phi,      # latitude
        #     phi_resolution=2,
        #     theta_resolution=2,
        #     radius_resolution=2,
        # )
        # edges = mesh.extract_all_edges()
        # plotter.add_mesh(
        #     edges,
        #     line_width=2,
        #     color="black",
        # )

    plotter.add_axes()

    # mean_phi, mean_theta = (min_phi + max_phi)/2, (lon_min + lon_max)/2

    # # set camera to look at equator from prime meridian
    # cam_pos = pv.grid_from_sph_coords(
    #     # theta = [0], # longitude
    #     phi = [90], # latitude
    #     theta = [mean_theta+20], # longitude
    #     # phi = [mean_phi+20], # latitude
    #     r = [r_outer*10], # distance from the origin
    # )

    # plotter.set_position(cam_pos.points)
    # plotter.set_viewup((0,0,1))
    plotter.show()

Which is used like this:

tesseroids = [
    [-70, -65, -40, -35, mean_radius - 100e3, mean_radius],
    [-55, -50, -40, -35, mean_radius - 100e3, mean_radius],
    [-70, -65, -25, -20, mean_radius - 100e3, mean_radius],
    [-55, -50, -25, -20, mean_radius - 100e3, mean_radius],
]
plot_tesseroids(tesseroids)

or like this:

# Create a synthetic relief
longitude = np.linspace(0, 10, 5)
latitude = np.linspace(2, 8, 4)
mesh_longitude, mesh_latitude = np.meshgrid(longitude, latitude)
surface_heights = np.arange(20, dtype=float).reshape((4, 5))
# Convert heights to radii
reference_radii = bl.WGS84.geocentric_radius(mesh_latitude)
surface_radii = surface_heights + reference_radii
# Define constant densities
density = 2670.0 * np.ones_like(surface_radii)
# Define a layer of tesseroid
tesseroids = hm.tesseroid_layer(
    (longitude, latitude),
    surface_radii,
    reference=reference_radii,
    properties={"density": density},
)

# extract tesseroids from the layer
tesseroids = tesseroid_layer.tesseroid_layer._to_tesseroids()

plot_tesseroids(tesseroids)

Plot layer of tesseroids:

`plot_tesseroid` function

def plot_tesseroid_layer(tesseroid_layer, add_sphere=True):
    """Plot a tesseroid using PyVista. This plots 1 large tesseroid, then adds some
    grid lines to denote each separate tesseroid.
    """

    # extract tesseroids from the layer
    tesseroids = tesseroid_layer.tesseroid_layer._to_tesseroids()
    density = tesseroid_layer.tesseroid_layer._obj.density.values

    # discard thin tesseroids
    mask = tesseroid_layer.tesseroid_layer._get_nonans_mask(property_name="density")
    density = density[mask]
    # tesseroids = _discard_thin_tesseroids(
    #     tesseroids,
    #     density,
    #     1,
    # )[0]

    # get arrays of boundaries
    lon_mins = [i[0] for i in tesseroids]
    lon_maxs = [i[1] for i in tesseroids]
    lat_mins = [i[2] for i in tesseroids]
    lat_maxs = [i[3] for i in tesseroids]
    r_inners = [i[4] for i in tesseroids]
    r_outers = [i[5] for i in tesseroids]

    # # transform latitudes to be centered at north pole
    # mean_lat = np.mean(np.array([lat_mins, lat_maxs]))
    # print(mean_lat)
    # lat_mins = [90-np.abs(x-mean_lat) for x in lat_mins]
    # lat_maxs = [90-np.abs(x-mean_lat) for x in lat_maxs]
    # print(min(lat_mins), max(lat_maxs))

    # phi (latitude) is defined as 0 at north pole and 180 at south pole in PyVista
    min_phis = [90 - x for x in lat_maxs]
    max_phis = [90 - x for x in lat_mins]

    # transform latitudes to be centered at north pole
    # mean_phi = np.mean(np.array([min_phis, max_phis]))
    # min_phis = [np.abs(x-mean_phi) for x in min_phis]
    # max_phis = [np.abs(x-mean_phi) for x in max_phis]
    # print(min(min_phis), max(max_phis))

    # # transform longitudes to be centered at 0deg
    # mean_lon = np.mean(np.array([lon_mins, lon_maxs]))
    # lon_mins = [x-(mean_lon) for x in lon_mins]
    # lon_maxs = [x-(mean_lon) for x in lon_maxs]

    mean_theta = np.mean(np.array([lon_mins, lon_maxs]))
    mean_phi = np.mean(np.array([min_phis, max_phis]))

    plotter = pv.Plotter()

    if add_sphere:
        # Create and plot the semi-transparent shell of a sphere
        sphere_mesh = pv.Sphere(radius=np.mean([r_outers, r_inners]))
        plotter.add_mesh(sphere_mesh, color="white", style="wireframe", opacity=0.1)

    # loop through each tesseroid, generate its mesh, and append to a list
    meshes = []
    edges = []
    for i, _j in enumerate(tesseroids):
        # Create a solid sphere sector (tesseroid)
        tesseroid_mesh = pv.SolidSphere(
            outer_radius=r_outers[i],
            inner_radius=r_inners[i],
            start_theta=lon_mins[i],  # longitude
            end_theta=lon_maxs[i],  # longitude
            start_phi=min_phis[i],  # latitude
            end_phi=max_phis[i],  # latitude
            phi_resolution=5,  # resolutions of 2 will be flat faces
            theta_resolution=5,
            radius_resolution=5,
        )
        plotter.add_mesh(
            tesseroid_mesh,
            # color="lightblue",
            scalars=[density[i]] * tesseroid_mesh.n_cells,
            cmap="viridis",
            show_edges=False,
            opacity=1,
        )

        meshes.append(tesseroid_mesh)
        tesseroid_edge = pv.SolidSphere(
            outer_radius=r_outers[i],
            inner_radius=r_inners[i],
            start_theta=lon_mins[i],  # longitude
            end_theta=lon_maxs[i],  # longitude
            start_phi=min_phis[i],  # latitude
            end_phi=max_phis[i],  # latitude
            phi_resolution=2,
            theta_resolution=2,
            radius_resolution=2,
        )
        edge = tesseroid_edge.extract_feature_edges()
        edges.append(edge)

    transform = pv.Transform()
    transform.rotate_x(-mean_phi)  # rotate so latitude is at north pole
    # transform.rotate_z(180-mean_theta) # rotate so longitude is at prime meridian (180)

    # # plot merged mesh of all tesseroids
    # mesh = pv.merge(meshes)
    # # mesh = mesh.transform(transform, inplace=False)
    # print(mesh)
    # plotter.add_mesh(
    #     mesh,
    #     color="lightblue",
    #     cmap="viridis",
    #     show_edges=False,
    #     opacity=1,
    # )

    # plot merged edges of all tesseroids
    edges = pv.merge(edges)
    # edges = edges.transform(transform, inplace=False)
    plotter.add_mesh(edges, line_width=2, color="black")

    # r = r_outer_max*10
    # phi = 90
    # theta = np.mean(np.array([lon_min, lon_max]))
    # print(r, phi, theta)
    # print(pv.spherical_to_cartesian(r, phi, theta))

    # set camera to look at equator from prime meridian

    # cam_pos = pv.grid_from_sph_coords(
    #     phi = 0, # latitude
    #     theta = 90, # longitude
    #     # phi = [mean_phi], # latitude
    #     r = [np.max(r_outers)*3], # distance from the origin
    # )
    # plotter.set_position(cam_pos.points)

    # plotter.camera_position = [
    #     (0.0, 0.0, 1.0),
    #     (0.0, 0.0, 0.0),
    #     (0.0, 1.0, 0.0)
    # ]
    # plotter.camera_position = 'yz' # looking directly at prime meridian and equator
    mean_radius = np.mean([r_outers, r_inners])

    plotter.camera_position = "yz"  # looking directly at prime meridian and equator
    # plotter.camera.distance=12e6
    plotter.camera.distance = mean_radius * 1.5
    plotter.camera.clipping_range = (mean_radius * 0.5, mean_radius * 3)
    plotter.camera.azimuth = mean_theta - 90
    plotter.camera.roll = mean_phi + 90
    plotter.camera.elevation = 10

    plotter.add_axes()
    # plotter.show_grid()
    # plotter.show()

    return plotter

Which is used like this:

plotter = plot_tesseroid_layer(
    tesseroid_layer,
    add_sphere=False,
)
plotter.show()

[mdtanker]

Here is a notebooks with the above functions and figures:
test.ipynb

[mdtanker]

Other links which may be helpful:
https://s3dlib.org/docs/clss_mod/surface.sph.html
https://community.plotly.com/t/spherical-wedges-with-volume-plot/78760/5

[santisoler]

Woow! Yes, I've been wanting to add this for a long time! Do you want to start a PR for it? I think we could mimic the API for prisms, i.e. having hm.visualizations.tesseroid_to_pyvista and DataAccessorTesseroidLayer.to_pyvista. What do you think?

[mdtanker]

Nice, glad you're on board! I can open a PR in a few weeks, but if you want to get started go ahead and I can contribute when I have time! And yeah I was thinking following the API for prisms would be good.