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     inputFileName_removeSpaces=inputFileName.replace(" ", "")
  41     field.writeVTK(inputFileName_removeSpaces)
  42
  43     numpy_array = Extract_VTK_data_over_line_to_numpyArray(inputFileName_removeSpaces+"_0.vtu", point1, point2, resolution)
  44
  45     os.remove(inputFileName_removeSpaces+"_0.vtu")
  46     return numpy_array
  47
  48 def Extract_field_data_over_line_to_txt_file(field, point1, point2, resolution, outputFileName):
  49
  50     numpy_array = Extract_field_data_over_line_to_numpyArray(field, point1, point2, resolution)
  51
  52     np.savetxt(outputFileName, numpy_array, delimiter=" ")
  53
  54 def Slice_VTK_data_to_numpyArray(inputFileName,
  55                                  point, normal,
  56                                  resolution
  57                                            ):
  58     reader = vtk.vtkXMLUnstructuredGridReader()
  59     reader.SetFileName(inputFileName)
  60     reader.Update()
  61
  62     plane = vtk.vtkPlane()
  63     plane.SetOrigin(point)
  64     plane.SetNormal(normal)
  65
  66     cutter = vtk.vtkCutter()
  67     cutter.SetCutFunction(plane)
  68     cutter.SetInputConnection(reader.GetOutputPort())
  69     cutter.Update()
  70
  71     vtkarray = cutter.GetOutput().GetPointData().GetArray(0)
  72     numpy_array = npvtk.vtk_to_numpy(vtkarray)
  73
  74     return numpy_array
  75
  76
  77 def Slice_VTK_data_to_txt_file(inputFileName, outputFileName,
  78                                            point, normal,
  79                                            resolution
  80                                            ):
  81     numpy_array =   Slice_VTK_data_to_numpyArray(inputFileName, point, normal, resolution )
  82
  83     np.savetxt(outputFileName, numpy_array, delimiter=" ")
  84
  85
  86 def Slice_field_data_to_numpyArray(field,
  87                                    point, normal,
  88                                    resolution
  89                                    ):
  90     inputFileName = field.getName()
  91     field.writeVTK(inputFileName)
  92
  93     numpy_array = Slice_VTK_data_to_numpyArray(inputFileName+"_0.vtu", point, normal, resolution)
  94
  95     os.remove(inputFileName+"_0.vtu")
  96     return numpy_array
  97
  98 def Slice_field_data_to_txt_file(field, outputFileName,
  99                                         point, normal,
 100                                         resolution):
 101     numpy_array = Slice_field_data_to_numpyArray(field, point, normal, resolution)
 102
 103     np.savetxt(outputFileName, numpy_array, delimiter=" ")
 104
 105 def Slice_VTK_data_to_VTK(inputFileName,
 106                              outputFileName,
 107                                  point, normal,
 108                                  resolution
 109                                            ):
 110     reader = vtk.vtkXMLUnstructuredGridReader()
 111     reader.SetFileName(inputFileName)
 112     reader.Update()
 113
 114     plane = vtk.vtkPlane()
 115     plane.SetOrigin(point)
 116     plane.SetNormal(normal)
 117
 118     cutter = vtk.vtkCutter()
 119     cutter.SetCutFunction(plane)
 120     cutter.SetInputConnection(reader.GetOutputPort())
 121     cutter.Update()
 122
 123     #Convert tht polydata structure générated by cutter into unstructured grid by triangulation
 124     triFilter = vtk.vtkDataSetTriangleFilter()
 125     triFilter.SetInputConnection(cutter.GetOutputPort())
 126     triFilter.Update()
 127
 128     writer = vtk.vtkXMLUnstructuredGridWriter()
 129     writer.SetInputData(triFilter.GetOutput())
 130     writer.SetFileName(outputFileName)
 131     writer.Write()
 132
 133 def Clip_VTK_data_to_VTK(inputFileName,
 134                              outputFileName,
 135                                  point, normal,
 136                                  resolution
 137                                            ):
 138     reader = vtk.vtkXMLUnstructuredGridReader()
 139     reader.SetFileName(inputFileName)
 140     reader.Update()
 141
 142     plane = vtk.vtkPlane()
 143     plane.SetOrigin(point)
 144     plane.SetNormal(normal)
 145
 146     clipper = vtk.vtkClipDataSet()
 147     clipper.SetClipFunction(plane)
 148     clipper.SetInputConnection(reader.GetOutputPort())
 149     clipper.Update()
 150
 151     writer = vtk.vtkXMLUnstructuredGridWriter()
 152     writer.SetInputData(clipper.GetOutput())
 153     writer.SetFileName(outputFileName)
 154     writer.Write()
 155
 156 def Save_VTK_data_to_picture_file(inputFileName, field_name,
 157                              node_or_cell, outputFileName
 158                              ):
 159     reader = vtk.vtkXMLUnstructuredGridReader()
 160     reader.SetFileName(inputFileName)
 161     reader.Update()
 162
 163     if node_or_cell== 'CELLS':
 164         reader.CellArrayStatus = [field_name]
 165         reader.GetOutput().GetCellData().SetActiveScalars(field_name)
 166     elif node_or_cell== 'NODES':
 167         reader.PointArrayStatus = [field_name]
 168         reader.GetOutput().GetPointData().SetActiveScalars(field_name)
 169     else:
 170         raise ValueError("unknown type : should be CELLS or NODES")
 171
 172 #-------------------------------------------------------------------------------
 173     mapper = vtk.vtkDataSetMapper()
 174     mapper.SetInputConnection(reader.GetOutputPort())
 175     actor = vtk.vtkActor()
 176     actor.SetMapper(mapper)
 177
 178     scalarBar=vtk.vtkScalarBarActor()
 179     scalarBar.SetLookupTable(mapper.GetLookupTable())
 180     scalarBar.SetTitle(field_name)
 181
 182     mapper.SetScalarRange(reader.GetOutput().GetScalarRange())
 183
 184     ren = vtk.vtkRenderer()
 185     renWin = vtk.vtkRenderWindow()
 186     renWin.AddRenderer(ren)
 187
 188     ren.AddViewProp(actor)
 189     ren.AddActor2D(scalarBar);
 190
 191     renWin.Render()
 192
 193     image = vtk.vtkWindowToImageFilter()
 194     image.SetInput(renWin)
 195     image.ReadFrontBufferOff()
 196     image.Update()
 197     writer = vtk.vtkPNGWriter()
 198     writer.SetFileName(outputFileName+".png")
 199     writer.SetInputConnection(image.GetOutputPort())
 200     writer.Write()
 201