FireFlowDemo

#!/usr/bin/env python3

import time
from dataclasses import dataclass
from pathlib import Path

# noinspection PyUnresolvedReferences
import vtkmodules.vtkInteractionStyle
# noinspection PyUnresolvedReferences
import vtkmodules.vtkRenderingFreeType
# noinspection PyUnresolvedReferences
import vtkmodules.vtkRenderingOpenGL2
from vtkmodules.vtkCommonColor import vtkNamedColors
from vtkmodules.vtkCommonCore import vtkCommand
from vtkmodules.vtkFiltersCore import (
    vtkContourFilter,
    vtkTubeFilter
)
from vtkmodules.vtkFiltersFlowPaths import vtkStreamTracer
from vtkmodules.vtkFiltersGeneric import vtkGenericOutlineFilter
from vtkmodules.vtkFiltersSources import (
    vtkPointSource,
    vtkSphereSource)
from vtkmodules.vtkIOImport import vtkVRMLImporter
from vtkmodules.vtkIOXML import vtkXMLUnstructuredGridReader
from vtkmodules.vtkInteractionStyle import vtkInteractorStyleTrackballCamera
from vtkmodules.vtkInteractionWidgets import (
    vtkSliderRepresentation2D,
    vtkSliderWidget)
from vtkmodules.vtkRenderingCore import (
    vtkActor,
    vtkPolyDataMapper,
    vtkRenderer,
    vtkRenderWindow,
    vtkRenderWindowInteractor
)


def get_program_parameters():
    import argparse
    description = 'Fire Flow Demonstration.'
    epilogue = '''
    The example illustrates how to combine a geometric description of a scene with a fluid flow solution.
    '''
    parser = argparse.ArgumentParser(description=description, epilog=epilogue,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('geometry',
                        help='The path to the geometry wrl file, e.g. room_vis.wrl.')
    parser.add_argument('velocity',
                        help='The path to the velocity vtu file, e.g. fire_ug.vtu.')
    parser.add_argument('-d', '--delay', type=int, default=100,
                        help='The delay in milliseconds, default=500ms.')
    args = parser.parse_args()
    return args.geometry, args.velocity, args.delay


def main():
    geometry_fn, velocity_fn, delay = get_program_parameters()
    geometry_path = Path(geometry_fn)
    velocity_path = Path(velocity_fn)
    file_check = True
    if not geometry_path.is_file():
        print(f'Missing geometry file: {geometry_path}.')
        file_check = False
    else:
        if geometry_path.suffix.lower() != '.wrl':
            print(f'The geometry file : {geometry_path} must have a .wrl suffix.')
            file_check = False

    if not Path(velocity_path).is_file():
        print(f'Missing velocity file: {velocity_path}.')
        file_check = False
    else:
        if velocity_path.suffix.lower() != '.vtu':
            print(f'The velocity file : {velocity_path} must have a .vtu suffix.')
            file_check = False
    if not file_check:
        return

    colors = vtkNamedColors()
    iso_surface_color = colors.GetColor3d('WhiteSmoke')
    sphere_color = colors.GetColor3d('HotPink')
    background_color = colors.GetColor3d('SlateGray')

    renderer = vtkRenderer(use_hidden_line_removal=True, background=background_color)

    render_window = vtkRenderWindow(size=(640, 512), window_name='FireFlowDemo')
    render_window.AddRenderer(renderer)

    render_window_interactor = vtkRenderWindowInteractor()
    render_window_interactor.render_window = render_window

    style = vtkInteractorStyleTrackballCamera()
    render_window_interactor.interactor_style = style

    # Import the VRML Files that define the geometry.
    vrml_import = vtkVRMLImporter(file_name=geometry_path, render_window=render_window)
    vrml_import.Update()

    # Read the UnstructuredGrid define the solution.
    solution = vtkXMLUnstructuredGridReader(file_name=velocity_path)
    solution.update()

    bounds = solution.output.bounds
    # center = list()
    # for i in range(0, 6, 2):
    #     print(i)
    #     center.append((bounds[i+1] - bounds[i])/2.0)
    center = (3.0, 1.6, 1.25)
    scalar_range = solution.output.scalar_range

    # Create an outline.
    outline = vtkGenericOutlineFilter()
    solution >> outline

    # Create Seeds.
    # seeds = vtkPointSource(radius=0.2, center=center, number_of_points=50)
    seeds = vtkPointSource(radius=0.2, center=center, number_of_points=50)

    # Create streamlines.
    stream_tracer = vtkStreamTracer(source_connection=seeds.output_port,
                                    maximum_propagation=50, initial_integration_step=0.2, minimum_integration_step=0.01,
                                    integrator_type=2, compute_vorticity=True)
    solution >> stream_tracer
    stream_tracer.SetIntegrationDirectionToBoth()

    tubes = vtkTubeFilter(number_of_sides=8, radius=0.02, vary_radius=False)
    stream_tracer >> tubes

    map_tubes = vtkPolyDataMapper(scalar_range=scalar_range)
    stream_tracer >> tubes >> map_tubes

    tubes_actor = vtkActor(mapper=map_tubes)

    # Create an Isosurface.
    iso_surface = vtkContourFilter()
    iso_surface.SetValue(0, 500.0)
    solution >> iso_surface

    iso_surface_mapper = vtkPolyDataMapper(scalar_visibility=False)
    solution >> iso_surface >> iso_surface_mapper

    iso_surface_actor = vtkActor(mapper=iso_surface_mapper)
    iso_surface_actor.property.opacity = 0.5
    iso_surface_actor.property.diffuse_color = iso_surface_color

    sphere = vtkSphereSource(center=seeds.center, radius=seeds.radius,
                             theta_resolution=20, phi_resolution=11)
    sphere_mapper = vtkPolyDataMapper()
    sphere >> sphere_mapper

    sphere_actor = vtkActor(mapper=sphere_mapper)
    sphere_actor.property.opacity = 1.0
    sphere_actor.property.specular = 0.4
    sphere_actor.property.specular_power = 80
    sphere_actor.property.diffuse_color = sphere_color

    renderer.AddActor(tubes_actor)
    renderer.AddActor(sphere_actor)
    renderer.AddActor(iso_surface_actor)

    sp = SliderProperties()
    sp.title_text = 'X'
    sp.range['minimum_value'] = bounds[0]
    sp.range['maximum_value'] = bounds[1]
    sp.range['value'] = seeds.center[0]
    sp.position['point1'] = (0.1, 0.08)
    sp.position['point2'] = (0.3, 0.08)
    widget_x = make_slider_widget(sp, renderer, render_window_interactor)
    cb_x = SliderCallback(0, seeds, sphere, delay)
    widget_x.AddObserver(vtkCommand.InteractionEvent, cb_x)

    sp.title_text = 'Y'
    sp.range['minimum_value'] = bounds[2]
    sp.range['maximum_value'] = bounds[3]
    sp.range['value'] = seeds.center[1]
    sp.position['point1'] = (0.4, 0.08)
    sp.position['point2'] = (0.6, 0.08)
    widget_y = make_slider_widget(sp, renderer, render_window_interactor)
    cb_y = SliderCallback(1, seeds, sphere, delay)
    widget_y.AddObserver(vtkCommand.InteractionEvent, cb_y)

    sp.title_text = 'Z'
    sp.range['minimum_value'] = bounds[4]
    sp.range['maximum_value'] = bounds[5]
    sp.range['value'] = seeds.center[2]
    sp.position['point1'] = (0.7, 0.08)
    sp.position['point2'] = (0.9, 0.08)
    widget_z = make_slider_widget(sp, renderer, render_window_interactor)
    cb_z = SliderCallback(2, seeds, sphere, delay)
    widget_z.AddObserver(vtkCommand.InteractionEvent, cb_z)

    render_window.Render()

    renderer.active_camera.Azimuth(20.0)
    renderer.active_camera.Elevation(10.0)
    renderer.active_camera.Dolly(1.25)
    renderer.ResetCameraClippingRange()

    render_window_interactor.Start()


class SliderProperties:
    dimensions = {
        'tube_width': 0.01,
        'slider_length': 0.02, 'slider_width': 0.05,
        'end_cap_length': 0.015, 'end_cap_width': 0.05,
        'title_height': 0.03, 'label_height': 0.025,
    }
    colors = {
        'title_color': 'AliceBlue', 'label_color': 'AliceBlue', 'slider_color': 'BurlyWood',
        'selected_color': 'Lime', 'bar_color': 'Black', 'bar_ends_color': 'Indigo',
        'value_color': 'DarkSlateGray'
    }
    range = {'minimum_value': 0.0, 'maximum_value': 1.0, 'value': 1.0}
    title_text = '',
    position = {'point1': (0.1, 0.05), 'point2': (0.2, 0.05)}


def make_slider_widget(slider_properties, renderer, interactor):
    """
    Make a slider widget.
    :param slider_properties: range, title name, dimensions, colors, and position.
    :param renderer: The renderer.
    :param interactor: The interactor.
    :return: The slider widget.
    """
    colors = vtkNamedColors()

    slider_rep = vtkSliderRepresentation2D(minimum_value=slider_properties.range['minimum_value'],
                                           maximum_value=slider_properties.range['maximum_value'],
                                           value=slider_properties.range['value'],
                                           title_text=slider_properties.title_text,
                                           tube_width=slider_properties.dimensions['tube_width'],
                                           slider_length=slider_properties.dimensions['slider_length'],
                                           slider_width=slider_properties.dimensions['slider_width'],
                                           end_cap_length=slider_properties.dimensions['end_cap_length'],
                                           end_cap_width=slider_properties.dimensions['end_cap_width'],
                                           title_height=slider_properties.dimensions['title_height'],
                                           label_height=slider_properties.dimensions['label_height'],
                                           )

    # Set the color properties.
    slider_rep.title_property.color = colors.GetColor3d(slider_properties.colors['title_color'])
    slider_rep.label_property.color = colors.GetColor3d(slider_properties.colors['label_color'])
    slider_rep.tube_property.color = colors.GetColor3d(slider_properties.colors['bar_color'])
    slider_rep.cap_property.color = colors.GetColor3d(slider_properties.colors['bar_ends_color'])
    slider_rep.slider_property.color = colors.GetColor3d(slider_properties.colors['slider_color'])
    slider_rep.selected_property.color = colors.GetColor3d(slider_properties.colors['selected_color'])

    # Set the position.
    slider_rep.point1_coordinate.coordinate_system = Coordinate.CoordinateSystem.VTK_NORMALIZED_VIEWPORT
    slider_rep.point1_coordinate.value = slider_properties.position['point1']
    slider_rep.point2_coordinate.coordinate_system = Coordinate.CoordinateSystem.VTK_NORMALIZED_VIEWPORT
    slider_rep.point2_coordinate.value = slider_properties.position['point2']

    slider_rep.renderer = renderer

    widget = vtkSliderWidget(representation=slider_rep, interactor=interactor, enabled=True)
    widget.SetNumberOfAnimationSteps(10)
    widget.SetAnimationModeToAnimate()

    return widget


class SliderCallback:
    def __init__(self, axis, seeds, sphere, delay=500):
        """
        The delay is used to make the animation between selections consistent.

        :paran: axis: The axis [0..3) corresponding to X, Y, Z.
        :param seeds: The vtkPointSource seeds.
        :param sphere: The vtkSphereSource.
        :param delay: Delay in milliseconds.
        """
        self.axis = axis
        self.seeds = seeds
        self.sphere = sphere
        self.delay = delay

    def __call__(self, caller, ev):
        slider_widget = caller
        value = slider_widget.representation.value
        center = list(self.seeds.center)
        center[self.axis] = value
        self.seeds.center = center
        self.sphere.center = center
        # Sleep for delay miliseconds.
        time.sleep(self.delay / 1000)


@dataclass(frozen=True)
class Coordinate:
    @dataclass(frozen=True)
    class CoordinateSystem:
        VTK_DISPLAY: int = 0
        VTK_NORMALIZED_DISPLAY: int = 1
        VTK_VIEWPORT: int = 2
        VTK_NORMALIZED_VIEWPORT: int = 3
        VTK_VIEW: int = 4
        VTK_POSE: int = 5
        VTK_WORLD: int = 6
        VTK_USERDEFINED: int = 7


if __name__ == '__main__':
    main()