CDMATH/postprocessing/VTK_routines.py

   1 #!/usr/bin/env python
   2 # -*-coding:utf-8 -*-
   3
   4 import os
   5 import numpy as np
   6 import vtk
   7 from vtk.util import numpy_support as npvtk
   8 # do I need to kill the pipeline?
   9
  10 def Extract_VTK_data_over_line_to_numpyArray(inputFileName, point1, point2, resolution):
  11
  12     reader = vtk.vtkXMLUnstructuredGridReader()
  13     reader.SetFileName(inputFileName)
  14     reader.Update()
  15
  16     probeLine = vtk.vtkLineSource()
  17     probeLine.SetPoint1(point1)
  18     probeLine.SetPoint2(point2)
  19     probeLine.SetResolution(resolution)
  20
  21     probe = vtk.vtkProbeFilter()
  22     probe.SetInputConnection(probeLine.GetOutputPort())
  23     probe.SetSourceData(reader.GetOutput())
  24     probe.Update()
  25
  26     vtkarray = probe.GetOutput().GetPointData().GetArray(0) # or Slice1.GetCellData() # or Clip1.GetCellData()
  27     numpy_array = npvtk.vtk_to_numpy(vtkarray)
  28
  29     return numpy_array
  30
  31 def Extract_VTK_data_over_line_to_txt_file(inputFileName, outputFileName, point1, point2, resolution):
  32
  33     numpy_array = Extract_VTK_data_over_line_to_numpyArray(inputFileName, point1, point2, resolution)
  34
  35     np.savetxt(outputFileName, numpy_array, delimiter=" ")
  36
  37 def Extract_field_data_over_line_to_numpyArray(field, point1, point2, resolution):
  38
  39     inputFileName = field.getName()#os.getcwd()+field.get_name()
  40     field.writeVTK(inputFileName)
  41
  42     numpy_array = Extract_VTK_data_over_line_to_numpyArray(inputFileName+"_1.vtu", point1, point2, resolution)
  43
  44     os.remove(inputFileName+"_1.vtu")
  45     return numpy_array
  46
  47 def Extract_field_data_over_line_to_txt_file(field, point1, point2, resolution, outputFileName):
  48
  49     numpy_array = Extract_field_data_over_line_to_numpyArray(field, point1, point2, resolution)
  50
  51     np.savetxt(outputFileName, numpy_array, delimiter=" ")
  52
  53 def Slice_VTK_data_to_numpyArray(inputFileName,
  54                                  point, normal,
  55                                  resolution
  56                                            ):
  57     reader = vtk.vtkXMLUnstructuredGridReader()
  58     reader.SetFileName(inputFileName)
  59     reader.Update()
  60
  61     plane = vtk.vtkPlane()
  62     plane.SetOrigin(point)
  63     plane.SetNormal(normal)
  64
  65     cutter = vtk.vtkFiltersCorePython.vtkCutter()
  66     cutter.SetCutFunction(plane)
  67     cutter.SetInputConnection(reader.GetOutputPort())
  68     cutter.Update()
  69
  70     vtkarray = cutter.GetOutput().GetPointData().GetArray(0)
  71     numpy_array = npvtk.vtk_to_numpy(vtkarray)
  72
  73     return numpy_array
  74
  75
  76 def Slice_VTK_data_to_txt_file(inputFileName, outputFileName,
  77                                            point, normal,
  78                                            resolution
  79                                            ):
  80     numpy_array =   Slice_VTK_data_to_numpyArray(inputFileName, point, normal, resolution )
  81
  82     np.savetxt(outputFileName, numpy_array, delimiter=" ")
  83
  84
  85 def Slice_field_data_to_numpyArray(field,
  86                                    point, normal,
  87                                    resolution
  88                                    ):
  89     inputFileName = field.getName()
  90     field.writeVTK(inputFileName)
  91
  92     numpy_array = Slice_VTK_data_to_numpyArray(inputFileName+"_1.vtu", point, normal, resolution)
  93
  94     os.remove(inputFileName+"_1.vtu")
  95     return numpy_array
  96
  97 def Slice_field_data_to_txt_file(field, outputFileName,
  98                                         point, normal,
  99                                         resolution):
 100     numpy_array = Slice_field_data_to_numpyArray(field, point, normal, resolution)
 101
 102     np.savetxt(outputFileName, numpy_array, delimiter=" ")
 103
 104 def Slice_VTK_data_to_VTK(inputFileName,
 105                              outputFileName,
 106                                  point, normal,
 107                                  resolution
 108                                            ):
 109     reader = vtk.vtkXMLUnstructuredGridReader()
 110     reader.SetFileName(inputFileName)
 111     reader.Update()
 112
 113     plane = vtk.vtkPlane()
 114     plane.SetOrigin(point)
 115     plane.SetNormal(normal)
 116
 117     cutter = vtk.vtkFiltersCorePython.vtkCutter()
 118     cutter.SetCutFunction(plane)
 119     cutter.SetInputConnection(reader.GetOutputPort())
 120     cutter.Update()
 121
 122     #Convert tht polydata structure générated by cutter into unstructured grid by triangulation
 123     triFilter = vtk.vtkDataSetTriangleFilter()
 124     triFilter.SetInputConnection(cutter.GetOutputPort())
 125     triFilter.Update()
 126
 127     writer = vtk.vtkXMLUnstructuredGridWriter()
 128     writer.SetInputData(triFilter.GetOutput())
 129     writer.SetFileName(outputFileName)
 130     writer.Write()
 131
 132 def Clip_VTK_data_to_VTK(inputFileName,
 133                              outputFileName,
 134                                  point, normal,
 135                                  resolution
 136                                            ):
 137     reader = vtk.vtkXMLUnstructuredGridReader()
 138     reader.SetFileName(inputFileName)
 139     reader.Update()
 140
 141     plane = vtk.vtkPlane()
 142     plane.SetOrigin(point)
 143     plane.SetNormal(normal)
 144
 145     clipper = vtk.vtkClipDataSet()
 146     clipper.SetClipFunction(plane)
 147     clipper.SetInputConnection(reader.GetOutputPort())
 148     clipper.Update()
 149
 150     writer = vtk.vtkXMLUnstructuredGridWriter()
 151     writer.SetInputData(clipper.GetOutput())
 152     writer.SetFileName(outputFileName)
 153     writer.Write()
 154
 155 def Save_VTK_data_to_picture_file(inputFileName, field_name,
 156                              node_or_cell, outputFileName
 157                              ):
 158     reader = vtk.vtkXMLUnstructuredGridReader()
 159     reader.SetFileName(inputFileName)
 160     reader.Update()
 161
 162     if node_or_cell== 'CELLS':
 163         reader.CellArrayStatus = [field_name]
 164         reader.GetOutput().GetCellData().SetActiveScalars(field_name)
 165     elif node_or_cell== 'NODES':
 166         reader.PointArrayStatus = [field_name]
 167         reader.GetOutput().GetPointData().SetActiveScalars(field_name)
 168     else:
 169         raise ValueError("unknown type : should be CELLS or NODES")
 170
 171 #-------------------------------------------------------------------------------
 172     mapper = vtk.vtkDataSetMapper()
 173     mapper.SetInputConnection(reader.GetOutputPort())
 174     actor = vtk.vtkActor()
 175     actor.SetMapper(mapper)
 176
 177     scalarBar=vtk.vtkScalarBarActor()
 178     scalarBar.SetLookupTable(mapper.GetLookupTable())
 179     scalarBar.SetTitle(field_name)
 180
 181     mapper.SetScalarRange(reader.GetOutput().GetScalarRange())
 182
 183     ren = vtk.vtkRenderer()
 184     renWin = vtk.vtkRenderWindow()
 185     renWin.AddRenderer(ren)
 186
 187     ren.AddViewProp(actor)
 188     ren.AddActor2D(scalarBar);
 189
 190     renWin.Render()
 191
 192     image = vtk.vtkWindowToImageFilter()
 193     image.SetInput(renWin)
 194     image.ReadFrontBufferOff()
 195     image.Update()
 196     writer = vtk.vtkPNGWriter()
 197     writer.SetFileName(outputFileName+".png")
 198     writer.SetInputConnection(image.GetOutputPort())
 199     writer.Write()
 200