Merge pull request #32 from anlavandier/master

Enhancement: extending `pyevtk` to allow for usage of Parallel VTK files

Author: renefritze

examples/rectilinear.py

@@ -30,9 +30,13 @@
 # * This example shows how to export a rectilinear grid.       *
 # **************************************************************
 
-from pyevtk.hl import gridToVTK
+from pyevtk.hl import gridToVTK, writeParallelVTKGrid
 import numpy as np
 
+# ==================
+# Serial example
+# ==================
+
 # Dimensions
 nx, ny, nz = 6, 6, 2
 lx, ly, lz = 1.0, 1.0, 1.0
@@ -58,3 +62,26 @@
     cellData={"pressure": pressure},
     pointData={"temp": temp},
 )
+
+
+# ==================
+# Parallel example
+# ==================
+
+# Dimensions
+x1 = np.linspace(0, 1, 10)
+x2 = np.linspace(1, 2, 20)
+
+y = np.linspace(0, 1, 10)
+z = np.linspace(0, 1, 10)
+
+gridToVTK("test.0", x1, y, z, start=(0, 0, 0))
+gridToVTK("test.1", x2, y, z, start=(9, 0, 0))
+
+writeParallelVTKGrid(
+    "test_full",
+    coordsData=((29, 10, 10), x.dtype),
+    starts=[(0, 0, 0), (9, 0, 0)],
+    ends=[(9, 9, 9), (28, 9, 9)],
+    sources=["test.0.vtr", "test.1.vtr"],
+)

examples/structured.py

@@ -30,10 +30,14 @@
 # * This example shows how to export a structured grid.        *
 # **************************************************************
 
-from pyevtk.hl import gridToVTK
+from pyevtk.hl import gridToVTK, writeParallelVTKGrid
 import numpy as np
 import random as rnd
 
+# ===================
+# Serial Example
+# ===================
+
 # Dimensions
 nx, ny, nz = 6, 6, 2
 lx, ly, lz = 1.0, 1.0, 1.0
@@ -72,3 +76,55 @@
     cellData={"pressure": pressure},
     pointData={"temp": temp},
 )
+
+
+# ===================
+# Parallel example
+# ===================
+
+# Dimensions
+x1 = np.linspace(0, 1, 20)
+
+x2_1 = np.linspace(0, 0.5, 10)
+x2_2 = np.linspace(0.5, 1, 10)
+
+x3 = np.linspace(0, 2, 30)
+
+XX1_1, XX2_1, XX3_1 = np.meshgrid(x1, x2_1, x3, indexing="ij")
+XX1_2, XX2_2, XX3_2 = np.meshgrid(x1, x2_2, x3, indexing="ij")
+
+pi = np.pi
+sphere = lambda R, Th, Ph: (
+    R * np.cos(pi * Ph) * np.sin(pi * Th),
+    R * np.sin(pi * Ph) * np.sin(pi * Th),
+    R * np.cos(pi * Th),
+)
+
+# First Half sphere
+gridToVTK(
+    "sphere.0",
+    *sphere(XX1_1, XX2_1, XX3_1),
+    start=(0, 0, 0),
+    pointData={"R": XX1_1, "Theta": XX2_1, "Phi": XX3_1}
+)
+# Second Half sphere
+gridToVTK(
+    "sphere.1",
+    *sphere(XX1_2, XX2_2, XX3_2),
+    start=(0, 9, 0),
+    pointData={"R": XX1_2, "Theta": XX2_2, "Phi": XX3_2}
+)
+
+# Write parallel file
+writeParallelVTKGrid(
+    "sphere_full",
+    coordsData=((20, 19, 30), XX1_1.dtype),
+    starts=[(0, 0, 0), (0, 9, 0)],
+    ends=[(19, 9, 29), (19, 18, 29)],
+    sources=["sphere.0.vts", "sphere.1.vts"],
+    pointData={
+        "R": (XX1_1.dtype, 1),
+        "Theta": (XX2_1.dtype, 1),
+        "Phi": (XX3_1.dtype, 1),
+    },
+)

pyevtk/hl.py

@@ -27,10 +27,15 @@
 import numpy as np
 from .vtk import (
     VtkFile,
+    VtkParallelFile,
     VtkUnstructuredGrid,
     VtkImageData,
     VtkRectilinearGrid,
     VtkStructuredGrid,
+    VtkPImageData,
+    VtkPRectilinearGrid,
+    VtkPStructuredGrid,
+    VtkUnstructuredGrid,
     VtkVertex,
     VtkLine,
     VtkPolyLine,
@@ -81,6 +86,33 @@ def _addDataToFile(vtkFile, cellData, pointData, fieldData=None):
         vtkFile.closeData("Field")
 
 
+def _addDataToParallelFile(vtkParallelFile, cellData, pointData):
+    assert isinstance(vtkParallelFile, VtkParallelFile)
+    # Point data
+    if pointData:
+        keys = list(pointData.keys())
+        # find first scalar and vector data key to set it as attribute
+        scalars = next((key for key in keys if pointData[key][1] == 1), None)
+        vectors = next((key for key in keys if pointData[key][1] == 3), None)
+        vtkParallelFile.openData("PPoint", scalars=scalars, vectors=vectors)
+        for key in keys:
+            dtype, ncomp = pointData[key]
+            vtkParallelFile.addHeader(key, dtype=dtype, ncomp=ncomp)
+        vtkParallelFile.closeData("PPoint")
+
+    # Cell data
+    if cellData:
+        keys = list(cellData.keys())
+        # find first scalar and vector data key to set it as attribute
+        scalars = next((key for key in keys if cellData[key][1] == 1), None)
+        vectors = next((key for key in keys if cellData[key][1] == 3), None)
+        vtkParallelFile.openData("PCell", scalars=scalars, vectors=vectors)
+        for key in keys:
+            dtype, ncomp = cellData[key]
+            vtkParallelFile.addHeader(key, dtype=dtype, ncomp=ncomp)
+        vtkParallelFile.closeData("PCell")
+
+
 def _appendDataToFile(vtkFile, cellData, pointData, fieldData=None):
     # Append data to binary section
     if pointData is not None:
@@ -112,6 +144,7 @@ def imageToVTK(
     cellData=None,
     pointData=None,
     fieldData=None,
+    start=(0, 0, 0),
 ):
     """
     Export data values as a rectangular image.
@@ -120,6 +153,10 @@ def imageToVTK(
     ----------
     path : str
         name of the file without extension where data should be saved.
+    start : tuple, optional
+        start of the coordinates.
+        Used in the distributed context where each process
+        writes its own vtk file. Default is (0, 0, 0).
     origin : tuple, optional
         grid origin.
         The default is (0.0, 0.0, 0.0).
@@ -138,7 +175,7 @@ def imageToVTK(
         Arrays must have same dimension in each direction and
         they should be equal to the dimensions of the cell data plus one and
         can contain scalar data ([n+1,n+1,n+1]) or
-        +1,n+1,n+1],[n+1,n+1,n+1],[n+1,n+1,n+1]).
+        ([n+1,n+1,n+1],[n+1,n+1,n+1],[n+1,n+1,n+1]).
         The default is None.
     fieldData : dict, optional
         dictionary with variables associated with the field.
@@ -157,7 +194,6 @@ def imageToVTK(
     assert cellData is not None or pointData is not None
 
     # Extract dimensions
-    start = (0, 0, 0)
     end = None
     if cellData is not None:
         keys = list(cellData.keys())
@@ -188,9 +224,11 @@ def imageToVTK(
 
 
 # ==============================================================================
-def gridToVTK(path, x, y, z, cellData=None, pointData=None, fieldData=None):
+def gridToVTK(
+    path, x, y, z, cellData=None, pointData=None, fieldData=None, start=(0, 0, 0)
+):
     """
-    Write data values as a rectilinear or rectangular grid.
+    Write data values as a rectilinear or structured grid.
 
     Parameters
     ----------
@@ -202,6 +240,10 @@ def gridToVTK(path, x, y, z, cellData=None, pointData=None, fieldData=None):
         y coordinate axis.
     z : array-like
         z coordinate axis.
+    start : tuple, optional
+        start of the coordinates.
+        Used in the distributed context where each process
+        writes its own vtk file. Default is (0, 0, 0).
     cellData : dict, optional
         dictionary containing arrays with cell centered data.
         Keys should be the names of the data arrays.
@@ -216,12 +258,10 @@ def gridToVTK(path, x, y, z, cellData=None, pointData=None, fieldData=None):
     fieldData : dict, optional
         dictionary with variables associated with the field.
         Keys should be the names of the variable stored in each array.
-
     Returns
     -------
     str
         Full path to saved file.
-
     Notes
     -----
     coordinates of the nodes of the grid. They can be 1D or 3D depending if
@@ -235,8 +275,6 @@ def gridToVTK(path, x, y, z, cellData=None, pointData=None, fieldData=None):
     In both cases the arrays dimensions should be
     equal to the number of nodes of the grid.
     """
-    # Extract dimensions
-    start = (0, 0, 0)
     nx = ny = nz = 0
 
     if x.ndim == 1 and y.ndim == 1 and z.ndim == 1:
@@ -249,13 +287,22 @@ def gridToVTK(path, x, y, z, cellData=None, pointData=None, fieldData=None):
         isRect = False
         ftype = VtkStructuredGrid
     else:
-        assert False
-    end = (nx, ny, nz)
+        raise ValueError(
+            f"x, y and z should have ndim == 3 but they have ndim of {x.ndim}, {y.ndim}"
+            f" and {z.ndim} respectively"
+        )
 
+    # Write extent
+    end = (start[0] + nx, start[1] + ny, start[2] + nz)
+
+    # Open File
     w = VtkFile(path, ftype)
+
+    # Open Grid part
     w.openGrid(start=start, end=end)
     w.openPiece(start=start, end=end)
 
+    # Add coordinates description
     if isRect:
         w.openElement("Coordinates")
         w.addData("x_coordinates", x)
@@ -267,16 +314,101 @@ def gridToVTK(path, x, y, z, cellData=None, pointData=None, fieldData=None):
         w.addData("points", (x, y, z))
         w.closeElement("Points")
 
+    # Add data description
     _addDataToFile(w, cellData, pointData, fieldData)
+
+    # Close Grid part
     w.closePiece()
     w.closeGrid()
+
     # Write coordinates
     if isRect:
         w.appendData(x).appendData(y).appendData(z)
     else:
         w.appendData((x, y, z))
     # Write data
     _appendDataToFile(w, cellData, pointData, fieldData)
+
+    # Close file
+    w.save()
+
+    return w.getFileName()
+
+
+def writeParallelVTKGrid(
+    path, coordsData, starts, ends, sources, ghostlevel=0, cellData=None, pointData=None
+):
+    """
+    Writes a parallel vtk file from grid-like data:
+    VTKStructuredGrid or VTKRectilinearGrid
+
+    Parameters
+    ----------
+    path : str
+        name of the file without extension.
+    coordsData : tuple
+        2-tuple (shape, dtype) where shape is the
+        shape of the coordinates of the full mesh
+        and dtype is the dtype of the coordinates.
+    starts : list
+        list of 3-tuple representing where each source file starts
+        in each dimension
+    source : list
+        list of the relative paths of the source files where the actual data is found
+    ghostlevel : int, optional
+        Number of ghost-levels by which
+        the extents in the individual source files overlap.
+    pointData : dict
+        dictionnary containing the information about the arrays
+        containing node centered data.
+        Keys shoud be the names of the arrays.
+        Values are (dtype, number of components)
+    cellData :
+        dictionnary containing the information about the arrays
+        containing cell centered data.
+        Keys shoud be the names of the arrays.
+        Values are (dtype, number of components)
+    """
+    # Check that every source as a start and an end
+    assert len(starts) == len(ends) == len(sources)
+
+    # Get the extension + check that it's consistent accros all source files
+    common_ext = sources[0].split(".")[-1]
+    assert all(s.split(".")[-1] == common_ext for s in sources)
+
+    if common_ext == "vts":
+        ftype = VtkPStructuredGrid
+        is_Rect = False
+    elif common_ext == "vtr":
+        ftype = VtkPRectilinearGrid
+        is_Rect = True
+    else:
+        raise ValueError("This functions is meant to work only with ")
+
+    w = VtkParallelFile(path, ftype)
+    start = (0, 0, 0)
+    (s_x, s_y, s_z), dtype = coordsData
+    end = s_x - 1, s_y - 1, s_z - 1
+
+    w.openGrid(start=start, end=end, ghostlevel=ghostlevel)
+
+    _addDataToParallelFile(w, cellData=cellData, pointData=pointData)
+
+    if is_Rect:
+        w.openElement("PCoordinates")
+        w.addHeader("x_coordinates", dtype=dtype, ncomp=1)
+        w.addHeader("y_coordinates", dtype=dtype, ncomp=1)
+        w.addHeader("z_coordinates", dtype=dtype, ncomp=1)
+        w.closeElement("PCoordinates")
+    else:
+        w.openElement("PPoints")
+        w.addHeader("points", dtype=dtype, ncomp=3)
+        w.closeElement("PPoints")
+
+    for start_source, end_source, source in zip(starts, ends, sources):
+        w.addPiece(start_source, end_source, source)
+
+    w.closeGrid()
     w.save()
     return w.getFileName()