Superimposed display of 3D data (discrete)

# sphinx_gallery_thumbnail_path = '_static/cigvis/more_demos/020.png'

import numpy as np
import cigvis
from cigvis import colormap
from pathlib import Path

root = Path(__file__).resolve().parent.parent.parent


def show(bg, fg):
    """
    Overlay seismic data and facies labels with alpha=0.5.
    """
    fg[fg == 1] = 100
    values = np.unique(fg)
    colors = ['red', 'green', 'yellow', 'blue', (0, 0.5, 0.5)]
    cmap = colormap.custom_disc_cmap(values, colors)
    cmap = colormap.set_alpha(cmap, 0.5)  # Note:

    nodes1 = cigvis.create_slices(bg)
    nodes1 = cigvis.add_mask(nodes1, fg, cmap=cmap, interpolation='nearest')
    nodes1 += cigvis.create_colorbar_from_nodes(
        nodes1,
        label_str='Facies',
        select='mask',
        discrete=True,
        disc_ticks=[values.astype(int)],
    )
    """
    Mask the smallest class, but keep it visible in the colorbar (white).
    """
    values = np.unique(fg)
    colors = ['red', 'green', 'yellow', 'blue', (0, 0.5, 0.5)]
    cmap = colormap.custom_disc_cmap(values, colors)
    cmap = colormap.set_alpha_except_min(cmap, 0.5)  # Note:

    nodes2 = cigvis.create_slices(bg)
    nodes2 = cigvis.add_mask(nodes2, fg, cmap=cmap, interpolation='nearest')
    nodes2 += cigvis.create_colorbar_from_nodes(
        nodes2,
        label_str='Facies',
        select='mask',
        discrete=True,
        disc_ticks=[values.astype(int)],
    )
    """
    Mask the smallest class and remove it from the colorbar.
    """
    values = np.unique(fg)
    colors = ['red', 'green', 'yellow', 'blue', (0, 0.5, 0.5)]
    cmap = colormap.custom_disc_cmap(values, colors)
    cmap = colormap.set_alpha_except_min(cmap, 0.5)  # Note:

    nodes3 = cigvis.create_slices(bg)
    nodes3 = cigvis.add_mask(nodes3, fg, cmap=cmap, interpolation='nearest')

    values = values[1:]

    nodes3 += cigvis.create_colorbar_from_nodes(
        nodes3,
        label_str='Facies',
        select='mask',
        discrete=True,
        disc_ticks=[values.astype(int)],
    )
    """
    Mask the largest class.
    """
    values = np.unique(fg)
    colors = ['red', 'green', 'yellow', 'blue', (0, 0.5, 0.5)]
    cmap = colormap.custom_disc_cmap(values, colors)
    cmap = colormap.set_alpha_except_max(cmap, 0.5)  # Note:

    nodes4 = cigvis.create_slices(bg)
    nodes4 = cigvis.add_mask(nodes4, fg, cmap=cmap, interpolation='nearest')

    values = values[:-1]
    nodes4 += cigvis.create_colorbar_from_nodes(
        nodes4,
        label_str='Facies',
        select='mask',
        discrete=True,
        disc_ticks=[values.astype(int)],
    )
    """
    Mask selected values, not necessarily the minimum or maximum; multiple values are allowed.
    """
    values = np.unique(fg)
    colors = ['red', 'green', 'yellow', 'blue', (0, 0.5, 0.5)]
    cmap = colormap.custom_disc_cmap(values, colors)
    cmap = colormap.set_alpha_except_values(cmap,
                                            0.5,
                                            clim=[fg.min(), fg.max()],
                                            values=[0, 100])

    nodes5 = cigvis.create_slices(bg)
    nodes5 = cigvis.add_mask(nodes5, fg, cmap=cmap, interpolation='nearest')

    values = values[values != 0]
    values = values[values != 100]
    nodes5 += cigvis.create_colorbar_from_nodes(
        nodes5,
        label_str='Facies',
        select='mask',
        discrete=True,
        disc_ticks=[values],
    )

    cigvis.plot3D([nodes1, nodes2, nodes3, nodes4, nodes5],
                  view=cigvis.Plot3DView(
                      grid=(2, 3),
                      size=(1300, 800),
                      cbar_region_ratio=0.24,
                  ),
                  save=cigvis.Plot3DSave(path='example.png'))


if __name__ == '__main__':
    sxp = root / 'data/seis_h360x600x400.dat'
    lxp = root / 'data/label_h360x600x400.dat'
    ni, nx, nt = 400, 600, 360

    sx = np.fromfile(sxp, np.float32).reshape(ni, nx, nt)
    lx = np.fromfile(lxp, np.float32).reshape(ni, nx, nt)

    show(sx, lx)