Improve performance of converting `xtgeo.RegularSurface` (→ `np.array`?) → RGBA image

Description

In webviz we use the DeckGLMap component to render surface grids as images. This component is a wrapper on deck.gl including e.g. additional layers targeting subsurface data, such as surface grids and well data.

To render surface grids we use a hillshading layer that is based on the DeckGL bitmap layer.

We send the surface grid to the hillshading layer as a bitmap image which is generated on-the-fly in the python backend.

Our current workflow for converting surface grids to images are:

1. Read surfaces stored on RMS Binary format using Xtgeo.
2. Extract surface values as numpy arrays and scale to RGB range
3. Encode the numpy array to RGBA png
4. Send png as base64 data from the python backend to the React frontend

)

Challenge

The execution time of converting a surface grid to base64 image is currently too slow. We should investigate how to improve the workflow by refactoring the code or looking at other options of generating these images from surface grids.

Example

Below is a working example that will start a Dash app with a map visualization using the workflow above.

import requests
import io
import base64
import numpy as np
from PIL import Image
import xtgeo
import dash

import webviz_subsurface_components as wsc


def xtgeo_surface_to_array(surface: xtgeo.RegularSurface) -> np.array:
    """Converts an xtgeo.RegularSurface to a numpy 2darray.
    Removes any rotation and flip"""
    # Unrotate the grid
    surface.unrotate()
    # Set masked values to undefined
    surface.fill(np.nan)
    # Flip order of values
    values = np.flip(surface.values.transpose(), axis=0)
    return values


def array2d_to_png(z_array):
    """The DeckGL map dash component takes in pictures as base64 data
    (or as a link to an existing hosted image). I.e. for containers wanting
    to create pictures on-the-fly from numpy arrays, they have to be converted
    to base64. This is an example function of how that can be done.

    This function encodes the numpy array to a RGBA png.
    The array is encoded as a heightmap, in a format similar to Mapbox TerrainRGB
    (https://docs.mapbox.com/help/troubleshooting/access-elevation-data/),
    but without the -10000 offset and the 0.1 scale.
    The undefined values are set as having alpha = 0. The height values are
    shifted to start from 0.
    """

    # Shift the values to start from 0 and scale them to cover
    # the whole RGB range for increased precision.
    scale_factor = (256 * 256 * 256 - 1) / (np.nanmax(z_array) - np.nanmin(z_array))
    z_array = (z_array - np.nanmin(z_array)) * scale_factor

    shape = z_array.shape
    z_array = np.repeat(z_array, 4)  # This will flatten the array
    z_array[0::4][np.isnan(z_array[0::4])] = 0  # Red
    z_array[1::4][np.isnan(z_array[1::4])] = 0  # Green
    z_array[2::4][np.isnan(z_array[2::4])] = 0  # Blue
    z_array[0::4] = np.floor((z_array[0::4] / (256 * 256)) % 256)  # Red
    z_array[1::4] = np.floor((z_array[1::4] / 256) % 256)  # Green
    z_array[2::4] = np.floor(z_array[2::4] % 256)  # Blue
    z_array[3::4] = np.where(np.isnan(z_array[3::4]), 0, 255)  # Alpha

    # Back to 2d shape + 1 dimension for the rgba values.

    z_array = z_array.reshape((shape[0], shape[1], 4))
    image = Image.fromarray(np.uint8(z_array), "RGBA")
    byte_io = io.BytesIO()
    image.save(byte_io, format="png")
    byte_io.seek(0)
    base64_data = base64.b64encode(byte_io.read()).decode("ascii")
    return f"data:image/png;base64,{base64_data}"


if __name__ == "__main__":

    # Download an example surface grid and load as xtgeo.RegularSurface
    # Note that this surface is small and is quick to convert.
    # A larger grid can easily be generated by refining the downloaded surface:
    # surface.refine(8)
    # surface.to_file("refined.gri")

    surface_stream = requests.get(
        "https://raw.githubusercontent.com/equinor/webviz-subsurface-testdata/master/01_drogon_ahm/realization-0/iter-0/share/results/maps/topvolon--ds_extract_postprocess.gri",
        stream=True,
    )
    surface = xtgeo.surface_from_file(io.BytesIO(surface_stream.content))
    values = xtgeo_surface_to_array(surface)
    min_value = np.nanmin(values)
    max_value = np.nanmax(values)
    map_image = array2d_to_png(values)
    map_bounds = [surface.xmin, surface.ymin, surface.xmax, surface.ymax]
    map_width = surface.xmax - surface.xmin
    map_height = surface.ymax - surface.ymin
    map_target = [surface.xmin + map_width / 2, surface.ymin + map_height / 2, 0]

    app = dash.Dash(__name__)
    app.layout = dash.html.Div(
        [
            wsc.DeckGLMap(
                id="map",
                coords={"visible": True, "multiPicking": True, "pickDepth": 10},
                scale={"visible": True},
                coordinateUnit="m",
                bounds=map_bounds,
                layers=[
                    {
                        "@@type": "ColormapLayer",
                        "colormap": "https://cdn.jsdelivr.net/gh/kylebarron/deck.gl-raster/assets/colormaps/viridis_r.png",
                        "valueRange": [min_value, max_value],
                        "image": map_image,
                        "bounds": map_bounds,
                    },
                    {
                        "@@type": "Hillshading2DLayer",
                        "bounds": map_bounds,
                        "valueRange": [min_value, max_value],
                        "image": map_image,
                    },
                ],
                editedData={
                    "data": {"type": "FeatureCollection", "features": []},
                },
            )
        ]
    )

    app.run_server(debug=True)