[gui] Fluid visualization utilities (#7682)

Issue: #

## Brief Summary
This PR implemented several visualization utilities to help simulation
users to quickly visualize their Scalar/Vector fields.

- `contour()` under the `Canvas` class to plot a scalar field in GGUI
- `vector_field()` method under the `Canvas` class to plot a 2D vector field in GGUI
- `write_vtk()` method that can output a 2D/3D scalar field into a `.vtr` file for visualization
- `contour()` and `vector_field()` methods in `ti.GUI` just for completeness; the user should be aware that `ti.GUI` will be deprecated in the future in favor of GGUI.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

python/taichi/tools/__init__.py

@@ -9,3 +9,4 @@
 from taichi.tools.image import *
 from taichi.tools.np2ply import *
 from taichi.tools.video import *
+from taichi.tools.vtk import *

python/taichi/tools/vtk.py

@@ -0,0 +1,31 @@
+import numpy as np
+
+
+def write_vtk(scalar_field, filename):
+    try:
+        from pyevtk.hl import \
+            gridToVTK  # pylint: disable=import-outside-toplevel
+    except ImportError:
+        raise RuntimeError('Failed to import pyevtk. Please install it via /\
+        `pip install pyevtk` first. ')
+
+    scalar_field_np = scalar_field.to_numpy()
+    field_shape = scalar_field_np.shape
+    dimensions = len(field_shape)
+
+    if dimensions not in (2, 3):
+        raise ValueError("The input field must be a 2D or 3D scalar field.")
+
+    if dimensions == 2:
+        scalar_field_np = scalar_field_np[np.newaxis, :, :]
+        zcoords = np.array([0, 1])
+    elif dimensions == 3:
+        zcoords = np.arange(0, field_shape[2])
+    gridToVTK(filename,
+              x=np.arange(0, field_shape[0]),
+              y=np.arange(0, field_shape[1]),
+              z=zcoords,
+              cellData={filename: scalar_field_np})
+
+
+__all__ = ['write_vtk']

python/taichi/ui/canvas.py

@@ -42,6 +42,46 @@ def set_image(self, img):
             info = get_field_info(staging_img)
             self.canvas.set_image(info)
 
+    def contour(self, scalar_field, cmap_name='plasma', normalize=False):
+        """Plot a contour view of a scalar field.
+
+        The input scalar_field will be converted to a Numpy array first, and then plotted
+        using Matplotlib's colormap. Users can specify the color map through the cmap_name
+        argument.
+
+        Args:
+            scalar_field (ti.field): The scalar field being plotted. Must be 2D.
+            cmap_name (str, Optional): The name of the color map in Matplotlib.
+            normalize (bool, Optional): Display the normalized scalar field if set to True.
+            Default is False.
+        """
+        try:
+            import numpy as np  # pylint: disable=import-outside-toplevel
+            from matplotlib import \
+                cm  # pylint: disable=import-outside-toplevel
+        except ImportError:
+            raise RuntimeError('Failed to import Numpy and Matplotlib. /\
+            Please install Numpy and Matplotlib before using contour().')
+
+        scalar_field_np = scalar_field.to_numpy()
+        field_shape = scalar_field_np.shape
+        ndim = len(field_shape)
+        if ndim != 2:
+            raise ValueError(\
+                'contour() can only be used on a 2D scalar field.')
+
+        if normalize:
+            scalar_max = np.max(scalar_field_np)
+            scalar_min = np.min(scalar_field_np)
+            scalar_field_np = (scalar_field_np - scalar_min) /\
+            (scalar_max - scalar_min + 1e-16)
+
+        cmap = cm.get_cmap(cmap_name)
+        output_rgba = cmap(scalar_field_np)
+        output_rgb = output_rgba.astype(np.float32)[:, :, :3]
+        output = np.ascontiguousarray(output_rgb)
+        self.set_image(output)
+
     def triangles(self,
                   vertices,
                   color=(0.5, 0.5, 0.5),
@@ -128,6 +168,73 @@ def circles(self,
         vbo_info = get_field_info(vbo)
         self.canvas.circles(vbo_info, has_per_vertex_color, color, radius)
 
+    def vector_field(self,
+                     vector_field,
+                     arrow_spacing=5,
+                     scale=0.1,
+                     width=0.002,
+                     color=(0, 0, 0)):
+        """Draw a vector field on this canvas.
+
+        Args:
+            vector_field: The vector field to be plotted on the canvas.
+            arrow_spacing (int): Spacing used when sample on the vector field.
+            scale (float): Maximum vector length proportional to the canvas.
+            width (float): Line width when drawing the arrow.
+            color (tuple[float]): The RGB color of arrows.
+        """
+        try:
+            import numpy as np  # pylint: disable=import-outside-toplevel
+
+            import taichi as ti  # pylint: disable=import-outside-toplevel
+        except ImportError:
+            raise RuntimeError("Can't import taichi and/or numpy.")
+        v_np = vector_field.to_numpy()
+        v_norm = np.linalg.norm(v_np, axis=-1)
+        v_max = np.max(v_norm)
+        nx, ny = vector_field.shape
+        N = 6 * nx * ny
+        points = ti.Vector.field(3, dtype=ti.f32)  # End points for arrows
+        fb = ti.FieldsBuilder()  # Prepare to destroy the points in memory
+        fb.dense(ti.i, N).place(points)
+        fb_snode_tree = fb.finalize()
+
+        @ti.kernel
+        def cal_lines_points():
+            for i, j in ti.ndrange(nx // arrow_spacing + 1,
+                                   ny // arrow_spacing + 1):
+                i = i * arrow_spacing if i < nx // arrow_spacing else nx - 1
+                j = j * arrow_spacing if j < ny // arrow_spacing else ny - 1
+                linear_id = (
+                    i + j * nx) * 6  # 6 because an arrow needs 6 end points
+                # Begin point of the arrow
+                x = i / (nx - 1)
+                y = j / (ny - 1)
+                points[linear_id] = ti.Vector([x, y, 0])
+                # End point of the arrow
+                scale_factor = scale / v_max
+                dx, dy = scale_factor * vector_field[i, j]
+                points[linear_id + 1] = ti.Vector([x + dx, y + dy, 0])
+                # The tip segments
+                line_angle = ti.atan2(dy, dx)
+                tip_angle_radians = 30 / 180 * 3.14
+                line_length = ti.sqrt(dx**2 + dy**2)
+                tip_length = line_length * 0.2
+                angle1 = line_angle + 3.14 - tip_angle_radians
+                angle2 = line_angle - 3.14 + tip_angle_radians
+                points[linear_id + 2] = ti.Vector([x + dx, y + dy, 0])
+                points[linear_id + 3] = ti.Vector([x + dx + tip_length * ti.cos(angle1),\
+                                                   y + dy + tip_length * ti.sin(angle1), \
+                                                   0])
+                points[linear_id + 4] = ti.Vector([x + dx, y + dy, 0])
+                points[linear_id + 5] = ti.Vector([x + dx + tip_length * ti.cos(angle2),\
+                                                   y + dy + tip_length * ti.sin(angle2),\
+                                                   0])
+
+        cal_lines_points()
+        self.lines(points, color=color, width=width)
+        fb_snode_tree.destroy()
+
     def scene(self, scene):
         """Draw a 3D scene on the canvas
 

python/taichi/ui/gui.py

@@ -342,6 +342,48 @@ def set_image(self, img):
 
         self.core.set_img(self.img.ctypes.data)
 
+    def contour(self, scalar_field, normalize=False):
+        """Plot a contour view of a scalar field.
+
+        The input scalar_field will be converted to a Numpy array first, and then plotted
+        by the Matplotlib colormap 'Plasma'. Notice this method will automatically perform
+        a bilinear interpolation on the field if the size of the field does not match with
+        the GUI window size.
+
+        Args:
+            scalar_field (ti.field): The scalar field being plotted.
+            normalize (bool, Optional): Display the normalized scalar field if set to True.
+            Default is False.
+        """
+        try:
+            from matplotlib import \
+                cm  # pylint: disable=import-outside-toplevel
+        except ImportError:
+            raise RuntimeError(
+                'Failed to import Matplotlib. Please install it via /\
+            `pip install matplotlib` first. ')
+        scalar_field_np = scalar_field.to_numpy()
+        if self.res != scalar_field_np.shape:
+            x, y = np.meshgrid(np.linspace(0, 1, self.res[1]),
+                               np.linspace(0, 1, self.res[0]))
+            x_idx = x * (scalar_field_np.shape[1] - 1)
+            y_idx = y * (scalar_field_np.shape[0] - 1)
+            x1 = x_idx.astype(int)
+            x2 = np.minimum(x1 + 1, scalar_field_np.shape[1] - 1)
+            y1 = y_idx.astype(int)
+            y2 = np.minimum(y1 + 1, scalar_field_np.shape[0] - 1)
+            array_y1 = scalar_field_np[y1, x1] * (1 - (x_idx - x1)) * (1 - (y_idx - y1)) +\
+                scalar_field_np[y1, x2] * (x_idx - x1) * (1 - (y_idx - y1))
+            array_y2 = scalar_field_np[y2, x1] * (1 - (x_idx - x1)) * (y_idx - y1) + \
+                scalar_field_np[y2, x2] * (x_idx - x1) * (y_idx - y1)
+            output = array_y1 + array_y2
+        else:
+            output = scalar_field_np
+        if normalize:
+            scalar_max, scalar_min = np.max(output), np.min(output)
+            output = (output - scalar_min) / (scalar_max - scalar_min)
+        self.set_image(cm.plasma(output))
+
     def circle(self, pos, color=0xFFFFFF, radius=1):
         """Draws a circle on canvas.
 
@@ -699,6 +741,31 @@ def arrow_field(self,
                                       2)
         self.arrows(base, direction, radius=radius, color=color, **kwargs)
 
+    def vector_field(self, vector_field, arrow_spacing=5, color=0xFFFFFF):
+        """Display a vector field on canvas.
+
+        Args:
+            vector_field (ti.Vector.field): The vector field being displayed.
+            arrow_spacing (int, optional): The spacing between vectors.
+            color (Union[int, np.array], optional): The color of vectors.
+
+        """
+        v_np = vector_field.to_numpy()
+        v_norm = np.linalg.norm(v_np, axis=-1)
+        nx, ny, ndim = v_np.shape
+        max_magnitude = np.max(v_norm)  # Find the largest vector magnitude
+
+        # The largest vector should occupy 10% of the window
+        scale_factor = 0.1 / (max_magnitude + 1e-16)
+
+        x = np.arange(0, 1, arrow_spacing / nx)
+        y = np.arange(0, 1, arrow_spacing / ny)
+        X, Y = np.meshgrid(x, y)
+        begin = np.dstack((X, Y)).reshape(-1, 2, order='F')
+        incre = (v_np[::arrow_spacing,::arrow_spacing] \
+                 * scale_factor).reshape(-1, 2, order='C')
+        self.arrows(orig=begin, direction=incre, radius=1, color=color)
+
     def show(self, file=None):
         """Shows the frame content in the gui window, or save the content to an
         image file.

requirements_test.txt

@@ -14,3 +14,4 @@ setproctitle
 nbmake
 marko
 PyYAML
+pyevtk