KochSnowflake

#!/usr/bin/env python

from math import pi, cos, sin, sqrt

# noinspection PyUnresolvedReferences
import vtkmodules.vtkInteractionStyle
# noinspection PyUnresolvedReferences
import vtkmodules.vtkRenderingOpenGL2
from vtkmodules.vtkCommonColor import vtkNamedColors
from vtkmodules.vtkCommonCore import (
    vtkIntArray,
    vtkLookupTable,
    vtkPoints
)
from vtkmodules.vtkCommonDataModel import (
    vtkCellArray,
    vtkPolyData,
    vtkPolyLine,
    vtkTriangle
)
from vtkmodules.vtkRenderingCore import (
    vtkActor,
    vtkPolyDataMapper,
    vtkRenderWindow,
    vtkRenderWindowInteractor,
    vtkRenderer
)

LEVEL = 6


def as_polyline(points, level):
    '''
    Koch Snowflake as a vtkPolyLine
    '''
    # Use the points from the previous iteration to create the points of the next
    # level. There is an assumption on my part that the curve is traversed in a
    # counterclockwise fashion. If the initial triangle above is written to
    # describe clockwise motion, the points will face inward instead of outward.
    for i in range(level):
        temp = vtkPoints()
        # The first point of the previous vtkPoints is the first point of the next vtkPoints.
        temp.InsertNextPoint(*points.GetPoint(0))

        # Iterate over 'edges' in the vtkPoints
        for i in range(1, points.GetNumberOfPoints()):
            x0, y0, z0 = points.GetPoint(i - 1)
            x1, y1, z1 = points.GetPoint(i)
            t = sqrt((x1 - x0) ** 2 + (y1 - y0) ** 2)
            nx = (x1 - x0) / t  # x-component of edge unit tangent
            ny = (y1 - y0) / t  # y-component of edge unit tangent

            # the points describing the Koch snowflake edge
            temp.InsertNextPoint(x0 + nx * t / 3, y0 + ny * t / 3, 0.)
            temp.InsertNextPoint(x0 + nx * t / 2 + ny * t * sqrt(3) / 6, y0 + ny * t / 2 - nx * t * sqrt(3) / 6, 0.)
            temp.InsertNextPoint(x0 + nx * 2 * t / 3, y0 + ny * 2 * t / 3, 0.)
            temp.InsertNextPoint(x0 + nx * t, y0 + ny * t, 0.)

        points = temp

    # draw the outline
    lines = vtkCellArray()
    pl = vtkPolyLine()
    pl.GetPointIds().SetNumberOfIds(points.GetNumberOfPoints())
    for i in range(points.GetNumberOfPoints()):
        pl.GetPointIds().SetId(i, i)
    lines.InsertNextCell(pl)

    # complete the polydata
    polydata = vtkPolyData()
    polydata.SetLines(lines)
    polydata.SetPoints(points)

    return polydata


def as_triangles(indices, cellarray, level, data):
    '''
    Koch Snowflake as a collection of vtkTriangles
    '''
    if len(indices) >= 3:
        stride = len(indices) // 4
        indices.append(indices[-1] + 1)

        triangle = vtkTriangle()
        triangle.GetPointIds().SetId(0, indices[stride])
        triangle.GetPointIds().SetId(1, indices[2 * stride])
        triangle.GetPointIds().SetId(2, indices[3 * stride])

        cellarray.InsertNextCell(triangle)
        data.InsertNextValue(level)

        as_triangles(indices[0: stride], cellarray, level + 1, data)
        as_triangles(indices[stride: 2 * stride], cellarray, level + 1, data)
        as_triangles(indices[2 * stride: 3 * stride], cellarray, level + 1, data)
        as_triangles(indices[3 * stride: -1], cellarray, level + 1, data)


def main():
    colors = vtkNamedColors()

    # Initially, set up the points to be an equilateral triangle. Note that the
    # first point is the same as the last point to make this a closed curve when
    # I create the vtkPolyLine.
    points = vtkPoints()
    for i in range(4):
        points.InsertNextPoint(cos(2.0 * pi * i / 3), sin(2 * pi * i / 3.0), 0.0)

    outline_pd = as_polyline(points, LEVEL)
    # You have already gone through the trouble of putting the points in the
    # right places - so 'all' you need to do now is to create polygons from the
    # points that are in the vtkPoints.

    # The points that are passed in, have an overlap of the beginning and the
    # end. For this next trick, I will need a list of the indices in the
    # vtkPoints. They're consecutive, so thats pretty straightforward.

    indices = [i for i in range(outline_pd.GetPoints().GetNumberOfPoints() + 1)]
    triangles = vtkCellArray()

    # Set this up for each of the initial sides, then call the recursive function.
    stride = (len(indices) - 1) // 3

    # The cell data will allow us to color the triangles based on the level of
    # the iteration of the Koch snowflake.
    data = vtkIntArray()
    data.SetNumberOfComponents(0)
    data.SetName('Iteration Level')

    # This is the starting triangle.
    t = vtkTriangle()
    t.GetPointIds().SetId(0, 0)
    t.GetPointIds().SetId(1, stride)
    t.GetPointIds().SetId(2, 2 * stride)
    triangles.InsertNextCell(t)
    data.InsertNextValue(0)

    as_triangles(indices[0: stride + 1], triangles, 1, data)
    as_triangles(indices[stride: 2 * stride + 1], triangles, 1, data)
    as_triangles(indices[2 * stride: -1], triangles, 1, data)

    triangle_pd = vtkPolyData()
    triangle_pd.SetPoints(outline_pd.GetPoints())
    triangle_pd.SetPolys(triangles)
    triangle_pd.GetCellData().SetScalars(data)

    # ---------------- #
    # rendering stuff  #
    # ---------------- #
    outline_mapper = vtkPolyDataMapper()
    outline_mapper.SetInputData(outline_pd)

    lut = vtkLookupTable()
    lut.SetNumberOfTableValues(256)
    lut.SetHueRange(0.6, 0.6)
    lut.SetSaturationRange(0.0, 1.0)
    lut.Build()

    triangle_mapper = vtkPolyDataMapper()
    triangle_mapper.SetInputData(triangle_pd)
    triangle_mapper.SetScalarRange(0.0, LEVEL)
    triangle_mapper.SetLookupTable(lut)

    outline_actor = vtkActor()
    outline_actor.SetMapper(outline_mapper)

    triangle_actor = vtkActor()
    triangle_actor.SetMapper(triangle_mapper)

    outline_ren = vtkRenderer()
    outline_ren.AddActor(outline_actor)
    outline_ren.SetViewport(0.0, 0.0, 0.5, 1.0)

    triangle_ren = vtkRenderer()
    triangle_ren.AddActor(triangle_actor)
    triangle_ren.SetViewport(0.5, 0.0, 1.0, 1.0)
    triangle_ren.SetActiveCamera(outline_ren.GetActiveCamera())

    renw = vtkRenderWindow()
    renw.AddRenderer(outline_ren)
    renw.AddRenderer(triangle_ren)
    renw.SetSize(800, 400)
    renw.SetWindowName('KochSnowflake')

    outline_ren.SetBackground(colors.GetColor3d('CornFlowerBLue'))
    triangle_ren.SetBackground(colors.GetColor3d('MistyRose'))

    iren = vtkRenderWindowInteractor()
    iren.SetRenderWindow(renw)

    outline_ren.ResetCamera()
    renw.Render()
    iren.Start()


if __name__ == '__main__':
    main()