Taichi Python version of "Lava New Version" by MelisaHot on Shadertoy https://www.shadertoy.com/view/ssKBzm

lava_new_version.py

"""
"Lava New Version" by MelisaHot on Shadertoy
https://www.shadertoy.com/view/ssKBzm
"""

import taichi as ti

ti.init(arch=ti.vulkan)

@ti.func
def rot(a):
    return ti.Matrix([
        [ti.cos(a), ti.sin(a)],
        [-ti.sin(a), ti.cos(a)]
    ])

@ti.func
def hash21(n):
    return ti.math.fract(ti.cos(n.dot(ti.Vector([5.9898, 4.1414]))) * 65899.89956)

@ti.func
def noise(n):
    d = ti.Vector([0.0, 1.0])
    b = ti.floor(n)
    f = ti.math.smoothstep(ti.Vector([0.0, 0.0]), ti.Vector([1.0, 1.0]), ti.math.fract(n))
    return ti.math.mix(
        ti.math.mix(hash21(b), hash21(b + d.yx), f.x),
        ti.math.mix(hash21(b + d.xy), hash21(b + d.yy), f.x),
        f.y
    )

@ti.func
def mixNoise(p):
    epsilon = 0.968785675
    noiseX = noise(ti.Vector([p.x + epsilon, p.y])) - noise(ti.Vector([p.x - epsilon, p.y]))
    noiseY = noise(ti.Vector([p.x, p.y + epsilon])) - noise(ti.Vector([p.x, p.y - epsilon]))
    return ti.Vector([noiseX, noiseY])

@ti.func
def fbm(p, iTime):
    amplitude = 3.0
    total = 0.0
    pom = p
    for i in range(7):
        p += iTime * 0.05
        pom += iTime * 0.09
        n = mixNoise(i * p * 0.3244243 + iTime * 0.131321)
        n = rot(iTime * 0.5 - (0.03456 * p.x + 0.0342322 * p.y) * 50.0) @ n
        p += n * 0.5
        total += (ti.sin(noise(p) * 8.5) * 0.55 + 0.4566) / amplitude
        p = ti.math.mix(pom, p, 0.5)
        amplitude *= 1.3
        p *= 2.007556
        pom *= 1.6895367
    return total

res = (800, 600)
pixels = ti.Vector.field(3, dtype=ti.f32, shape=res)

@ti.kernel
def render(iTime: ti.f32):
    for i, j in pixels:
        uv = ti.Vector([i / res[0], j / res[1]])
        uv.x *= res[0] / res[1]
        uv *= 2.2
        fbm_value = fbm(uv, iTime)
        col = ti.Vector([0.212, 0.08, 0.03]) / max(fbm_value, 0.0001)
        col = col ** 1.5
        pixels[i, j] = col

gui = ti.GUI("Lava Shader", res)
iTime = 0.0
while gui.running:
    render(iTime)
    gui.set_image(pixels)
    gui.show()
    iTime += 0.03