FireFlow

#!/usr/bin/env python3

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.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.vtkRenderingCore import (
    vtkActor,
    vtkPolyDataMapper,
    vtkRenderer,
    vtkRenderWindow,
    vtkRenderWindowInteractor
)


# !/usr/bin/env python3


def get_program_parameters():
    import argparse
    description = 'Fire Flow.'
    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.')
    args = parser.parse_args()
    return args.geometry, args.velocity


def main():
    geometry_fn, velocity_fn = 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='FireFlow')
    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)

    # 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)

    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()


if __name__ == '__main__':
    main()