Added support for vector field in numpy -> paraview conversion

This commit is contained in:
Guilhem Lavaux 2016-11-04 18:25:08 +01:00
parent aec35ded1c
commit 8070594b46

View File

@ -33,6 +33,9 @@ class ProgrammableDensityLoad(object):
import vtk import vtk
from vtk.util import numpy_support from vtk.util import numpy_support
if len(density.shape) > 3:
_, Nx, Ny, Nz = density.shape
else:
Nx, Ny, Nz = density.shape Nx, Ny, Nz = density.shape
ido = source.GetOutput() ido = source.GetOutput()
@ -47,6 +50,11 @@ class ProgrammableDensityLoad(object):
ido.SetOrigin(*origin) ido.SetOrigin(*origin)
ido.SetSpacing(*spacing) ido.SetSpacing(*spacing)
ido.SetExtent([0,Nx-1,0,Ny-1,0,Nz-1]) ido.SetExtent([0,Nx-1,0,Ny-1,0,Nz-1])
if len(density.shape) > 3 and density.shape[0] == 3:
N = Nx*Ny*Nz
density = density.transpose().astype(np.float64).reshape((N,3))
arr = numpy_support.numpy_to_vtk(density, deep=1)
else:
arr = numpy_support.numpy_to_vtk(density.transpose().astype(np.float64).ravel(), deep=1) arr = numpy_support.numpy_to_vtk(density.transpose().astype(np.float64).ravel(), deep=1)
arr.SetName(aname) arr.SetName(aname)
ido.GetPointData().AddArray(arr) ido.GetPointData().AddArray(arr)