hyperbolic_tiling_taichi.py

"""
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Simple Hyperbolic tiling animation using taichi
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
"""
import taichi as ti
ti.init(arch=ti.cpu)

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

pqr = (3, 3, 7)
resolution = (600, 360)  # window resolution
inf = -1
inf_cos = 1.1  # =1 for paracompact tiling and >1 for noncompact tiling
pi = 3.141592654
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, ti.f32, 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

def vec2(*args):
    assert 1 <= len(args) <= 2
    if len(args) == 1:
        x = float(args[0])
        result = ti.Vector([x, x])
    else:
        x, y = args
        result = ti.Vector([float(x), float(y)])

    return result


def vec3(*args):
    assert len(args) == 1 or len(args) == 3
    if len(args) == 1:
        x = float(args[0])
        result = ti.Vector([x, x, x])
    else:
        x, y, z = args
        result = ti.Vector([float(x), float(y), float(z)])

    return result


@ti.func
def clamp(v, v_min, v_max):
    return ti.max(ti.min(v, v_max), v_min)


@ti.func
def fract(x):
    return x - ti.floor(x)


@ti.func
def mix(a, b, t):
    return a * (1.0 - t) + b * t


@ti.func
def smoothstep(edge1, edge2, v):
    assert(edge1 != edge2)
    t = (v - edge1) / float(edge2 - edge1)
    t = clamp(t, 0.0, 1.0)
    return (3.0 - 2.0 * t) * t * t


@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[1] - ti.sqrt(delta))

    n = vec2(-mB[1], mB[0])
    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 min(max(d[0], d[1]), 0.) + ti.max(d, 0.).norm()


@ti.func
def lBox(p, a, b, ew):
    ang = ti.atan2(b[1] - a[1], b[0] - a[0])
    p -= (a + b) / 2
    p = rotate2d(p, ang)
    l = ti.Vector([(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[0] *= W / H
    if mouse.norm() < 1e-3:
        mouse += 1e-3

    if abs(mouse[0]) > 0.7 or abs(mouse[1]) > 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[0] *= -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 = ti.Vector([(2. * ix / W - 1) * W / H, 2. * iy / H - 1]) * 1.05
        p = uv

        if ti.static(rotate_pattern):
            p = mouse_inversion(p, ti.Vector([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 = min(ln2, lBox(p, vec2(0), m0, .007))
        pnt = min(pnt, (p - v0).norm())
        pnt = 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 = float(count) * pi / 4.0
        oCol = 0.55 + 0.45 * ti.cos(ti.Vector([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[1] * 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 = 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 *= 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)

    for i in range(1000000):
        gui.get_event()
        if gui.is_pressed(ti.GUI.ESCAPE):
            exit()

        mouse_x, mouse_y = 0, 0
        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()