initial commit

.gitignore

···
+.git
+__pycache__
+out

abelian_sandpile.py

···
+import math
+import random
+
+import torch
+
+from util import *
+
+@settings
+def _settings():
+    scale = 2 ** 10
+    w = scale + 1
+    h = scale + 1
+
+    iterations = 500000
+    debug_modulus = 10000
+
+    out_file = "hng"
+
+    device="cuda"
+    ctype=torch.cdouble
+    dtype=torch.double
+
+@ifmain(__name__)
+@timed
+def _main():
+    _ = _settings()
+    torch.set_default_device(_.device)
+    torch.set_default_dtype(_.dtype)
+    dim = [_.h, _.w]
+
+    grid = torch.zeros(dim, dtype=torch.int)
+
+    #offsets = [torch.tensor(x) for x in [[0,1],[0,-1],[1,0],[-1,0]]]
+    #offsets = [[1,1],[1,-1],[-1,1],[-1,-1],[0,2],[0,-2],[2,0],[-2,0]]
+    #offsets = [[0,10],[10,0],[5,10],[10,5],[5,5],[-1,0],[0,-1]]
+    #offsets = [[0,5],[5,0],[0,-5],[-5,0],[1,4],[4,1],[-1,4],[4,-1],[1,-4],[-4,1],[-4,-1],[-1,-4]]
+    offsets = [[0,3],[3,0],[1,2],[2,1],[13,3],[3,13],[1,0],[0,1],[1,1],[0,0]]
+    #offsets = [x for x in [
+    #    [1,1],
+    #    [0,1],
+    #    [1,0],
+    #    [0,-5],
+    #    [-5,-5],
+    #    [-5,0]
+    #    ]]
+
+    max_c = 100
+    c_total = 0
+    last_c_total = 100
+
+    n = len(offsets)
+
+    for iteration in range(_.iterations):
+        #spots = torch.randn(dim) > 2.5
+        #grid += spots
+        #p = [_.h // 2, _.w // 2]
+        p = [0,0]
+        #p = [random.randint(0,_.h-1), random.randint(0,_.w-1)]
+        last = torch.clone(grid)
+        grid[p[0],p[1]] = n
+        peaks = (grid >= n).int()
+        #peaks = torch.nonzero(grid >= 4)
+        c = 0
+        while torch.count_nonzero(peaks):
+        #while len(peaks) != 0:
+            c += 1
+            #p0 = peaks[0]
+            #pts = [p0 + o for o in offsets]
+            #pts = [p for p in pts if 0 < p[0] < _.h and 0 < p[1] < _.w]
+            #for p in pts:
+            #    grid[p[0],p[1]] += 1
+            #grid[p0[0],p0[1]] -= 4
+            shift = -(n * peaks)
+            for offset in offsets:
+                p_shift = torch.roll(peaks, offset, dims=[0,1])
+                if offset[0] < 0:
+                    p_shift[offset[0]:] = 0
+                else:
+                    p_shift[:offset[0]] = 0
+                if offset[1] < 0:
+                    p_shift[:,offset[1]:] = 0
+                else:
+                    p_shift[:,:offset[1]] = 0
+                shift += p_shift
+            grid += shift
+
+            #pa = torch.roll(peaks, 1, dims=0)
+            #pb = torch.roll(peaks, -1, dims=0)
+            #pc = torch.roll(peaks, 1, dims=1)
+            #pd = torch.roll(peaks, -1, dims=1)
+            #pa[0] = 0
+            #pb[-1] = 0
+            #pc[:,0] = 0
+            #pd[:,-1] = 0
+            #shift = pa + pb + pc + pd - (4 * peaks)
+            #grid += shift
+
+            peaks = (grid >= n).int()
+            #peaks = torch.nonzero(grid >= 4)
+        c_total += c
+        if c > 20:
+            s = "*" if c >= max_c else ""
+            print(f"{iteration}: {c}{s}")
+            if c > (0.9 * max_c):
+                msave(grid / (n-1), f"{_.out_file}_{iteration}_{c}")
+                msave(last / (n-1), f"{_.out_file}_{iteration}_prev")
+            if c > max_c:
+                max_c = c
+        if iteration % _.debug_modulus == 0:
+            msave(grid / (n-1), f"{_.out_file}_{iteration}_")
+        if c_total > (last_c_total * 2):
+            msave(grid / (n-1), f"{_.out_file}_{iteration}_ct{c_total}")
+            last_c_total = c_total
+
+    msave(grid / (n-1), _.out_file)

fluid.py

···
+import time
+from pathlib import Path
+
+import torch
+from torch.nn.functional import conv2d as convolve
+
+from util import *
+
+@settings
+def _s():
+    scale = 2**8
+    w = scale
+    h = scale
+
+    max_iterations = 10000
+    save_every = 10
+
+    v_diffusion_rate = 2000
+    m_diffusion_rate = 0.5
+    m_timescale = 0.1
+    timescale = 0.01
+
+    name = "fluid"
+
+    device = "cuda"
+
+
+diffuse_kernel = torch.tensor([[[[0,1,0],[1,0,1],[0,1,0]]]], dtype=torch.float, device="cuda")
+project_kernel_x = torch.tensor([[[[0,0,0],[-1,0,1],[0,0,0]]]], dtype=torch.float, device="cuda")
+project_kernel_y = torch.tensor([[[[0,-1,0],[0,0,0],[0,1,0]]]], dtype=torch.float, device="cuda")
+
+def continuous_boundary(field):
+    field[0] = field[1]
+    field[-1] = field[-2]
+    field[:,0] = field[:,1]
+    field[:,-1] = field[:,-2]
+    field[(0,0,-1,-1),(0,-1,0,-1)] = 0.5 * (field[(0,0,-2,-2),(1,-2,0,-1)] + field[(1,1,-1,-1),(0,-1,1,-2)])
+
+def opposed_v_boundary(field):
+    field[0] = field[1]
+    field[-1] = field[-2]
+    field[:,0] = -field[:,1]
+    field[:,-1] = -field[:,-2]
+    field[(0,0,-1,-1),(0,-1,0,-1)] = 0.5 * (field[(0,0,-2,-2),(1,-2,0,-1)] + field[(1,1,-1,-1),(0,-1,1,-2)])
+
+def opposed_h_boundary(field):
+    field[0] = -field[1]
+    field[-1] = -field[-2]
+    field[:,0] = field[:,1]
+    field[:,-1] = field[:,-2]
+    field[(0,0,-1,-1),(0,-1,0,-1)] = 0.5 * (field[(0,0,-2,-2),(1,-2,0,-1)] + field[(1,1,-1,-1),(0,-1,1,-2)])
+
+
+def diffuse(field, rate, set_boundary, dt, h, w):
+    a = dt * rate
+    result = torch.clone(field)
+    if field.shape != (h, w):
+        print("bad field shape in diffuse")
+    for n in range(20):
+        convolution = a * convolve(result.unsqueeze(0), diffuse_kernel, bias=None, padding=[0], stride=[1])[0]
+        result[1:h-1,1:w-1] = field[1:h-1,1:w-1] + convolution
+        result /= 1 + 4 * a
+        set_boundary(result)
+        #result = result * ~border_mask[0] + field * border_mask[0]
+    return result
+
+def advect(field, velocities, dt, h, w):
+    dth, dtw = dt, dt
+    inds_x = torch.arange(1,w-1).repeat(h-2,1).float()
+    inds_y = torch.arange(1,h-1).repeat(w-2,1).t().float()
+    inds_x += dtw * velocities[1,1:h-1,1:w-1]
+    inds_y += dth * velocities[0,1:h-1,1:w-1]
+    inds_x = inds_x.clamp(1.5, w - 2.5)
+    inds_y = inds_y.clamp(1.5, h - 2.5)
+    inds_x_i = inds_x.int()
+    inds_y_i = inds_y.int()
+    inds_x -= inds_x_i
+    inds_y -= inds_y_i
+    inds_x_inv = 1 - inds_x
+    inds_y_inv = 1 - inds_y
+    inds_x_i_next = inds_x_i + 1
+    inds_y_i_next = inds_y_i + 1
+    inds_x_all = torch.stack([inds_x_i, inds_x_i_next, inds_x_i, inds_x_i_next])
+    inds_y_all = torch.stack([inds_y_i, inds_y_i, inds_y_i_next, inds_y_i_next])
+    if field.shape[0] == 1:
+        values = torch.cat([field[:,1:h-1,1:w-1] * inds_x_inv * inds_y_inv,
+                            field[:,1:h-1,1:w-1] * inds_x * inds_y_inv,
+                            field[:,1:h-1,1:w-1] * inds_x_inv * inds_y,
+                            field[:,1:h-1,1:w-1] * inds_x * inds_y])
+        res = torch.zeros_like(field[0])
+        res.index_put_((inds_y_all, inds_x_all), values, accumulate=True)
+        continuous_boundary(res)
+        return res.unsqueeze(0)
+    else:
+        values = torch.stack([field[:,1:h-1,1:w-1] * inds_x_inv * inds_y_inv,
+                              field[:,1:h-1,1:w-1] * inds_x * inds_y_inv,
+                              field[:,1:h-1,1:w-1] * inds_x_inv * inds_y,
+                              field[:,1:h-1,1:w-1] * inds_x * inds_y])
+        res = torch.zeros_like(field)
+        res[0].index_put_((inds_y_all, inds_x_all), values[:,0,:,:], accumulate=True)
+        res[1].index_put_((inds_y_all, inds_x_all), values[:,1,:,:], accumulate=True)
+        opposed_h_boundary(res[1])
+        opposed_v_boundary(res[0])
+        return res
+
+def project(field, h, w):
+    hx = -1# / w
+    hy = -1# / h
+    divergence = convolve(field[1].unsqueeze(0), project_kernel_x, bias=None, stride=[1], padding=[0])[0] * hx
+    divergence += convolve(field[0].unsqueeze(0), project_kernel_y, bias=None, stride=[1], padding=[0])[0] * hy
+    divergence *= 0.5
+    continuous_boundary(divergence)
+    p = torch.zeros_like(field[0])
+    for i in range(40):
+        p[1:h-1,1:w-1] = (divergence + convolve(p.unsqueeze(0), diffuse_kernel, bias=None, stride=[1], padding=[0])[0]) / 4
+        continuous_boundary(p)
+    field[1,1:h-1,1:w-1] += 0.5 * convolve(p.unsqueeze(0), project_kernel_x, bias=None, stride=[1], padding=[0])[0] / hx
+    field[0,1:h-1,1:w-1] += 0.5 * convolve(p.unsqueeze(0), project_kernel_y, bias=None, stride=[1], padding=[0])[0] / hy
+    opposed_h_boundary(field[1])
+    opposed_v_boundary(field[0])
+
+@ifmain(__name__, _s)
+@timed
+def _main(settings):
+    globals().update(settings.get())
+    run_dir = time.strftime(f"%d.%m.%Y/{name}_t%H.%M.%S")
+    Path("out/" + run_dir).mkdir(parents=True, exist_ok=True)
+    with open(f"out/{run_dir}/settings.py", "w") as f:
+        f.write(settings.src)
+    torch.set_default_device(device)
+
+
+    #velocities = torch.zeros([2, h, w])
+    upsample = torch.nn.Upsample(size=[h, w], mode='bilinear')
+    velocities = torch.randn([2, h//32, w//32]) * 100
+    densities = torch.ones([1, h, w]) * 0.1
+
+    cg = cgrid(h, w, 0 + 0j, 4 + 4j, dtype=torch.float, ctype=torch.cfloat)
+
+    velocities = upsample(velocities.unsqueeze(0))[0]#cg.real - cg.imag * torch.sin(cg.real)
+    #velocities[1] = 0#-cg.imag - cg.real * torch.cos(cg.imag)
+
+    #for i in range(3):
+    #    for j in range(3):
+    #        w7 = w // 7
+    #        h7 = h // 7
+    #        densities[0][(2*j+1)*h7:(2*j+2)*h7][(2*i+1)*w7:(2*i+2)*w7] = 1
+    #        print((2*j+1)*h7)
+    #        print((2*j+2)*h7)
+
+
+    #w7 = w//7
+    #h7 = h//7
+    #for i in range(w7):
+    #    for j in range(h7):
+    #        for a in range(3):
+    #            for b in range(3):
+    #                densities[0][j + h7 * (2 * a + 1)][w7 * (2 * b + 1) + i] = 1.0
+
+    initial = torch.cat((velocities * 0.5 + 0.5, densities), dim=0)
+    save(initial, f"{run_dir}/initial")
+    del initial
+
+    border_mask = torch.ones([2,h,w], dtype=torch.int)
+    border_mask[:,1:h-1,1:w-1] = 0;
+
+    for iteration in range(max_iterations):
+        velocities[0,3 * h//4 - 10:(3*h//4),w//2-10:w//2+10] = -100
+        velocities[1,3 * h//4:(3*h//4)+10,w//2-10:w//2+10] = -90
+        velocities[0,h//4:(h//4)+10,w//2-3:w//2+3] = 100
+        velocities[1,h//4:(h//4)+10,w//2-3:w//2+3] = 200
+        velocities[0] += 0.01
+        densities[0,3*h//4-10:(3*h//4),w//2-9:w//2+9] += 0.1
+        #densities[0,h//4:(h//4)+10,w//2+30:w//2+55] += 0.1
+        densities += 0.001
+        velocities[0] = diffuse(velocities[0], v_diffusion_rate, opposed_h_boundary, timescale, h, w)
+        velocities[1] = diffuse(velocities[1], v_diffusion_rate, opposed_v_boundary, timescale, h, w)
+        project(velocities, h, w)
+        velocities = advect(velocities, velocities, timescale, h, w)
+        project(velocities, h, w)
+        densities = diffuse(densities[0], m_diffusion_rate, continuous_boundary, m_timescale, h, w).unsqueeze(0)
+        densities = advect(densities, velocities, m_timescale, h, w)
+        densities *= 0.999
+        if iteration % save_every == 0:
+            #res = torch.cat((velocities * 0.5 + 0.5, densities), dim=0)
+            #save(res, f"{run_dir}/{iteration}")
+            msave(densities[0], f"{run_dir}/_{iteration}")
+
+    final = torch.cat((velocities * 0.5 + 0.5, densities), dim=0)
+    save(final, f"{run_dir}/final")
+
+
+
+
+
+
+
+
+
+
+
+

fractal.py

···
+import time
+import torch
+from PIL import Image
+from util import *
+
+import math
+
+# Settings
+
+w = 512 * 2
+h = w
+
+#x_range = [-1.8, -1.7]
+#y_range = [-0.095, 0.005]
+
+#ship
+#x_start = -1.8
+#x_span = 0.1
+#y_start = -0.095
+#y_span = 0.1
+
+x_start = -math.e / 2
+x_span = 1.54
+y_start = -2
+y_span = 2
+
+alpha = 2.5
+
+x_range = [x_start, x_start + x_span]
+y_range = [y_start, y_start + y_span]
+
+final_file = "gauss_1"
+
+device = "cuda"
+ctype = torch.cfloat
+dtype = torch.float
+
+# how many random samples in each point
+sample_ct = 40
+noise_scale = x_span / (2 * w)
+
+do_jitter = False
+jitter_scale = 2.5 / w
+jitter_chunk_size = 128
+
+# how many iterations of the z_(n+1) = f(z_n) recurrence
+iterations = 120
+
+# threshold for defining escape
+limit = 2.0
+
+if __name__ == "__main__":
+    t0 = time.perf_counter()
+
+    z_init = torch.zeros([h, w], dtype=ctype, device=device)
+    c = torch.zeros([h, w], dtype=ctype, device=device)
+    j = torch.zeros([h, w], dtype=ctype, device=device)
+    j.imag = torch.ones([h, w], dtype=dtype, device=device)
+
+    yspace = torch.linspace(y_range[0], y_range[1], h, dtype=dtype, device=device)
+    xspace = torch.linspace(x_range[0], x_range[1], w, dtype=dtype, device=device)
+
+    for x in range(h):
+        c[x] += xspace
+
+    for y in range(w):
+        c[:, y] += yspace * 1j
+
+    #save(c, "c")
+    #save(z_init, "z0")
+
+    #z = z_init
+
+    totals = torch.zeros([h, w], dtype=torch.int, device=device)
+    jitter = torch.randn([h // jitter_chunk_size, w // jitter_chunk_size], dtype=ctype, device=device) * jitter_scale
+    if do_jitter:
+        upsample = torch.nn.Upsample(size=[h, w], mode='nearest')
+        jitter = torch.view_as_real(jitter).permute((2,0,1))[None, ...]
+        jitter = torch.view_as_complex(upsample(jitter)[0].permute((1,2,0)).contiguous())
+
+    #csave(jitter, "jitter")
+    #exit()
+
+    for sample_idx in range(sample_ct):
+        noise = torch.randn([h, w], dtype=ctype, device=device)
+        z = z_init# + noise * noise_scale
+        c_n = c + noise * noise_scale
+        mask = torch.abs(z) < limit
+        iters = torch.zeros_like(mask, dtype=torch.uint8)
+        for i in range(iterations):
+            #im = torch.abs(z.imag) * j
+            #z = (torch.abs(z.real) + im)**2 + c_n
+            z = torch.exp(-alpha * z) + c_n
+
+            if do_jitter:
+                z += jitter
+            mask &= (torch.abs(z) < limit)
+            iters += mask
+        #for i in range(iterations):
+        #    im = torch.abs(z.imag) * j
+        #    z = (torch.abs(z.real) + im)**2 + c_n
+        #    totals +=
+
+        #totals += ~mask
+        totals += iters
+            #csave(mask, f"mask")
+
+            #csave(z, f"z{i}")
+
+        #msave(mask, f"mfinal_s{sample_idx}")
+        #csave(z, f"zfinal_s{sample_idx}")
+
+        #im = torch.log(iters)
+    #im = iters
+        #im = im / torch.max(im)
+
+        #msave(im, f"ifinal_s{sample_idx}")
+    #im = im * ~mask
+
+    #totals = torch.log(1. - totals)
+    #totals *= totals
+
+    #m = totals == 0. #torch.max(totals)
+    #totals = totals + torch.max(totals) * 0.5 * m
+
+    totals = totals / sample_ct / iterations#torch.max(totals)
+    msave(totals, final_file + "_")
+
+    #msave(totals**2, final_file + "_sq")
+    #msave(torch.sqrt(totals), final_file + "_rt")
+    totals = 1. - totals
+    #msave(totals, final_file + "_inv_")
+    #msave(totals**2, final_file + "_inv_sq")
+    msave(torch.sqrt(totals), final_file + "_inv_rt")
+
+
+    #csave(torch.nan_to_num(z, 0) * (mask), "zmfinal")
+
+    #zm_r = torch.nan_to_num(z.real, 0) * mask
+    #zm_i = torch.nan_to_num(z.imag, 0) * mask
+
+    #print(torch.max(torch.abs(torch.nan_to_num(z, 0) * mask)))
+
+    #msave(im, "imfinal")
+
+    #final = torch.stack([zm_r / 4 + 0.5, zm_i / 4 + 0.5, im])
+
+    #save(final, "final")
+
+    print(f"done in {time.perf_counter() - t0}s")
+
+
+
+
+
+

lyapunov.py

···
+import torch
+import time
+import itertools
+
+from PIL import Image
+# monochrome, 0 to 1
+def mpilify(z):
+    z_3d = torch.stack([z, z, z])
+    z_norm = z_3d.clamp(0, 1)
+    z_np = z_norm.detach().cpu().permute(1, 2, 0).numpy()
+    z_bytes = (z_np * 255).round().astype("uint8")
+    return Image.fromarray(z_bytes)
+
+def msave(x, f):
+    mpilify(x).save(f"out/{f}.png")
+
+# config
+
+final_file = "abaaa_squish_tall2"
+
+# 2 ** 9 = 512; 10 => 1024; 11 => 2048
+scale = 2 ** 9
+h = scale * 2
+w = scale
+
+origin = 3 + (1 / 3), 2
+span = 1 / 3, 2
+
+zooms = []
+
+dev = "cuda"
+t_fp = torch.double
+t_cfp = torch.cdouble
+
+seq = "ABAAAAAAAA"
+c = 3.5
+
+iterations = 100
+
+
+# funcs
+
+def logistic_map(r, x):
+    x_inf_mask = torch.isinf(x)
+    _x = r * x * (1 - x)
+    return torch.nan_to_num(_x) * ~x_inf_mask - x * x_inf_mask
+
+def lyapunov_term(r, x):
+    x_inf_mask = torch.isinf(x)
+    term = torch.log(torch.abs(r * (1 - 2 * x)))
+    return torch.nan_to_num(term) * ~x_inf_mask - x * x_inf_mask
+
+def repeat(seq):
+    i = 0
+    l = len(seq)
+    while True:
+        yield seq[i%l]
+        i += 1
+
+def c_avg(prev, val, n):
+    """cumulative average
+    prev: previous result (should be zero for n=1)
+    val: next value
+    n: how many values given so far, including this value"""
+    prev_inf_mask = torch.isinf(prev)
+    _next = prev + (val - prev) / n
+    return torch.nan_to_num(_next) * ~prev_inf_mask + prev * prev_inf_mask
+
+# init
+
+for z in zooms:
+    # TODO shift & scale
+    pass
+
+x_min = origin[0]
+y_min = origin[1]
+x_max = origin[0] + span[0]
+y_max = origin[1] + span[1]
+
+
+
+
+
+def main():
+    x = torch.ones([h, w], device=dev, dtype=t_fp) / 2
+    ab = torch.zeros([h, w], device=dev, dtype=t_cfp)
+
+    lyapunov_avg = torch.zeros([h, w], device=dev, dtype=t_fp)
+
+    yspace = torch.linspace(y_min, y_max, h, dtype=t_fp, device=dev)
+    xspace = torch.linspace(x_min, x_max, w, dtype=t_fp, device=dev)
+    for _x in range(h):
+        ab[_x] += xspace
+    for _y in range(w):
+        ab[:, _y] += yspace * 1j
+
+    gen = itertools.islice(repeat(seq), iterations)
+
+    for idx, seq_elem in enumerate(gen):
+        if seq_elem == "C":
+            r = c
+        else:
+            r = ab.real if seq_elem == "A" else ab.imag
+
+        if idx > 0:
+            lyapunov_avg = c_avg(lyapunov_avg, lyapunov_term(r, x), idx)
+
+        if idx < iterations - 1:
+            x = logistic_map(r, x)
+
+        #msave(torch.tanh(lyapunov_avg / 4) / 2 + 0.5, f"avg_{idx}")
+
+    result = torch.tanh(lyapunov_avg) / 2 + 0.5
+    #msave(result, final_file)
+    #msave(1 - result**2, final_file + "sq")
+    #msave(1 - torch.sqrt(result), final_file + "sqrt")
+    #msave(result ** 2, final_file + "sq_p")
+    msave(torch.sqrt(result), final_file + "sqrt_p")
+
+
+if __name__ == "__main__":
+    t0 = time.perf_counter()
+    main()
+    print(f"done in {time.perf_counter() - t0}s")
+

raytrace.py

···
+import torch
+import time
+from math import sin, cos
+
+from util import msave
+
+out_file = "qjulia_11_40"
+
+# 2**10 = 1024
+scale = 2 ** 11
+w = 2 * scale // 3
+h = scale
+
+max_steps = 500
+mandel_iters = 20#50
+max_samples = 4
+
+noise_scale = 1. / (2 * scale)
+
+zoom = 1
+#frame_center = -0.5, -0.5, 0
+frame_center = 0, 0, -1 * zoom
+step_dir = 0, 0, 1
+#frame_span = 1, 1
+frame_span = 2 * zoom, 3 * zoom
+max_depth = 2 * zoom
+
+dev = "cuda"
+t_cfp = torch.cfloat
+t_fp = torch.float
+
+pi = 3.14159265358979323
+a, b, c = pi / 3, pi + 0.2834298, pi / 4
+
+sina, sinb, sinc = sin(a), sin(b), sin(c)
+cosa, cosb, cosc = cos(a), cos(b), cos(c)
+
+rot_x = torch.tensor([
+    [1,0,0],
+    [0,cosc,-sinc],
+    [0,sinc,cosc]], dtype=t_fp, device=dev)
+rot_y = torch.tensor([
+    [cosb,0,sinb],
+    [0,1,0],
+    [-sinb,0,cosb]], dtype=t_fp, device=dev)
+rot_z = torch.tensor([
+    [cosa,-sina,0],
+    [sina,cosa,0],
+    [0,0,1]], dtype=t_fp, device=dev)
+
+rot = rot_z @ rot_y @ rot_x
+
+quat_last = 0
+quat_c = [-1, 0.2, 0, 0]
+
+def hit_0(pos):
+    z = 0 #pos[2] * (0.2 * pos[2] + 0.5j)
+    c = torch.view_as_complex(pos[0:2].permute(1,2,0).contiguous())
+    for i in range(mandel_iters):
+        z = z**2 + pos[2]*z + c
+    return torch.abs(z) < 2
+
+def hit_1(pos):
+    #v = rot @ pos #torch.clone(pos)
+    v = (pos.permute(1,2,0) @ rot).permute(2,0,1)
+    for i in range(mandel_iters):
+        x = v[0]
+        y = v[1]
+        z = v[2]
+        x_term = 1j * torch.sqrt(y**2 + z**2)
+        y_term = 1j * torch.sqrt(x**2 + z**2)
+        z_term = 1j * torch.sqrt(y**2 + x**2)
+        v[0] = (x + x_term)**9 / 2 + (x - x_term)**9 / 2 + pos[0]
+        v[1] = (y + y_term)**9 / 2 + (y - y_term)**9 / 2 + pos[1]
+        v[2] = (z + z_term)**9 / 2 + (z - z_term)**9 / 2 + pos[2]
+    n = torch.norm(v, dim=0)
+    return ~(n < 2)
+
+def hit_3(pos):
+    v = (pos.permute(1,2,0) @ rot).permute(2,0,1)
+    v = torch.cat((v, torch.ones([1,h,w], dtype=t_fp, device=dev) * quat_last))
+    for i in range(mandel_iters):
+        r, a, b, c = v[0], v[1], v[2], v[3]
+        _r = r*r - a*a - b*b - c*c
+        _a = 2*r*a
+        _b = 2*r*b
+        _c = 2*r*c
+        v[0] = _r + quat_c[0]
+        v[1] = _a + quat_c[1]
+        v[2] = _b + quat_c[2]
+        v[3] = _c + quat_c[3]
+    n = torch.norm(v, dim=0)
+    return n < 4
+
+def hit_2(pos):
+    return torch.norm(pos, dim=0) > 0# 0.9
+
+def main():
+    start = torch.zeros([h, w], device=dev, dtype=t_cfp)
+    #res = torch.zeros([h,w], device=dev, dtype=t_fp)
+    res = torch.ones([h,w], device=dev, dtype=t_fp)
+
+    half_x = frame_span[0] / 2
+    half_y = frame_span[1] / 2
+    yspace = torch.linspace(-half_y, half_y, h, dtype=t_fp, device=dev)
+    xspace = torch.linspace(-half_x, half_x, w, dtype=t_fp, device=dev)
+    for _x in range(h):
+        start[_x] += xspace
+    for _y in range(w):
+        start[:, _y] += yspace * 1j
+
+    start = torch.stack((start.real - frame_center[0], start.imag - frame_center[1], torch.ones([h, w], dtype=t_fp, device=dev) * frame_center[2]))
+
+    step = torch.stack([torch.ones([h,w], device=dev, dtype=t_fp)*x for x in step_dir]) * max_depth / max_steps
+
+    s_res = []
+    v_res = []
+    for sample_idx in range(max_samples):
+        pos = start + torch.randn([3,h,w], device=dev, dtype=t_fp) * noise_scale
+        s_res.append(torch.ones_like(res))
+        v_res.append(torch.zeros_like(res))
+        for step_idx in range(max_steps):
+            pos += step
+            hit = hit_3(pos) #* hit_2(pos)
+            hit_adj = hit * (s_res[sample_idx] > (step_idx / max_steps))
+            hit_val = hit_adj * (step_idx / max_steps)
+            #hit_val = hit * (1 / max_steps)
+            v_res[sample_idx] = v_res[sample_idx] + hit * (1. / max_steps)
+            s_res[sample_idx] = (s_res[sample_idx] * ~hit_adj) + hit_val #hit / max_samples / max_steps
+            if step_idx % 10 == 999:
+                print(step_idx)
+                print(hit.sum())
+                print(hit_adj.sum())
+                print(s_res[sample_idx].sum())
+
+    s_res = torch.stack(s_res).mean(dim=0)
+    v_res = torch.stack(v_res).mean(dim=0)
+
+    #res = res % (1/5)
+    #res = torch.sqrt(res)# * 5)
+    msave(1-s_res, f"{out_file}_s")
+    msave(v_res, f"{out_file}_v")
+    #msave(1-res, out_file)
+    #res_mv = torch.cat([res, torch.flip(res, dims=(0,))])
+    #res_mh = torch.cat([res_mv, torch.flip(res_mv, dims=(1,))], dim=1)
+    #msave(res_mh, out_file)
+
+if __name__ == "__main__":
+    t0 = time.perf_counter()
+    main()
+    print(f"done in {time.perf_counter() - t0}")

slime.py

···
+import time
+import inspect
+from random import randint
+from pathlib import Path
+
+import torch
+from torchvision.transforms import GaussianBlur
+
+from util import *
+
+@settings
+def _s():
+    from math import tau
+
+    name = "slime"
+
+    scale = 2 ** 11
+    w = scale
+    h = scale
+
+    max_iterations = 1000
+    save_mod = 100
+
+    particle_count = 1000000 // 1
+    decay_rate = 0.9
+
+    view_angle = tau / 8
+    view_distance = 1.414 * 100
+
+    direction_count = 4
+    turn_amount = tau / direction_count * 1.2
+    move_distance = 1.414 * 3
+
+    blur_size = 5
+    blur_sigma = 1
+
+    dtype = torch.double
+    ctype = torch.cdouble
+    device = "cuda"
+
+def offset_coord(position, angle, distance):
+    res = torch.clone(position)
+    res[:,0] += (torch.sin(angle) * distance)
+    res[:,1] += (torch.cos(angle) * distance)
+    return res
+
+def clamp_coords(t, h, w):
+    t[:,0] = (t[:,0] + (h-1)) % (h-1)
+    t[:,1] = (t[:,1] + (w-1)) % (w-1)
+
+def idxput(tensor, indices, value):
+    tensor.index_put_((indices[:,0],indices[:,1]), value)
+
+@ifmain(__name__, _s)
+@timed
+def _main(settings):
+    globals().update(settings.get())
+    run_dir = time.strftime(f"%d.%m.%Y/{name}_t%H.%M.%S")
+    Path("out/" + run_dir).mkdir(parents=True, exist_ok=True)
+    with open(f"out/{run_dir}/settings.py", "w") as f:
+        f.write(settings.src)
+    torch.set_default_device(device)
+
+    blur = GaussianBlur(blur_size, blur_sigma)
+
+    # p_ denotes particle data
+    p_positions = torch.rand([particle_count,2])
+    p_positions[:,0] *= h
+    p_positions[:,1] *= w
+
+    p_directions = ((torch.rand(particle_count)*direction_count).floor() / direction_count) * 2 * tau
+
+    world = torch.zeros([3,h,w], dtype=dtype)
+    scratch = torch.zeros([3,h,w], dtype=dtype)
+    t0 = torch.tensor([0], dtype=dtype)
+    t1 = torch.tensor([1], dtype=dtype)
+
+    for iteration in range(max_iterations):
+
+        # sense
+        angles = (-view_angle, 0, view_angle)
+        angles = (p_directions + a for a in angles)
+        sensor_coords = (offset_coord(p_positions, a, view_distance) for a in angles)
+        sc = [sc for sc in sensor_coords]
+        for c in sc:
+            clamp_coords(c, h, w)
+
+        sci = [c.long() for c in sc]
+
+        senses = [world[0].take(c[:,0] * w + c[:,1]) for c in sci]
+        sense_neg = senses[0] > senses[1]
+        sense_pos = senses[2] > senses[1]
+        #sense_mid = ~sense_neg * ~sense_pos # forward highest
+        sense_out = sense_neg * sense_pos # forward lowest
+        sense_neg = sense_neg * ~sense_out # negative lowest
+        sense_pos = sense_pos * ~sense_out # negative highest
+
+        # rotate
+            #p_directions += torch.rand(particle_count) * sense_out * turn_amount
+        mask = torch.rand(particle_count) > 0.5
+        p_directions += mask * sense_out * turn_amount
+        p_directions -= ~mask * sense_out * turn_amount
+        p_directions += sense_pos * turn_amount
+        p_directions -= sense_neg * turn_amount
+
+        # move
+        p_positions = offset_coord(p_positions, p_directions, move_distance)
+        clamp_coords(p_positions, h, w)
+
+        # deposit
+        idxput(world[0], p_positions.long(), t1)
+
+        # diffuse
+        world = blur(world)
+
+        # render
+        if iteration % save_mod == 0 and iteration != 0:
+            msave(world[0], f"{run_dir}/{iteration}")
+
+        # decay
+        if iteration < max_iterations - 1:
+            world *= decay_rate
+
+    msave(world[0], f"{run_dir}/final")

util.py

···
+import torch
+from PIL import Image
+
+class Settings(dict):
+    __getattr__ = dict.__getitem__
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__
+
+    def __init__(self, _vars):
+        for key in _vars:
+            self[key] = _vars[key]
+
+class _Settings():
+    def __init__(self, _get, _src):
+        self.get = _get
+        self.src = _src
+
+def settings(f):
+    import inspect
+    f_src = inspect.getsource(f)
+    f_lines = f_src.split('\n')
+    f_lines = [l for l in f_lines if f"@settings" not in l]
+    f_lines.append("    return Settings(vars())")
+    new_src = "\n".join(f_lines)
+    exec(new_src)
+    return _Settings(eval(f.__name__),new_src)
+
+def first(l, p):
+    return next((idx,value) for idx,value in enumerate(l) if p(value))
+
+def timed(f):
+    import time
+    def _f(*args, **kwargs):
+        t0 = time.perf_counter()
+        f(*args, **kwargs)
+        print(f"{f.__name__}: {time.perf_counter() - t0}s")
+    return _f
+
+def ifmain(name, provide_arg=None):
+    if name == "__main__":
+        if provide_arg is None:
+            return lambda f: f()
+        return lambda f: f(provide_arg)
+    return lambda f: f
+
+def cpilify(z):
+    z_3d = torch.stack([z.real, z.imag, torch.zeros_like(z.real)])
+    z_norm = (z_3d / 2 + 0.5).clamp(0, 1)
+    z_np = z_norm.detach().cpu().permute(1, 2, 0).numpy()
+    z_bytes = (z_np * 255).round().astype("uint8")
+    return Image.fromarray(z_bytes)
+
+# complex, from -1 to 1 & -i to i
+def csave(x, f):
+    cpilify(x).save(f"out/{f}.png")
+
+# monochrome, 0 to 1
+def mpilify(z):
+    z_3d = torch.stack([z, z, z])
+    z_norm = z_3d.clamp(0, 1)
+    z_np = z_norm.detach().cpu().permute(1, 2, 0).numpy()
+    z_bytes = (z_np * 255).round().astype("uint8")
+    return Image.fromarray(z_bytes)
+
+def msave(x, f):
+    mpilify(x).save(f"out/{f}.png")
+
+# 3 channels
+def pilify(z):
+    z_norm = z.clamp(0, 1)
+    z_np = z_norm.detach().cpu().permute(1, 2, 0).numpy()
+    z_bytes = (z_np * 255).round().astype("uint8")
+    return Image.fromarray(z_bytes)
+
+def save(x, f):
+    pilify(x).save(f"out/{f}.png")
+
+# grid of complex numbers
+def cgrid(h,w,center,span,ctype=torch.cdouble,dtype=torch.double,**_):
+    g = torch.zeros([h, w], dtype=ctype)
+
+    low = center - span / 2
+    hi = center + span / 2
+
+    yspace = torch.linspace(low.imag, hi.imag, h, dtype=dtype)
+    xspace = torch.linspace(low.real, hi.real, w, dtype=dtype)
+
+    for _x in range(h):
+        g[_x] += xspace
+    for _y in range(w):
+        g[:, _y] += yspace * 1j
+
+    return g
+
+# result, iterations; iterations == -1 if no convergence before limit
+def gauss_seidel(a, b):
+    x = torch.zeros_like(b)
+    itlim = 1000
+    for it in range(1, itlim):
+        xn = torch.zeros_like(x)
+        for i in range(a.shape[0]):
+            s1 = a[i, :i].dot(xn[:i])
+            s2 = a[i, i+1:].dot(x[i+1:])
+            xn[i] = (b[i] - s1 - s2) / a[i, i]
+        if torch.allclose(x, xn, rtol=1e-8):
+            return xn, it
+        x = xn
+    return x, -1