Streamlines


In [1]:
import k3d
import vtk
from vtk.util import numpy_support
import numpy as np

In [2]:
plot = k3d.plot(camera_auto_fit=False)
plot.display()



In [3]:
reader = vtk.vtkSTLReader()
reader.SetFileName('./cfd/c0006.stl')
reader.Update()
geometry_wss = reader.GetOutput()

In [4]:
reader = vtk.vtkXMLImageDataReader()
reader.SetFileName('./cfd/000005595.vti')
reader.Update()
vti = reader.GetOutput()
bounds = vti.GetBounds()
shape = vti.GetDimensions()
spacing = vti.GetSpacing()
seeds_points = vtk.vtkPoints()
velocity = numpy_support.vtk_to_numpy(vti.GetPointData().GetArray('v [m/s]'))
velocity = velocity.flatten().reshape(shape[2], shape[1], shape[0], 3)

speed = np.linalg.norm(velocity, axis=3)    
speed[np.isnan(speed)] = 0.0

In [5]:
points = vtk.vtkPoints()
points.SetNumberOfPoints(vti.GetNumberOfPoints())

for i in range(points.GetNumberOfPoints()):
    points.SetPoint(i, vti.GetPoint(i))
        
grid = vtk.vtkStructuredGrid()
grid.SetDimensions(shape)
grid.SetPoints(points)
grid.GetPointData().SetVectors(
    vti.GetPointData().GetVectors('v [m/s]')
)

seeds = vtk.vtkPolyData()
points = vtk.vtkPoints()

for (z,x) in zip(*np.where(speed[:,0,:] > 0.0)):
    y = 0
    
    if x % 4 == 0 and z % 4 == 0:
        points.InsertNextPoint((x * spacing[0] + bounds[0], 
                                y * spacing[1] + bounds[2], 
                                z * spacing[2] + bounds[4]))

seeds.SetPoints(points)

In [6]:
streamer = vtk.vtkStreamTracer()
streamer.SetInputData(grid)
streamer.SetSourceData(seeds)
streamer.SetIntegrationDirectionToForward()
streamer.SetComputeVorticity(0)
streamer.SetIntegrator(vtk.vtkRungeKutta4())
streamer.Update()

streamline = streamer.GetOutput()

In [8]:
streamlines_points = numpy_support.vtk_to_numpy(streamline.GetPoints().GetData())
streamlines_velocity = numpy_support.vtk_to_numpy(streamline.GetPointData().GetArray('v [m/s]'))
streamlines_speed = np.linalg.norm(streamlines_velocity, axis=1)


C:\Tools\Anaconda\lib\site-packages\numpy\linalg\linalg.py:2286: RuntimeWarning: overflow encountered in multiply
  s = (x.conj() * x).real

In [9]:
vtkLines = streamline.GetLines()
vtkLines.InitTraversal()
point_list = vtk.vtkIdList()

lines = []
lines_attributes = []

while vtkLines.GetNextCell(point_list):
    start_id = point_list.GetId(0)
    end_id = point_list.GetId(point_list.GetNumberOfIds() - 1)
    l= []
    v= []
    
    for i in range(start_id, end_id):
        l.append(streamlines_points[i])
        v.append(streamlines_speed[i])
    
    lines.append(np.array(l))
    lines_attributes.append(np.array(v))

In [10]:
cm = k3d.matplotlib_color_maps.Inferno

for l, a in zip(lines, lines_attributes):
    plot += k3d.line(l, attribute=a, width=0.0001, 
                     color_map=cm, color_range=[0,0.5], 
                     shader='mesh', 
                     compression_level=9)

In [11]:
plot += k3d.vtk_poly_data(geometry_wss, opacity=0.1, wireframe=True, color=0xaaaaaa)

In [ ]: