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,7 +33,10 @@ class ProgrammableDensityLoad(object):
import vtk
from vtk.util import numpy_support
Nx, Ny, Nz = density.shape
if len(density.shape) > 3:
_, Nx, Ny, Nz = density.shape
else:
Nx, Ny, Nz = density.shape
ido = source.GetOutput()
ido.SetDimensions(Nx, Ny, Nz)
@ -47,7 +50,12 @@ class ProgrammableDensityLoad(object):
ido.SetOrigin(*origin)
ido.SetSpacing(*spacing)
ido.SetExtent([0,Nx-1,0,Ny-1,0,Nz-1])
arr = numpy_support.numpy_to_vtk(density.transpose().astype(np.float64).ravel(), deep=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.SetName(aname)
ido.GetPointData().AddArray(arr)