CurvatureBandsWithGlyphs.py

#!/usr/bin/env python

import math
from collections import OrderedDict

import numpy as np
import vtk
from vtk.util import numpy_support


def main():
    # ------------------------------------------------------------
    # Create the surface, lookup tables, contour filter etc.
    # ------------------------------------------------------------
    # desired_surface = 'ParametricTorus'
    # desired_surface = 'Plane'
    desired_surface = 'RandomHills'
    # desired_surface = 'Sphere'
    # desired_surface = 'Torus'
    surface = desired_surface.lower()
    available_surfaces = ['parametrictorus', 'plane', 'randomhills', 'sphere', 'torus']
    if surface not in available_surfaces:
        print('No surface specified.')
        return
    if surface == 'parametrictorus':
        src = make_parametric_torus()
    elif surface == 'plane':
        src = make_elevations(make_plane())
    elif surface == 'randomhills':
        src = make_parametric_hills()
    elif surface == 'sphere':
        src = make_elevations(make_sphere())
    elif surface == 'torus':
        src = make_elevations(make_torus())
    else:
        print('No surface specified.')
        return
    print(desired_surface)

    curvatures = ComputeCurvatures(src)
    curvatures.update()

    src.GetPointData().SetActiveScalars('Gauss_Curvature')
    scalar_range_curvatures = src.GetPointData().GetScalars('Gauss_Curvature').GetRange()
    scalar_range_elevation = src.GetPointData().GetScalars('Elevation').GetRange()

    lut = make_categorical_lut()
    lut1 = make_diverging_lut()
    lut.SetTableRange(scalar_range_curvatures)
    lut1.SetTableRange(scalar_range_elevation)
    number_of_bands = lut.GetNumberOfTableValues()
    bands = make_bands(scalar_range_curvatures, number_of_bands, False)
    if surface == 'randomhills':
        # These are my custom bands.
        # Generated by first running:
        # bands = make_bands(scalar_range_curvatures, number_of_bands, False)
        # then:
        #  freq = frequencies(bands, src)
        #  print_bands_frequencies(bands, freq)
        # Finally using the output to create this table:
        # my_bands = [
        #     [-0.630, -0.190], [-0.190, -0.043], [-0.043, -0.0136],
        #     [-0.0136, 0.0158], [0.0158, 0.0452], [0.0452, 0.0746],
        #     [0.0746, 0.104], [0.104, 0.251], [0.251, 1.131]]
        #  This demonstrates that the gaussian curvature of the surface
        #   is mostly planar with some hyperbolic regions (saddle points)
        #   and some spherical regions.
        my_bands = [
            [-0.630, -0.190], [-0.190, -0.043], [-0.043, 0.0452], [0.0452, 0.0746],
            [0.0746, 0.104], [0.104, 0.251], [0.251, 1.131]]
        # Comment this out if you want to see how allocating
        # equally spaced bands works.
        bands = make_custom_bands(scalar_range_curvatures, number_of_bands, my_bands)
        # bands = make_bands(scalar_range_curvatures, number_of_bands, False)
        # Adjust the number of table values
        lut.SetNumberOfTableValues(len(bands))

    # Let's do a frequency table.
    # The number of scalars in each band.
    freq = frequencies(bands, src)

    min_key = min(freq.keys())
    max_key = max(freq.keys())
    if surface == 'sphere':
        freq[0] = 0
    first, last = adjust_frequency_ranges(freq)
    for idx in range(min_key, first):
        freq.pop(idx)
        bands.pop(idx)
    for idx in range(last + 1, max_key + 1):
        freq.popitem()
        bands.popitem()
    old_keys = freq.keys()
    adj_freq = OrderedDict()
    adj_bands = OrderedDict()

    for idx, k in enumerate(old_keys):
        adj_freq[idx] = freq[k]
        adj_bands[idx] = bands[k]
    # print_bands_frequencies(bands, freq)
    print_bands_frequencies(adj_bands, adj_freq)

    min_key = min(adj_freq.keys())
    max_key = max(adj_freq.keys())
    scalar_range_curvatures = (adj_bands[min_key][0], adj_bands[max_key][2])
    lut.SetTableRange(scalar_range_curvatures)
    lut.SetNumberOfTableValues(len(adj_bands))

    # We will use the midpoint of the band as the label.
    labels = []
    for k in adj_bands:
        labels.append('{:4.2f}'.format(adj_bands[k][1]))

    # Annotate
    values = vtk.vtkVariantArray()
    for i in range(len(labels)):
        values.InsertNextValue(vtk.vtkVariant(labels[i]))
    for i in range(values.GetNumberOfTuples()):
        lut.SetAnnotation(i, values.GetValue(i).ToString())

    # Create a lookup table with the colors reversed.
    lutr = reverse_lut(lut)

    # Create the contour bands.
    bcf = vtk.vtkBandedPolyDataContourFilter()
    bcf.SetInputData(src)
    # Use either the minimum or maximum value for each band.
    for k in adj_bands:
        bcf.SetValue(k, adj_bands[k][2])
    # We will use an indexed lookup table.
    bcf.SetScalarModeToIndex()
    bcf.GenerateContourEdgesOn()

    # Generate the glyphs on the original surface.
    glyph = make_glyphs(src, False)

    # ------------------------------------------------------------
    # Create the mappers and actors
    # ------------------------------------------------------------

    colors = vtk.vtkNamedColors()

    # Set the background color.
    colors.SetColor('BkgColor', [179, 204, 255, 255])
    colors.SetColor("ParaViewBkg", [82, 87, 110, 255])

    src_mapper = vtk.vtkPolyDataMapper()
    src_mapper.SetInputConnection(bcf.GetOutputPort())
    src_mapper.SetScalarRange(scalar_range_curvatures)
    src_mapper.SetLookupTable(lut)
    src_mapper.SetScalarModeToUseCellData()

    src_actor = vtk.vtkActor()
    src_actor.SetMapper(src_mapper)

    # Create contour edges
    edge_mapper = vtk.vtkPolyDataMapper()
    edge_mapper.SetInputData(bcf.GetContourEdgesOutput())
    edge_mapper.SetResolveCoincidentTopologyToPolygonOffset()

    edge_actor = vtk.vtkActor()
    edge_actor.SetMapper(edge_mapper)
    edge_actor.GetProperty().SetColor(colors.GetColor3d('Black'))

    glyph_mapper = vtk.vtkPolyDataMapper()
    glyph_mapper.SetInputConnection(glyph.GetOutputPort())
    glyph_mapper.SetScalarModeToUsePointFieldData()
    glyph_mapper.SetColorModeToMapScalars()
    glyph_mapper.ScalarVisibilityOn()
    glyph_mapper.SelectColorArray('Elevation')
    # Colour by scalars.
    glyph_mapper.SetLookupTable(lut1)
    glyph_mapper.SetScalarRange(scalar_range_elevation)

    glyph_actor = vtk.vtkActor()
    glyph_actor.SetMapper(glyph_mapper)

    window_width = 800
    window_height = 800

    # Add scalar bars.
    scalar_bar = vtk.vtkScalarBarActor()
    # This LUT puts the lowest value at the top of the scalar bar.
    # scalar_bar->SetLookupTable(lut);
    # Use this LUT if you want the highest value at the top.
    scalar_bar.SetLookupTable(lutr)
    scalar_bar.SetTitle('Gaussian\nCurvature')
    scalar_bar.GetTitleTextProperty().SetColor(
        colors.GetColor3d('AliceBlue'))
    scalar_bar.GetLabelTextProperty().SetColor(
        colors.GetColor3d('AliceBlue'))
    scalar_bar.GetAnnotationTextProperty().SetColor(
        colors.GetColor3d('AliceBlue'))
    scalar_bar.UnconstrainedFontSizeOn()
    scalar_bar.SetMaximumWidthInPixels(window_width // 8)
    scalar_bar.SetMaximumHeightInPixels(window_height // 3)
    scalar_bar.SetPosition(0.85, 0.05)

    scalar_bar_elev = vtk.vtkScalarBarActor()
    # This LUT puts the lowest value at the top of the scalar bar.
    # scalar_bar_elev->SetLookupTable(lut);
    # Use this LUT if you want the highest value at the top.
    scalar_bar_elev.SetLookupTable(lut1)
    scalar_bar_elev.SetTitle('Elevation')
    scalar_bar_elev.GetTitleTextProperty().SetColor(
        colors.GetColor3d('AliceBlue'))
    scalar_bar_elev.GetLabelTextProperty().SetColor(
        colors.GetColor3d('AliceBlue'))
    scalar_bar_elev.GetAnnotationTextProperty().SetColor(
        colors.GetColor3d('AliceBlue'))
    scalar_bar_elev.UnconstrainedFontSizeOn()
    scalar_bar_elev.SetNumberOfLabels(5)
    scalar_bar_elev.SetMaximumWidthInPixels(window_width // 8)
    scalar_bar_elev.SetMaximumHeightInPixels(window_height // 3)
    # scalar_bar_elev.SetBarRatio(scalar_bar_elev.GetBarRatio() * 0.5)
    scalar_bar_elev.SetPosition(0.85, 0.4)

    # ------------------------------------------------------------
    # Create the RenderWindow, Renderer and Interactor
    # ------------------------------------------------------------
    ren = vtk.vtkRenderer()
    ren_win = vtk.vtkRenderWindow()
    iren = vtk.vtkRenderWindowInteractor()
    style = vtk.vtkInteractorStyleTrackballCamera()
    iren.SetInteractorStyle(style)

    ren_win.AddRenderer(ren)
    # Important: The interactor must be set prior to enabling the widget.
    iren.SetRenderWindow(ren_win)
    cam_orient_manipulator = vtk.vtkCameraOrientationWidget()
    cam_orient_manipulator.SetParentRenderer(ren)
    # Enable the widget.
    cam_orient_manipulator.On()

    # add actors
    ren.AddViewProp(src_actor)
    ren.AddViewProp(edge_actor)
    ren.AddViewProp(glyph_actor)
    ren.AddActor2D(scalar_bar)
    ren.AddActor2D(scalar_bar_elev)

    ren.SetBackground(colors.GetColor3d('ParaViewBkg'))
    ren_win.SetSize(window_width, window_height)
    ren_win.SetWindowName('CurvatureBandsWithGlyphs')

    if desired_surface == "RandomHills":
        camera = ren.GetActiveCamera()
        camera.SetPosition(10.9299, 59.1505, 24.9823)
        camera.SetFocalPoint(2.21692, 7.97545, 7.75135)
        camera.SetViewUp(-0.230136, 0.345504, -0.909761)
        camera.SetDistance(54.6966)
        camera.SetClippingRange(36.3006, 77.9852)
        ren_win.Render()

    iren.Start()


def make_bands(d_r, number_of_bands, nearest_integer):
    """
    Divide a range into bands
    :param: d_r - [min, max] the range that is to be covered by the bands.
    :param: number_of_bands - the number of bands, a positive integer.
    :param: nearest_integer - if True then [floor(min), ceil(max)] is used.
    :return: A dictionary consisting of the band number and [min, midpoint, max] for each band.
    """
    bands = OrderedDict()
    if (d_r[1] < d_r[0]) or (number_of_bands <= 0):
        return bands
    x = list(d_r)
    if nearest_integer:
        x[0] = math.floor(x[0])
        x[1] = math.ceil(x[1])
    dx = (x[1] - x[0]) / float(number_of_bands)
    b = [x[0], x[0] + dx / 2.0, x[0] + dx]
    i = 0
    while i < number_of_bands:
        bands[i] = b
        b = [b[0] + dx, b[1] + dx, b[2] + dx]
        i += 1
    return bands


def make_custom_bands(d_r, number_of_bands, my_bands):
    """
    Divide a range into custom bands.

    You need to specify each band as an list [r1, r2] where r1 < r2 and
    append these to a list.
    The list should ultimately look
    like this: [[r1, r2], [r2, r3], [r3, r4]...]

    :param: d_r - [min, max] the range that is to be covered by the bands.
    :param: number_of_bands - the number of bands, a positive integer.
    :return: A dixtionary consisting of band number and [min, midpoint, max] for each band.
    """
    bands = OrderedDict()
    if (d_r[1] < d_r[0]) or (number_of_bands <= 0):
        return bands
    x = my_bands
    # Determine the index of the range minimum and range maximum.
    idx_min = 0
    for idx in range(0, len(my_bands)):
        if my_bands[idx][1] > d_r[0] >= my_bands[idx][0]:
            idx_min = idx
            break

    idx_max = len(my_bands) - 1
    for idx in range(len(my_bands) - 1, -1, -1):
        if my_bands[idx][1] > d_r[1] >= my_bands[idx][0]:
            idx_max = idx
            break

    # Set the minimum to match the range minimum.
    x[idx_min][0] = d_r[0]
    x[idx_max][1] = d_r[1]
    x = x[idx_min: idx_max + 1]
    for idx, e in enumerate(x):
        bands[idx] = [e[0], e[0] + (e[1] - e[0]) / 2, e[1]]
    return bands


def frequencies(bands, src):
    """
    Count the number of scalars in each band.
    :param: bands - the bands.
    :param: src - the vtkPolyData source.
    :return: The frequencies of the scalars in each band.
    """
    freq = OrderedDict()
    for i in range(len(bands)):
        freq[i] = 0
    tuples = src.GetPointData().GetScalars().GetNumberOfTuples()
    for i in range(tuples):
        x = src.GetPointData().GetScalars().GetTuple1(i)
        for j in range(len(bands)):
            if x <= bands[j][2]:
                freq[j] = freq[j] + 1
                break
    return freq


def adjust_frequency_ranges(freq):
    """
    Get the indices of the first and last non-zero elements.
    :param freq: The frequency dictionary.
    :return: The indices of the first and last non-zero elements.
    """
    first = 0
    for k, v in freq.items():
        if v != 0:
            first = k
            break
    rev_keys = list(freq.keys())[::-1]
    last = rev_keys[0]
    for idx in list(freq.keys())[::-1]:
        if freq[idx] != 0:
            last = idx
            break
    return first, last


def make_elevations(src):
    """
    Generate elevations over the surface.
    :param: src - the vtkPolyData source.
    :return: - vtkPolyData source with elevations.
    """
    bounds = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    src.GetBounds(bounds)
    if abs(bounds[2]) < 1.0e-8 and abs(bounds[3]) < 1.0e-8:
        bounds[3] = bounds[2] + 1
    elev_filter = vtk.vtkElevationFilter()
    elev_filter.SetInputData(src)
    elev_filter.SetLowPoint(0, bounds[2], 0)
    elev_filter.SetHighPoint(0, bounds[3], 0)
    elev_filter.SetScalarRange(bounds[2], bounds[3])
    elev_filter.Update()
    return elev_filter.GetPolyDataOutput()


def make_parametric_hills():
    """
    Make a parametric hills surface as the source.
    :return: vtkPolyData with normal and scalar data.
    """
    fn = vtk.vtkParametricRandomHills()
    fn.AllowRandomGenerationOn()
    fn.SetRandomSeed(1)
    fn.SetNumberOfHills(30)
    # Make the normals face out of the surface.
    # Not needed with VTK 8.0 or later.
    # if fn.GetClassName() == 'vtkParametricRandomHills':
    #    fn.ClockwiseOrderingOff()

    source = vtk.vtkParametricFunctionSource()
    source.SetParametricFunction(fn)
    source.SetUResolution(50)
    source.SetVResolution(50)
    source.SetScalarModeToZ()
    source.Update()
    # Name the arrays (not needed in VTK 6.2+ for vtkParametricFunctionSource).
    # source.GetOutput().GetPointData().GetNormals().SetName('Normals')
    # source.GetOutput().GetPointData().GetScalars().SetName('Scalars')
    # Rename the scalars to 'Elevation' since we are using the Z-scalars as elevations.
    source.GetOutput().GetPointData().GetScalars().SetName('Elevation')

    transform = vtk.vtkTransform()
    transform.Translate(0.0, 5.0, 15.0)
    transform.RotateX(-90.0)
    transform_filter = vtk.vtkTransformPolyDataFilter()
    transform_filter.SetInputConnection(source.GetOutputPort())
    transform_filter.SetTransform(transform)
    transform_filter.Update()

    return transform_filter.GetOutput()


def make_parametric_torus():
    """
    Make a parametric torus as the source.
    :return: vtkPolyData with normal and scalar data.
    """

    fn = vtk.vtkParametricTorus()
    fn.SetRingRadius(5)
    fn.SetCrossSectionRadius(2)

    source = vtk.vtkParametricFunctionSource()
    source.SetParametricFunction(fn)
    source.SetUResolution(50)
    source.SetVResolution(50)
    source.SetScalarModeToZ()
    source.Update()

    # Name the arrays (not needed in VTK 6.2+ for vtkParametricFunctionSource).
    # source.GetOutput().GetPointData().GetNormals().SetName('Normals')
    # source.GetOutput().GetPointData().GetScalars().SetName('Scalars')
    # Rename the scalars to 'Elevation' since we are using the Z-scalars as elevations.
    source.GetOutput().GetPointData().GetScalars().SetName('Elevation')

    transform = vtk.vtkTransform()
    transform.RotateX(-90.0)
    transform_filter = vtk.vtkTransformPolyDataFilter()
    transform_filter.SetInputConnection(source.GetOutputPort())
    transform_filter.SetTransform(transform)
    transform_filter.Update()

    return transform_filter.GetOutput()


def make_plane():
    """
    Make a plane as the source.
    :return: vtkPolyData with normal and scalar data.
    """

    source = vtk.vtkPlaneSource()
    source.SetOrigin(-10.0, -10.0, 0.0)
    source.SetPoint2(-10.0, 10.0, 0.0)
    source.SetPoint1(10.0, -10.0, 0.0)
    source.SetXResolution(20)
    source.SetYResolution(20)
    source.Update()

    transform = vtk.vtkTransform()
    transform.Translate(0.0, 0.0, 0.0)
    transform.RotateX(-90.0)
    transform_filter = vtk.vtkTransformPolyDataFilter()
    transform_filter.SetInputConnection(source.GetOutputPort())
    transform_filter.SetTransform(transform)
    transform_filter.Update()

    # We have a m x n array of quadrilaterals arranged as a regular tiling in a
    # plane. So pass it through a triangle filter since the curvature filter only
    # operates on polys.
    tri = vtk.vtkTriangleFilter()
    tri.SetInputConnection(transform_filter.GetOutputPort())

    # Pass it though a CleanPolyDataFilter and merge any points which
    # are coincident, or very close
    cleaner = vtk.vtkCleanPolyData()
    cleaner.SetInputConnection(tri.GetOutputPort())
    cleaner.SetTolerance(0.005)
    cleaner.Update()

    return cleaner.GetOutput()


def make_sphere():
    source = vtk.vtkSphereSource()
    source.SetCenter(0.0, 0.0, 0.0)
    source.SetRadius(10.0)
    source.SetThetaResolution(32)
    source.SetPhiResolution(32)
    source.Update()

    return source.GetOutput()


def make_torus():
    """
    Make a torus as the source.
    :return: vtkPolyData with normal and scalar data.
    """
    source = vtk.vtkSuperquadricSource()
    source.SetCenter(0.0, 0.0, 0.0)
    source.SetScale(1.0, 1.0, 1.0)
    source.SetPhiResolution(64)
    source.SetThetaResolution(64)
    source.SetThetaRoundness(1)
    source.SetThickness(0.5)
    source.SetSize(10)
    source.SetToroidal(1)

    # The quadric is made of strips, so pass it through a triangle filter as
    # the curvature filter only operates on polys
    tri = vtk.vtkTriangleFilter()
    tri.SetInputConnection(source.GetOutputPort())

    # The quadric has nasty discontinuities from the way the edges are generated
    # so let's pass it though a CleanPolyDataFilter and merge any points which
    # are coincident, or very close
    cleaner = vtk.vtkCleanPolyData()
    cleaner.SetInputConnection(tri.GetOutputPort())
    cleaner.SetTolerance(0.005)
    cleaner.Update()

    return cleaner.GetOutput()


def clipper(src, dx, dy, dz):
    """
    Clip a vtkPolyData source.
    A cube is made whose size corresponds the the bounds of the source.
    Then each side is shrunk by the appropriate dx, dy or dz. After
    this operation the source is clipped by the cube.
    :param: src - the vtkPolyData source
    :param: dx - the amount to clip in the x-direction
    :param: dy - the amount to clip in the y-direction
    :param: dz - the amount to clip in the z-direction
    :return: vtkPolyData.
    """
    bounds = [0, 0, 0, 0, 0, 0]
    src.GetBounds(bounds)

    plane1 = vtk.vtkPlane()
    plane1.SetOrigin(bounds[0] + dx, 0, 0)
    plane1.SetNormal(1, 0, 0)

    plane2 = vtk.vtkPlane()
    plane2.SetOrigin(bounds[1] - dx, 0, 0)
    plane2.SetNormal(-1, 0, 0)

    plane3 = vtk.vtkPlane()
    plane3.SetOrigin(0, bounds[2] + dy, 0)
    plane3.SetNormal(0, 1, 0)

    plane4 = vtk.vtkPlane()
    plane4.SetOrigin(0, bounds[3] - dy, 0)
    plane4.SetNormal(0, -1, 0)

    plane5 = vtk.vtkPlane()
    plane5.SetOrigin(0, 0, bounds[4] + dz)
    plane5.SetNormal(0, 0, 1)

    plane6 = vtk.vtkPlane()
    plane6.SetOrigin(0, 0, bounds[5] - dz)
    plane6.SetNormal(0, 0, -1)

    clip_function = vtk.vtkImplicitBoolean()
    clip_function.SetOperationTypeToUnion()
    clip_function.AddFunction(plane1)
    clip_function.AddFunction(plane2)
    clip_function.AddFunction(plane3)
    clip_function.AddFunction(plane4)
    clip_function.AddFunction(plane5)
    clip_function.AddFunction(plane6)

    # Clip it.
    pd_clipper = vtk.vtkClipPolyData()
    pd_clipper.SetClipFunction(clip_function)
    pd_clipper.SetInputData(src)
    pd_clipper.GenerateClipScalarsOff()
    pd_clipper.GenerateClippedOutputOff()
    # pd_clipper.GenerateClippedOutputOn()
    pd_clipper.Update()
    return pd_clipper.GetOutput()


def calculate_curvatures(src):
    """
    The source must be triangulated.
    :param: src - the source.
    :return: vtkPolyData with normal and scalar data representing curvatures.
    """
    curvature = vtk.vtkCurvatures()
    curvature.SetCurvatureTypeToGaussian()
    curvature.SetInputData(src)
    curvature.Update()
    return curvature.GetOutput()


def get_color_series():
    color_series = vtk.vtkColorSeries()
    # Select a color scheme.
    # color_series_enum = color_series.BREWER_DIVERGING_BROWN_BLUE_GREEN_9
    # color_series_enum = color_series.BREWER_DIVERGING_SPECTRAL_10
    # color_series_enum = color_series.BREWER_DIVERGING_SPECTRAL_3
    # color_series_enum = color_series.BREWER_DIVERGING_PURPLE_ORANGE_9
    # color_series_enum = color_series.BREWER_SEQUENTIAL_BLUE_PURPLE_9
    # color_series_enum = color_series.BREWER_SEQUENTIAL_BLUE_GREEN_9
    color_series_enum = color_series.BREWER_QUALITATIVE_SET3
    # color_series_enum = color_series.CITRUS
    color_series.SetColorScheme(color_series_enum)
    return color_series


def make_categorical_lut():
    """
    Make a lookup table using vtkColorSeries.
    :return: An indexed (categorical) lookup table.
    """
    color_series = get_color_series()
    # Make the lookup table.
    lut = vtk.vtkLookupTable()
    color_series.BuildLookupTable(lut, color_series.CATEGORICAL)
    lut.SetNanColor(0, 0, 0, 1)
    return lut


def make_ordinal_lut():
    """
    Make a lookup table using vtkColorSeries.
    :return: An ordinal (not indexed) lookup table.
    """
    color_series = get_color_series()
    # Make the lookup table.
    lut = vtk.vtkLookupTable()
    color_series.BuildLookupTable(lut, color_series.ORDINAL)
    lut.SetNanColor(0, 0, 0, 1)
    return lut


def make_diverging_lut():
    """
    See: [Diverging Color Maps for Scientific Visualization](https://www.kennethmoreland.com/color-maps/)
                       start point         midPoint            end point
     cool to warm:     0.230, 0.299, 0.754 0.865, 0.865, 0.865 0.706, 0.016, 0.150
     purple to orange: 0.436, 0.308, 0.631 0.865, 0.865, 0.865 0.759, 0.334, 0.046
     green to purple:  0.085, 0.532, 0.201 0.865, 0.865, 0.865 0.436, 0.308, 0.631
     blue to brown:    0.217, 0.525, 0.910 0.865, 0.865, 0.865 0.677, 0.492, 0.093
     green to red:     0.085, 0.532, 0.201 0.865, 0.865, 0.865 0.758, 0.214, 0.233

    :return:
    """
    ctf = vtk.vtkColorTransferFunction()
    ctf.SetColorSpaceToDiverging()
    # Cool to warm.
    ctf.AddRGBPoint(0.0, 0.085, 0.532, 0.201)
    ctf.AddRGBPoint(0.5, 0.865, 0.865, 0.865)
    ctf.AddRGBPoint(1.0, 0.758, 0.214, 0.233)

    table_size = 256
    lut = vtk.vtkLookupTable()
    lut.SetNumberOfTableValues(table_size)
    lut.Build()

    for i in range(0, table_size):
        rgba = list(ctf.GetColor(float(i) / table_size))
        rgba.append(1)
        lut.SetTableValue(i, rgba)

    return lut


def reverse_lut(lut):
    """
    Create a lookup table with the colors reversed.
    :param: lut - An indexed lookup table.
    :return: The reversed indexed lookup table.
    """
    lutr = vtk.vtkLookupTable()
    lutr.DeepCopy(lut)
    t = lut.GetNumberOfTableValues() - 1
    rev_range = reversed(list(range(t + 1)))
    for i in rev_range:
        rgba = [0.0] * 3
        v = float(i)
        lut.GetColor(v, rgba)
        rgba.append(lut.GetOpacity(v))
        lutr.SetTableValue(t - i, rgba)
    t = lut.GetNumberOfAnnotatedValues() - 1
    rev_range = reversed(list(range(t + 1)))
    for i in rev_range:
        lutr.SetAnnotation(t - i, lut.GetAnnotation(i))
    return lutr


def make_glyphs(src, reverse_normals):
    """
    Glyph the normals on the surface.

    You may need to adjust the parameters for mask_pts, arrow and glyph for a
    nice appearance.

    :param: src - the surface to glyph.
    :param: reverse_normals - if True the normals on the surface are reversed.
    :return: The glyph object.

    """
    # Sometimes the contouring algorithm can create a volume whose gradient
    # vector and ordering of polygon (using the right hand rule) are
    # inconsistent. vtkReverseSense cures this problem.
    reverse = vtk.vtkReverseSense()

    # Choose a random subset of points.
    mask_pts = vtk.vtkMaskPoints()
    mask_pts.SetOnRatio(5)
    mask_pts.RandomModeOn()
    if reverse_normals:
        reverse.SetInputData(src)
        reverse.ReverseCellsOn()
        reverse.ReverseNormalsOn()
        mask_pts.SetInputConnection(reverse.GetOutputPort())
    else:
        mask_pts.SetInputData(src)

    # Source for the glyph filter
    arrow = vtk.vtkArrowSource()
    arrow.SetTipResolution(16)
    arrow.SetTipLength(0.3)
    arrow.SetTipRadius(0.1)

    glyph = vtk.vtkGlyph3D()
    glyph.SetSourceConnection(arrow.GetOutputPort())
    glyph.SetInputConnection(mask_pts.GetOutputPort())
    glyph.SetVectorModeToUseNormal()
    glyph.SetScaleFactor(1.0)
    glyph.SetColorModeToColorByVector()
    glyph.SetScaleModeToScaleByVector()
    glyph.OrientOn()
    glyph.Update()
    return glyph


def print_bands(bands):
    s = f'Bands:\n'
    for k, v in bands.items():
        for j, q in enumerate(v):
            if j == 0:
                s += f'{k:4d} ['
            if j == len(v) - 1:
                s += f'{q:8.3f}]\n'
            else:
                s += f'{q:8.3f}, '
    print(s)


def print_frequencies(freq):
    s = ''
    for i, p in freq.items():
        if i == 0:
            s += f'Frequencies: ['
        if i == len(freq) - 1:
            s += f'{i}: {p} ]'
        else:
            s += f'{i}: {p}, '
    print(s)


def print_bands_frequencies(bands, freq):
    if len(bands) != len(freq):
        print('Bands and frequencies must be the same size.')
        return
    s = f'Bands & frequencies:\n'
    for k, v in bands.items():
        for j, q in enumerate(v):
            if j == 0:
                s += f'{k:4d} ['
            if j == len(v) - 1:
                s += f'{q:8.3f}]: {freq[k]:6d}\n'
            else:
                s += f'{q:8.3f}, '
    print(s)


class ComputeCurvatures:
    """
    This class takes a vtkPolyData source and:
     - calculates Gaussian and Mean curvatures,
     - adjusts curvatures along the edges using a weighted average,
     - inserts the adjusted curvatures into the vtkPolyData source.

     Additional methods are provided for setting bounds and precision.
    """

    def __init__(self, polydata_source, gauss_eps=1.0e-08, mean_eps=1.0e-08):
        """

        :param polydata_source: The polydata source.
        :param gauss_eps: Gaussian curvatures less than this will be set to zero.
        :param mean_eps: Mean curvatures less than this will be set to zero.
        """
        self.source = polydata_source
        self.curvature_type = ['Gauss_Curvature', 'Mean_Curvature']
        self.adjusted_curvatures = dict()
        self.bounds = {'Gauss_Curvature': [0.0, 0.0], 'Mean_Curvature': [0.0, 0.0]}
        self.bounds_state = {'Gauss_Curvature': False, 'Mean_Curvature': False}
        self.epsilons = {'Gauss_Curvature': gauss_eps, 'Mean_Curvature': mean_eps}

    def update(self):
        for curvature_name in self.curvature_type:
            self.compute_curvature_and_fix_up_boundary(curvature_name)
            #  Set small values to zero.
            if self.epsilons[curvature_name] != 0.0:
                eps = abs(self.epsilons[curvature_name])
                self.adjusted_curvatures[curvature_name] = np.where(
                    abs(self.adjusted_curvatures[curvature_name]) < eps, 0,
                    self.adjusted_curvatures[curvature_name])
            # Set upper and lower bounds.
            if self.bounds_state[curvature_name]:
                lower_bound = self.bounds[curvature_name][0]
                self.adjusted_curvatures[curvature_name] = np.where(
                    self.adjusted_curvatures[curvature_name] < lower_bound, lower_bound,
                    self.adjusted_curvatures[curvature_name])
                upper_bound = self.bounds[curvature_name][1]
                self.adjusted_curvatures[curvature_name] = np.where(
                    self.adjusted_curvatures[curvature_name] > upper_bound, upper_bound,
                    self.adjusted_curvatures[curvature_name])
            self.update_curvatures(curvature_name)

    def compute_curvature_and_fix_up_boundary(self, curvature_name):
        # Curvature as vtkPolyData.
        curvature_data = self.compute_curvature(curvature_name)
        # Curvature as python list.
        curvature = self.extract_data(curvature_data, curvature_name)
        # Ids of the boundary points.
        p_ids = self.extract_boundary_ids()
        # Remove duplicate Ids.
        p_ids_set = set(p_ids)

        # Iterate over the edge points and compute the curvature as the weighted
        # average of the neighbors.
        count_invalid = 0
        for p_id in p_ids:
            p_ids_neighbors = self.point_neighborhood(p_id)
            # Keep only interior points.
            p_ids_neighbors -= p_ids_set
            # Compute distances and extract curvature values.
            curvs = [curvature[p_id_n] for p_id_n in p_ids_neighbors]
            dists = [self.compute_distance(p_id_n, p_id) for p_id_n in p_ids_neighbors]
            curvs = np.array(curvs)
            dists = np.array(dists)
            curvs = curvs[dists > 0]
            dists = dists[dists > 0]
            if len(curvs) > 0:
                weights = 1 / np.array(dists)
                weights /= weights.sum()
                new_curv = np.dot(curvs, weights)
            else:
                # Corner case.
                count_invalid += 1
                new_curv = 0
            # Set the new curvature value.
            curvature[p_id] = new_curv
        self.adjusted_curvatures[curvature_name] = np.array(curvature)

    def compute_curvature(self, curvature_name):
        curvature_filter = vtk.vtkCurvatures()
        curvature_filter.SetInputData(self.source)
        if 'gaus' in curvature_name.lower():
            curvature_filter.SetCurvatureTypeToGaussian()
        else:
            curvature_filter.SetCurvatureTypeToMean()
        curvature_filter.Update()
        return curvature_filter.GetOutput()

    @staticmethod
    def extract_data(curvature_data, curvature_name):
        array = curvature_data.GetPointData().GetAbstractArray(curvature_name)
        n = curvature_data.GetNumberOfPoints()
        data = []
        for i in range(n):
            data.append(array.GetVariantValue(i).ToDouble())
        return data

    def extract_boundary_ids(self):
        """
        See here: https://discourse.vtk.org/t/2530/3
        """
        array_name = 'ids'
        id_filter = vtk.vtkIdFilter()
        id_filter.SetInputData(self.source)
        id_filter.SetPointIds(True)
        id_filter.SetCellIds(False)
        id_filter.SetPointIdsArrayName(array_name)
        id_filter.SetCellIdsArrayName(array_name)
        id_filter.Update()

        edges = vtk.vtkFeatureEdges()
        edges.SetInputConnection(id_filter.GetOutputPort())
        edges.BoundaryEdgesOn()
        edges.ManifoldEdgesOff()
        edges.NonManifoldEdgesOff()
        edges.FeatureEdgesOff()
        edges.Update()

        array = edges.GetOutput().GetPointData().GetArray(array_name)
        n = edges.GetOutput().GetNumberOfPoints()
        boundary_ids = []
        for i in range(n):
            boundary_ids.append(array.GetValue(i))
        return boundary_ids

    def point_neighborhood(self, p_id):
        """
        Extract the topological neighbors for point pId. In two steps:
        1) self.source.GetPointCells(pId, cell_ids)
        2) self.source.GetCellPoints(c_id, point_ids) for all c_id in cell_ids
        """
        cell_ids = vtk.vtkIdList()
        self.source.GetPointCells(p_id, cell_ids)
        neighbors = set()
        for i in range(0, cell_ids.GetNumberOfIds()):
            cell_id = cell_ids.GetId(i)
            cell_point_ids = vtk.vtkIdList()
            self.source.GetCellPoints(cell_id, cell_point_ids)
            for j in range(0, cell_point_ids.GetNumberOfIds()):
                neighbors.add(cell_point_ids.GetId(j))
        return neighbors

    def compute_distance(self, pt_id_a, pt_id_b):
        pt_a = np.array(self.source.GetPoint(pt_id_a))
        pt_b = np.array(self.source.GetPoint(pt_id_b))
        return np.linalg.norm(pt_a - pt_b)

    def update_curvatures(self, curvature_name):
        """
        Add the adjusted curvatures into the self.source.
         :return:
        """
        if self.source.GetNumberOfPoints() != len(self.adjusted_curvatures[curvature_name]):
            print(curvature_name, ':\nCannot add the adjusted curvatures to the source.\n'
                                  ' The number of points in source does not equal the\n'
                                  ' number of point ids in the adjusted curvature array.')
            return
        curvatures = numpy_support.numpy_to_vtk(num_array=self.adjusted_curvatures[curvature_name].ravel(), deep=True,
                                                array_type=vtk.VTK_DOUBLE)
        curvatures.SetName(curvature_name)
        self.source.GetPointData().AddArray(curvatures)
        self.source.GetPointData().SetActiveScalars(curvature_name)

    # Remember to run Update after these set and on/off methods.
    def set_gauss_curvature_bounds(self, lower=0.0, upper=0.0):
        self.bounds['Gauss_Curvature'] = [lower, upper]

    def gauss_bounds_on(self):
        self.bounds_state['Gauss_Curvature'] = True

    def gauss_bounds_off(self):
        self.bounds_state['Gauss_Curvature'] = False

    def set_mean_curvature_bounds(self, lower=0.0, upper=0.0):
        self.bounds['Mean_Curvature'] = [lower, upper]

    def mean_bounds_on(self):
        self.bounds_state['Mean_Curvature'] = True

    def mean_bounds_off(self):
        self.bounds_state['Mean_Curvature'] = False

    def set_epsilons(self, gauss_eps=1.0e-08, mean_eps=1.0e-08):
        self.epsilons = {'Gauss_Curvature': gauss_eps, 'Mean_Curvature': mean_eps}


if __name__ == '__main__':
    main()