hyperbolic_tiling.py

"""
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Simple Hyperbolic tiling animation using taichi
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
"""
import taichi as ti
from taichi.math import *

ti.init(arch=ti.cpu)

# ----------------------------------------------
# global settings

pqr = (3, 3, 7)
resolution = (640, 400)  # window resolution
inf = -1
inf_cos = 1.1  # =1 for paracompact tiling and >1 for noncompact tiling
max_iterations = 15
black_region_size = 0.065  # control size of the black region

klein_model = False  # if you want to draw the tiling in Klein model
rotate_pattern = True

# ----------------------------------------------
# the pixels array we will draw on
pixels = ti.Vector.field(3, float, shape=resolution)

# mA, mB, (cen, rad) are the 3 reflection mirrors, where mA, mB are lines
# and (cen, rad) represents a circle

# v0, m0 are two vertices of the fundamental triangle, (0, 0) is the 3rd one

# ----------------------------------------------
# some glsl-style helper functions


@ti.func
def rotate2d(p, ang):
    ca = ti.cos(ang)
    sa = ti.sin(ang)
    return ti.Matrix([[ca, sa], [-sa, ca]]) @ p

# ----------------------------------------------
# routines for compute the tiling

def dihedral(x):
    return inf_cos if x == inf else ti.cos(pi / float(x))


def init_tiling_data(pqr):
    """
    :param pqr: the hyperbolic triangle group (p, q, r) of the tiling.
    Here the first entry `p` must be finite, i.e. 2 <= p < inf.
    """
    p, q, r = pqr
    cos_AB = dihedral(p)
    sin_AB = ti.sqrt(1 - cos_AB * cos_AB)
    cos_AC = dihedral(q)
    cos_BC = dihedral(r)

    mA = vec2(1, 0)
    mB = vec2(-cos_AB, sin_AB)

    k1 = cos_AC
    k2 = (cos_BC + cos_AB * cos_AC) / sin_AB
    rad = 1 / ti.sqrt(k1 * k1 + k2 * k2 - 1)
    cen = vec2(k1 * rad, k2 * rad)

    delta = rad * rad - cen[0] * cen[0]
    v0 = vec2(0, 1)
    if delta >= 0:
        v0 = vec2(0, cen.y - ti.sqrt(delta))

    n = vec2(-mB.y, mB.x)
    b = cen.dot(n)
    c = cen.norm_sqr() - rad * rad
    k = -1
    if b * b >= c:
        k = b + ti.sqrt(b * b - c)

    m0 = k * n

    return mA, mB, cen, rad, v0, m0


@ti.func
def try_reflect_line(p, n, count: ti.template()):
    """Try to reflect a point `p` about a line with normal `n` and update `count`.
    """
    k = p.dot(n)
    if k < 0:
        p -= 2. * k  * n
        count += 1
    return p


@ti.func
def try_reflect_circ(p, center, radius, count: ti.template()):
    """Try to invert a point `p` about a circle `(center, radius)` and update `count`.
    """
    dist = (p - center).norm() - radius
    if dist < 0:
        p -= center
        p *= (radius * radius) / p.norm_sqr()
        p += center
        count += 1

    return p


@ti.func
def fold(p, count: ti.template()):
    """Iteratively map a point `p` into the fundamental triangle.
    """
    for _ in ti.static(range(max_iterations)):
        p = try_reflect_line(p, mA, count)
        p = try_reflect_line(p, mB, count)
        p = try_reflect_circ(p, cen, rad, count)

    return p


# ----------------------------------------------
# render functions

@ti.func
def sBox(p, b):
    d = abs(p) - b
    return ti.min(ti.max(d.x, d.y), 0.) + ti.max(d, 0.).norm()


@ti.func
def lBox(p, a, b, ew):
    x, y = b - a
    ang = ti.atan2(y, x)
    p -= (a + b) / 2
    p = rotate2d(p, ang)
    l = vec2((b - a).norm(), ew)
    return sBox(p, (l + ew) / 2.)


@ti.func
def mouse_inversion(p, mouse):
    W, H = resolution
    mouse = 2 * mouse - 1
    mouse.x *= W / H
    if mouse.norm() < 1e-3:
        mouse += 1e-3

    if abs(mouse.x) > 0.7 or abs(mouse.y) > 0.7:
        mouse *= 0.98

    k = 1.0 / mouse.norm_sqr()
    invCtr = mouse * k
    t = (k - 1.0) / (p - invCtr).norm_sqr()
    p = mix(invCtr, p, t)
    p.x *= -1.0
    return p


@ti.kernel
def render(iTime: ti.f32, mouse_x: ti.f32, mouse_y: ti.f32):
    W, H = resolution
    for ix, iy in pixels:
        # map screen coordinates to complex plane points
        uv = vec2((2. * ix / W - 1) * W / H, 2. * iy / H - 1) * 1.05
        p = uv

        if ti.static(rotate_pattern):
            p = mouse_inversion(p, vec2(mouse_x, mouse_y))
            p = rotate2d(p, iTime / 16)

        if p.norm() > 1:  # invert if p is outside of the unit disk
            p /= p.norm_sqr()

        if ti.static(klein_model):
            p = p / (1 + ti.sqrt(1 - p.norm_sqr()))

        count = 0
        p = fold(p, count)

        ln = ln2 = pnt = 1e5
        ln = ti.min(ln, lBox(p, vec2(0), v0, 0.007))
        ln = ti.min(ln, lBox(p, vec2(0), m0, 0.007))
        ln = ti.min(ln, (p - cen).norm() - rad - 0.007)

        ln2 = ti.min(ln2, lBox(p, vec2(0), m0, .007))
        pnt = ti.min(pnt, (p - v0).norm())
        pnt = ti.min(pnt, (p - m0).norm())

        cir = uv.norm()
        # smoothing factor
        ssf = (2 - smoothstep(0, 0.25, abs(cir - 1.) - 0.25))
        sf = 2.0 * ssf / H
        # color map
        cm = count * pi / 4.0
        oCol = 0.55 + 0.45 * ti.cos(vec3(cm, cm + 1, cm + 2))
        pat = smoothstep(0, 0.25, abs(fract(ln2 * 50.0 - 0.2) - 0.5) * 2.0 - 0.2)
        sh = clamp(0.65 + ln / v0.y * 4, 0, 1)

        col = ti.min(oCol * (pat * 0.2 + 0.9) * sh, 1.0)
        col = mix(col, vec3(0), 1 - smoothstep(0, sf, ln))
        col = mix(col, vec3(0), 1 - smoothstep(0, sf, -(ln - black_region_size)))

        pnt -= .032
        pnt = ti.min(pnt, p.norm() - .032)
        col = mix(col, vec3(0), 1 - smoothstep(0, sf, pnt))
        col = mix(col, vec3(1, .8, .3), 1 - smoothstep(0, sf, pnt + .02))

        bg = vec3(1, .2, .4)
        bg *= 0.7 * (mix(col, vec3(1) * (col.dot(vec3(.299, .587, .114))), .5) * 0.5 + .5)
        pat = smoothstep(0, 0.25, abs(fract((uv.x - uv.y) * 43. - .25) - .5) * 2. - .5)
        bg *= ti.max(1 - cir * 0.5, 0.) * (pat * .2 + .9)

        col = mix(col, vec3(0), (1 - smoothstep(0, sf * 10., abs(cir - 1.) - .05)) * .7)
        col = mix(col, vec3(0), (1 - smoothstep(0, sf * 2., abs(cir - 1.) - .05)))
        col = mix(col, vec3(0.9) + bg, (1 - smoothstep(0, sf, abs(cir - 1.) - .03)))
        col = mix(col, col * max(1 - cir * 0.5, 0), (1 - smoothstep(0, sf, -cir + 1.05)))
        col = mix(col, bg, (1 - smoothstep(0, sf, -cir + 1.05)))

        col = mix(col, vec3(0), (1 - smoothstep(0, sf, abs(cir - 1.035) - .03)) * .8)
        col = mix(col, 1 - ti.exp(-col), .35)
        col = clamp(col, 0, 1)
        col = ti.sqrt(col)
        pixels[ix, iy] = col


if __name__ == "__main__":
    mA, mB, cen, rad, v0, m0 = init_tiling_data(pqr)
    gui = ti.GUI("Hyperbolic Tiling", res=resolution)
    mouse_x, mouse_y = 0.0, 0.0
    for i in range(1000000):
        gui.get_event(ti.GUI.PRESS)
        if gui.is_pressed(ti.GUI.ESCAPE):
            exit()

        if gui.is_pressed(ti.GUI.LMB):
            mouse_x, mouse_y = gui.get_cursor_pos()

        render(i * 0.03, mouse_x, mouse_y)
        gui.set_image(pixels)
        gui.show()