TL: added vtk support for 3D VTR files

Author: tlunet

etc/environment-base.yml

@@ -10,6 +10,7 @@ dependencies:
   - sympy>=1.0
   - numba>=0.35
   - dill>=0.2.6
+  - vtk
   - pip
   - pip:
     - qmat>=0.1.8

pySDC/helpers/fieldsIO.py

@@ -43,7 +43,7 @@
 To use MPI collective writing, you need to call first the class methods :class:`Rectilinear.initMPI` (cf their docstring).
 Also, `Rectilinear.setHeader` **must be given the global grids coordinates**, wether the code is run in parallel or not.
 
-Also : this feature is only available with Python 3.11 or higher ...
+> ⚠️ Also : this module can only be imported with **Python 3.11 or higher** !
 """
 import os
 import numpy as np
@@ -640,7 +640,9 @@ def readField(self, idx):
         return t, field
 
 
-# Utility function used for testing
+# -----------------------------------------------------------------------------------------------
+# Utility functions used for testing
+# -----------------------------------------------------------------------------------------------
 def initGrid(nVar, gridSizes):
     dim = len(gridSizes)
     coords = [np.linspace(0, 1, num=n, endpoint=False) for n in gridSizes]

pySDC/helpers/vtk.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Helper functions to write and read fields from VTK files (to be used with Paraview or PyVista)
+"""
+import os
+import vtk
+from vtkmodules.util import numpy_support
+import numpy as np
+
+
+def writeToVTR(fileName: str, data, coords, varNames):
+    """
+    Write a data array containing variables from a 3D rectilinear grid into a VTR file.
+
+    Parameters
+    ----------
+    fileName : str
+        Name of the VTR file, can be with or without the .vtr extension.
+    data : np.4darray
+        Array containing all the variables with [nVar, nX, nY, nZ] shape.
+    coords : list[np.1darray]
+        Coordinates in each direction.
+    varNames : list[str]
+        Variable names.
+
+    Returns
+    -------
+    fileName : str
+        Name of the VTR file.
+    """
+    data = np.asarray(data)
+    nVar, *gridSizes = data.shape
+
+    assert len(gridSizes) == 3, "function can be used only for 3D grid data"
+    assert nVar == len(varNames), "varNames must have as many variable as data"
+    assert [len(np.ravel(coord)) for coord in coords] == gridSizes, "coordinate size incompatible with data shape"
+
+    nX, nY, nZ = gridSizes
+    vtr = vtk.vtkRectilinearGrid()
+    vtr.SetDimensions(nX, nY, nZ)
+
+    vect = lambda x: numpy_support.numpy_to_vtk(num_array=x, deep=True, array_type=vtk.VTK_FLOAT)
+    x, y, z = coords
+    vtr.SetXCoordinates(vect(x))
+    vtr.SetYCoordinates(vect(y))
+    vtr.SetZCoordinates(vect(z))
+
+    field = lambda u: numpy_support.numpy_to_vtk(num_array=u.ravel(order='F'), deep=True, array_type=vtk.VTK_FLOAT)
+    pointData = vtr.GetPointData()
+    for name, u in zip(varNames, data):
+        uVTK = field(u)
+        uVTK.SetName(name)
+        pointData.AddArray(uVTK)
+
+    writer = vtk.vtkXMLRectilinearGridWriter()
+    if not fileName.endswith(".vtr"):
+        fileName += ".vtr"
+    writer.SetFileName(fileName)
+    writer.SetInputData(vtr)
+    writer.Write()
+
+    return fileName
+
+
+def readFromVTR(fileName: str):
+    """
+    Read a VTR file into a numpy 4darray
+
+    Parameters
+    ----------
+    fileName : str
+        Name of the VTR file, can be with or without the .vtr extension.
+
+    Returns
+    -------
+    data : np.4darray
+        Array containing all the variables with [nVar, nX, nY, nZ] shape.
+    coords : list[np.1darray]
+        Coordinates in each direction.
+    varNames : list[str]
+        Variable names.
+    """
+    if not fileName.endswith(".vtr"):
+        fileName += ".vtr"
+    assert os.path.isfile(fileName), f"{fileName} does not exist"
+
+    reader = vtk.vtkXMLRectilinearGridReader()
+    reader.SetFileName(fileName)
+    reader.Update()
+
+    vtr = reader.GetOutput()
+    dims = vtr.GetDimensions()
+    assert len(dims) == 3, "can only read 3D data"
+
+    vect = lambda x: numpy_support.vtk_to_numpy(x)
+    coords = [vect(vtr.GetXCoordinates()), vect(vtr.GetYCoordinates()), vect(vtr.GetZCoordinates())]
+    pointData = vtr.GetPointData()
+    varNames = [pointData.GetArrayName(i) for i in range(pointData.GetNumberOfArrays())]
+    data = [numpy_support.vtk_to_numpy(pointData.GetArray(name)).reshape(dims, order="F") for name in varNames]
+    data = np.array(data)
+
+    return data, coords, varNames

pySDC/tests/test_helpers/test_vtk.py

@@ -0,0 +1,23 @@
+import pytest
+import numpy as np
+
+
+@pytest.mark.parametrize("nZ", [1, 5, 16])
+@pytest.mark.parametrize("nY", [1, 5, 16])
+@pytest.mark.parametrize("nX", [1, 5, 16])
+@pytest.mark.parametrize("nVar", [1, 2, 3])
+def testVTR(nVar, nX, nY, nZ):
+    from pySDC.helpers.vtk import writeToVTR, readFromVTR
+
+    data1 = np.random.rand(nVar, nX, nY, nZ)
+    coords1 = [np.sort(np.random.rand(n)) for n in [nX, nY, nZ]]
+    varNames1 = [f"var{i}" for i in range(nVar)]
+
+    data2, coords2, varNames2 = readFromVTR(writeToVTR("testVTR", data1, coords1, varNames1))
+
+    for i, (x1, x2) in enumerate(zip(coords1, coords2)):
+        print(x1, x2)
+        assert np.allclose(x1, x2), f"coordinate mismatch in dir. {i}"
+    assert varNames1 == varNames2, f"varNames mismatch"
+    assert data1.shape == data2.shape, f"data shape mismatch"
+    assert np.allclose(data1, data2), f"data values mismatch"