From 01859ef5a11881a5e556bfb48f9de561ccb226f8 Mon Sep 17 00:00:00 2001
From: Martin Bauer <martin.bauer@fau.de>
Date: Wed, 1 Nov 2017 20:21:56 +0100
Subject: [PATCH] VTK output for serial scenarios, now with vector field
 support

---
 vtk.py | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 vtk.py

diff --git a/vtk.py b/vtk.py
new file mode 100644
index 000000000..c256f0019
--- /dev/null
+++ b/vtk.py
@@ -0,0 +1,41 @@
+from pyevtk.vtk import VtkFile, VtkImageData
+from pyevtk.hl import _addDataToFile, _appendDataToFile
+
+
+def imageToVTK(path, origin=(0.0, 0.0, 0.0), spacing=(1.0, 1.0, 1.0), cellData=None, pointData=None):
+    """Patched version of same pyevtk function that also support vector data"""
+    assert cellData != None or pointData != None
+    # Extract dimensions
+    start = (0, 0, 0)
+    end = None
+    if cellData:
+        keys = list(cellData.keys())
+        data = cellData[keys[0]]
+        if hasattr(data, 'shape'):
+            end = data.shape[:-1]
+        elif data[0].ndim == 3 and data[1].ndim == 3 and data[0].ndim == 3:
+            keys = list(cellData.keys())
+            data = cellData[keys[0]]
+            end = data[0].shape
+    elif pointData:
+        keys = list(pointData.keys())
+        data = pointData[keys[0]]
+        if hasattr(data, 'shape'):
+            end = data.shape
+            end = (end[0] - 1, end[1] - 1, end[2] - 1)
+        # Added for vector support...
+        elif data[0].ndim == 3 and data[1].ndim == 3 and data[0].ndim == 3:
+            keys = list(pointData.keys())
+            data = pointData[keys[0]]
+            end = data[0].shape
+            end = (end[0] - 1, end[1] - 1, end[2] - 1)
+    # Write data to file
+    w = VtkFile(path, VtkImageData)
+    w.openGrid(start=start, end=end, origin=origin, spacing=spacing)
+    w.openPiece(start=start, end=end)
+    _addDataToFile(w, cellData, pointData)
+    w.closePiece()
+    w.closeGrid()
+    _appendDataToFile(w, cellData, pointData)
+    w.save()
+    return w.getFileName()
-- 
GitLab