[ENH] Converting PyVista Mesh to Raster

[AlexanderJuestel]

Is your feature request related to a problem? Please describe.
GemPy/PyVista Meshes should be converted to rasters to be read into GIS software

Describe the solution you'd like
We need a function that reads in the mesh and that returns a raster

Describe alternatives you've considered
Converting the triangles of the mesh to polygons and saving them as shape file

Additional context
I am currently working on a solution using barycentric coordinates.

[AlexanderJuestel]

PyVista Mesh

GeoDataFrame with Raster Cells

Numpy Array which can be saved as tiff (raster), conversion from gdf to raster not shown yet

import shapely
import pyvista as pv
import geopandas as gpd
from shapely.geometry import Polygon
import numpy as np
import matplotlib.pyplot as plt


def convert_mesh_to_raster(mesh: pv.core.pointset.UnstructuredGrid,
                           crs: str,
                           resolution: int):
    """Convert PyVista Mesh to Raster.
    
    """

    # Extracting faces from mesh
    faces = mesh.cells.reshape(int(len(mesh_clipped.cells)/4), 4)[:,1:]

    # Creating polygons (with x,y,z vertices) from faces
    polygons = [Polygon(mesh.points[face][:,:3]) for face in faces]

    # Extracting mean z value
    z = [np.mean(mesh.points[face][:,2]) for face in faces]

    # Creating GeoDataFrame from polygons
    gdf = gpd.GeoDataFrame(geometry=polygons, crs=crs)

    # Appending z value to GeoDataFrame
    gdf['Z'] = z

    # Creating bounding box from GeoDataFrame
    bbox = shapely.box(*gdf.total_bounds)

    # Creating raster mask
    mask = create_polygon_mask(gdf=bbox, 
                               step_size=resolution, 
                               crs=crs)

    # Creating centroids from raster mask
    centroids = gpd.GeoDataFrame(geometry=mask.centroid, crs=crs)
    
    # Spatially joining polygons with centroids
    gdf_centroids_sjoined = gdf.sjoin(centroids).drop_duplicates(subset='index_right').reset_index(drop=True)

    # Selecting points based on index
    points = [centroids.iloc[i].geometry for i in gdf_centroids_sjoined['index_right']]

    # Appending points to joined GeoDataFrame
    gdf_centroids_sjoined['centroids'] = points

    # Extracting z coordinates of centroids on plane
    z = [point_on_triangle_plane(list(gdf_centroids_sjoined.iloc[i]['centroids'].coords)[0], 
                   list(set(list(gdf_centroids_sjoined.iloc[i].geometry.exterior.coords)))) for i in range(len(gdf_centroids_sjoined))]

    # Calculating weights 
    l_values_list = [barycentric_coords_triangle(np.append(np.array(list(gdf_centroids_sjoined.iloc[i]['centroids'].coords)), z[i]), 
                   list(set(list(gdf_centroids_sjoined.iloc[i].geometry.exterior.coords)))) for i in range(len(z))]

    # Extracting z values of vertices
    z_vertices_list = [[list(set(list(gdf_centroids_sjoined.iloc[j].geometry.exterior.coords)))[i][2] for i in range(len(list(set(list(gdf_centroids_sjoined.iloc[j].geometry.exterior.coords)))))] for j in range(len(gdf_centroids_sjoined))]

    # Calculating weighted z values 
    z_raster_values = [np.sum(z_vertices*l_values) for z_vertices, l_values in zip(z_vertices_list, l_values_list)]

    # Replacing polygon geometry with centroids
    gdf_centroids_sjoined['geometry'] = gdf_centroids_sjoined['centroids']

    # Dropping centroids column
    gdf_centroids_sjoined = gdf_centroids_sjoined.drop('centroids', axis=1)

    # Appending raster z values
    gdf_centroids_sjoined['Z'] = z_raster_values

    # Sorting values
    gdf_centroids_sjoined = gdf_centroids_sjoined.sort_values(by='index_right', ascending=True).reset_index(drop=True)

    # Joining mask with centroids and raster z values
    mask_joined = mask.join(gdf_centroids_sjoined, rsuffix='_centroids')

    # Dropping columns
    mask_joined = mask_joined.drop(['geometry_centroids', 'index_right'], axis=1)

    return mask_joined
    
def point_on_triangle_plane(p, coords):
    """
    Given x,y coordinates, compute z on the plane defined by triangle (v0,v1,v2).
    """
    # Normal vector of the triangle plane
    v0, v1, v2 = np.array(coords)
    x, y = p
    u = v1 - v0
    v = v2 - v0
    n = np.cross(u, v)

    if abs(n[2]) < 1e-14:
        raise ValueError("Triangle is vertical in XY-projection (n_z≈0).")

    x0, y0, z0 = v0
    # Solve plane equation for z
    z = z0 - (n[0]*(x - x0) + n[1]*(y - y0)) / n[2]
    return z

def barycentric_coords_triangle(p, coords):
    """
    Compute barycentric coordinates (λ0, λ1, λ2) of point p w.r.t triangle (v0,v1,v2).
    """
    v0, v1, v2 = np.array(coords)
    u = v1 - v0
    v = v2 - v0
    w = p - v0

    d_uu = np.dot(u, u)
    d_uv = np.dot(u, v)
    d_vv = np.dot(v, v)
    d_wu = np.dot(w, u)
    d_wv = np.dot(w, v)

    D = d_uu * d_vv - d_uv * d_uv
    if abs(D) < 1e-14:
        raise ValueError("Degenerate triangle")

    l1 = (d_vv * d_wu - d_uv * d_wv) / D
    l2 = (d_uu * d_wv - d_uv * d_wu) / D
    l0 = 1.0 - l1 - l2
    return np.array([l0, l1, l2])

from typing import Union
import pyproj
from itertools import product

def create_polygon_mask(
    gdf: Union[gpd.GeoDataFrame, Polygon],
    step_size: int,
    crop_gdf: bool = False,
    crs: Union[str, pyproj.crs.crs.CRS] = None,
) -> gpd.GeoDataFrame:
    
    # Converting Polygon to GeoDataFrame
    if isinstance(gdf, Polygon):
        if not isinstance(crs, (str, pyproj.crs.crs.CRS)):
            raise TypeError("CRS must be provided as string or pyproj CRS object")
        gdf = gpd.GeoDataFrame(geometry=[gdf], crs=crs)

    # Checking that the gdf is of type GeoDataFrame
    if not isinstance(gdf, gpd.GeoDataFrame):
        raise TypeError("gdf must be provided as GeoDataFrame")

    # Checking that the step_size is of type int
    if not isinstance(step_size, int):
        raise TypeError("The step_size must be provided as int")

    # Checking that crop_gdf is of type bool
    if not isinstance(crop_gdf, bool):
        raise TypeError("crop_gdf must be provided as bool")

    # Creating arrays
    x = np.arange(gdf.total_bounds[0], gdf.total_bounds[2], step_size)

    y = np.arange(gdf.total_bounds[1], gdf.total_bounds[3], step_size)

    # Creating polygons
    polygons = [
        Polygon(
            [
                (a, b),
                (a + step_size, b),
                (a + step_size, b + step_size),
                (a, b + step_size),
            ]
        )
        for a, b in product(x, y)
    ]

    # Converting polygons to GeoDataFrame
    gdf_mask = gpd.GeoDataFrame(geometry=polygons, crs=gdf.crs)

    # Dropping duplicate cells
    gdf_mask = gdf_mask.drop_duplicates(ignore_index=True)

    # Cropping the gdf if crop_gdf is True
    if crop_gdf:
        gdf_mask = gdf_mask.sjoin(gdf).reset_index()[["geometry"]]

    return gdf_mask

[AlexanderJuestel]

Converting GeoDataframe (raster) to numpy array

raster = convert_mesh_to_raster(mesh=mesh_clipped,
                           crs='EPSG:4647',
                           resolution=100)

width = raster.total_bounds[2]-raster.total_bounds[0]
height = raster.total_bounds[3]-raster.total_bounds[1]


resolution = 100 #input from above
width_pixels = width/resolution 
height_pixels = height/resolution

import numpy as np
from rasterio import features
from affine import Affine


xmin, ymin, xmax, ymax = raster.total_bounds  # extent of all polygons

# Define affine transform (maps pixel coordinates -> spatial coordinates)
xres = (xmax - xmin) / width_pixels
yres = (ymax - ymin) / height_pixels
transform = Affine.translation(xmin, ymax) * Affine.scale(xres, -yres)

# Prepare shapes iterable: (geometry, value)
shapes = [(geom, val) for geom, val in zip(raster.geometry, raster['Z'])]

# Rasterize
raster_out = features.rasterize(
    shapes=shapes,
    out_shape=(int(height_pixels), int(width_pixels)),
    transform=transform,
    fill=0,            # background value
    dtype="float32"    # or "int16"
)

plt.imshow(raster_out, extent=[raster.total_bounds[0], 
                               raster.total_bounds[2], 
                               raster.total_bounds[1], 
                               raster.total_bounds[3]] )

[javoha]

Hi @AlexanderJuestel,
thanks for the idea and the code. As we are trying to keep GemPy lean and slim and not blow it up to much with convenience functions I think this would be good content for a GemPy plugin (which are now hosted in their own repo within the GemPy framework). Would this not also fit the functionality of GemGIS?
Cheers,
Jan

[AlexanderJuestel]

Yep it is going to be implemented in GemGIS, just as a reference for future GemPy users