From bc2e126075a295d1dc9f5881b92ee731b7fd25c2 Mon Sep 17 00:00:00 2001 From: Guilhem Lavaux Date: Fri, 6 Mar 2020 08:23:55 +0100 Subject: [PATCH] Add support for multiple attributes in CTPV --- python/cosmotool/ctpv.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/cosmotool/ctpv.py b/python/cosmotool/ctpv.py index 22a9d22..4d90f1a 100644 --- a/python/cosmotool/ctpv.py +++ b/python/cosmotool/ctpv.py @@ -87,7 +87,7 @@ class ProgrammableParticleHistoryLoad(object): class ProgrammableDensityLoad(object): @staticmethod - def main_script(source, density, extents=None, aname="default"): + def main_script(source, density, extents=None, aname="default", extra_arrays={}): import vtk from vtk.util import numpy_support @@ -116,6 +116,11 @@ class ProgrammableDensityLoad(object): arr = numpy_support.numpy_to_vtk(density.transpose().astype(np.float64).ravel(), deep=1) arr.SetName(aname) ido.GetPointData().AddArray(arr) + + for k in extra_arrays.keys(): + arr = numpy_support.numpy_to_vtk(extra_arrays[k].transpose().astype(np.float64).ravel(), deep=1) + arr.SetName(k) + ido.GetPointData().AddArray(arr) @staticmethod def request_information(source, density=None, dims=None):