Vtk output depending on Format

src/darsia/__init__.py

@@ -27,6 +27,7 @@
 from darsia.utils.linear_solvers.mg import *
 from darsia.utils.andersonacceleration import *
 from darsia.utils.dtype import *
+from darsia.utils.formats import *
 from darsia.utils.grid import *
 from darsia.utils.fv import *
 from darsia.utils.kernels import *

src/darsia/image/image.py

@@ -1637,7 +1637,16 @@ def to_vtk(self, path: Union[str, Path], name: Optional[str] = None) -> None:
             name = self.name
         if name is None:
             name = "data"
-        darsia.plotting.to_vtk(path, [(name, self)])
+        darsia.plotting.to_vtk(
+            path,
+            [
+                (
+                    name,
+                    self,
+                    darsia.Format.SCALAR if self.scalar else darsia.Format.TENSOR,
+                )
+            ],
+        )
 
     # ! ---- Auxiliary routines
 

src/darsia/measure/wasserstein.py

@@ -1837,17 +1837,17 @@ def wasserstein_distance_to_vtk(
 
     """
     data = [
-        (key, info[key])
-        for key in [
-            "src",
-            "dst",
-            "mass_diff",
-            "flux",
-            "weighted_flux",
-            "pressure",
-            "transport_density",
-            "weight",
-            "weight_inv",
+        (key, info[key], format)
+        for key, format in [
+            ("src", darsia.Format.SCALAR),
+            ("dst", darsia.Format.SCALAR),
+            ("mass_diff", darsia.Format.SCALAR),
+            ("flux", darsia.Format.VECTOR),
+            ("weighted_flux", darsia.Format.VECTOR),
+            ("pressure", darsia.Format.SCALAR),
+            ("transport_density", darsia.Format.SCALAR),
+            ("weight", darsia.Format.TENSOR),
+            ("weight_inv", darsia.Format.TENSOR),
         ]
     ]
     darsia.plotting.to_vtk(path, data)

src/darsia/utils/formats.py

@@ -0,0 +1,12 @@
+"""Introduce concept of formats: Scalar, Vector, Tensor through enumeration."""
+
+# Introduce concept of formats: Scalar, Vector, Tensor through enumeration
+from enum import Enum
+
+
+class Format(Enum):
+    """Enumeration of formats for data."""
+
+    SCALAR = 0
+    VECTOR = 1
+    TENSOR = 2

src/darsia/utils/plotting.py

@@ -115,15 +115,15 @@ def plot_2d_wasserstein_distance(
 
 def to_vtk(
     path: Union[str, Path],
-    data: list[tuple[Union[darsia.Image, np.ndarray], str]],
+    data: list[tuple[Union[darsia.Image, np.ndarray, darsia.Format], str]],
 ) -> None:
     """Write data to a VTK file.
 
     Args:
         path (Union[str, Path]): path to the VTK file
-        data (list[tuple[Union[darsia.Image, np.ndarray], str]]): data to write, includes
-            the data and the name of the data. Require at least one data point to be an
-            image.
+        data (list[tuple[Union[darsia.Image, np.ndarray, darsia.Format], str]]): data to
+            write, includes the data and the name of the data. Require at least one data
+            point to be an image.
 
     NOTE: Requires pyevtk to be installed.
 
@@ -134,7 +134,7 @@ def to_vtk(
         # Check whether the data contains at least one image, and pick the first one
         image = None
         for d in data:
-            name, img = d
+            name, img, format = d
             if isinstance(img, darsia.Image):
                 image = img
                 break
@@ -179,7 +179,7 @@ def conditional_flip(x, cond):
         # Convert cell data to right format
         cellData = {}
         for d in data:
-            name, img = d
+            name, img, format = d
             if isinstance(img, darsia.Image):
                 img = img.img.copy()
             assert isinstance(img, np.ndarray), "Data must be of type np.ndarray."
@@ -219,9 +219,19 @@ def conditional_flip(x, cond):
                 if image.space_dim == 1:
                     img = tuple(img)
                 elif image.space_dim == 2:
-                    img = (-img[1], img[0], img[2])
+                    if format == darsia.Format.VECTOR:
+                        img = (-img[1], img[0], img[2])
+                    elif format == darsia.Format.TENSOR:
+                        img = (img[1], img[0], img[2])
+                    else:
+                        raise ValueError("Format not supported.")
                 elif image.space_dim == 3:
-                    img = (-img[1], img[2], img[0])
+                    if format == darsia.Format.VECTOR:
+                        img = (-img[1], img[2], img[0])
+                    elif format == darsia.Format.TENSOR:
+                        img = (img[1], img[2], img[0])
+                    else:
+                        raise ValueError("Format not supported.")
             cellData[name] = img
 
         # Make directory if necessary