[example] Add real func version of marching_square

ghstack-source-id: 7cb9b47f4aac401a90b238b6a650fc81512f6484
Pull Request resolved: #8195

python/taichi/examples/real_func/algorithm/marching_squares.py

@@ -0,0 +1,182 @@
+"""
+Marching squares algorithm in Taichi.
+See "https://en.wikipedia.org/wiki/Marching_squares"
+"""
+import time
+
+import numpy as np
+
+import taichi as ti
+import taichi.math as tm
+
+ti.init(arch=ti.cpu)
+
+W, H = 800, 600
+resolution = (W, H)
+grid_size = 8
+level = 0.15  # Draw contour of isofunc=level
+pixels = ti.Vector.field(3, float, shape=resolution)
+
+# Cases 0-15 in the wikipedia page
+_edges_np = np.array(
+    [
+        [[-1, -1], [-1, -1]],
+        [[3, 0], [-1, -1]],
+        [[0, 1], [-1, -1]],
+        [[1, 3], [-1, -1]],
+        [[1, 2], [-1, -1]],
+        [[0, 1], [2, 3]],
+        [[0, 2], [-1, -1]],
+        [[2, 3], [-1, -1]],
+        [[2, 3], [-1, -1]],
+        [[0, 2], [-1, -1]],
+        [[0, 3], [1, 2]],
+        [[1, 2], [-1, -1]],
+        [[1, 3], [-1, -1]],
+        [[0, 1], [-1, -1]],
+        [[3, 0], [-1, -1]],
+        [[-1, -1], [-1, -1]],
+    ],
+    dtype=np.int32,
+)
+edge_table = ti.Matrix.field(2, 2, int, 16)
+edge_table.from_numpy(_edges_np)
+
+Edge = ti.types.struct(p0=tm.vec2, p1=tm.vec2)
+edges = Edge.field()
+ti.root.dynamic(ti.i, 1024, chunk_size=32).place(edges)
+iTime = ti.field(float, shape=())
+
+
+@ti.experimental.real_func
+def hash22(p_: tm.vec2) -> tm.vec2:
+    n = tm.sin(tm.dot(p_, tm.vec2(41, 289)))
+    p = tm.fract(tm.vec2(262144, 32768) * n)
+    return tm.sin(p * 6.28 + iTime[None])
+
+
+@ti.experimental.real_func
+def noise(p_: tm.vec2) -> float:
+    ip = tm.floor(p_)
+    p = p_ - ip
+    v = tm.vec4(
+        tm.dot(hash22(ip), p),
+        tm.dot(hash22(ip + tm.vec2(1, 0)), p - tm.vec2(1, 0)),
+        tm.dot(hash22(ip + tm.vec2(0, 1)), p - tm.vec2(0, 1)),
+        tm.dot(hash22(ip + tm.vec2(1, 1)), p - tm.vec2(1, 1)),
+    )
+    p = p * p * p * (p * (p * 6 - 15) + 10)
+    return tm.mix(tm.mix(v.x, v.y, p.x), tm.mix(v.z, v.w, p.x), p.y)
+
+
+@ti.experimental.real_func
+def isofunc(p: tm.vec2) -> float:
+    return noise(p / 4 + 1)
+
+
+@ti.experimental.real_func
+def interp(p1: tm.vec2, p2: tm.vec2, v1: float, v2: float, isovalue: float) -> tm.vec2:
+    return tm.mix(p1, p2, (isovalue - v1) / (v2 - v1))
+
+
+@ti.experimental.real_func
+def get_vertex(vertex_id: int, values: tm.vec4, isovalue: float) -> tm.vec2:
+    v = tm.vec2(0)
+    square = [tm.vec2(0), tm.vec2(1, 0), tm.vec2(1, 1), tm.vec2(0, 1)]
+    if vertex_id == 0:
+        v = interp(square[0], square[1], values.x, values.y, isovalue)
+    elif vertex_id == 1:
+        v = interp(square[1], square[2], values.y, values.z, isovalue)
+    elif vertex_id == 2:
+        v = interp(square[2], square[3], values.z, values.w, isovalue)
+    else:
+        v = interp(square[3], square[0], values.w, values.x, isovalue)
+    return v
+
+
+@ti.kernel
+def march_squares() -> int:
+    edges.deactivate()
+    x_range = tm.ceil(grid_size * W / H, int)
+    y_range = grid_size
+    for i, j in ti.ndrange((-x_range, x_range + 1), (-y_range, y_range + 1)):
+        case_id = 0
+        values = tm.vec4(
+            isofunc(tm.vec2(i, j)),
+            isofunc(tm.vec2(i + 1, j)),
+            isofunc(tm.vec2(i + 1, j + 1)),
+            isofunc(tm.vec2(i, j + 1)),
+        )
+        if values.x > level:
+            case_id |= 1
+        if values.y > level:
+            case_id |= 2
+        if values.z > level:
+            case_id |= 4
+        if values.w > level:
+            case_id |= 8
+
+        # Fix the ambiguity for case 5 and 10
+        if case_id == 5 or case_id == 10:
+            center = isofunc(tm.vec2(i + 0.5, j + 0.5))
+            if center < level:  # A valley, switch the case_id
+                case_id = 15 - case_id
+
+        for k in ti.static(range(2)):
+            if edge_table[case_id][k, 0] != -1:
+                ind1 = edge_table[case_id][k, 0]
+                ind2 = edge_table[case_id][k, 1]
+                p0 = tm.vec2(i, j) + get_vertex(ind1, values, level)
+                p1 = tm.vec2(i, j) + get_vertex(ind2, values, level)
+                edges.append(Edge(p0, p1))
+
+    return edges.length()
+
+
+@ti.experimental.real_func
+def dseg(p_: tm.vec2, a: tm.vec2, b_: tm.vec2) -> ti.f32:
+    p = p_ - a
+    b = b_ - a
+    h = tm.clamp(tm.dot(p, b) / tm.dot(b, b), 0, 1)
+    return tm.length(p - h * b)
+
+
+@ti.kernel
+def render():
+    for i, j in pixels:
+        p = (2 * tm.vec2(i, j) - tm.vec2(resolution)) / H
+        p *= grid_size
+        p.y += 1.2
+
+        q = tm.fract(p) - 0.5
+        d2 = 0.5 - abs(q)
+        dgrid = ti.min(d2.x, d2.y)
+        dedge = 1e5
+        dv = 1e5
+        for k in range(edges.length()):
+            s, e = edges[k].p0, edges[k].p1
+            dedge = ti.min(dedge, dseg(p, s, e))
+            dv = ti.min(dv, tm.length(p - s), tm.length(p - e))
+
+        col = tm.vec3(0.3, 0.6, 0.8)
+        if isofunc(p) > level:
+            col = tm.vec3(1, 0.8, 0.3)
+
+        # Draw background grid
+        col = tm.mix(col, tm.vec3(0), 1 - tm.smoothstep(0, 0.04, dgrid - 0.02))
+        # Draw edges
+        col = tm.mix(col, tm.vec3(1, 0, 0), 1 - tm.smoothstep(0, 0.05, dedge - 0.02))
+        # Draw small circles at the vertices
+        col = tm.mix(col, tm.vec3(0), 1 - tm.smoothstep(0, 0.05, dv - 0.1))
+        col = tm.mix(col, tm.vec3(1), 1 - tm.smoothstep(0, 0.05, dv - 0.08))
+        pixels[i, j] = tm.sqrt(tm.clamp(col, 0, 1))
+
+
+t0 = time.perf_counter()
+gui = ti.GUI("2D Marching Squares", res=resolution, fast_gui=True)
+while gui.running and not gui.get_event(gui.ESCAPE):
+    iTime[None] = time.perf_counter() - t0
+    march_squares()
+    render()
+    gui.set_image(pixels)
+    gui.show()