From 8070594b4648b4f93ce9c16eec36e95ee15548e7 Mon Sep 17 00:00:00 2001 From: Guilhem Lavaux Date: Fri, 4 Nov 2016 18:25:08 +0100 Subject: [PATCH] Added support for vector field in numpy -> paraview conversion --- python/cosmotool/ctpv.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/python/cosmotool/ctpv.py b/python/cosmotool/ctpv.py index e9bbd7b..471eec9 100644 --- a/python/cosmotool/ctpv.py +++ b/python/cosmotool/ctpv.py @@ -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)