# Colormap Choices
# ================

# Use a Matplotlib, Colorcet, cmocean, or custom colormap when plotting scalar
# values with `pyvista.plot()` and `Plotter` methods.
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.colors import ListedColormap
import matplotlib.pyplot as plt
import numpy as np
import pyvista as pv
from pyvista import examples

# Any colormap built for `matplotlib`, `colorcet`, or `cmocean` is fully
# compatible with PyVista. Colormaps are typically specified by passing the
# string name of the colormap to the plotting routine via the `cmap`
# argument.
# See Matplotlib's complete list of available colormaps,
# Colorcet's complete list, and cmocean's complete list.

# Custom Made Colormaps
# ---------------------

# To get started using a custom colormap, download some data with scalar values to
# plot.
mesh = examples.download_st_helens().warp_by_scalar()
# Add scalar array with range (0, 100) that correlates with elevation
mesh['values'] = pv.plotting.tools.normalize(mesh['Elevation']) * 100

# Build a custom colormap - here we make a colormap with 5 discrete colors
# and we specify the ranges where those colors fall:
# Define the colors we want to use
blue = np.array([12 / 256, 238 / 256, 246 / 256, 1.0])
black = np.array([11 / 256, 11 / 256, 11 / 256, 1.0])
grey = np.array([189 / 256, 189 / 256, 189 / 256, 1.0])
yellow = np.array([255 / 256, 247 / 256, 0 / 256, 1.0])
red = np.array([1.0, 0.0, 0.0, 1.0])

mapping = np.linspace(mesh['values'].min(), mesh['values'].max(), 256)
newcolors = np.empty((256, 4))
newcolors[mapping >= 80] = red
newcolors[mapping < 80] = grey
newcolors[mapping < 55] = yellow
newcolors[mapping < 30] = blue
newcolors[mapping < 1] = black

# Make the colormap from the listed colors
my_colormap = ListedColormap(newcolors)

# Simply pass the colormap to the plotting routine.
mesh.plot(scalars='values', cmap=my_colormap)

# Or you could make a simple colormap… any Matplotlib colormap can be passed
# to PyVista.
boring_cmap = plt.get_cmap('viridis', 5)
mesh.plot(scalars='values', cmap=boring_cmap)

# You can also pass a list of color strings to the color map.  This
# approach divides up the colormap into 5 equal parts.
mesh.plot(scalars=mesh['values'], cmap=['black', 'blue', 'yellow', 'grey', 'red'])

# If you still wish to have control of the separation of values, you
# can do this by creating a scalar array and passing that to the
# plotter along with the colormap
scalars = np.empty(mesh.n_points)
scalars[mesh['values'] >= 80] = 4  # red
scalars[mesh['values'] < 80] = 3  # grey
scalars[mesh['values'] < 55] = 2  # yellow
scalars[mesh['values'] < 30] = 1  # blue
scalars[mesh['values'] < 1] = 0  # black

mesh.plot(scalars=scalars, cmap=['black', 'blue', 'yellow', 'grey', 'red'])

# Matplotlib vs. Colorcet
# -----------------------

# Let's compare Colorcet's perceptually uniform "fire" colormap to Matplotlib's
# "hot" colormap much like the example on the first page of Colorcet's docs.
# The "hot" version washes out detail at the high end, as if the image is
# overexposed, while "fire" makes detail visible throughout the data range.
# Please note that in order to use Colorcet's colormaps including "fire", you
# must have Colorcet installed in your Python environment:
# `pip install colorcet`
pl = pv.Plotter(shape=(2, 2), border=False)
pl.subplot(0, 0)
pl.add_mesh(
    mesh,
    scalars='Elevation',
    cmap='fire',
    lighting=True,
    scalar_bar_args={'title': 'Colorcet Fire'},
)

pl.subplot(0, 1)
pl.add_mesh(
    mesh,
    scalars='Elevation',
    cmap='fire',
    lighting=False,
    scalar_bar_args={'title': 'Colorcet Fire (No Lighting)'},
)

pl.subplot(1, 0)
pl.add_mesh(
    mesh,
    scalars='Elevation',
    cmap='hot',
    lighting=True,
    scalar_bar_args={'title': 'Matplotlib Hot'},
)

pl.subplot(1, 1)
pl.add_mesh(
    mesh,
    scalars='Elevation',
    cmap='hot',
    lighting=False,
    scalar_bar_args={'title': 'Matplotlib Hot (No Lighting)'},
)

pl.show()

# Recreating Native Solver Colormaps (CalculiX)
# ---------------------------------------------

# Sometimes you want to process results in PyVista but keep the visual
# style identical to the solver's native post-processing tool.
# Below is an example of how to recreate the default `cgx` (CalculiX GraphiX)
# colormap using Matplotlib's `matplotlib.colors.LinearSegmentedColormap`.
# It may be particular useful when visualizing `.frd` files with
# `FRDReader`.
# This specific colormap uses exact RGB nodes at specific scalar fractions
# (0.0, 0.1, 0.3, 0.5, 0.8, 0.9, 1.0) to create its distinct color bands.
# Define the exact RGB nodes from the cgx source code
cgx_nodes = [
    (0.0, (0.50, 0.0, 1.0)),  # Purple
    (0.1, (0.00, 0.0, 0.8)),  # Dark Blue
    (0.3, (0.00, 1.0, 0.4)),  # Greenish Blue
    (0.5, (0.00, 0.5, 0.0)),  # Dark Green
    (0.8, (1.00, 1.0, 0.0)),  # Yellow
    (0.9, (1.00, 0.5, 0.0)),  # Orange
    (1.0, (0.75, 0.0, 0.0)),  # Dark Red
]

# Create the colormap
cgx_cmap = LinearSegmentedColormap.from_list('cgx', cgx_nodes)

# Download a sample mesh to demonstrate the colormap
mesh = examples.download_notch_stress()

# Plot the mesh using the custom cgx colormap
cpos = pv.CameraPosition(
    position=(0.1249, -0.1242, 0.1698),
    focal_point=(0.1965, 0.0467, -0.0036),
    viewup=(0.2512, 0.6356, 0.7299),
)

pl = pv.Plotter()
pl.add_mesh(
    mesh,
    cmap=cgx_cmap,
    n_colors=21,
    show_edges=True,
)
pl.camera_position = cpos
pl.show()

# ----------------------------------------------------------------------
# Generated by sphinx-examples-as-code https://github.com/pyvista/sphinx-examples-as-code