merge old work into this repo, expose old sketch folder, begin work on webui

lib/diffusion/sdxl_vae.py

···
 import torch.nn as nn
 from torch.nn import functional as func
 from safetensors import safe_open as st_open
+from PIL import Image
 
 class Attention(nn.Module):
     def __init__(self):
···
             for key in sd.keys():
                 sd[key].copy_(file.get_tensor(key))
 
+approximation_matrix = [
+    [0.85, 0.85, 0.6], # seems to be mainly value
+    [-0.35, 0.2, 0.5], # mainly blue? maybe a little green, def not red
+    [0.15, 0.15, 0], # yellow. but mainly encoding texture not color, i think
+    [0.15, -0.35, -0.35] # inverted value? but also red
+]
+
+def save_approx_decode(latents, path):
+    lmin = latents.min()
+    l = latents - lmin
+    lmax = latents.max()
+    l = latents / lmax
+    l = l.float().mul_(0.5).add_(0.5)
+    ims = []
+    for lat in l:
+        apx_mat = torch.tensor(approximation_matrix).to("cuda")
+        approx_decode = torch.einsum("...lhw,lr -> ...rhw", lat, apx_mat).mul_(255).round()
+        #lat -= lat.min()
+        #lat /= lat.max()
+        im_data = approx_decode.permute(1,2,0).detach().cpu().numpy().astype("uint8")
+        #im_data = im_data.round().astype("uint8")
+        im = Image.fromarray(im_data).resize(size=(im_data.shape[1]*8,im_data.shape[0]*8), resample=Image.NEAREST)
+        ims += [im]
+
+    #clear_output()
+    for im in ims:
+        #im.save(f"out/tmp_approx_decode/{index:06d}.bmp")
+        im.save(path)
+        #display(im)
+

lib/internal/witchery.py

···
+import dis
+
+class Errs(type):
+    def __iter__(_class):
+        return iter((k,v) for (k,v) in vars(_class).items() if not k.startswith("_"))
+
+    def __repr__(_class):
+        return str(list(k for (k,v) in _class))
+
+def errs(errs):
+    def _errs(fn):
+        fn.errs = errs
+        return fn
+    return _errs
+
+# stuff above this line i probs move to another file
+
+class _Errs(metaclass=Errs):
+    FOUND_RETURN = "encountered return instruction"
+    IN_DEFINITION = "error in fn definition"
+    IN_EXECUTION = "error in fn execution"
+    NON_DICT_RETURN = "fn dumped a non-dict"
+
+def _modify_to_dump_locals(fn, fn_source, deftime_globals, log):
+    instructions = dis.get_instructions(fn)
+    for instruction in instructions:
+        if instruction.opcode == dis.opmap["RETURN_VALUE"]:
+            return (False, _Errs.FOUND_RETURN)
+    fn_source = fn_source
+    fn_source_modified = f"{fn_source}\n    return locals()"
+    sandbox = {}
+    try:
+        exec(fn_source_modified, globals=deftime_globals, locals=sandbox)
+    except KeyboardInterrupt:
+        raise
+    except:
+        log.indented().trace(source=fn)
+        return (False, _Errs.IN_DEFINITION)
+    return (True, sandbox[fn.__name__])
+
+@errs(_Errs)
+def try_dump_locals(fn, fn_source, args, kwargs, deftime_globals, log):
+    (success, fn_or_err) = _modify_to_dump_locals(fn, fn_source, deftime_globals, log)
+    if not success:
+        return (False, fn_or_err)
+    try:
+        res = fn_or_err(*args, **kwargs)
+    except KeyboardInterrupt:
+        raise
+    except:
+        log.indented().trace(source=fn)
+        return (False, _Errs.IN_EXECUTION)
+    if not isinstance(res, dict):
+        return (False, _Errs.NON_DICT_RETURN)
+    return (True, res)

lib/iter.py

···
 def first(l, p):
     return next((idx,value) for idx,value in enumerate(l) if p(value))
 
+# just use itertools.pairwise
 def pairs(l):
     return zip(l, l[1:])
 

lib/log.py

···
 import sys
 import inspect
 import traceback
+import time
 
 from dataclasses import dataclass
 from typing import Optional
···
         tag_color = "\033[31m"
     elif mode == "success":
         tag_color = "\033[92m"
+    elif mode == "warning":
+        tag_color = "\033[33m"
     else:
         tag_color = "\033[96m"
     print(f"{tag_color}{' '*indent}[{tag}]{reset_color} {content}{final_color}")

main.py

···
 from lib import util
 from lib.log import Logger, inner_log
 from lib.internal.parse import lsnap
+from lib.internal.witchery import try_dump_locals
 
 parser = ArgParser("snakepyt")
 #parser.add_argument("sketch", help="the sketch to run", type=str)
···
         schedule.append((fn, args))
     return (schedule, _schedule)
 
-def modify_to_dump_locals(fn, deftime_globals, log, sources):
-    instructions = dis.get_instructions(fn)
-    for instruction in instructions:
-        if instruction.opcode == dis.opmap["RETURN_VALUE"]:
-            return (False, "encountered return instruction")
-    fn_source = sources[fn.__name__]
-    fn_source_modified = f"{fn_source}\n    return locals()"
-    sandbox = {}
-    try:
-        exec(fn_source_modified, globals=deftime_globals, locals=sandbox)
-    except KeyboardInterrupt:
-        raise
-    except:
-        log.indented().trace(source=fn)
-        return (False, "error in definition")
-    return (True, sandbox[fn.__name__])
-
-
 def run(fn, arg, partial_id, outer_scope, log, sources, finalizer=None):
     scope = dict(outer_scope)
     schedule, schedule_fn = establish_scheduler()
     scope["schedule"] = schedule_fn
     scope["print"] = inner_log(source=fn, indent=4)
-    #fn_name = fn.__name__
     run_id = partial_id + (arg,)
-    #instructions = dis.get_instructions(fn)
-    #for instruction in instructions:
-    #    if instruction.opcode == dis.opmap["RETURN_VALUE"]:
-    #        log(f"{fn_name} cannot be scheduled because it has a return instruction", mode="error")
-    #        return False
-    #fn_source = sources[fn.__name__]
-    #fn_source_modified = f"{fn_source}\n    return locals()"
-    #sandbox = {}
 
-    (succ, fn_or_err) = modify_to_dump_locals(fn, scope, log, sources)
-    if not succ:
-        log.indented().log(fn_or_err, mode="error")
-        return False
-    fn_modified = fn_or_err
-    try:
-        #exec(fn_source_modified, globals=scope, locals=sandbox)
-        #fn_modified = sandbox[fn_name]
-        if arg is not None:
-            fn_locals = fn_modified(arg)
+    args = [] if arg is None else [arg]
+    (success, locals_or_err) = try_dump_locals(fn, sources[fn.__name__], args, {}, scope, log)
+    if success:
+        scope.update(locals_or_err)
+    else:
+        err = locals_or_err
+        errs = try_dump_locals.errs
+        if err == errs.NON_DICT_RETURN:
+            log.indented().log("could not extract locals. check for an early return", mode="warning")
         else:
-            fn_locals = fn_modified()
-        scope.update(fn_locals)
-    except KeyboardInterrupt:
-        raise
-    except:
-        log.indented().trace(source=fn)
-        return False
+            log.indented().log(locals_or_err, mode="error")
+            return False
+
+
     for (fn, args) in schedule:
         if args is not None:
             for arg in args:
···
     if sketch_name in persistent_hashes:
         if bytecode_hash == persistent_hashes[sketch_name]:
             return True
-    (success, fn_or_err) = modify_to_dump_locals(persistent_fn, module_globals, log, sources)
+    (success, locals_or_err) = try_dump_locals(persistent_fn, sources[persistent_fn.__name__], [], {}, module_globals, log)
     if success:
-        try:
-            fn_locals = fn_or_err()
-        except KeyboardInterrupt:
-            raise
-        except:
-            log.indented().trace(source=persistent_fn)
-            return False
         persistent_hashes[sketch_name] = bytecode_hash
-        persistent_state.update(fn_locals)
+        persistent_state.update(locals_or_err)
     else:
-        log(f"failed to run persistent function: {fn_or_err}", mode="error")
+        log(f"failed to run persistent function: {locals_or_err}", mode="error")
+
     return success
 
 
···
     if "prefix" in persistent_state:
         message = " ".join([persistent_state["prefix"], message])
 
+    message = message.lstrip()
+
+    if message.startswith("."):
+        if message.rstrip() == ".":
+            pyt_print(persistent_state)
+        else:
+            segments = [segment.strip() for segment in message.split(".")][1:]
+            selection = ("base scope", persistent_state)
+            for segment in segments:
+                if segment == "":
+                    pyt_print("repeated dots (..) are redundant", mode="warning")
+                    continue
+                try:
+                    selection = (segment, selection[1][segment])
+                    pyt_print(f"{selection[0]}: {selection[1]}")
+                except KeyError:
+                    pyt_print(f"no \"{segment}\" in {selection[0]}", mode="error")
+                    continue
+                except TypeError:
+                    pyt_print(f"{selection[0]} is not a scope", mode="error")
+                    pyt_print.indented()(f"{selection[0]}: {selection[1]}", mode="info")
+
+        continue
+
+
     (command, remainder) = lsnap(message)
 
-    if command == "state":
-        pyt_print(persistent_state)
     if command == "flush":
         persistent_state = {}
         persistent_hashes = {}
         pyt_print("state flushed")
+        continue
     if command in ["exit", "quit", ":q", ",q"]:
         pyt_print.blank().log("goodbye <3").blank()
         repl_continue = False
         continue
+    if command in ["hello", "hi"]:
+        pyt_print("hiii :3")
+        continue
     if command == "crash":
         raise Exception("crashing on purpose :3")
+        continue
     if command == "run":
         sketch_name, remainder = lsnap(remainder)
 
···
         with open(f"out/{run_dir}/.snakepyt", "w") as metadata:
             metadata.write(f"snakepyt version {snakepyt_version[0]}.{snakepyt_version[1]}\n")
 
-        if hasattr(sketch, "init"):
+        if hasattr(sketch, "main"):
             if hasattr(sketch, "final"):
                 finalizer = sketch.final.__code__
             else:
                 finalizer = None
             try:
-                run(sketch.init, None, (), sketch.__dict__, log, sources, finalizer)
+                run(sketch.main, None, (), sketch.__dict__, log, sources, finalizer)
             except KeyboardInterrupt:
                 pyt_print.blank().log("aborted", mode="info").blank()
                 continue
         else:
             log("sketch has no init function", mode="error", indent=4)
-            #run(None, settings_functions, (), sketch.__dict__)
 
         log(f"finished all runs in {perf_counter() - t0:.3f}s")
+        continue
+
+    pyt_print(f"command: {command}", mode="info")
 

original.py

···
-
-import inspect
-import dis
-import traceback
-import hashlib
-import sys
-import os
-import time
-import shutil
-from pathlib import Path
-from argparse import ArgumentParser as ArgParser
-from importlib import import_module, reload
-from time import perf_counter
-
-from lib import util
-from lib.log import Logger, inner_log
-from lib.internal.parse import lsnap
-
-parser = ArgParser("snakepyt")
-#parser.add_argument("sketch", help="the sketch to run", type=str)
-args = parser.parse_args()
-
-snakepyt_version = (0, 0)
-
-def establish_scheduler():
-    schedule = []
-    def _schedule(fn, args):
-        schedule.append((fn, args))
-    return (schedule, _schedule)
-
-def modify_to_dump_locals(fn, deftime_globals, log, sources):
-    instructions = dis.get_instructions(fn)
-    for instruction in instructions:
-        if instruction.opcode == dis.opmap["RETURN_VALUE"]:
-            return (False, "encountered return instruction")
-    fn_source = sources[fn.__name__]
-    fn_source_modified = f"{fn_source}\n    return locals()"
-    sandbox = {}
-    try:
-        exec(fn_source_modified, globals=deftime_globals, locals=sandbox)
-    except KeyboardInterrupt:
-        raise
-    except:
-        log.indented().trace(source=fn)
-        return (False, "error in definition")
-    return (True, sandbox[fn.__name__])
-
-
-def run(fn, arg, partial_id, outer_scope, log, sources, finalizer=None):
-    scope = dict(outer_scope)
-    schedule, schedule_fn = establish_scheduler()
-    scope["schedule"] = schedule_fn
-    scope["print"] = inner_log(source=fn, indent=4)
-    fn_name = fn.__name__
-    run_id = partial_id + (arg,)
-    instructions = dis.get_instructions(fn)
-    for instruction in instructions:
-        if instruction.opcode == dis.opmap["RETURN_VALUE"]:
-            log(f"{fn_name} cannot be scheduled because it has a return instruction", mode="error")
-            return False
-    fn_source = sources[fn.__name__]
-    fn_source_modified = f"{fn_source}\n    return locals()"
-    sandbox = {}
-    try:
-        exec(fn_source_modified, globals=scope, locals=sandbox)
-        fn_modified = sandbox[fn_name]
-        if arg is not None:
-            fn_locals = fn_modified(arg)
-        else:
-            fn_locals = fn_modified()
-        scope.update(fn_locals)
-    except KeyboardInterrupt:
-        raise
-    except:
-        log.indented().trace(source=fn)
-        return False
-    for (fn, args) in schedule:
-        if args is not None:
-            for arg in args:
-                run(fn, arg, run_id, scope, log.indented(), sources)
-        else:
-            t0 = perf_counter()
-            log(f"begin {fn.__name__} {run_id}")
-            success = run(fn, None, run_id, scope, log.indented(), sources)
-            log(f"done in {perf_counter() - t0:.3f}s", mode="success" if success else "error")
-            log.blank()
-    if finalizer is not None:
-        try:
-            exec(finalizer, globals=scope)
-        except KeyboardInterrupt:
-            raise
-        except:
-            log.indented().trace(source=fn)
-            return False
-    return True
-
-persistent_state = {}
-persistent_hashes = {}
-def handle_persistent(persistent_fn, module_globals, log, sources):
-    bytecode_hash = hashlib.sha256(persistent_fn.__code__.co_code).hexdigest()
-    if sketch_name in persistent_hashes:
-        if bytecode_hash == persistent_hashes[sketch_name]:
-            return True
-    (success, fn_or_err) = modify_to_dump_locals(persistent_fn, module_globals, log, sources)
-    if success:
-        try:
-            fn_locals = fn_or_err()
-        except KeyboardInterrupt:
-            raise
-        except:
-            log.indented().trace(source=persistent_fn)
-            return False
-        persistent_hashes[sketch_name] = bytecode_hash
-        persistent_state.update(fn_locals)
-    else:
-        log(f"failed to run persistent function: {fn_or_err}", mode="error")
-    return success
-
-
-try:
-    username = os.getlogin()
-except:
-    username = None
-
-pyt_print = Logger().mode("success").tag("snakepyt")
-pyt_print(f"hello {username}! <3" if username else "hello! <3")
-pyt_print.blank()
-
-repl_continue = True
-while repl_continue:
-    try:
-        message = pyt_print.input(username)
-    except (KeyboardInterrupt, EOFError):
-        pyt_print.blank().log("goodbye <3").blank()
-        repl_continue = False
-        continue
-
-    if "prefix" in persistent_state:
-        message = " ".join([persistent_state["prefix"], message])
-
-    (command, remainder) = lsnap(message)
-
-    if command == "state":
-        pyt_print(persistent_state)
-    if command == "flush":
-        persistent_state = {}
-        persistent_hashes = {}
-        pyt_print("state flushed")
-    if command in ["exit", "quit", ":q", ",q"]:
-        pyt_print.blank().log("goodbye <3").blank()
-        repl_continue = False
-        continue
-    if command == "crash":
-        raise Exception("crashing on purpose :3")
-    if command == "run":
-        sketch_name, remainder = lsnap(remainder)
-
-        try:
-            pyt_print(f"loading sketch \"{sketch_name}\"", mode="info")
-            module_name = f"sketch.{sketch_name}"
-            if module_name in sys.modules:
-                del sys.modules[module_name]
-            sketch = import_module(module_name)
-        except ModuleNotFoundError:
-            pyt_print.indented().trace()
-            pyt_print("no such sketch", mode="error", indent=4).blank()
-            continue
-        except KeyboardInterrupt:
-            pyt_print("aborted", mode="info").blank()
-            continue
-        except:
-            pyt_print.indented().trace()
-            continue
-
-        sources = { name : inspect.getsource(member)
-                    for name, member in inspect.getmembers(sketch)
-                    if inspect.isfunction(member) and member.__module__ == module_name
-                   }
-
-        log = pyt_print.tag(sketch_name).mode("info")
-
-        t0 = perf_counter()
-        if hasattr(sketch, "persistent"):
-            if not handle_persistent(sketch.persistent, sketch.__dict__, log, sources):
-                log.blank()
-                continue
-
-        sketch.__dict__.update(persistent_state)
-
-        run_dir = time.strftime(f"%d.%m.%Y/{sketch_name}_t%H.%M.%S")
-        sketch.__dict__["run_dir"] = run_dir
-        Path("out/" + run_dir).mkdir(parents=True, exist_ok=True)
-
-        shutil.copy(sketch.__file__, f"out/{run_dir}/{sketch_name}.py")
-
-        with open(f"out/{run_dir}/.snakepyt", "w") as metadata:
-            metadata.write(f"snakepyt version {snakepyt_version[0]}.{snakepyt_version[1]}\n")
-
-        if hasattr(sketch, "init"):
-            if hasattr(sketch, "final"):
-                finalizer = sketch.final.__code__
-            else:
-                finalizer = None
-            try:
-                run(sketch.init, None, (), sketch.__dict__, log, sources, finalizer)
-            except KeyboardInterrupt:
-                pyt_print.blank().log("aborted", mode="info").blank()
-                continue
-        else:
-            log("sketch has no init function", mode="error", indent=4)
-            #run(None, settings_functions, (), sketch.__dict__)
-
-        log(f"finished all runs in {perf_counter() - t0:.3f}s")
-

sketch/old/pre_snakepyt/abelian_sandpile.py

···
+import math
+import random
+import time
+from pathlib import Path
+
+import torch
+
+from util import *
+
+@settings
+def _s():
+    scale = 2 ** 10
+    w = scale + 1
+    h = scale + 1
+
+    iterations = 500000
+    debug_modulus = 10000
+
+    offsets = [[-n,0] for n in range(17)] + [[n,0] for n in range(17)] + [[0,17]]
+    drop_points = [[h//2,w//2+n] for n in range(17)]
+
+    wrap_vertical = False
+    wrap_horizontal = True
+
+    name = "sand"
+    out_file = ""
+
+    device="cuda"
+    ctype=torch.cdouble
+    dtype=torch.double
+
+@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)
+    torch.set_default_dtype(dtype)
+    dim = [h, w]
+
+    grid = torch.zeros(dim, dtype=torch.int)
+
+    max_c = 500
+    c_total = 0
+    last_c_total = 500
+
+    n = len(offsets)
+
+    for iteration in range(iterations):
+        last = torch.clone(grid)
+        for p in drop_points:
+            #grid[p[0],p[1]] += 1
+            grid[p[0],p[1]] = n
+            # todo: make this index_put_
+        peaks = (grid >= n).int()
+        c = 0
+        while torch.count_nonzero(peaks):
+            c += 1
+
+            shift = -(n * peaks)
+            for offset in offsets:
+                p_shift = torch.roll(peaks, offset, dims=[0,1])
+                if offset[0] < 0 and not wrap_vertical:
+                    p_shift[offset[0]:] = 0
+                elif not wrap_vertical:
+                    p_shift[:offset[0]] = 0
+                if offset[1] < 0 and not wrap_horizontal:
+                    p_shift[:,offset[1]:] = 0
+                elif not wrap_horizontal:
+                    p_shift[:,:offset[1]] = 0
+                shift += p_shift
+            grid += shift
+
+
+            peaks = (grid >= n).int()
+        c_total += c
+        if c > 20:
+            s = "*" if c >= max_c else ""
+            print(f"{iteration}: {c}{s}")
+            if c > (0.98 * max_c):
+                msave(grid / (n-1), f"{run_dir}/{out_file}_{iteration}_{c}")
+                msave(last / (n-1), f"{run_dir}/{out_file}_{iteration}_prev")
+            if c > max_c:
+                max_c = c
+        if iteration % debug_modulus == 0:
+            msave(grid / (n-1), f"{run_dir}/{out_file}_{iteration}_")
+        if c_total > (last_c_total * 2):
+            msave(grid / (n-1), f"{run_dir}/{out_file}_{iteration}_ct{c_total}")
+            last_c_total = c_total
+
+    msave(grid / (n-1), out_file)
+

sketch/old/pre_snakepyt/blus.py

···
+import time
+from pathlib import Path
+from random import random
+
+import torch
+
+from util import *
+
+type fp_range = tuple[float, float]
+type fp_region2 = tuple[fp_range, fp_range]
+type fp_coords2 = tuple[float, float]
+type hw = tuple[int, int]
+type region_mapping = tuple[fp_region2, hw]
+
+@settings
+def _s():
+    import math
+    name = "blus"
+    device = "cuda"
+    t_real = torch.double
+    t_complex = torch.cdouble
+
+    tau = 3.14159265358979323 * 2
+
+    flow = 0.01
+    ebb = 0 #1 - (1 / 1000)
+    randomize = False
+    random_range = tau #/ 50
+    rescaling = False
+    show_gridlines = False
+
+    particle_count = 1000000
+
+    iterations = 5000
+
+    # 2 ** 9 = 512; 10 => 1024; 11 => 2048
+    scale_power = 11
+    scale = 2 ** scale_power
+
+    origin = 0, 0
+    span = 5, 5
+
+    stretch = 1, 1
+
+    zooms = [
+            ]
+
+def blus(a, b):
+    theta_a = a.angle()
+    return torch.polar(a.abs() + b.abs() * torch.cos(b.angle() - theta_a), theta_a)
+
+@ifmain(__name__, _s)
+@timed
+def run(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)
+    Path("out/" + run_dir + "/aggregate").mkdir(parents=True, exist_ok=True)
+    Path("out/" + run_dir + "/frame").mkdir(parents=True, exist_ok=True)
+
+    global span
+
+    x_min = origin[0] - (span[0] / 2)
+    y_min = origin[1] - (span[1] / 2)
+
+    for ((xa, xb), (ya, yb)) in zooms:
+        x_min += span[0] * xa
+        y_min += span[1] * ya
+        span = span[0] * (xb - xa), span[1] * (yb - ya)
+
+    x_max = x_min + span[0]
+    y_max = y_min + span[1]
+
+    aspect = span[0] * stretch[0] / (span[1] * stretch[1])
+
+    if aspect < 1:
+        h = scale
+        w = int(scale * aspect)
+    else:
+        w = scale
+        h = int(scale / aspect)
+
+    x_range = (x_min, x_max)
+    y_range = (y_min, y_max)
+    region = (x_range, y_range)
+    mapping = (region, (h,w))
+
+    p_positions = (torch.rand([particle_count], device=device, dtype=t_complex) - (0.5 + 0.5j)) * 4
+    #p_positions.imag = torch.linspace(0, tau, particle_count)
+    #p_positions.real = torch.linspace(0, 1, particle_count)
+    #p_positions = torch.polar(p_positions.real, p_positions.imag)
+
+    p_colors = torch.ones([particle_count,3], device=device, dtype=t_real)
+    color_rotation = torch.linspace(0, tau / 4, particle_count)
+
+    p_colors[:,0] = p_positions.real / 4 + 0.5
+    p_colors[:,2] = p_positions.imag / 4 + 0.5
+
+    #p_colors[:,0] = torch.cos(color_rotation)
+    p_colors[:,1] *= 0
+    #p_colors[:,2] *= 0
+    #p_colors[:,2] = torch.sin(color_rotation)
+
+    canvas = torch.zeros([3, h, w], device=device, dtype=t_real)
+    scratch = torch.zeros([h, w, 3], device=device, dtype=t_real)
+    gridlines = torch.zeros([h,w], device=device)
+
+    def project(p):
+        return torch.view_as_real(p).permute(1,0)
+
+    ones = torch.ones_like(p_positions)
+    global direction
+    direction = flow * torch.polar(ones.real, 1 * tau * ones.real)
+    def next_positions(p):
+        global direction
+        #return blus(p, ones) - ones
+        if randomize:
+            direction = flow * torch.polar(ones.real, (random() * random_range - random_range / 2) * ones.real)
+        result = blus(p, direction) - direction * ebb
+        res_abs = result.abs()
+        if rescaling:
+            result = torch.polar(res_abs / res_abs.max(), result.angle())
+        return result
+
+    gridlines[(h-1)//2,:] = 1
+    gridlines[:,(w-1)//2] = 1
+
+    for line in range(1,6):
+        pos = line - y_min
+        pos *= (h-1) / (y_max - y_min)
+        try:
+            gridlines[math.floor(pos), :] = 0.8
+        except:
+            pass
+    for line in range(1,6):
+        pos = -line - y_min
+        pos *= (h-1) / (y_max - y_min)
+        try:
+            gridlines[math.floor(pos), :] = 0.8
+        except:
+            pass
+    for line in range(1,6):
+        pos = line - x_min
+        pos *= (w-1) / (x_max - x_min)
+        try:
+            gridlines[:,math.floor(pos)] = 0.8
+        except:
+            pass
+    for line in range(1,6):
+        pos = -line - x_min
+        pos *= (w-1) / (x_max - x_min)
+        try:
+            gridlines[:,math.floor(pos)] = 0.8
+        except:
+            pass
+    for line in range(1,10):
+        pos = -(line / 10) - y_min
+        pos *= (h-1) / (y_max - y_min)
+        try:
+            gridlines[math.floor(pos), :] = 0.3
+        except:
+            pass
+    for line in range(1,10):
+        pos = (line / 10) - y_min
+        pos *= (h-1) / (y_max - y_min)
+        try:
+            gridlines[math.floor(pos), :] = 0.3
+        except:
+            pass
+
+    def insert_at_coords(coords, values, target, mapping: region_mapping):
+        (region, hw) = mapping
+        (xrange, yrange) = region
+        (h,w) = hw
+        (x_min, x_max) = xrange
+        (y_min, y_max) = yrange
+
+        mask = torch.ones([particle_count], device=device)
+        mask *= (coords[1] >= x_min) * (coords[1] <= x_max)
+        mask *= (coords[0] >= y_min) * (coords[0] <= y_max)
+        in_range = mask.nonzero().squeeze()
+
+        # TODO: combine coord & value tensors so there's only one index_select necessary
+        coords_filtered = torch.index_select(coords.permute(1,0), 0, in_range)
+        values_filtered = torch.index_select(values, 0, in_range)
+
+        coords_filtered[:,1] -= x_min
+        coords_filtered[:,1] *= (w-1) / (x_max - x_min)
+        coords_filtered[:,0] -= y_min
+        coords_filtered[:,0] *= (h-1) / (y_max - y_min)
+        indices = coords_filtered.long()
+
+        target.index_put_((indices[:,0],indices[:,1]), values_filtered, accumulate=True)
+
+
+    for iteration in range(iterations):
+        p_projected = project(p_positions).clone()
+
+        if iteration % 1 == 0:
+
+            scratch *= 0
+            insert_at_coords(p_projected, p_colors, scratch, mapping)
+            canvas += scratch.permute(2,0,1)
+
+
+            temp = canvas.clone()
+            #for d in range(3):
+            #    temp[d] -= temp[d].mean()
+            #    temp[d] /= 8 * temp[d].std()
+            #    temp[d] -= temp[d].min()
+
+            temp /= temp.max()
+
+            save(1 - temp.sqrt().sqrt(), f"{run_dir}/aggregate/{iteration:06d}")
+            #scratch /= scratch.max()
+            #scratch = scratch.sqrt().sqrt()
+            if show_gridlines:
+                scratch[:,:,1] += gridlines
+                scratch[:,:,2] += gridlines
+            save(scratch.permute(2,0,1), f"{run_dir}/frame/{iteration:06d}")
+
+        p_positions = next_positions(p_positions)
+
+    #for d in range(3):
+    #    canvas[d] -= canvas[d].mean()
+    #    canvas[d] /= 8 * canvas[d].std()
+    #    canvas[d] -= canvas[d].min()
+
+    canvas /= canvas.max()
+    save(1 - canvas, f"{run_dir}/final")
+
+    with open(f"out/{run_dir}/settings.py", "w") as f:
+        f.write(settings.src)
+    torch.set_default_device(device)
+

sketch/old/pre_snakepyt/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**11
+    w = scale
+    h = scale
+
+    max_iterations = 100000
+    save_every = 100
+
+    v_diffusion_rate = 2000
+    m_diffusion_rate = 0.5
+    m_timescale = 0.01
+    timescale = 0.001
+
+    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]) * 1 #* 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] = -10 * 40
+        velocities[1,3 * h//4:(3*h//4)+10,w//2-10:w//2+10] = -9 * 40
+        velocities[0,h//4:(h//4)+10,w//2-3:w//2+3] = 10 * 40
+        velocities[1,h//4:(h//4)+10,w//2-3:w//2+3] = 20 * 40
+        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.99
+        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")
+
+
+
+
+
+
+
+
+
+
+
+

sketch/old/pre_snakepyt/fluid_clean.py

···
+# pytorch adaptation of the methods described here:
+# https://www.dgp.toronto.edu/public_user/stam/reality/Research/pdf/GDC03.pdf (overall design)
+# https://www.karlsims.com/fluid-flow.html (divergence clearing step)
+
+# only dependencies are pillow and pytorch
+
+import math
+import time
+from pathlib import Path
+from PIL import Image
+
+import torch
+from torch.nn.functional import conv2d
+
+# settings
+
+result_directory = "out"
+Path(result_directory).mkdir(parents=True, exist_ok=True)
+Path(result_directory + "/velocities").mkdir(parents=True, exist_ok=True)
+Path(result_directory + "/densities").mkdir(parents=True, exist_ok=True)
+
+scale = 10
+w = 2**scale
+h = 2**scale
+
+max_iterations = 100000
+save_every = 1
+
+velocity_diffusion_rate = 20
+density_diffusion_rate = None
+density_timescale = 0.005
+timescale = 0.01
+
+torch.set_default_device("cuda")
+
+diffusion_solver_steps = 10
+divergence_clearing_steps = 20
+
+
+# save images
+
+def monochrome_save(tensor, filename):
+    """save a 2d tensor of shape (h,w) as a monochrome image"""
+    scaled = torch.clone(tensor).clamp_(0,1).mul_(255).round().type(torch.uint8)
+    # copy to make 3 color channels
+    rearranged = scaled.unsqueeze(2).expand(-1, -1, 3).cpu().numpy()
+    Image.fromarray(rearranged).save(f"{result_directory}/{filename}.png")
+
+def save(tensor, filename):
+    """save a 3d tensor with shape (3,h,w) as an rgb image"""
+    scaled = torch.clone(tensor).clamp_(0, 1).mul_(255).round().type(torch.uint8)
+    rearranged = scaled.permute(1, 2, 0).cpu().numpy()
+    Image.fromarray(rearranged).save(f"{result_directory}/{filename}.png")
+
+
+
+
+
+# convolution
+
+# the backbone of this simulation is a handful of discrete convolutions
+# detailed definition here: https://en.wikipedia.org/wiki/Convolution#Discrete_convolution
+
+# roughly: we have a discrete field and a "kernel".
+# we look at a small window of values in the field, multiply each value there by a number,
+# and sum the results. the kernel describes what number to multiply each value in the window by.
+# we do this for every possible window, centered on every point we can fit the window around,
+# skipping points where the window would hang off the edge of the field. thus, we end up with
+# a result that's slightly smaller than the field
+
+def convolve(field, kernel):
+    # conv2d works in batches & we only ever need to do a 1-element batch
+    # "unsqueeze" is pytorch's absolutely bizarre term for wrapping a tensor with another dimension. like going from [1,2] to [[1,2]]
+    return conv2d(field.unsqueeze(0), kernel, bias=None, padding=[0], stride=[1])[0]
+
+# convolution kernels:
+
+# total of values of immediate neighbors in all 4 cardinal directions
+diffusion_kernel = torch.tensor([[[
+    [0, 1, 0],
+    [1, 0, 1],
+    [0, 1, 0]
+]]], dtype=torch.float)
+
+# gradient of divergence in the x direction
+x_div_kernel = torch.tensor([[[
+    [ 1, 2, 1],
+    [-2,-4,-2],
+    [ 1, 2, 1]
+]]], dtype=torch.float)
+
+div_opp_kernel = torch.tensor([[[
+    [ 1, 0,-1],
+    [ 0, 0, 0],
+    [-1, 0, 1]
+]]], dtype=torch.float)
+
+# gradient of divergence in the y direction
+y_div_kernel = torch.tensor([[[
+    [ 1,-2, 1],
+    [ 2,-4, 2],
+    [ 1,-2, 1]
+]]], dtype=torch.float)
+
+# various boundary conditions
+
+def continuous_boundary(field):
+    """set the border values of the provided discretized 2d scalar field (that is, a 2d array of numbers) to be equal to their inner neighbors (& average the corner points)"""
+    field[0] = field[1]
+    field[-1] = field[-2]
+    field[:,0] = field[:,1]
+    field[:,-1] = field[:,-2]
+    # this is doing four indexing operations at once, getting all four corners in one go
+    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_vertical_boundary(field):
+    """set the border values of a discretized 2d field to be equal to their inner neighbors horizontally but opposite their neighbors vertically. average at corners"""
+    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_horizontal_boundary(field):
+    """set border values equal vertically, opposite horizontally, average at corners"""
+    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, boundary_condition, timescale):
+    """diffuse the scalar field. basically means repeatedly applying a blur to it"""
+    a = timescale * rate
+    result = torch.clone(field)
+    for n in range(diffusion_solver_steps):
+        # scaled sum of surrounding points
+        convolution = a * convolve(result, diffusion_kernel)
+        # convolution operation doesn't produce output for the border rows & columns, thus the "1:n-1" indexing
+        result[1:h-1,1:w-1] = field[1:h-1,1:w-1] + convolution
+        result /= 1 + 4 * a
+        boundary_condition(field)
+    return result
+
+# generate indices outside the function for a tiny performance improvement
+indices_x = torch.arange(1,w-1,dtype=torch.float).repeat(h-2,1)
+indices_y = torch.arange(1,h-1,dtype=torch.float).repeat(w-2,1).t()
+indices = torch.stack((indices_y, indices_x))
+
+# defining these here to reuse the same gpu memory for each advect() call
+# TODO do the same for the rest of the tensors
+offsets = indices.clone()
+inverse_offsets = indices.clone()
+indices_int = indices.int()
+next_indices = indices_int.clone()
+
+
+# advection
+# this is a point where i deviate from the original paper.
+# i take the velocity at every point, and add a scaled version of that to the index at each point
+# to find where that velocity will carry whatever is being advected along it in one timestep.
+# that will be a point that i decompose into the integer component & fractional component. like (3.4, 1.2) -> (3,1) + (0.4, 0.2)
+# the point itself will be somewhere between two indices on each axis. the integer component determines the first of these two
+# indices, and the fractional component determines which of the points it's closer to
+
+# the paper i based this on did this weird backward version of that, where they use the negative velocity at each point to
+# select a point & interpolate the values of the field at the surrounding points
+
+def advect(field, velocities, timescale):
+    """given any field, and a velocity field of the same height & width, move the values in the field a small distance in the direction of the velocity field"""
+    offsets[...] = indices
+    offsets.add_(timescale * velocities[:,1:h-1,1:w-1])
+    offsets[1].clamp_(1.5, w - 2.5)
+    offsets[0].clamp_(1.5, h - 2.5)
+    indices_int[...] = offsets.int()
+    offsets.sub_(indices_int)
+    inverse_offsets[...] = 1 - offsets
+    next_indices[...] = indices_int + 1
+    inds_x_all = torch.stack([indices_int[1], next_indices[1], indices_int[1], next_indices[1]])
+    inds_y_all = torch.stack([indices_int[0], indices_int[0], next_indices[0], next_indices[0]])
+    values = torch.stack([field[:,1:h-1,1:w-1] * inverse_offsets[1] * inverse_offsets[0],
+                        field[:,1:h-1,1:w-1] * offsets[1] * inverse_offsets[0],
+                        field[:,1:h-1,1:w-1] * inverse_offsets[1] * offsets[0],
+                        field[:,1:h-1,1:w-1] * offsets[1] * offsets[0]])
+    res = torch.zeros_like(field)
+    res[0].index_put_((inds_y_all, inds_x_all), values[:,0,:,:], accumulate=True)
+    if field.shape[0] == 1:
+        continuous_boundary(res[0])
+    else:
+        res[1].index_put_((inds_y_all, inds_x_all), values[:,1,:,:], accumulate=True)
+        opposed_horizontal_boundary(res[1])
+        opposed_vertical_boundary(res[0])
+    return res
+
+def clear_divergence(field):
+    opposed_horizontal_boundary(field[0])
+    opposed_vertical_boundary(field[1])
+    for i in range(divergence_clearing_steps):
+        x_op = convolve(field[0], div_opp_kernel)
+        y_op = convolve(field[1], div_opp_kernel)
+        field[1,1:h-1,1:w-1] += (convolve(field[1], y_div_kernel) + x_op) / 8
+        field[0,1:h-1,1:w-1] += (convolve(field[0], x_div_kernel) + y_op) / 8
+        opposed_horizontal_boundary(field[0])
+        opposed_vertical_boundary(field[1])
+
+
+
+
+# initialize a small field of random velocities, then upscale it interpolating between those velocities,
+# so that the variation in velocity is somewhat smooth instead of pure per-pixel noise,
+# which would lead to a less interesting simulation
+velocities = torch.randn([2, h//32, w//32]) * 120
+upsample = torch.nn.Upsample(size=[h, w], mode='bilinear')
+velocities = upsample(velocities.unsqueeze(0))[0]
+
+# initialize "dye" densities to a uniform field
+densities = torch.ones([1, h, w]) * 0.3
+# add lines
+#for x in range(20):
+#    densities[0,:,x*w//20] += 0.8
+#for y in range(20):
+#    densities[0,y*h//20,:] += 0.8
+
+frame_index = 0
+
+last_frame = time.perf_counter()
+
+limit = 1 / (timescale * 2 * math.sqrt(2))
+
+# core loop of the simulation
+for iteration in range(max_iterations):
+    # add a tiny bit of "dye" to the fluid at all points
+    densities.add_(0.003)
+    #if iteration % (30 * save_every) == 0:
+    #    for x in range(20):
+    #        densities[0,:,x*w//20] += 0.3
+    #    for y in range(20):
+    #        densities[0,y*h//20,:] += 0.3
+
+
+    #velocities[1,h//12+h//3:h//12+2*h//3,w//8] += 50 + 50 * math.sin(iteration * 0.005)
+    #velocities[1,h//3-h//12:2*h//3-h//12,7*w//8] -= 20 + 70 * math.sin(iteration * 0.01)
+    #velocities[0,7*h//8,4*w//6:5*w//6] -= 60 + 50 * math.sin(iteration * 0.03)
+
+    # diffuse the velocities in both directions, enforcing a boundary
+    # condition that maintains a constant inward velocity matching the outward velocity along the edges. that is, a wall
+    velocities[0] = diffuse(velocities[0], velocity_diffusion_rate, opposed_horizontal_boundary, timescale)
+    velocities[1] = diffuse(velocities[1], velocity_diffusion_rate, opposed_vertical_boundary, timescale)
+    clear_divergence(velocities)
+    velocities.clamp_(-limit, limit)
+
+    # let the velocity field flow along itself
+    velocities = advect(velocities, velocities, timescale)
+    clear_divergence(velocities)
+
+    # diffuse the densities & let them flow along the velocity field
+    if density_diffusion_rate is not None:
+        densities = diffuse(densities[0], density_diffusion_rate, continuous_boundary, density_timescale).unsqueeze(0)
+    densities = advect(densities, velocities, density_timescale)
+
+    # remove a little density
+    densities.sub_(0.003)
+    densities.clamp_(0,1)
+
+    if iteration % save_every == 0:
+        frame_index += 1
+        frame_time = time.perf_counter() - last_frame
+        image = torch.cat((0.5 + 0.5 * velocities / torch.sqrt(velocities[0]**2 + velocities[1]**2), torch.zeros((1,h,w))), dim=0)
+        save(image, f"velocities/{frame_index:06d}")
+        monochrome_save(densities[0], f"densities/{frame_index:06d}")
+
+        print(f"frame {frame_index:06d}: {frame_time:06f}")
+        print(f"[{frame_time/save_every:06f}/it for {save_every} iterations]")
+        last_frame = time.perf_counter()
+

sketch/old/pre_snakepyt/fluid_minimal.py

···
+import torch as t;from PIL import Image as I;w,h,iters,save_every,r,dt,conv=2**9,2**9,100000,1,10,0.005,lambda f,k:t.nn.functional.conv2d(f.unsqueeze(0),k,bias=None,padding=[0],stride=[1])[0]
+t.set_default_device("cuda");kd,kph,kpv=t.tensor([[[[0,1.0,0],[1,0,1],[0,1,0]]]]),t.tensor([[[[0,0,0],[1.0,0,-1],[0,0,0]]]]),t.tensor([[[[0,1.0,0],[0,0,0],[0,-1,0]]]])
+def bd(f,lt=lambda t:(t[0],t[1])):f[(0,-1)],f[:,(0,-1)]=lt((f[(1,-2)],f[:,(1,-2)]));f[(0,0,-1,-1),(0,-1,0,-1)]=0.5*(f[(0,0,-2,-2),(1,-2,0,-1)]+f[(1,1,-1,-1),(0,-1,1,-2)])
+ind,lim=t.stack((t.arange(1,h-1,dtype=t.float).repeat(w-2,1).t(),t.arange(1,w-1,dtype=t.float).repeat(h-2,1))),1/(dt*1.4142)
+def adv(f, v):
+    off=ind.clone().add_(dt*v[:,1:h-1,1:w-1]);off[1].clamp_(1.5,w-2.5);off[0].clamp_(1.5,h-2.5);ind_int=off.int()
+    inv_off,next_ind=1-off.sub_(ind_int),ind_int+1;i=(t.stack([ind_int[0],ind_int[0],next_ind[0],next_ind[0]]),t.stack([ind_int[1],next_ind[1],ind_int[1],next_ind[1]]))
+    res,values=t.zeros_like(f),t.stack([f[:,1:h-1,1:w-1]*inv_off[1]*inv_off[0],f[:,1:h-1,1:w-1]*off[1]*inv_off[0],f[:,1:h-1,1:w-1]*inv_off[1]*off[0],f[:,1:h-1,1:w-1]*off[1]*off[0]])
+    res[0].index_put_(i,values[:,0,:,:],accumulate=True)
+    if f.shape[0]==1:bd(res[0])
+    else:res[1].index_put_(i,values[:,1,:,:],accumulate=True);bd(res[1],lambda t:(-t[0],t[1]));bd(res[0],lambda t:(t[0],-t[1]))
+    return res
+def proj(f):
+    div,p=(conv(f[1],kph)+conv(f[0],kpv))*0.5,t.zeros_like(f[0]);bd(div)
+    for i in range(80):p[1:h-1,1:w-1]=(div+conv(p,kd))/4;bd(p)
+    f[1,1:h-1,1:w-1]+=0.5*conv(p,kph);f[0,1:h-1,1:w-1]+=0.5*conv(p,kpv);bd(f[1],lambda t:(-t[0],t[1]));bd(f[0],lambda t:(t[0],-t[1]))
+v,d,i=t.nn.Upsample(size=[h,w],mode='bilinear')((t.randn([2,h//128,w//128])*50).unsqueeze(0))[0],t.ones([1,h,w])*0.1,1
+d[0,:,tuple(x*w//20 for x in range(20))]+=0.8;d[0,tuple(x*h//20 for x in range(20)),:]+=0.8
+for it in range(iters):
+    v.nan_to_num_(0);
+    d.add_(0.003);v[1,h//3:2*h//3,w//8].add_(40).clamp_(-lim,lim);v[0,1:h-1,1:w-1]=(v[0,1:h-1,1:w-1]+dt*r*conv(v[0],kd))/(1+4*dt*r)
+    v[1,1:h-1,1:w-1]=(v[1,1:h-1,1:w-1]+dt*r*conv(v[1],kd))/(1+4*dt*r);bd(v);proj(v);v=adv(v,v);proj(v);d=adv(d,v).sub_(0.003).clamp_(0,1)
+    print(f"{it}: {t.count_nonzero(t.isnan(v.view(-1))).item()} v nans")
+    if it%save_every == 0:
+        i+=1
+        I.fromarray(d[0].detach().clone().clamp_(0,1).mul_(255).round().type(t.uint8).unsqueeze(2).expand(-1,-1,3).cpu().numpy()).save(f"out/densities/{i:06d}.png")
+        #I.fromarray(t.cat((v.isnan(),t.zeros_like(v[0]).unsqueeze(0))).detach().clone().clamp_(0,1).mul_(255).round().type(t.uint8).permute(1,2,0).cpu().numpy()).save(f"out/velocities/nan_{i:06d}.png")
+        I.fromarray(t.cat(((0.5 + 0.5 * v / t.sqrt(v[0]**2 + v[1]**2)), t.zeros_like(v[0]).unsqueeze(0))).detach().clone().clamp_(0,1).mul_(255).round().type(t.uint8).permute(1,2,0).cpu().numpy()).save(f"out/velocities/dir_{i:06d}.png")
+        #I.fromarray((0.5 + 0.5 * (v[0]**2+v[1]**2) / (v[0]**2 + v[1]**2).abs().max()).detach().clone().clamp_(0,1).mul_(255).round().type(t.uint8).unsqueeze(2).expand(-1,-1,3).cpu().numpy()).save(f"out/velocities/mag_{i:06d}.png")
+

sketch/old/pre_snakepyt/fluid_vids

···
+ffmpeg -i out/velocities/%06d.png -framerate 24 -c:v libx264 -pix_fmt yuv420p out/velocity.mp4
+ffmpeg -i out/densities/%06d.png -framerate 24 -c:v libx264 -pix_fmt yuv420p out/density.mp4

sketch/old/pre_snakepyt/fractal.py

···
+import time
+import torch
+from PIL import Image
+from util import *
+
+import math
+
+# Settings
+
+w = 2**13
+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 = "ship_big"
+
+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")
+
+
+
+
+
+

sketch/old/pre_snakepyt/ising.py

···
+import time
+from pathlib import Path
+
+import torch
+
+from util import *
+
+@settings
+def _s():
+    name = "template"
+    device = "cuda"
+
+    t_real = torch.double
+    t_complex = torch.cdouble
+
+    # 9 -> 512 | 10 -> 1024 | 11 -> 2048
+    scale_power = 9
+    scale = 2 ** scale_power
+
+    w, h = scale, scale
+
+
+@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)
+
+    lattice = torch.randn([h,w], dtype=t_complex, device=device).round(0, 1)
+
+    with open(f"out/{run_dir}/settings.py", "w") as f:
+        f.write(settings.src)
+    torch.set_default_device(device)
+

sketch/old/pre_snakepyt/lorenz.py

···
+import time
+from pathlib import Path
+
+import torch
+
+from util import *
+
+@settings
+def _s():
+    name = "lorenz"
+    device = "cuda"
+    t_fp = torch.double
+
+    sigma = 10
+    rho = 28
+    beta = 8/3
+
+    tau = 3.14159265358979323 * 2
+
+    particle_count = 5000
+    iterations = 1000
+
+    # 2 ** 9 = 512; 10 => 1024; 11 => 2048
+    scale_power = 10
+    scale = 2 ** scale_power
+
+    origin = 0, 0
+    span = 60, 40
+
+    zooms = [
+            ]
+
+    dt = 0.005
+
+
+def rk4_step(get_derivative, take_step, position_0):
+    direction_0 = get_derivative(position_0)
+    position_1 = take_step(position_0, direction_0, 0.5)
+    direction_1 = get_derivative(position_1)
+    position_2 = take_step(position_0, direction_1, 0.5)
+    direction_2 = get_derivative(position_2)
+    position_3 = take_step(position_0, direction_2, 1)
+    direction_3 = get_derivative(position_3)
+    final_direction = (direction_0 + 2 * (direction_1 + direction_2) + direction_3) / 6
+    return take_step(position_0, final_direction, 1)
+
+def euler_step(get_derivative, take_step, position_0):
+    return take_step(position_0, get_derivative(position_0), 1)
+
+
+def idxput(tensor, indices, values):
+    tensor.index_put_((indices[:,0],indices[:,1]), values, accumulate=False)
+
+@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)
+
+    global span
+
+    x_min = origin[0] - (span[0] / 2)
+    y_min = origin[1] - (span[1] / 2)
+
+    for ((xa, xb), (ya, yb)) in zooms:
+        x_min += span[0] * xa
+        y_min += span[1] * ya
+        span = span[0] * (xb - xa), span[1] * (yb - ya)
+
+    x_max = x_min + span[0]
+    y_max = y_min + span[1]
+
+    aspect = span[0] / span[1]
+
+    if aspect < 1:
+        h = scale
+        w = int(scale * aspect)
+    else:
+        w = scale
+        h = int(scale / aspect)
+
+
+    dx = lambda x,y,z: sigma * (y - x)
+    dy = lambda x,y,z: x * (rho - z) - y
+    dz = lambda x,y,z: x * y - beta * z
+
+    p_positions = (torch.rand([3, particle_count], device=device, dtype=t_fp) - 0.5) * 40
+    p_positions[0] = torch.linspace(-0.1, 0.1, particle_count)
+    p_positions[1] = torch.linspace(-0.1, 0.1, particle_count)
+    p_positions[2] = torch.linspace(-0.1, 0.1, particle_count)
+
+
+    p_colors = torch.rand([particle_count,3], device=device, dtype=t_fp)
+
+    color_rotation = torch.linspace(0, tau / 4, particle_count)
+
+    p_colors[:,0] = torch.cos(color_rotation)
+    p_colors[:,1] *= 0
+    p_colors[:,2] = torch.sin(color_rotation)
+
+    canvas = torch.zeros([3, h, w], device=device, dtype=t_fp)
+    scratch = torch.zeros([h, w, 3], device=device, dtype=t_fp)
+
+    step = lambda p, dp, h: p + dp * h * dt
+
+    def derivative(p):
+        dp = torch.zeros_like(p)
+        dp[0] = dx(*p)
+        dp[1] = dy(*p)
+        dp[2] = dz(*p)
+        return dp
+
+    rk4_curried = lambda p: rk4_step(derivative, step, p)
+
+    def project(p):
+        return p[0:2]
+
+    for iteration in range(iterations):
+        p_projected = project(p_positions).clone()
+
+        if iteration % 500 == 0:
+            p_mask = torch.ones([particle_count], device=device)
+            p_mask *= (p_projected[1] >= x_min) * (p_projected[1] <= x_max)
+            p_mask *= (p_projected[0] >= y_min) * (p_projected[0] <= y_max)
+
+            in_range = p_mask.nonzero().squeeze()
+
+            p_projected_filtered = torch.index_select(p_projected.permute(1,0), 0, in_range)
+            p_indices = p_projected_filtered
+            p_colors_filtered = torch.index_select(p_colors, 0, in_range)
+
+            p_indices[:,1] -= x_min
+            p_indices[:,1] *= (w-1) / (x_max - x_min)
+            p_indices[:,0] -= y_min
+            p_indices[:,0] *= (h-1) / (y_max - y_min)
+
+            p_indices = p_indices.long()
+
+            scratch *= 0
+            idxput(scratch, p_indices, p_colors_filtered)
+            #canvas += scratch.permute(2,0,1)
+
+
+            #temp = canvas.clone()
+            #temp[0] -= temp[0].mean()
+            #temp[1] -= temp[1].mean()
+            #temp[2] -= temp[2].mean()
+
+            #temp[0] /= 8 * temp[0].std()
+            #temp[1] /= 8 * temp[1].std()
+            #temp[2] /= 8 * temp[2].std()
+
+            #temp += 0.5
+            #save(temp, f"{run_dir}/{iteration}")
+            save(scratch.permute(2,0,1), f"{run_dir}/{iteration}_b")
+
+        p_positions = rk4_curried(p_positions)
+
+    #canvas[0] -= canvas[0].mean()
+    #canvas[1] -= canvas[1].mean()
+    #canvas[2] -= canvas[2].mean()
+
+    #canvas[0] /= 8 * canvas[0].std()
+    #canvas[1] /= 8 * canvas[1].std()
+    #canvas[2] /= 8 * canvas[2].std()
+
+    #canvas += 0.5
+
+    #save(canvas, f"{run_dir}/final")
+
+    with open(f"out/{run_dir}/settings.py", "w") as f:
+        f.write(settings.src)
+    torch.set_default_device(device)
+

sketch/old/pre_snakepyt/lyapunov.py

···
+import torch
+import time
+import itertools
+from random import random
+
+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
+
+
+# 2 ** 9 = 512; 10 => 1024; 11 => 2048
+scale = 2 ** 14
+
+origin = 3.80576, 3.8096
+span = 0.01309, 0.00892
+
+zooms = [
+        ]
+
+dev = "cuda"
+t_fp = torch.double
+t_cfp = torch.cdouble
+
+seq = "BBABABA"
+c = 3.5
+x_init = 0.499
+alpha = 0.907
+
+iterations = 4000
+
+discontinuity = True
+
+final_file = f"{seq}_sq_e14_dis_rngb"
+
+# 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
+
+x_min = origin[0]
+y_min = origin[1]
+
+for ((xa, xb), (ya, yb)) in zooms:
+    x_min += span[0] * xa
+    y_min += span[1] * ya
+    span = span[0] * (xb - xa), span[1] * (yb - ya)
+
+x_max = x_min + span[0]
+y_max = y_min + span[1]
+
+aspect = span[0] / span[1]
+
+if aspect < 1:
+    h = scale
+    w = int(scale * aspect)
+else:
+    w = scale
+    h = int(scale / aspect)
+
+
+def main():
+    x = torch.ones([h, w], device=dev, dtype=t_fp) * x_init
+    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)
+
+    flip = False
+    for idx, seq_elem in enumerate(gen):
+        if random() > 0.99:
+            flip = not flip
+        if seq_elem == "C":
+            r = c
+        else:
+            if flip:
+                if seq_elem == "A":
+                    seq_elem = "B"
+                else:
+                    seq_elem = "A"
+            r = ab.real if seq_elem == "A" else ab.imag
+
+        if idx > 600:
+            lyapunov_avg = c_avg(lyapunov_avg, lyapunov_term(r, x), idx)
+
+        if idx < iterations - 1:
+            if discontinuity:
+                mask = x <= 0.5
+            x = logistic_map(r, x)
+            if discontinuity:
+                x += mask * (alpha - 1) * (r - 2) / 4
+
+        #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")
+

sketch/old/pre_snakepyt/raytrace.py

···
+import torch
+import time
+
+from pathlib import Path
+
+from util import *
+
+@settings
+def _s():
+    name = "qjulia_11_40"
+
+    # 2**10 = 1024
+    scale = 2 ** 8
+    w = 2 * scale // 3
+    h = scale
+
+    max_steps = 500
+    mandel_iters = 40#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
+
+    from math import sin, cos
+    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_mandel(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_power9_bulb(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_julia_q(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_ball(pos):
+    return torch.norm(pos, dim=0) > 0.9
+
+@timed
+@ifmain(__name__, _s)
+def _m(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(dev)
+
+    start = torch.zeros([h, w], dtype=t_cfp)
+    #res = torch.zeros([h,w], device=dev, dtype=t_fp)
+    res = torch.ones([h,w], 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)
+    xspace = torch.linspace(-half_x, half_x, w, dtype=t_fp)
+    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) * frame_center[2]))
+
+    step = torch.stack([torch.ones([h,w], 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], 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"{run_dir}/res_s")
+    msave(v_res, f"{run_dir}/res_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)
+

sketch/old/pre_snakepyt/scratch.py

···
+from util import *
+import torch
+import torchvision
+
+@timed
+def a(z, n):
+    for i in range(n):
+        #msave_gpu(z, "test/f")
+        #torchvision.utils.save_image(z, f"out/test/new_{i}.png", "png")
+        msave(z, f"test/old_{i}")
+        #mpilify(z)
+
+@timed
+def b(z, n):
+    for i in range(n):
+        #torch.save(z.type(torch.uint8), f"out/test/new_{i}.pt")
+        msave_alt(z, f"test/new_{i}")
+        #mpilify_gpu(z)
+
+@ifmain(__name__)
+def _main():
+    s = 2048
+    h,w = s, s
+
+    z = torch.randn((h, w), device="cuda", dtype=torch.double);
+
+    z.mul_(3.14159)
+
+    n = 20
+
+    a(z, n)
+
+    b(z, n)
+
+    a(z, n)
+
+    b(z, n)

sketch/old/pre_snakepyt/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
+    also_save = [10, 50, 150]
+
+    particle_count = 10000000 // 4
+    decay_rate = 0.1
+
+    view_angle = tau / 5
+    view_distance = 1.414 * 140
+
+    direction_count = 12
+    turn_amount = tau / 6#direction_count
+    move_distance = 1.414 * 2
+
+    blur_size = 1
+    blur_sigma = 1.5
+
+    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}")
+
+        if iteration in also_save:
+            msave(world[0], f"{run_dir}/{iteration}")
+
+        # decay
+        if iteration < max_iterations - 1:
+            world *= decay_rate
+
+    msave(world[0], f"{run_dir}/final")

sketch/old/pre_snakepyt/stdmap.py

···
+# "standard map" aka Chirikov-Taylor map
+
+# p_next = p + K * sin(theta)
+# theta_next = theta + p_next
+
+import torch
+import time
+import itertools
+from random import random
+
+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 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 msave(x, f):
+    mpilify(x).save(f"out/{f}.png")
+
+def save(x, f):
+    pilify(x).save(f"out/{f}.png")
+
+# config
+
+tau = 3.14159265358979323 * 2
+
+# 2 ** 9 = 512; 10 => 1024; 11 => 2048
+scale_power = 12
+scale = 2 ** scale_power
+
+origin = tau / 2, tau / 2
+span = tau, tau
+
+zooms = [
+            #((0.333, 0.667), (0.333, 0.667))
+        ]
+
+dev = "cuda"
+t_fp = torch.double
+t_cfp = torch.cdouble
+
+k = 0.971635406
+
+particle_count = 500000
+
+iterations = 1000
+
+randomness = 0
+
+growth = 0.00001
+
+final_file = f"k{k}_i{iterations}_p{particle_count}_s{scale_power}_r{randomness}_b"
+
+# funcs
+
+def offset_mod(x):
+    return torch.remainder(x, tau)
+
+def idxput(tensor, indices, values):
+    tensor.index_put_((indices[:,0],indices[:,1]), values, accumulate=False)
+
+# init
+
+x_min = origin[0] - (span[0] / 2)
+y_min = origin[1] - (span[1] / 2)
+
+for ((xa, xb), (ya, yb)) in zooms:
+    x_min += span[0] * xa
+    y_min += span[1] * ya
+    span = span[0] * (xb - xa), span[1] * (yb - ya)
+
+x_max = x_min + span[0]
+y_max = y_min + span[1]
+
+aspect = span[0] / span[1]
+
+if aspect < 1:
+    h = scale
+    w = int(scale * aspect)
+else:
+    w = scale
+    h = int(scale / aspect)
+
+print(f"x {x_min} -> {x_max}")
+print(f"y {y_min} -> {y_max}")
+
+def main():
+    world = torch.zeros([3, h, w], device=dev, dtype=t_fp)
+    scratch = torch.zeros([h, w, 3], device=dev, dtype=t_fp)
+    p_positions = torch.rand([particle_count,2], device=dev, dtype=t_fp) * tau
+
+    p_positions[:,0] = torch.linspace(0, tau, particle_count)
+    p_positions[:,1] = torch.linspace(0, tau, particle_count)
+
+    p_colors = torch.rand([particle_count,3], device=dev, dtype=t_fp)
+
+    color_rotation = torch.linspace(0, tau / 4, particle_count)
+
+    p_colors[:,0] = torch.cos(color_rotation)
+    p_colors[:,1] *= 0
+    p_colors[:,2] = torch.sin(color_rotation)
+
+    for iteration in range(iterations):
+        k_r = k + (random() - 0.5) * randomness + growth * iteration
+        p_positions[:,0] = p_positions[:,0] - k_r * torch.sin(p_positions[:,1])
+        p_positions[:,1] += p_positions[:,0]
+        p_positions[:,0] = offset_mod(p_positions[:,0])
+        p_positions[:,1] = offset_mod(p_positions[:,1])
+
+        if len(zooms) > 0:
+            p_mask = torch.ones([particle_count], device=dev)
+            p_mask *= (p_positions[:,0] >= x_min) * (p_positions[:,0] <= x_max)
+            p_mask *= (p_positions[:,1] >= y_min) * (p_positions[:,1] <= y_max)
+
+            in_range = p_mask.nonzero().squeeze()
+
+            p_positions_filtered = torch.index_select(p_positions, 0, in_range)
+            p_indices = p_positions_filtered
+            p_colors_filtered = torch.index_select(p_colors, 0, in_range)
+        else:
+            p_positions_filtered = p_positions
+            p_indices = p_positions.clone()
+            p_colors_filtered = p_colors
+
+        p_indices[:,0] -= x_min
+        p_indices[:,0] *= (w-1) / (x_max - x_min)
+        p_indices[:,1] -= y_min
+        p_indices[:,1] *= (h-1) / (y_max - y_min)
+
+        p_indices = p_indices.long()
+
+        scratch *= 0
+        idxput(scratch, p_indices, p_colors_filtered)
+        world += scratch.permute(2,0,1)
+
+        if iteration % 1000 == 0:
+            print(f"iteration {iteration}")
+
+    world[0] -= world[0].mean()
+    world[1] -= world[1].mean()
+    world[2] -= world[2].mean()
+
+    world[0] /= 8 * world[0].std()
+    world[1] /= 8 * world[1].std()
+    world[2] /= 8 * world[2].std()
+
+    world += 0.5
+
+    save(world, final_file)
+
+if __name__ == "__main__":
+    t0 = time.perf_counter()
+    main()
+    print(f"done in {time.perf_counter() - t0}s")
+

sketch/old/pre_snakepyt/template.py

···
+import time
+from pathlib import Path
+
+import torch
+
+from util import *
+
+@settings
+def _s():
+    name = "template"
+    device = "cuda"
+
+    t_real = torch.double
+    t_complex = torch.cdouble
+
+@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)
+
+    # { code }
+
+    with open(f"out/{run_dir}/settings.py", "w") as f:
+        f.write(settings.src)
+    torch.set_default_device(device)
+

sketch/old/pre_snakepyt/worley.py

···
+import time
+from pathlib import Path
+
+import torch
+import math
+
+from util import *
+
+@settings
+def _s():
+    name = "worley"
+    keep_records = False
+
+    subcell_size = (16, 16, 16)
+    subcell_counts = (8, 8, 8)
+    max_feature_points = 9
+    feature_point_density = 4
+
+    device = "cuda"
+    dtype = torch.double
+
+    assert len(subcell_size) == len(subcell_counts), "dimensional consistency"
+
+def tuple_mul(a,b):
+    return tuple(a*b for a,b in zip(a,b))
+
+def factorial(tensor):
+    return torch.exp(torch.lgamma(tensor + 1))
+
+def prob_n_points(n, density):
+    return 1 / (math.pow(density, -n) * math.exp(density) * math.factorial(n))
+
+
+@ifmain(__name__, _s)
+@timed
+def _main(settings):
+    globals().update(settings.get())
+    run_dir = ""
+    if keep_records:
+        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)
+
+    dim = len(subcell_size)
+
+    seed_points_shape = subcell_counts + (max_feature_points, dim)
+
+    seed_points = torch.rand(seed_points_shape, dtype=dtype).mul_(2).add_(-1)
+
+    result_shape = tuple_mul(subcell_size, subcell_counts)
+
+    result = torch.zeros(result_shape, dtype=dtype)
+
+    rng = torch.rand(subcell_counts)
+    mask = torch.zeros(subcell_counts + (max_feature_points,), dtype=torch.bool)
+    prob = 0
+    for n in range(max_feature_points):
+        prob += prob_n_points(n, feature_point_density)
+        mask[...,n] = rng > prob
+    #infs = ~mask * torch.inf
+    #for dim_index in range(dim):
+    #    seed_points[...,dim_index].mul_(infs)
+
+    linspaces = [torch.arange(0, subcell_counts[d], 1 / subcell_size[d]) for d in range(dim)]
+    grid_coords = torch.meshgrid(linspaces, indexing="ij")
+
+    grid_indices = grid_coords[0].long().clone()
+    for d in range(dim-1):
+        grid_indices.add_(grid_coords[d+1].long() * subcell_counts[d])
+
+    for d in range(dim):
+        print(torch.take(seed_points[...,d], grid_indices))
+
+    #breakpoint()
+
+    # MxNxO points
+    # AxBxC subcells = M//subdiv N//subdiv O//subdiv
+    # one point randomly placed in each cell
+    # AxBxCx
+    # every point is closest to a point in one of its neighbor cells
+    # figure out distance to that point
+

sketch/old/snakepyt_0_0/autograd/image_approx.py

···
+
+import torch
+
+from lib.util import load_image_tensor, save, msave
+from lib.spaces import grid, xrange_yrange
+
+def init():
+    t_real = torch.double
+
+    torch.set_default_device("cuda")
+    torch.set_default_dtype(t_real)
+
+    im = load_image_tensor("in/flower.jpg").to(t_real).to("cuda")
+
+    (c, h, w) = im.shape
+
+
+    n_a = 4
+
+    a = torch.rand([n_a,c,h]) - 0.5 #* 0.1 - 0.05
+    b = torch.rand([n_a,c,w]) - 0.5 #* 0.1 - 0.05
+    a *= 0.01
+    b *= 0.01
+    shifts = torch.rand([12])
+    coefs = torch.ones([12]) / 12
+
+    a = gradify(a)
+    b = gradify(b)
+    coefs = gradify(coefs)
+    shifts = gradify(shifts)
+
+    train_iters = 100000
+    def get_lr(i):
+        if i < 10:
+            return 1
+        if i < 100:
+            return 1
+        if i < 1000:
+            return 1e-1
+        if i < 10000:
+            return 1e-1
+        return 1e-1
+
+    save_on = lambda i: (i % 100 == 0) if i > 100 else (i % 10 == 0)
+
+    schedule(train, None)
+
+def gradify(t):
+    return t.clone().detach().requires_grad_(True)
+
+def train():
+    global a, b, coefs, shifts
+
+    im_bw = im.norm(p=2,dim=0)
+
+    g = grid((xrange_yrange((0,0), (7,7)), (h,w))).permute(2,0,1)
+
+
+    for n in range(train_iters):
+        copies = a.view(n_a,c,h,1) * b.view(n_a,c,1,w)
+
+        _g = g * 5
+        _g[0] += shifts[0]
+        _g[1] += shifts[1]
+        s_0 = torch.sin(_g[0]) * copies[0] + torch.cos(_g[0]) * copies[1]
+        s_1 = torch.sin(_g[1]) * copies[2] + torch.cos(_g[1]) * copies[3]
+        _g = g * 11
+        _g[0] += shifts[2]
+        _g[1] += shifts[3]
+        s_2 = torch.sin(_g[0]) * copies[0] + torch.cos(_g[0]) * copies[2]
+        s_3 = torch.sin(_g[1]) * copies[1] + torch.cos(_g[1]) * copies[3]
+        _g = g * 23
+        _g[0] += shifts[4]
+        _g[1] += shifts[5]
+        s_4 = torch.sin(_g[0]) * copies[1] + torch.cos(_g[0]) * copies[0]
+        s_5 = torch.sin(_g[1]) * copies[3] + torch.cos(_g[1]) * copies[2]
+        _g = g * 57
+        _g[0] += shifts[6]
+        _g[1] += shifts[7]
+        s_6 = torch.sin(_g[0]) * copies[0] + torch.cos(_g[0]) * copies[1]
+        s_7 = torch.sin(_g[1]) * copies[2] + torch.cos(_g[1]) * copies[3]
+        _g = g * 101
+        _g[0] += shifts[8]
+        _g[1] += shifts[9]
+        s_8 = torch.sin(_g[0] + _g[1]) * copies[2] + torch.cos(_g[0] * _g[0] + _g[1] * _g[1]) * copies[3]
+        s_9 = torch.sin(_g[1] - _g[0]) * copies[0] + torch.cos(_g[1] * _g[1] + _g[0] * _g[0]) * copies[1]
+
+        s_10 = (g[0] + shifts[10]) * copies[2] - (g[0] + shifts[10]) * copies[3]
+        s_11 = (g[1] + shifts[11]) * copies[0] - (g[1] + shifts[11]) * copies[1]
+
+        srcs = torch.stack((s_0, s_1, s_2, s_3, s_4, s_5, s_6, s_7, s_8, s_9, s_10, s_11))
+
+        coefs_ = coefs #+ torch.rand_like(coefs) * 0.1
+
+        srcs_scaled = (srcs * coefs_.view(12, 1, 1, 1))
+        copy = srcs_scaled.sum(dim=0)
+
+        diff = copy - im
+        mse = (diff ** 2).sum()
+
+        mse.backward()
+
+        torch.nn.utils.clip_grad_norm_([a,b,coefs], max_norm=1.0)
+        with torch.no_grad():
+            learn_rate = get_lr(n)
+            a -= (a.grad) * learn_rate
+            b -= (b.grad) * learn_rate
+            coefs -= (coefs.grad) * learn_rate
+            shifts -= (shifts.grad) * learn_rate * 0.01
+            if save_on(n):
+                print(a.grad.norm())
+                print(b.grad.norm())
+                print(coefs.grad.norm())
+                print(shifts.grad.norm())
+
+        a.grad.zero_()
+        b.grad.zero_()
+        coefs.grad.zero_()
+
+
+        if save_on(n):
+            print(f"mse {mse.item()}")
+            with torch.no_grad():
+                out = copy - copy.min()
+                out /= out.max()
+            #print(coefs)
+            save(out, f"{run_dir}/{n:06d}")
+            #with torch.no_grad():
+            #    c0 = copies[0] - copies[0].min()
+            #    c1 = copies[1] - copies[1].min()
+            #    c2 = copies[2] - copies[2].min()
+            #    c3 = copies[3] - copies[3].min()
+            #    c0 /= c0.max()
+            #    c1 /= c1.max()
+            #    c2 /= c2.max()
+            #    c3 /= c3.max()
+            print(coefs)
+            print(shifts)
+            #save(c0, f"{run_dir}/c0_{n:06d}")
+            #save(c1, f"{run_dir}/c1_{n:06d}")
+            #save(c2, f"{run_dir}/c2_{n:06d}")
+            #save(c3, f"{run_dir}/c3_{n:06d}")
+            #for s in range(12):
+            #    with torch.no_grad():
+            #        src = srcs_scaled[s]
+            #        src -= src.min()
+            #        src /= src.max()
+            #        save(src, f"{run_dir}/s{s}_{n:06d}")
+

sketch/old/snakepyt_0_0/labyrinth_chaos.py

···
+# Thomas' cyclically symmetric attractor / "Labyrinth Chaos"
+# see "Deterministic chaos seen in terms of feedback circuits: Analysis, synthesis, 'labyrinth chaos'" by René Thomas
+
+import time
+import math
+from pathlib import Path
+
+import torch
+
+from lib.spaces import insert_at_coords, map_space
+from lib.ode import rk4_step
+from lib.util import *
+
+def init():
+    schedule(run, None)
+
+name = "labyrinth"
+device = "cuda"
+torch.set_default_device(device)
+t_fp = torch.double
+torch.set_default_dtype(t_fp)
+
+dissipation = 0.04 # "b" parameter in the literature
+
+tau = 3.14159265358979323 * 2
+
+particle_count = 50000
+iterations = 1000
+save_if = lambda i: i > 150 and i % 5 == 0
+rotation_rate = tau / 2000
+
+# 2 ** 9 = 512; 10 => 1024; 11 => 2048
+scale_power = 10
+
+scale = 2560#2 ** scale_power
+origin = 0, 0
+s = 2
+span = 16 * s, 9 * s
+zooms = [
+        ]
+stretch = 1, 1
+
+mapping = map_space(origin, span, zooms, stretch, scale)
+(_, (h,w)) = mapping
+
+dt = 0.01
+
+
+def derivative(p):
+    return torch.sin(p[[1,2,0]]) - p * dissipation
+
+step = lambda p, dp, h: p + dp * h * dt
+rk4_curried = lambda p: rk4_step(derivative, step, p)
+
+p_positions = (torch.rand([3, particle_count], device=device, dtype=t_fp) - 0.5) * 20
+
+p_colors = torch.rand([particle_count,3], device=device, dtype=t_fp)
+
+color_rotation = torch.linspace(0, tau / 4, particle_count)
+
+p_colors[:,0] = (p_positions[0,:] / 20) + 0.5
+p_colors[:,1] = (p_positions[1,:] / 20) + 0.5
+p_colors[:,2] = (p_positions[2,:] / 20) + 0.5
+
+def project(p, colors, i):
+    c = math.cos(i * rotation_rate)
+    s = math.sin(i * rotation_rate)
+    rotation = torch.tensor([
+        [1, 0, 0],
+        [0, c,-s],
+        [0, s, c]])
+    #rotation = torch.tensor([
+    #    [ c, 0, s],
+    #    [ 0, 1, 0],
+    #    [-s, 0, c]])
+    #rotation = torch.tensor([
+    #    [c, -s, 0],
+    #    [s,  c, 0],
+    #    [0,  0, 1]])
+    alt_colors = colors.clone()
+    res = (rotation @ p)
+    color_filter = res[2].abs() < 0.7
+    alt_colors[:,0] *= color_filter
+    alt_colors[:,1] *= color_filter
+    alt_colors[:,2] *= color_filter
+    return (res[0:2], alt_colors)
+
+frame_index = [0]
+
+def run():
+    run_dir = time.strftime(f"%d.%m.%Y/{name}_t%H.%M.%S")
+    Path("out/" + run_dir).mkdir(parents=True, exist_ok=True)
+
+    scratch = torch.zeros([h, w, 3], device=device, dtype=t_fp)
+
+    for iteration in range(iterations):
+        if save_if(iteration) or iteration == iterations - 1:
+            (p_projected, alt_colors) = project(p_positions, p_colors, iteration)
+            frame_index[0] += 1
+            scratch *= 0
+            insert_at_coords(p_projected, alt_colors, scratch, mapping)
+            #save(scratch.permute(2,0,1), f"{run_dir}/{frame_index[0]:06d}")
+
+        p_positions.copy_(rk4_curried(p_positions))
+
+

sketch/old/snakepyt_0_0/laplace.py

···
+# laplacian growth
+# see "Fast Simulation of Laplacian Growth" by Theodore Kim, Jason Sewall, Avneesh Sud and Ming C. Lin
+
+import itertools
+from math import tau, log
+
+import torch
+
+from lib.util import load_image_tensor, save, msave
+
+
+
+def init():
+    torch.set_default_device("cuda")
+
+    dims = 2
+    #scale = 2 ** 10
+    #h = scale
+    #w = scale
+    margin = 2**4
+
+    physics_scale = 0.5
+    eta = 10.7 # fractal dimension or something
+
+    max_points = 2 ** 13
+
+    t_real = torch.half
+
+    point_positions = torch.zeros([max_points, dims], dtype=t_real)
+    point_values = torch.zeros([max_points], dtype=t_real)
+    is_charge = torch.zeros([max_points], dtype=torch.bool)
+    do_neighbors = torch.zeros([max_points], dtype=torch.bool)
+    is_candidate = torch.zeros_like(is_charge)
+    is_free = torch.ones_like(is_charge)
+
+    neighbors = torch.tensor([v for v in itertools.product([-1,0,1], repeat=dims) if any(v)], dtype=t_real)
+
+    inner_radius = physics_scale / 2
+
+    def insert_charge(position, value, was_candidate=True, no_neighbors=False):
+        if was_candidate: # index could probably be provided by the caller for efficiency
+            match_coords = point_positions == position.expand(point_positions.shape)
+            match_pos = match_coords.prod(dim=1)
+            index = torch.nonzero(is_candidate & match_pos)[0]
+            is_candidate[index] = False
+            is_charge[index] = True
+            do_neighbors[index] = not no_neighbors
+            point_values[index] = value
+        else:
+            # create the charge
+            free_index = torch.nonzero(is_free)[0] # maybe not ideal
+            is_charge[free_index] = True
+            is_free[free_index] = False
+            do_neighbors[free_index] = not no_neighbors
+            point_positions[free_index] = position
+            point_values[free_index] = value
+
+        # update existing candidate potentials (eqn 11)
+        candidate_pos = point_positions[is_candidate]
+        dist = (candidate_pos - position.expand(candidate_pos.shape)).norm(p=2, dim=1)
+        point_values[is_candidate] += 1 - inner_radius/dist
+
+        if no_neighbors:
+            return
+
+        # add new candidates
+        charge_pos = point_positions[is_charge & do_neighbors]
+        for charge_index in range(charge_pos.shape[0]):
+            neighborhood = neighbors + charge_pos[charge_index].expand(neighbors.shape)
+            for neighbor_index in range(neighborhood.shape[0]):
+                maybe_insert_candidate(neighborhood[neighbor_index])
+
+    def maybe_insert_candidate(position):
+        if torch.any((point_positions == position.expand(point_positions.shape)).prod(dim=1)):
+            return
+
+        free_index = torch.nonzero(is_free)[0] # maybe not ideal
+        is_candidate[free_index] = True
+        is_free[free_index] = False
+        point_positions[free_index] = position
+        charge_pos = point_positions[is_charge]
+        charge_val = point_values[is_charge]
+        dist = (charge_pos - position.expand(charge_pos.shape)).norm(p=2, dim=1)
+        point_values[free_index] = (1 - inner_radius/dist).sum() # eqn 10
+
+    def select_candidate(scale):
+        try:
+            candidate_val = point_values[is_candidate]
+            candidate_pos = point_positions[is_candidate]
+            max_val = candidate_val.max()
+            min_val = candidate_val.min()
+            normalized_val = (candidate_val - min_val) / (max_val - min_val) # eqn 13
+            val_pow = normalized_val ** eta
+            selection_prob = val_pow / val_pow.sum()
+
+            prob_sum = selection_prob.cumsum(dim=0)
+
+            choices = []
+            for _ in range(scale):
+                r = torch.rand(1, dtype=t_real)
+                choice = torch.nonzero(prob_sum > r)[0]
+                if choice not in choices:
+                    insert_charge(candidate_pos[choice], 1)
+                    choices.append(choice)
+        except KeyboardInterrupt:
+            raise
+        except:
+            print(f"selection failed")
+
+
+    schedule(run, None)
+
+
+def run():
+
+    # set initial conditions
+    insert_charge(torch.tensor([0]*dims, dtype=t_real), 1, was_candidate=False)
+    for n in range(-40,40):
+        insert_charge(torch.tensor([2, n], dtype=t_real), 1, was_candidate=False, no_neighbors=True)
+
+    schedule(main_loop, None)
+
+def main_loop():
+    for iteration in range(10000):
+        select_candidate(int(log(iteration * 100 + 10)))
+        if iteration % 10 == 0:
+            final_charges = point_positions[is_charge].clone()
+
+            for d in range(dims):
+                final_charges[:,d] -= final_charges[:,d].min()
+                final_charges[:,d] += margin
+
+            h = int(final_charges[:,0].max().item() + margin)
+            w = int(final_charges[:,1].max().item() + margin)
+
+            indices = final_charges.long()
+
+            canvas = torch.ones([h,w], dtype=t_real)
+            canvas[(indices[:,0], indices[:,1])] = 0
+
+            msave(canvas, f"{run_dir}/{iteration:06d}")
+

sketch/old/snakepyt_0_0/laplace_b.py

···
+
+from random import randint
+from math import cos, sin, tau
+
+import torch
+from torch.nn.functional import conv2d, softmax
+
+from lib.spaces import xrange_yrange, grid
+from lib.util import msave, save
+
+def init():
+    device = "cuda"
+    torch.set_default_device(device)
+
+    t_real = torch.double
+    torch.set_default_dtype(t_real)
+
+
+    scale = 2**10
+    h = scale
+    w = scale
+
+    charge_size = 50
+
+    iterations = 100001
+    selection_range = 0.05
+    growth_limit = 100
+
+    seed_pts = 10
+    eta = 5
+
+    span = (10, 10)
+    region = xrange_yrange((0,0), span)
+
+    neighbor_kernel = torch.ones([1,1,3,3],dtype=t_real)
+    neighbor_kernel[0,0,1,1] = 0
+    def neighbors(positions):
+        conv = conv2d(positions.unsqueeze(0).unsqueeze(0), neighbor_kernel, bias=None, padding="same", stride=[1])[0,0]
+        return ((conv > 0) & (positions == 0))
+
+
+    schedule(growth, None)
+
+def point_charge(position, positions, potentials, single_charge, no_seed=False, inverse=False, strength=1):
+    if not no_seed:
+        positions[position[0],position[1]] += 1
+    if inverse:
+        potentials.add_(torch.roll((1-single_charge)*strength, shifts=position, dims=(0,1)))
+    else:
+        potentials.add_(torch.roll(single_charge*strength, shifts=position, dims=(0,1)))
+
+
+def growth():
+    potentials = torch.zeros([h,w], dtype=t_real)
+    positions = grid((region, (h,w)))
+    #positions[:,:,0] -= (2*span[0]) / ((2 * h))
+    #positions[:,:,1] -= (2*span[1]) / ((2 * w))
+    single_charge_potential = 1 - (charge_size / (positions[:,:,0]**2 + positions[:,:,1]**2)**0.5)
+    #msave(single_charge_potential / single_charge_potential.max(), f"{run_dir}/scp")
+    single_charge_potential = torch.roll(single_charge_potential, shifts=(h//2, w//2), dims=(0,1))
+
+
+    positions = torch.zeros([h,w], dtype=t_real)
+
+    #r = torch.rand([seed_pts,2]) * (h) #+ ((h//2))
+    #r = r.long()
+    #for i in range(seed_pts):
+    #    point_charge((r[i,0].item(), r[i,1].item()), positions, potentials, single_charge_potential)
+
+    pts = [
+        #(h//2 + 1, w//2),
+        #(h//2 - 1, w//2)
+        #(h//2,w//2)
+        (0,w//2)
+    ]
+
+    for p in pts:
+        point_charge((p[0], p[1]), positions, potentials, single_charge_potential)
+
+    schedule(loop, None)
+
+
+def loop():
+    canvas = torch.zeros([3,h,w], dtype=t_real)
+
+    #U = 80
+    #for u in range(U//2):
+        #rad = u*h//U
+        #n_pts = int(tau*rad)
+        #for p in range(n_pts):
+        #    x = int(cos(tau*p/n_pts) * (rad) + (h//2))
+        #    y = int(sin(tau*p/n_pts) * (rad) + (h//2))
+        #    point_charge((x,y), positions, potentials, single_charge_potential, True, False, 5)
+
+        #rad = int((u+0.5)*(h//U)*10)
+        #n_pts = int(tau*rad)
+        #for p in range(n_pts):
+        #    x = int(cos(tau*p/n_pts) * (rad) + (h//4))
+        #    y = int(sin(tau*p/n_pts) * (rad) + (h//4))
+        #    point_charge((x,y), positions, potentials, single_charge_potential, True, True, 1)
+
+
+    for p in range(h):
+        point_charge((h-1, p), positions, potentials, single_charge_potential, True, False, 2)
+    #for p in range(h):
+    #    point_charge((p, 4 * h//10 - 1), positions, potentials, single_charge_potential, True, False, 8)
+    #    point_charge((p, 6 * h//10 + 1), positions, potentials, single_charge_potential, True, False, 8)
+
+    rad = h//10
+    n_pts = int(tau*rad)
+    for p in range(n_pts):
+        x = int(cos(tau*p/n_pts) * rad + h//2 - h//5)
+        x2 = int(cos(tau*p/n_pts) * rad + h//2)
+        x3 = int(cos(tau*p/n_pts) * rad + h//2 + h//5)
+        y = int(sin(tau*p/n_pts) * rad + h//2)
+        point_charge((x,y), positions, potentials, single_charge_potential, True, True, 2)
+        point_charge((x2,y), positions, potentials, single_charge_potential, True, True, 2)
+        point_charge((x3,y), positions, potentials, single_charge_potential, True, True, 2)
+
+
+    for iteration in range(iterations):
+
+        candidates = neighbors(positions)
+
+        cand_values = (potentials * candidates)
+        cand_values[cand_values==0] = torch.inf
+
+        cand_values -= cand_values.min() # infs necessary to keep min from being zero
+        cand_values = torch.nan_to_num(cand_values, 0, 0, 0) # remove infs
+        #msave(cand_values, f"{run_dir}/candv_{iteration}")
+        cand_values /= cand_values.max()
+
+        #msave(cand_values, f"{run_dir}/cand_{iteration:06d}")
+        #cand_values *= torch.rand_like(cand_values)
+        #cand_probs = softmax(cand_values)
+        cand_probs = cand_values
+        cand_probs.pow_(eta)
+        #cand_probs.mul_(0.9 + 0.1 * torch.rand_like(cand_probs))
+
+        #selected = torch.multinomial(cand_probs, 3, replacement=False)
+
+        selected = torch.argwhere((cand_probs > ((1 - selection_range) * cand_probs.max())) & (cand_probs > 0))
+
+
+        print(selected.shape[0])
+        if selected.shape[0] > 0:
+            r = randint(0,selected.shape[0] - 1)
+            selected = torch.roll(selected, shifts=r, dims=0)
+
+        #print(selected.shape)
+
+        #ys = (selected // w)
+        #xs = (selected % h)
+        #point_charge((x,y), positions, potentials, single_charge_potential)
+
+        if iteration % 100 == 0:
+            canvas[2] = potentials
+            canvas[2] -= canvas[2].min()
+            canvas[2] /= canvas[2].max()
+            canvas[2] = 1 - canvas[2]
+            #canvas[0] = canvas[0] ** 5
+            #/ potentials.max())**10
+            #canvas[1] = cand_values
+            #canvas[2] = single_charge_potential
+            #canvas[0] = positions
+            #msave(single_charge_potential, f"{run_dir}/scp_{iteration:06d}")
+            #msave(candidates, f"{run_dir}/cand_{iteration}")
+            #msave(potentials / potentials.max(), f"{run_dir}/pot_{iteration}")
+            #msave(potentials / potentials.max(), f"{run_dir}/{iteration:06d}")
+            save(canvas, f"{run_dir}/{iteration:06d}")
+
+        ys = selected[:,0]
+        xs = selected[:,1]
+
+        n = selected.shape[0]
+        if n > growth_limit: n = growth_limit
+
+        for index in range(n):
+            inds = ys[index].item(), xs[index].item()
+            point_charge(inds, positions, potentials, single_charge_potential)
+

sketch/old/snakepyt_0_0/lyapunov.py

···
+import torch
+import time
+import math
+import itertools
+import sys
+import subprocess
+from random import random
+
+from PIL import Image
+
+from lib.util import *
+from lib.spaces import map_space, cgrid, center_span, apply_zooms
+from lib.io import VideoWriter
+
+zoom_plan = [
+    ((0.6, 1.0), (0.6, 1.0)),
+    ((0.7, 0.8), (0.6, 0.7)),
+    ((0.7, 0.8), (0.8, 0.9)),
+    ((0.7, 0.8), (0.2, 0.3)),
+    ((0.5, 0.6), (0.4, 0.5)),
+    ((0.3, 0.4), (0.6, 0.7)),
+    ((0.0, 0.1), (0.0, 0.1)),
+    ((0.4, 0.5), (0.6, 0.7)),
+    ((0.2, 0.3), (0.7, 0.8)),
+    ((0.8, 0.9), (0.7, 0.8)),
+    ((0.1, 0.2), (0.4, 0.5)),
+    ((0.6, 0.7), (0.7, 0.8)),
+    ((0.0, 0.1), (0.3, 0.4)),
+    #((0.7, 0.8), (0.4, 0.5)),
+    #((0.0, 0.1), (0.6, 0.7)),
+    #((0.6, 0.7), (0.3, 0.4)),
+    #((0.4, 0.5), (0.7, 0.8)),
+    #((0.6, 0.7), (0.3, 0.4)),
+    #((0.2, 0.3), (0.1, 0.2)),
+    #((0.7, 0.8), (0.1, 0.2)),
+    #((0.1, 0.2), (0.6, 0.7)),
+    #((0.7, 0.8), (0.3, 0.4)),
+    #((0.9, 1.0), (0.8, 0.9)),
+]
+
+planning_mode = False
+smoothed_path = True
+
+def init():
+    frame_index = [0]
+
+    _origin = 0, 0
+    _span = 9, 9
+    mapping = map_space(_origin, _span, [], (), scale)
+    (_, dims) = mapping
+
+
+    #video = VideoWriter(dims, 24,
+    #        f"out/{run_dir}/vid.mp4",
+    #        f"out/{run_dir}/ffmpeg_log")
+
+    if planning_mode:
+        zooms = zoom_plan
+        schedule(run, None)
+    else:
+        if smoothed_path:
+            frames = 120#720
+            schedule(smooth_zoom_interp, range(frames))
+        else:
+            schedule(zoom_level, range(len(zoom_plan)-1))
+
+def final():
+    pass
+    #video.close()
+
+def zoom_level():
+    zoom_plan_partial = zoom_plan[:index+2]
+
+    frames = 10
+
+    schedule(zoom_interp, range(frames))
+
+
+
+def zoom_interp(index):
+    ((xl, xr), (yl, yr)) = zoom_plan_partial[-1]
+    #t = index / frames
+    t = math.log(index + 1, frames+1)
+
+    zooms = list(zoom_plan_partial)
+    zooms[-1] = ((lerp(0, xl, t), lerp(1, xr, t)), (lerp(0, yl, t), lerp(1, yr, t)))
+
+
+    schedule(run, None)
+
+def eerp(a,b,t):
+    return math.pow(b/a, t)*a
+
+def lerp(a,b,t):
+    return (1-t) * a + t * b
+
+def smooth_zoom_interp(index):
+    final_center, final_span = center_span(*apply_zooms(_origin, _span, zoom_plan))
+    zooms = []
+
+    t = (index / frames)#math.log(index + 1, frames+1)
+
+    span = (eerp(_span[0], final_span[0], t), eerp(_span[1], final_span[1], t))
+
+    t = (span[0] - _span[0]) / (final_span[0] - _span[0])
+
+    origin = (lerp(_origin[0], final_center[0], t), lerp(_origin[1], final_center[1], t))
+
+    schedule(run, None)
+
+# 2 ** 9 = 512; 10 => 1024; 11 => 2048
+scale = 2 ** 11
+
+#origin = 3.80576, 3.8096
+#span = 0.01309, 0.00892
+
+#origin, span = center_span((3.795, 3.824), (3.80287, 3.82417))
+#origin, span = center_span((3.8097, 3.8175), (3.8138, 3.8181))
+
+
+stretch = 1, 1
+
+
+dev = "cuda"
+t_fp = torch.double
+t_cfp = torch.cdouble
+
+torch.set_default_device(dev)
+torch.set_default_dtype(t_fp)
+
+seq = "BBABABA"
+c = 3.5
+x_init = 0.499
+alpha = 0.907
+
+iterations = 1000
+
+discontinuity = True
+
+flip_mode = "individual"
+flip_chance = 0.00#1
+
+# 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 opposite(element):
+    if element == "A":
+        return "B"
+    else:
+        return "A"
+
+def repeat(seq, flip_chance, flip_mode):
+    i = 0
+    l = len(seq)
+    flipped = False
+    while True:
+        yield seq[i%l] if not flipped else opposite(seq[i%l])
+        i += 1
+        if i % l == 0 or flip_mode == "individual":
+            flipped = random() < flip_chance
+
+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
+
+
+
+def run():
+    mapping = map_space(origin, span, zooms, stretch, scale)
+    (_, (h,w)) = mapping
+
+    x = torch.ones([h, w]) * x_init
+
+    ab = cgrid(mapping)
+
+    tenths = torch.zeros([h, w])
+    for i in range(10):
+        frac = i / 10
+        tenths[math.floor(frac*h), :] = 1
+        tenths[:, math.floor(frac*w)] = 1
+
+    lyapunov_avg = torch.zeros([h, w])
+
+    gen = itertools.islice(repeat(seq, 0, flip_mode), iterations)
+
+    flip = False
+    for idx, seq_elem in enumerate(gen):
+        if seq_elem == "C":
+            r = c
+        else:
+            r_normal = ab.real if seq_elem == "A" else ab.imag
+            if flip_chance > 0:
+                r_flipped = ab.imag if seq_elem == "A" else ab.real
+                flip_mask = torch.rand([h,w]) < flip_chance
+                r = r_normal * ~flip_mask + r_flipped * flip_mask
+            else:
+                r = r_normal
+
+        if idx > 600:
+            lyapunov_avg = c_avg(lyapunov_avg, lyapunov_term(r, x), idx)
+
+        if idx < iterations - 1:
+            if discontinuity:
+                mask = x <= 0.5
+            x = logistic_map(r, x)
+            if discontinuity:
+                x += mask * (alpha - 1) * (r - 2) / 4
+
+        #msave(torch.tanh(lyapunov_avg / 4) / 2 + 0.5, f"avg_{idx}")
+
+    result = torch.tanh(lyapunov_avg)
+    result -= result.mean()
+    result /= 5 * result.std()
+    result += 0.5
+
+    if planning_mode:
+        result = torch.stack([result + tenths, result, result])
+        save(result, f"{run_dir}/{frame_index[0]:06d}")
+    else:
+        msave(result, f"{run_dir}/{frame_index[0]:06d}")
+        #video.mframe(result)
+    frame_index[0] += 1
+    #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")
+
+

sketch/old/snakepyt_0_0/ode.py

···
+# Thomas' cyclically symmetric attractor / "Labyrinth Chaos"
+# see "Deterministic chaos seen in terms of feedback circuits: Analysis, synthesis, 'labyrinth chaos'" by René Thomas
+
+import time
+import math
+from pathlib import Path
+
+import torch
+
+from lib.spaces import insert_at_coords, map_space
+from lib.ode import rk4_step
+from lib.util import *
+
+def init():
+    schedule(run, None)
+
+device = "cuda"
+torch.set_default_device(device)
+t_fp = torch.double
+torch.set_default_dtype(t_fp)
+
+dissipation = 0.04 # "b" parameter in the literature
+
+tau = 3.14159265358979323 * 2
+
+particle_count = 50000
+iterations = 1000
+save_if = lambda i: True #i > 150 and i % 5 == 0
+rotation_rate = tau / 2000
+
+# 2 ** 9 = 512; 10 => 1024; 11 => 2048
+scale_power = 10
+
+scale = 2 ** scale_power
+origin = 0, 0
+s = 4
+span = s, s
+zooms = [
+        ]
+stretch = 1, 1
+
+mapping = map_space(origin, span, zooms, stretch, scale)
+(_, (h,w)) = mapping
+
+dt = 1
+
+def proj(u, v):
+    dot = (u * v).sum(dim=1)
+    scale = (dot / (v.norm(p=2, dim=1) ** 2))
+    out = v.clone()
+    for dim in range(3):
+        out[:,dim] *= scale
+    return out
+
+def proj_shift(a, b):
+    return a + proj(b, a)
+
+def get_transformation(direction, flow, ebb, rescaling):
+    def transformation(p):
+        #return blus(p, ones) - ones
+        #if randomize:
+        #    direction = flow * torch.polar(ones.real, (random() * random_range - random_range / 2) * ones.real)
+        if flow == 0:
+            result = p - direction * ebb
+        else:
+            result = proj_shift(p, direction * flow) - direction * ebb
+        #res_abs = result.abs()
+        if rescaling:
+            result /= result.norm(p=2,dim=1).max()
+            #result = torch.polar(2 * res_abs / res_abs.max(), result.angle())
+        return result
+    return transformation
+
+def _deriv(b, c):
+    def derivative(p):
+        return p * ((b * p).sum(dim=0)) - c * b
+    return derivative
+
+def gpt_deriv(b, k=1.0, c=0.1, omega=0.5):
+    """
+    Returns a function deriv(points) -> [3,N] where points are 3D column vectors.
+    b: tensor of shape [3] (will be normalized internally)
+    k: scalar gain for tanh nonlinearity
+    c: scalar drift magnitude along -b
+    omega: scalar rotation strength around b
+    """
+    b = b / b.norm()  # normalize once
+
+    def deriv(points):
+        # <a, b> for each column
+        #dots = torch.einsum('ij,i->j', points, b)  # shape [N]
+        dots = (points * b[:]).sum(dim=0)  # shape [N]
+
+        # radial scaling term
+        radial = torch.tanh(k * dots).unsqueeze(0) * points  # [3,N]
+
+        # drift along -b
+        drift = -c * b[:].expand_as(points)  # [3,N]
+
+
+        # rotation around b: b × a
+        rotation = omega * torch.cross(b[:].expand_as(points), points, dim=0)
+
+        return radial + drift + rotation
+
+    return deriv
+
+p_positions = (torch.rand([3, particle_count], device=device, dtype=t_fp) - 0.5) * 3
+#p_positions[2] = 0
+direction = torch.zeros_like(p_positions) / math.sqrt(3)
+direction[1] += 1
+
+#derivative = _deriv(direction, 0.5)
+derivative = gpt_deriv(direction, 1.0, 0.1, 0.5)
+
+step = lambda p, dp, h: p + dp * h * dt
+rk4_curried = lambda p: rk4_step(derivative, step, p)
+
+
+p_colors = torch.rand([particle_count,3], device=device, dtype=t_fp)
+
+color_rotation = torch.linspace(0, tau / 4, particle_count)
+
+p_colors[:,0] = (p_positions[0,:] / 2) + 0.5
+p_colors[:,1] = (p_positions[1,:] / 2) + 0.5
+p_colors[:,2] = (p_positions[2,:] / 2) + 0.5
+
+def project(p, colors, i):
+    c = math.cos(i * rotation_rate)
+    s = math.sin(i * rotation_rate)
+    rotation = torch.tensor([
+        [1, 0, 0],
+        [0, c,-s],
+        [0, s, c]])
+    #rotation = torch.tensor([
+    #    [ c, 0, s],
+    #    [ 0, 1, 0],
+    #    [-s, 0, c]])
+    #rotation = torch.tensor([
+    #    [c, -s, 0],
+    #    [s,  c, 0],
+    #    [0,  0, 1]])
+    alt_colors = colors.clone()
+    res = (rotation @ p)
+    #color_filter = res[2].abs() < 0.7
+    #alt_colors[:,0] *= color_filter
+    #alt_colors[:,1] *= color_filter
+    #alt_colors[:,2] *= color_filter
+    return (res[0:2], alt_colors)
+
+frame_index = [0]
+
+def run():
+    scratch = torch.zeros([h, w, 3], device=device, dtype=t_fp)
+
+    for iteration in range(iterations):
+        if save_if(iteration) or iteration == iterations - 1:
+            (p_projected, alt_colors) = project(p_positions, p_colors, iteration)
+            frame_index[0] += 1
+            scratch *= 0
+            insert_at_coords(p_projected, alt_colors, scratch, mapping)
+            save(scratch.permute(2,0,1), f"{run_dir}/{frame_index[0]:06d}")
+
+        p_positions.copy_(rk4_curried(p_positions))
+
+

sketch/old/snakepyt_0_0/orbits.py

···
+
+import math
+
+import torch
+
+from lib.util import save
+
+def legendre_p(l, x):
+    P = (torch.ones_like(x), x)
+    if l < 2:
+        return P[l]
+    for n in range(2, l+1):
+        p1_term = (2*n - 1) * x * P[1]
+        p0_term = (n - 1) * P[0]
+        Pnext = (p1_term - p0_term) / n
+        P = (P[1], Pnext)
+    return P[1]
+
+def associated_legendre(l, m, x):
+    def derivative(f, degree, dx=1e-5):
+        for iteration in range(degree):
+            f = lambda x, f=f, dx=dx: (f(x+dx) - f(x-dx)) / (2 * dx)
+        return f
+
+    p = lambda x: legendre_p(l, x)
+    dp = derivative(p, m)
+    return (-1)**m * (1 - x**2)**(m / 2) * dp(x)
+
+def spherical_harmonic(l, m, theta, phi):
+    x = torch.cos(theta)
+    p = associated_legendre(l, abs(m), x)
+    if m == 0:
+        return p * (1 + 0j)
+    return p * torch.exp(1j * m * phi)
+
+def generalized_laguerre(k, a, x):
+    total = torch.zeros_like(x)
+    for i in range(k + 1):
+        numer = math.factorial(k + a)
+        denom = math.factorial(k - i) * math.factorial(a + i) * math.factorial(i)
+        total += ((-x)**i) * (numer / denom)
+    return total
+
+def radial_wavefunc(n, l, r):
+    rho = 2 * r / n
+    laguerre = generalized_laguerre(n - l - 1, 2 * l + 1, rho)
+    return r**l * torch.exp(-rho / 2) * laguerre
+
+def orbital_wavefunc(n, l, m, r, theta, phi):
+    R = radial_wavefunc(n, l, r)
+    Y = spherical_harmonic(l, m, theta, phi)
+    return R * Y
+
+def sphericalize(coords):
+    r = coords.norm(p=2,dim=0)
+    theta = torch.acos(coords[2]/r)#torch.clamp(coords[2] / r, -1.0, 0.0))
+    phi = torch.atan2(coords[1], coords[0])
+    return r, theta, phi
+
+
+def init():
+    torch.set_default_device("cuda")
+
+    h = 2**9
+    w = 2**9
+    d = 2**9
+
+    scale = 20#0
+
+    rx = -scale, scale
+    ry = -scale, scale
+    rz = -scale, scale
+
+    coords = torch.zeros([3, h, w, d])
+
+    xspace = torch.linspace(rx[0], rx[1], w, dtype=torch.double)
+    yspace = torch.linspace(ry[0], ry[1], h, dtype=torch.double)
+    zspace = torch.linspace(rz[0], rz[1], d, dtype=torch.double)
+
+    coords[0] = xspace.view([w,1,1]).expand([w,h,d]).permute(1,0,2)
+    coords[1] = yspace.view([h,1,1]).expand([h,w,d])
+    coords[2] = zspace.view([d,1,1]).expand([d,h,w]).permute(1,2,0)
+
+    schedule(per_n, [9])
+    #schedule(per_n, range(1,6))
+
+def per_n(n):
+    schedule(per_l, range(n))
+    #schedule(per_l, [3])
+    #schedule(per_l, range(-5,6))
+
+def per_l(l):
+    schedule(per_m, range(l+1))
+    #schedule(per_m, range(-l,l+1))
+    #schedule(per_m, [3,2])
+    #schedule(per_m, range(-5,6))
+
+def per_m(m):
+    #schedule(per_frame, range(60*4))
+    schedule(per_frame, [0])
+
+def per_frame(frame_index):
+    torch.cuda.empty_cache()
+
+    rotation_rate = 1 * math.tau / 60
+    rotation_rate_b = 1 * math.tau / 60
+    rotation_rate_c = 1 * math.tau / 60
+
+    rotation_rate   = 0 if frame_index > 60 and frame_index < 180 else rotation_rate
+    rotation_rate_b = 0 if frame_index < 60 or frame_index > 120 else rotation_rate_b
+    rotation_rate_c = 0 if frame_index < 120 else rotation_rate_c
+
+    c = math.cos(frame_index * rotation_rate)
+    s = math.sin(frame_index * rotation_rate)
+    rotation = torch.tensor([
+        [1, 0, 0],
+        [0, c,-s],
+        [0, s, c]])
+    c = math.cos(frame_index * rotation_rate_b)
+    s = math.sin(frame_index * rotation_rate_b)
+    rotation_b = torch.tensor([
+        [ c, 0, s],
+        [ 0, 1, 0],
+        [-s, 0, c]])
+    c = math.cos(frame_index * rotation_rate_c)
+    s = math.sin(frame_index * rotation_rate_c)
+    rotation_c = torch.tensor([
+        [c, -s, 0],
+        [s,  c, 0],
+        [0,  0, 1]])
+
+
+    _coords = (coords.permute(1,2,3,0) @ (rotation_c @ rotation_b @ rotation).T).permute(3,0,1,2)
+
+    _coords *= n
+
+    r, theta, phi = sphericalize(_coords)
+    res = orbital_wavefunc(n, l, m, r, theta, phi)
+    for dim in [0]:
+
+        if m > 0:
+            res_c = res.real.abs() * res.imag
+            res_pos = res_c.clamp(0,torch.inf).sum(dim=dim)
+            res_pos /= res_pos.abs().max()
+            res_neg = (-res_c).clamp(0,torch.inf).sum(dim=dim)
+            res_neg /= res_neg.abs().max()
+
+            res_r = res_pos#.pow(0.3)
+            res_b = res_neg#.pow(0.3)
+
+        else:
+            res_r = res.real.clamp(0,torch.inf).sum(dim=dim)
+            res_b = (-res.real).clamp(0,torch.inf).sum(dim=dim)
+
+            res_r /= res_r.max()
+            res_b /= res_b.max()
+
+        res_g = torch.zeros_like(res_r)
+
+        res_rgb = torch.stack((res_r, res_g, res_b))
+
+        save(res_rgb, f"{run_dir}/{n}_{l}_{m}_view{dim}_{frame_index:06d}")
+    del r, theta, phi, res, res_r, res_g, res_b
+

sketch/old/snakepyt_0_0/physarum.py

···
+import time
+from random import randint
+from pathlib import Path
+from math import tau
+import math
+
+import torch
+from torchvision.transforms import GaussianBlur
+
+from lib.util import *
+from lib.io import VideoWriter
+
+def init():
+    scale = 2 ** 12
+    w = scale
+    h = scale
+
+    max_iterations = 2000
+    save_mod = 100
+    also_save = []
+
+    particle_count = int((w * h) * 3)
+    decay_rate = 0.03
+
+    view_angle = tau / 10
+    view_distance = 1.414 * 10
+
+    direction_count = 40000
+    turn_amount = tau / 10#direction_count
+    move_distance = 1.414 * 1
+
+    blur_size = 3
+    blur_sigma = 1.5
+
+    r_strength = 0.8
+    b_strength = 50
+    g_strength = 0.03
+
+    dtype = torch.double
+    ctype = torch.cdouble
+    device = "cuda"
+    torch.set_default_device(device)
+    torch.set_default_dtype(dtype)
+
+    video = VideoWriter((h,w), 60,
+            f"out/{run_dir}/vid.mp4",
+            f"out/{run_dir}/ffmpeg_log")
+
+    schedule(run, None)
+
+def final():
+    video.close()
+
+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)
+    tensor[:,indices[:,0],indices[:,1]] += value
+
+def run():
+
+    blur = GaussianBlur(blur_size, blur_sigma)
+
+    # p_ denotes particle data
+    p_positions = torch.rand([particle_count,2]) - 0.5
+    p_positions[:,0] *= h / 10
+    p_positions[:,1] *= w / 10
+    p_positions[:,0] += h / 10
+    p_positions[:,1] += w / 2
+
+    #stimulus_positions_r = torch.tensor([
+    #    [h//3,w//3],
+    #], dtype=torch.long)
+
+    stimulus_positions_r = torch.rand([1000,2]) - 0.5
+    stimulus_positions_r[:,0] *= h #* 0.9
+    stimulus_positions_r[:,1] *= w #* 0.55
+    #stimulus_positions_r[0:100,1] *= 0.02
+    #stimulus_positions_r[0:100,0] *= 1 / 0.9
+    stimulus_positions_r[:,0] += h / 2
+    stimulus_positions_r[:,1] += w / 2
+
+    #stimulus_positions_b = torch.tensor([
+    #    [h//3,2*w//3],
+    #    [2*h//3,w//3],
+    #], dtype=torch.long)
+
+    stimulus_positions_b = torch.rand([100,2]) - 0.5
+    stimulus_positions_b[:,0] *= h * 0.6
+    stimulus_positions_b[:,1] *= w #* 0.5
+    stimulus_positions_b[:,0] += h / 2
+    stimulus_positions_b[:,1] += w / 2
+
+    p_directions = ((torch.rand(particle_count)*direction_count).floor() / direction_count) * tau
+
+    world = torch.zeros([3,h,w], dtype=dtype)
+    scratch = torch.zeros([3,h,w], dtype=dtype)
+    black = torch.tensor([0, 0, 0], dtype=dtype)[:,None]
+    white = torch.tensor([1, 1, 1], dtype=dtype)[:,None]
+    red = torch.tensor([1, 0, 0], dtype=dtype)[:, None]
+    green = torch.tensor([0, 1, 0], dtype=dtype)[:, None]
+    blue = torch.tensor([0, 0, 1], dtype=dtype)[:, None]
+
+    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]
+        senses = [world[:,c[:,0], c[:,1]] for c in sci]
+
+        senses = [30 * s[2] + s[1] - (500 * s[0]) for s in senses]
+
+        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
+
+        #sense_out = sense_out[1] & ~sense_out[0]
+        #sense_pos = sense_pos[1] & ~sense_neg[0]
+        #sense_neg = sense_neg[1] & ~sense_pos[0]
+
+        # 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
+
+        p_directions = ((p_directions * (direction_count / tau)).floor() / direction_count) * tau
+
+        # move
+        p_positions = offset_coord(p_positions, p_directions, move_distance)
+        clamp_coords(p_positions, h, w)
+
+        # deposit
+        idxput(world, p_positions.long(), green * g_strength)
+
+        idxput(world, stimulus_positions_r.long(), red * r_strength)
+        idxput(world, stimulus_positions_b.long(), blue * b_strength)
+
+        # diffuse
+        world = blur(world)
+        world[2][None] = blur(world[2][None])
+
+        crowding_mask = world[1] > 0.9
+        world[0] += 0.1 * crowding_mask
+        loneliness_mask = world[1] < 0.01
+        world[0] += 0.005 * loneliness_mask
+
+        world[2] *= (1 - world[1]).clamp(0,1) ** 0.4
+
+        # render
+        video.frame(world)
+
+        if iteration % save_mod == 0 and iteration != 0:
+            save(world, f"{run_dir}/{iteration:06d}")
+
+        if iteration in also_save:
+            save(world, f"{run_dir}/{iteration:06d}")
+
+        # decay
+        if iteration < max_iterations - 1:
+            world *= (1-decay_rate)
+
+    save(world, f"{run_dir}/final")

sketch/old/snakepyt_0_0/pix_sort.py

···
+
+import torch
+from torch.nn.functional import conv2d
+from math import tau
+
+from lib.util import load_image_tensor, save, msave
+from diffusers import AutoencoderKL
+
+t_fp = torch.double
+
+def persistent():
+    pass
+    #vae = AutoencoderKL.from_pretrained(
+    #    "/ssd/0/ml_models/ql/hf-diff/stable-diffusion-xl-base-0.9", subfolder="vae"
+    #)
+    #vae.to("cuda")
+    #vae.to(torch.float)
+
+def decode(l, vae):
+    l = l.to(vae.dtype)
+    with torch.no_grad():
+        ims = vae.decode(l).sample
+    return ims
+
+def encode(im, vae):
+    im = im.to(vae.dtype)
+    with torch.no_grad():
+        l = vae.encode(im)
+    return l
+
+def init():
+    torch.set_default_device("cuda")
+    torch.set_default_dtype(t_fp)
+
+    im = load_image_tensor("in/sunset.jpg").to(t_fp).to("cuda")
+    #im = encode(im.unsqueeze(0), vae)[0].mean[0].to(t_fp)
+
+    #im = im.transpose(1,2)
+
+    threshold = 0.4
+
+    min_strip_size = 2
+    max_strip_size = 400000
+
+    # 1 = vertical, 2 = horizontal
+    sort_dim = 1
+
+    ranges = [(0.15, 0.33), (0.44, 0.9)]
+    def in_ranges(x):
+        for r in ranges:
+            if x > r[0] and x < r[1]:
+                return True
+        return False
+
+    column_selection = lambda w: [x for x in range(w) if in_ranges(x/w)]
+
+
+    schedule(run, None)
+
+def convolve(field, kernel):
+    return conv2d(field.unsqueeze(0), kernel, bias=None, padding=[1, 0], stride=[1])[0]
+
+kernel3 = torch.tensor([
+    [[[0], [-1], [1]], [[0], [0], [0]], [[0], [0], [0]]],
+    [[[0], [0], [0]], [[0], [-1], [1]], [[0], [0], [0]]],
+    [[[0], [0], [0]], [[0], [0], [0]], [[0], [-1], [1]]],
+], dtype=t_fp, device="cuda")
+
+kernel4 = torch.tensor([
+    [[[0], [-1], [1]], [[0], [0], [0]], [[0], [0], [0]], [[0], [0], [0]]],
+    [[[0], [0], [0]], [[0], [-1], [1]], [[0], [0], [0]], [[0], [0], [0]]],
+    [[[0], [0], [0]], [[0], [0], [0]], [[0], [-1], [1]], [[0], [0], [0]]],
+    [[[0], [0], [0]], [[0], [0], [0]], [[0], [0], [0]], [[0], [-1], [1]]],
+], dtype=t_fp, device="cuda")
+
+
+def run():
+    (_, h, w) = im.shape
+
+    out = im.clone() #torch.zeros_like(im)
+    group_masks = torch.zeros([h,w], dtype=torch.long)
+
+    edges = convolve(im, kernel3)
+    edges = (edges.norm(p=2, dim=0) > threshold)
+
+    scanner_encounters = torch.zeros([w], dtype=torch.long)
+
+    n_edges = edges.sum(dim=0)
+    print(n_edges)
+
+    for row_index in range(h):
+        scanner_encounters += edges[row_index]
+
+        group_masks[row_index] = scanner_encounters
+
+    n_groups = group_masks.max() + 1
+
+    for group_index in range(n_groups):
+        group_mask = group_masks == group_index
+
+        for column in column_selection(w):
+            inds = torch.where(group_mask[:, column])[0]
+            if inds.shape[0] < min_strip_size or inds.shape[0] > max_strip_size:
+                continue
+            i_min, i_max = inds.min(), inds.max()
+            strip = im[:, i_min:i_max, column]
+            strip_vals = strip.norm(p=2, dim=0)
+            s_vals, s_inds = torch.sort(strip_vals)
+            out[:, i_min:i_max, column] = strip[:, s_inds]
+
+    #out = 0.8 * out + 0.2 * im
+    #out = out.transpose(1,2)
+    #out = decode(out.unsqueeze(0), vae)[0]
+    save(out, f"{run_dir}/out")
+    #save(out, f"{run_dir}/out")
+
+

sketch/old/snakepyt_0_0/plume/lines.py

···
+
+from math import tau
+
+import torch
+
+from lib.spaces import map_space, draw_points_2d, dotted_lines_2d, center_span, grid
+from lib.util import save
+
+def init():
+    t_real = torch.double
+    torch.set_default_device("cuda")
+
+    scale = 2 ** 12
+
+    #r = 1
+
+    #c,s = center_span((-r,r),(-r,r))
+    c = (0, -0.5)
+    s = (1,1)
+
+    mapping = map_space(c, s, [], (1,1), scale)
+    _, (h,w) = mapping
+
+    canvas = torch.zeros([3,h,w], dtype=t_real)
+
+    schedule(test0, None)
+
+
+def test0():
+    global canvas
+    inner_scale = 50
+    points = torch.rand([2, 2, inner_scale ** 2], dtype=t_real)
+    colors = torch.ones([3], dtype=t_real)
+    r = colors.clone()
+    r[1] = 0
+    r[2] = 0
+
+    points[0] = grid(map_space(c, s, [], (1,1), inner_scale)).permute(2,0,1).view(2, inner_scale**2)
+    points[1] = points[0]
+    points[1,0] += s[1] / inner_scale
+
+    #p = points[0].clone()
+
+    #draw_points_2d(points[0], r, canvas, mapping)
+    #save(canvas, f"{run_dir}/test")
+
+    _points = points.clone()
+
+    dotted_lines_2d(_points, colors, 100, canvas, mapping)
+    #draw_points_2d(_points[0], r, canvas, mapping)
+    save(canvas, f"{run_dir}/{0:06d}")
+
+    for n in range(24 * 6):
+        p = _points
+
+        mult = 1 + (p[:,0]) / (p[:,0] * p[:,0] + p[:,1] * p[:,1])
+        points[:,0] = p[:,0] * mult - 0.804
+        points[:,1] = p[:,1] * mult
+        #p = points[1]
+
+        _points[0] = points[0]
+        _points[1] = points[1]
+
+        #_points[1] /= _points[1].norm(p=2,dim=0)
+        #_points[1] /= inner_scale / s[0]
+        #points[1] *= 10
+        #_points[1] += _points[0]
+        canvas *= 0
+        dotted_lines_2d(_points, colors, 100, canvas, mapping)
+        #draw_points_2d(_points[0], r, canvas, mapping)
+        save(canvas, f"{run_dir}/{n+1:06d}")
+
+
+    #dotted_lines_2d(points, colors, 30, canvas, mapping)
+    #draw_points_2d(points[0], r, canvas, mapping)
+
+    #save(canvas, f"{run_dir}/canvas")
+

sketch/old/snakepyt_0_0/plume/plume.py

···
+
+import math
+import time
+from pathlib import Path
+from random import random
+
+import torch
+
+from lib.util import *
+from lib.spaces import insert_at_coords, map_space
+
+def init():
+    schedule(_settings_0, range(30))
+
+name = "blus"
+
+device = "cuda"
+torch.set_default_device(device)
+
+t_real = torch.double
+t_complex = torch.cdouble
+
+tau = 3.14159265358979323 * 2
+
+randomize = False
+random_range = tau #/ 50
+show_gridlines = False
+
+particle_count = 100000
+
+
+# 2 ** 9 = 512; 10 => 1024; 11 => 2048
+scale_power = 10
+scale = 2 ** scale_power
+
+origin = 0, 0
+span = 5, 5
+
+stretch = 1, 1
+
+zooms = [
+        ]
+
+
+p_positions = (torch.rand([particle_count], device=device, dtype=t_complex) - (0.5 + 0.5j)) * 4
+p_colors = torch.ones([particle_count,3], device=device, dtype=t_real)
+color_rotation = torch.linspace(0, tau / 4, particle_count)
+
+p_colors[:,0] = torch.frac((p_positions.real / 4 + 0.5) * 10)
+p_colors[:,2] = torch.frac((p_positions.imag / 4 + 0.5) * 10)
+p_colors[:,1] = (1 - (p_colors[:,0] + p_colors[:,2])).clamp(0,1)
+
+ones = torch.ones_like(p_positions)
+
+def blus(a, b):
+    theta_a = a.angle()
+    return torch.polar(a.abs() + b.abs() * torch.cos(b.angle() - theta_a), theta_a)
+
+def get_transformation(direction, flow, ebb, rescaling):
+    def transformation(p):
+        #return blus(p, ones) - ones
+        #if randomize:
+        #    direction = flow * torch.polar(ones.real, (random() * random_range - random_range / 2) * ones.real)
+        result = blus(p, direction * flow) - direction * ebb
+        res_abs = result.abs()
+        if rescaling:
+            result = torch.polar(2 * res_abs / res_abs.max(), result.angle())
+        return result
+    return transformation
+
+
+def _settings_0(seed):
+    if seed % 2 == 0:
+        ebb = 0
+        iterations = 1
+        flow = 2 / iterations
+        rescaling = False
+    elif seed % 2 == 1:
+        flow = 0
+        iterations = 1
+        ebb = 0.333 / iterations
+        rescaling = False
+    elif False:#seed % 3 == 2:
+        flow = 0
+        ebb = 0
+        iterations = 10
+        rescaling = True
+    angle = 0 #if seed < 2 else tau / 4
+    direction = torch.polar(ones.real, 1 * tau * ones.real)
+    next_positions = get_transformation(direction, flow, ebb, rescaling)
+    schedule(run, None)
+
+def settings_1(idx):
+    if idx % 2 == 0:
+        iterations = 60
+        ebb = 0
+        flow = 1 / iterations
+    else:
+        iterations = 30
+        ebb = 1 / iterations
+        flow = 0
+    rescaling = False
+    angle = (idx // 2) * (tau / 4)
+    direction = torch.polar(ones.real, angle * ones.real)
+    next_positions = get_transformation(direction, flow, ebb, rescaling)
+    schedule(run, None)
+
+
+
+
+frame_index = [0]
+
+mapping = map_space(origin, span, zooms, stretch, scale)
+(_, (h,w)) = mapping
+
+def run():
+    scratch = torch.zeros([h, w, 3], device=device, dtype=t_real)
+
+    def project(p):
+        return torch.view_as_real(p).permute(1,0)
+
+    for iteration in range(iterations):
+        frame_index[0] += 1
+        p_projected = project(p_positions).clone()
+
+        if iteration % 1 == 0:
+            scratch *= 0
+            insert_at_coords(p_projected, p_colors, scratch, mapping)
+            if show_gridlines:
+                scratch[:,:,1] += gridlines
+                scratch[:,:,2] += gridlines
+            save(scratch.permute(2,0,1).sqrt(), f"{run_dir}/{frame_index[0]:06d}")
+
+        p_positions.copy_(next_positions(p_positions))
+
+

sketch/old/snakepyt_0_0/plume/plume3d.py

···
+
+import math
+import time
+from pathlib import Path
+from random import random
+
+import torch
+
+from lib.util import *
+from lib.spaces import insert_at_coords, map_space
+
+def init():
+    schedule(_settings_0, range(10000))
+
+name = "plume3"
+
+device = "cuda"
+torch.set_default_device(device)
+
+t_real = torch.double
+t_complex = torch.cdouble
+
+tau = 3.14159265358979323 * 2
+
+randomize = False
+random_range = tau #/ 50
+
+particle_count = 1000000
+
+
+# 2 ** 9 = 512; 10 => 1024; 11 => 2048
+scale_power = 10
+scale = 2 ** scale_power
+
+origin = 0, 0
+span = 2, 2
+
+stretch = 1, 1
+
+zooms = [
+        ]
+
+
+p_positions = (torch.rand([particle_count,3], device=device, dtype=t_real) - (0.5))
+#p_positions[:,0] -= 0.5
+#p_positions[:,1] -= 0.5
+#p_positions[:,2] = 0# -= 0.5
+p_positions *= 1
+p_colors = torch.ones([particle_count,3], device=device, dtype=t_real)
+color_rotation = torch.linspace(0, tau / 4, particle_count)
+
+p_colors[:,0] = torch.frac((p_positions[:,0] / 4 + 0.5) * 10)
+#p_colors[:,1] = torch.frac((p_positions[:,1] / 2 + 0.5) * 10)
+p_colors[:,2] = torch.frac((p_positions[:,1] / 4 + 0.5) * 10)
+p_colors[:,1] = (1 - (p_colors[:,0] + p_colors[:,2])).clamp(0,1)
+
+ones = torch.ones_like(p_positions)
+
+def proj(u, v):
+    dot = (u * v).sum(dim=1)
+    scale = (dot / (v.norm(p=2, dim=1) ** 2))
+    out = v.clone()
+    # todo: find better way
+    for dim in range(3):
+        out[:,dim] *= scale
+    return out
+
+def proj_shift(a, b):
+    return a + proj(b, a)
+
+def get_transformation(direction, flow, ebb, rescaling):
+    def transformation(p):
+        #return blus(p, ones) - ones
+        #if randomize:
+        #    direction = flow * torch.polar(ones.real, (random() * random_range - random_range / 2) * ones.real)
+        if flow == 0:
+            result = p - direction * ebb
+        else:
+            result = proj_shift(p, direction * flow) - direction * ebb
+        #res_abs = result.abs()
+        if rescaling:
+            result /= result.norm(p=2,dim=1).max()
+            #result = torch.polar(2 * res_abs / res_abs.max(), result.angle())
+        return result
+    return transformation
+
+
+def _settings_0(seed):
+    direction = ones.clone()
+    direction[:,1] = 0
+    direction[:,2] = 0
+    show_frame = True
+    if seed % 3 == 0:
+        ebb = 0
+        iterations = 1
+        flow = 1 / iterations
+        rescaling = False
+        show_frame = seed == 9999
+    elif seed % 3 == 1:
+        flow = 0
+        iterations = 1
+        ebb = 0.87 / iterations
+        rescaling = False
+        show_frame = False
+    elif seed % 3 == 2:
+        flow = 0
+        ebb = 0
+        iterations = 1
+        rescaling = True
+        show_frame = False
+    next_positions = get_transformation(direction, flow, ebb, rescaling)
+    schedule(run, None)
+
+
+frame_index = [0]
+
+mapping = map_space(origin, span, zooms, stretch, scale)
+(_, (h,w)) = mapping
+
+rotation_rate = 0#.02
+rotation_rate_b = 0#.025
+
+def project(p, c, transform):
+    (offset, rotation) = transform
+    _p = p - offset
+    _p = _p / _p.norm(p=2,dim=1).max()
+    _p = (rotation @ _p.transpose(1,0)).transpose(1,0)
+    c_filter = _p[:,2] < 99990
+    _c = c.clone()
+    for d in range(3):
+        _c[:,d] *= c_filter
+    return (_p[:,0:2].permute(1,0), _c)
+
+scratch = torch.zeros([h, w, 3], device=device, dtype=t_real)
+
+def run():
+    global scratch
+    for iteration in range(iterations):
+        if show_frame:
+            c = math.cos(frame_index[0] * rotation_rate)
+            s = math.sin(frame_index[0] * rotation_rate)
+            rotation_a = torch.tensor([
+                [1, 0, 0],
+                [0, c,-s],
+                [0, s, c]], dtype=torch.double)
+            c = math.cos(frame_index[0] * rotation_rate_b)
+            s = math.sin(frame_index[0] * rotation_rate_b)
+            rotation_b = torch.tensor([
+                [c,-s, 0],
+                [s, c, 0],
+                [0, 0, 1]], dtype=torch.double)
+            dist = direction
+            offset = flow * dist * (iteration)
+            offset -= ebb * dist * iteration
+
+            transform = (offset, rotation_a @ rotation_b)
+
+            frame_index[0] += 1
+            p_projected, c_filtered = project(p_positions, p_colors, transform)
+
+            if iteration % 1 == 0:
+                scratch *= 0
+                insert_at_coords(p_projected, c_filtered, scratch, mapping)
+                save(scratch.permute(2,0,1).sqrt(), f"{run_dir}/{frame_index[0]:06d}")
+
+        p_positions.copy_(next_positions(p_positions))
+
+

sketch/old/snakepyt_0_0/plume/plume3d_b.py

···
+
+import math
+import time
+from pathlib import Path
+from random import random
+import gc
+
+import torch
+
+from lib.util import *
+from lib.spaces import insert_at_coords, map_space, cgrid
+
+
+name = "plume3"
+
+device = "cuda"
+torch.set_default_device(device)
+gc.collect()
+torch.cuda.empty_cache()
+
+t_real = torch.double
+t_complex = torch.cdouble
+
+tau = 3.14159265358979323 * 2
+
+randomize = False
+random_range = tau #/ 50
+
+particle_count = 100
+
+
+# 2 ** 9 = 512; 10 => 1024; 11 => 2048
+scale_power = 7
+scale = 2 ** scale_power
+
+origin = 0.75, 0.75
+span = 0.5, 0.5
+
+stretch = 1, 1
+
+zooms = [
+        ]
+
+#p_positions = (torch.rand([particle_count,2], device=device, dtype=t_real) - (0.5))
+
+#ones = torch.ones_like(p_positions)
+
+def proj(u, v):
+    dot = (u * v).sum(dim=2)
+    norm_sq = v.norm(p=2, dim=2) ** 2
+    zero_mask = norm_sq == 0
+    norm_sq += zero_mask
+    scale = (dot / norm_sq) * ~zero_mask
+    out = v.clone()
+    for dim in range(out.shape[2]):
+        out[:,:,dim] *= scale
+    return out
+
+def proj_shift(a,b):
+    return a + proj(b,a)
+
+def proj_shift_(a, b):
+    dot = (a * b).sum(dim=2)
+    scale = (dot / (b.norm(p=2, dim=2) ** 2))
+    for dim in range(a.shape[2]):
+        a[:,:,dim] *= 1 + scale
+    return a
+
+def get_transformation(direction, flow, ebb, rescaling):
+    def transformation(p):
+        #if flow == 0:
+        #    result = p - direction * ebb
+        #else:
+        _ebb = torch.outer(ebb, direction).unsqueeze(1).expand(-1, particle_count, -1)
+        _flow = torch.outer(flow, direction).unsqueeze(1).expand(-1, particle_count, -1)
+        result = proj_shift(p, _flow) - _ebb
+        if rescaling:
+            result /= result.norm(p=2,dim=1).max()
+        return result
+    return transformation
+
+def init():
+    mapping = map_space(origin, span, zooms, stretch, scale)
+    (_, (h,w)) = mapping
+
+    ab = cgrid(mapping)
+
+    image = torch.zeros([h, w, 3], device="cpu", dtype=t_real)
+
+    p_positions = torch.rand([w, particle_count, 2], device=device, dtype=t_real) - 0.5
+
+    schedule(rows, range(h))
+
+def rows(row_index):
+
+    schedule(columns, None)#range(w))
+
+    schedule(out, None)
+
+def out():
+    im = image.clone()
+    for c in range(3):
+        im[:,:,c] = im[:,:,c] > 0
+        #im[:,:,c] /= im[:,:,c].max()
+    save(image.permute(2,0,1), f"{run_dir}/{frame_index[0]}")
+    frame_index[0] += 1
+
+def columns():
+    direction = torch.tensor([1,0], dtype=t_real)#ones_like(p_positions)#.clone()
+    #direction[:,1] = 0
+
+    iterations = 1000
+    ebb = ab[row_index].real
+    flow = ab[row_index].imag
+
+    rescaling = True
+
+
+    next_positions = get_transformation(direction, flow, ebb, rescaling)
+
+    schedule(run, None)
+
+
+
+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
+
+
+frame_index = [0]
+
+def run():
+    global scratch
+
+    lyapunov_avg = torch.zeros([w,particle_count], dtype=t_real)
+    diverge_iter = torch.ones([w], dtype=t_real) * torch.inf
+
+    var_avg = torch.zeros([w], dtype=t_real)
+    converge_iter = torch.ones([w], dtype=t_real) * torch.inf
+
+    skip = 500
+
+    for iteration in range(iterations):
+        _next = next_positions(p_positions)
+        if (iteration > skip):
+            var = _next.var(dim=(1,2))
+            lyapunov_term = (_next - p_positions).norm(p=2,dim=2).log()
+            lyapunov_avg = c_avg(lyapunov_avg, lyapunov_term, iteration - skip)
+            converge = (var - var_avg) < -(10**5)
+            var_avg = c_avg(var_avg, var, iteration - 5)
+            if converge.any():
+                div_now = converge * (iteration - skip) / (iterations - skip) + torch.nan_to_num(~converge * torch.inf, 0, torch.inf)
+                converge_iter = torch.min(converge_iter, div_now)
+        p_positions.copy_(_next)
+
+    image[row_index, :, 0] = torch.nan_to_num(converge_iter, 0, 0, 0)
+    means = lyapunov_avg.mean(dim=1)
+    image[row_index, :, 2] = torch.nan_to_num(means, 0, 0, 0).abs()
+    #image[row_index, :, 1] = torch.nan_to_num(means == 0, 0, 0, 0).abs()
+
+
+    means = var_avg.abs()
+    image[row_index, :, 1] = torch.nan_to_num(means, 0, 0, 0).abs()
+
+
+

sketch/old/snakepyt_0_0/plume/plume_c.py

···
+
+import math
+import time
+from pathlib import Path
+from random import random
+
+import torch
+
+from lib.util import *
+from lib.spaces import insert_at_coords, map_space, cgrid
+
+
+device = "cuda"
+torch.set_default_device(device)
+
+t_real = torch.double
+t_complex = torch.cdouble
+
+tau = 3.14159265358979323 * 2
+
+randomize = False
+random_range = tau #/ 50
+
+# 2 ** 9 = 512; 10 => 1024; 11 => 2048
+scale_power = 13
+scale = 2 ** scale_power
+
+
+stretch = 1, 1
+
+zooms = [
+        ]
+
+def proj(u, v):
+    dot = (u * v).sum(dim=1)
+    scale = (dot / (v.norm(p=2, dim=1) ** 2))
+    out = v.clone()
+    for dim in range(out.shape[1]):
+        out[:,dim] *= scale
+    return out
+
+def bad_proj(u, v):
+    dot = (u * v).sum(dim=1)
+    scale = (dot / v.norm(p=2, dim=1))
+    out = v.clone()
+    for dim in range(out.shape[1]):
+        out[:,dim] *= scale
+    return out
+
+def proj_shift(a, b):
+    return a + proj(b, a)
+
+def bad_proj_shift(a, b):
+    return a + bad_proj(b, a)
+
+def get_transformation(direction, flow, ebb, rescaling):
+    def transformation(p):
+        d = direction
+        d[:,0] += torch.rand_like(direction[:,0]) * 0.0001
+        if flow == 0:
+            result = p - direction * ebb
+        else:
+            result = proj_shift(p, d * flow)
+            result -= direction * ebb #proj_shift(result, direction * ebb)
+        if rescaling:
+            result /= result.norm(p=2,dim=1).max()
+        return result
+    return transformation
+
+def init():
+    steps = 1
+    schedule(per_t, [i / steps for i in range(steps)])
+
+def eerp(a,b,t):
+    return math.pow(b/a, t)*a
+
+def per_t(t):
+    origin = 0.0625, -0.15
+
+    #s = eerp(1000, 0.001, t)
+    s = 0.25
+    x_range = 0.5 * s
+    y_range = 1 * s
+
+    span = x_range, y_range
+    mapping = map_space(origin, span, zooms, stretch, scale)
+    (_, (h,w)) = mapping
+
+    grid = cgrid(mapping)
+
+    scratch = torch.zeros([h, w, 3], device=device, dtype=t_real)
+    p_positions = torch.zeros([h*w,2], device=device, dtype=t_real)
+    p_positions[:,0] = grid.real.reshape((h*w))
+    p_positions[:,1] = grid.imag.reshape((h*w))
+
+    direction = torch.ones_like(p_positions)
+    direction[:,0] = 0
+
+    iterations = 1301
+    #show_frame = lambda i: i == iterations - 1
+    show_frame = lambda i: i % 100 == 0#True
+    #show_frame = lambda i: True
+
+    #ebb = 0.804#0.9000469530469531 - 0.001 + 0.002 * t
+    #ebb = 1 + 0.4 * t
+    #ebb = 0.4
+    ebb = 0.83314# + (0.001 * t)
+    #ebb = (0.5 + 0.5 * t)#0.804#482 / 600
+    #ebb = 0.0 + 0.01 * t #0.9000469530469531 #+ 0.00001 * t #0.9000595404595405
+    flow = 1
+    rescaling = False
+
+
+    next_positions = get_transformation(direction, flow, ebb, rescaling)
+    schedule(run, None)
+
+frame_index = [0]
+
+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
+
+
+def run():
+    global scratch
+
+    diff_avg = torch.zeros([h*w], device=device, dtype=t_real)
+
+    for iteration in range(iterations):
+        _next = next_positions(p_positions)
+        diff = (p_positions - 0).norm(p=2, dim=1)
+
+        diff_avg = c_avg(diff_avg, diff, iteration + 1)
+        #diff /= diff.max()
+
+        if show_frame(iteration):
+            #diff_avg = diff
+            frame_index[0] += 1
+            print(f"{frame_index[0]}: ebb={ebb}")
+
+            #pos_reshape = p_positions.reshape((h,w,2))
+            #scratch[:,:,0] = pos_reshape[:,:,0] + 0.5
+            #scratch[:,:,1] = pos_reshape[:,:,1] + 0.5
+
+            #scratch += 0.5
+            #scratch *= 2
+            scratch[:,:,2] = diff_avg.nan_to_num().reshape((h,w))
+            #scratch[:,:,2] = diff_avg.nan_to_num().reshape((h,w))
+
+            scratch[scratch < 0] = 0
+
+            #scratch[:,:,0] = scratch[:,:,0].pow(2)
+            #scratch[:,:,2] = scratch[:,:,2].pow(0.5)
+
+            scratch[:,:,2] -= scratch[:,:,2].mean()
+            scratch[:,:,2] /= scratch[:,:,2].std() * 6
+            #scratch[:,:,2] /= scratch[:,:,2].max()
+
+            scratch[:,:,2] += 0.5
+
+            #scratch[:,:,1] = diff.reshape((h,w)) / diff.max()
+
+            save(scratch.permute(2,0,1), f"{run_dir}/{frame_index[0]:06d}")
+
+        p_positions.copy_(_next)
+
+

sketch/old/snakepyt_0_0/seam_carve.py

···
+import torch
+from torch.nn.functional import conv2d
+from math import tau
+
+from diffusers import AutoencoderKL
+
+from lib.util import load_image_tensor, save, msave
+
+t_fp = torch.double
+
+def persistent():
+    vae = AutoencoderKL.from_pretrained(
+        "/ssd/0/ml_models/ql/hf-diff/stable-diffusion-xl-base-0.9", subfolder="vae"
+    )
+    vae.to("cuda")
+    vae.to(torch.float)
+
+def decode(l, vae):
+    l_dtype = l.dtype
+    l = l.unsqueeze(0).to(vae.dtype)
+    with torch.no_grad():
+        ims = vae.decode(l).sample
+    return ims.to(l_dtype)[0]
+
+def encode(im, vae):
+    im_dtype = im.dtype
+    im = im.unsqueeze(0).to(vae.dtype)
+    with torch.no_grad():
+        l = vae.encode(im)
+    return l[0].mean[0].to(im_dtype)
+
+
+def init():
+    torch.set_default_device("cuda")
+    torch.set_default_dtype(t_fp)
+    do_vae = False
+    shrink = False
+    do_horiz = True
+    remove = 1000
+
+    im = load_image_tensor("in/seam_code_b.png").to(t_fp).to("cuda")
+    if do_vae:
+        im = encode(im, vae)
+    if do_horiz:
+        im = im.transpose(1,2)
+    schedule(run, None)
+
+def convolve(field, kernel, padding):
+    return conv2d(field.unsqueeze(0), kernel, bias=None, padding=padding, stride=[1])[0]
+
+kernel3 = torch.tensor([
+    [[[-1], [0], [1]], [[0], [0], [0]], [[0], [0], [0]]],
+    [[[0], [0], [0]], [[-1], [0], [1]], [[0], [0], [0]]],
+    [[[0], [0], [0]], [[0], [0], [0]], [[-1], [0], [1]]],
+], dtype=t_fp, device="cuda")
+
+kernel4 = torch.tensor([
+    [[[-1], [0], [1]], [[0], [0], [0]], [[0], [0], [0]], [[0], [0], [0]]],
+    [[[0], [0], [0]], [[-1], [0], [1]], [[0], [0], [0]], [[0], [0], [0]]],
+    [[[0], [0], [0]], [[0], [0], [0]], [[-1], [0], [1]], [[0], [0], [0]]],
+    [[[0], [0], [0]], [[0], [0], [0]], [[0], [0], [0]], [[-1], [0], [1]]],
+], dtype=t_fp, device="cuda")
+
+def run():
+    prev_im = im
+    _shrink = shrink
+    try:
+        for iteration in range(remove):
+            prev_shrink = _shrink
+            _shrink = (iteration % 4) > 0
+
+            prev_im = prev_im.transpose(1,2) if prev_shrink != _shrink else prev_im
+            (c, h, w) = prev_im.shape
+
+            if _shrink:
+                out = torch.zeros([c, h, w-1])
+            else:
+                out = torch.zeros([c, h, w+1])
+
+
+
+            kernel = kernel3 if c == 3 else kernel4
+            edges = convolve(prev_im, kernel, [1,0])
+            edges += convolve(prev_im, kernel.transpose(2,3), [0,1])
+            energy = edges.norm(p=2, dim=0)
+
+            paths = energy.clone()
+            if not _shrink:
+                paths += torch.rand_like(energy) * (11 if not do_vae else 100)
+            paths_l = torch.roll(paths, 1, dims=1)
+            paths_r = torch.roll(paths, -1, dims=1)
+            paths_l[0] = torch.inf
+            paths_r[-1] = torch.inf
+            paths_stack = torch.stack((paths_l, paths, paths_r))
+            path_mins = torch.min(paths_stack, dim=0)
+
+            for row_index in range(1,h):
+                prev_vals = paths_stack[:, row_index - 1]
+                mins, inds = torch.min(prev_vals, dim=0)
+                paths_stack[:, row_index] += mins.expand((3,w))
+
+            paths = paths_stack[1]
+
+            vals, ind = torch.min(paths[-1], dim=0)
+            min_path = [ind.item()]
+            for row_index in range(h-1)[::-1]:
+                prev_ind = min_path[-1]
+                l_ind = ((prev_ind - 1) + w) % w
+                r_ind = (prev_ind + 2) % (w+1)
+                if prev_ind == 0 or prev_ind == (w - 1):
+                    v, ind = torch.min(paths[row_index,(l_ind,prev_ind,r_ind-1)], dim=0)
+                else:
+                    v, ind = torch.min(paths[row_index,l_ind:r_ind], dim=0)
+                min_path.append((prev_ind + ind.item() - 1 + w) % w)
+
+            for row_index in range(h):
+                split_at = min_path[-row_index]
+                if _shrink:
+                    out[:,row_index,:split_at] = prev_im[:,row_index,:split_at]
+                    out[:,row_index,split_at:] = prev_im[:,row_index,(split_at+1):]
+                else: # grow
+                    out[:,row_index,:split_at] = prev_im[:,row_index,:split_at]
+                    out[:,row_index,split_at] = prev_im[:,row_index,split_at]
+                    out[:,row_index,(split_at+1):] = prev_im[:,row_index,split_at:]
+
+            prev_im = out
+            _do_horiz = not prev_shrink
+
+            if iteration % 1 == 0:
+                to_save = decode(out, vae) if do_vae else out
+                to_save = to_save.transpose(1,2) if _do_horiz else to_save
+                save(to_save, f"{run_dir}/{iteration:06d}")
+                #save(out.transpose(1,2), f"{run_dir}/{iteration:06d}")
+    except:
+        to_save = decode(prev_im, vae) if do_vae else out
+        to_save = to_save.transpose(1,2) if do_horiz else to_save
+        save(to_save, f"{run_dir}/{iteration:06d}")
+        #save(prev_im.transpose(1,2), f"{run_dir}/{iteration:06d}")
+        raise
+
+    #out = out.transpose(1,2)
+    to_save = decode(out, vae) if do_vae else out
+    to_save = to_save.transpose(1,2) if do_horiz else to_save
+    save(to_save, f"{run_dir}/out")
+
+

sketch/webserver.py

···
     except Exception as e:
         return build_response(500, f"Server error: {e}".encode())
 
-def run():
+def main():
     ssl_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
     ssl_ctx.load_cert_chain(certfile=SSL_CERT, keyfile=SSL_KEY)
 
···
         print(f"Serving HTTPS on {HOST}:{PORT}")
         with ssl_ctx.wrap_socket(server, server_side=True) as ssock:
             while True:
-                conn, addr = ssock.accept()
+                try:
+                    conn, addr = ssock.accept()
+                except KeyboardInterrupt:
+                    raise
+                except:
+                    print("exc")
+                    continue
                 with conn:
                     request = conn.recv(4096)
                     response = handle_request(request)

webui/content/main.html

···
         <link id="themeLink" rel="stylesheet" type="text/css" href="style/theme/void.css">
         <link rel="stylesheet" type="text/css" href="style/main.css">
         <link rel="stylesheet" type="text/css" href="style/load.css">
-        <link rel="stylesheet" type="text/css" href="style/../../../foo.css">
     </head>
     <body>
         <script src="main.js"></script>

webui/js/col.js

···
+
+const sheet = new CSSStyleSheet();
+sheet.replaceSync(`
+    .hsplitter {
+        height: 1px;
+        background-color: var(--main-faded);
+        cursor: row-resize;
+        user-select: none;
+        -webkit-user-drag: none;
+        overflow: visible;
+    }
+
+    .hsplitter::before {
+        content: '';
+        position: relative;
+        display: inline-block;
+        top: -6px;
+        left: 0;
+        width: 100%;
+        height: 13px;
+        /*background-color: rgba(255, 0, 0, 0.1);*/
+    }
+
+    .target.top {
+        padding-bottom: 1rem;
+    }
+
+    .target.bottom {
+        padding-top: 1rem;
+    }
+
+    .target.middle {
+        padding-top: 1rem;
+        padding-bottom: 1rem;
+    }
+`);
+document.adoptedStyleSheets = [...document.adoptedStyleSheets, sheet];
+
+export function main(target, n = 2) {
+    n = parseInt(n, 10);
+    if (!Number.isInteger(n) || n < 2) {
+        n = 2;
+    }
+
+    const container = document.createElement('div');
+    container.style.cssText = 'display: flex; flex-direction: column; height: 100%; width: 100%;';
+
+    const targets = [];
+    const splitters = [];
+
+    function createDragHandler(splitter, i) {
+        splitter.onmousedown = (e) => {
+            e.preventDefault();
+
+            function resizeCallback(e) {
+                const containerRect = container.getBoundingClientRect();
+                const relativeY = e.clientY - containerRect.top;
+                const percent = (relativeY / containerRect.height) * 100;
+
+                // Calculate total height of targets above this splitter
+                let topTotal = 0;
+                for (let j = 0; j <= i; j++) {
+                    if (j === i) {
+                        const newHeight = Math.max(1, Math.min(99, percent - topTotal));
+                        targets[j].style.height = newHeight + '%';
+                    } else {
+                        topTotal += parseFloat(targets[j].style.height);
+                    }
+                }
+            }
+
+            function cleanup() {
+                document.removeEventListener('mousemove', resizeCallback);
+                document.removeEventListener('mouseup', cleanup);
+                document.removeEventListener('mouseleave', cleanup);
+            }
+
+            document.addEventListener('mousemove', resizeCallback);
+            document.addEventListener('mouseup', cleanup);
+            document.addEventListener('mouseleave', cleanup);
+        };
+    }
+
+    for (let i = 0; i < n; i++) {
+        const target = document.createElement('div');
+        target.style.height = `${100/n}%`;
+
+        if (i === 0) {
+            target.className = "target top";
+        } else if (i === n - 1) {
+            target.className = "target bottom";
+            target.style.flex = '1'; // Last one gets flex to handle rounding
+            target.style.height = 'auto';
+        } else {
+            target.className = "target middle";
+        }
+
+        targets.push(target);
+        container.appendChild(target);
+
+        if (i === n - 1) continue;
+
+        const splitter = document.createElement('div');
+        splitter.className = 'hsplitter';
+        splitters.push(splitter);
+        container.appendChild(splitter);
+
+        createDragHandler(splitter, i);
+    }
+
+    target.appendChild(container);
+
+    return {
+        targets: targets
+    };
+}
+

webui/js/exit.js

···
+
+export function main(target) {
+    const dummy = document.createElement("div");
+
+    return {
+        targets: [dummy]
+    };
+}
+

webui/js/main.js

···
 window.userPreferences.setLoadButton(localStorage.getItem('loadButton'));
 
 
+import("/prompt.js").then(
+    module => module.main(document.body)
+)
+

webui/js/prompt.js

···
+
+export function main(target) {
+    const container = document.createElement('div');
+
+    const input = document.createElement('input');
+    input.type = 'text';
+    input.placeholder = '...';
+
+    async function handleLoad() {
+        const inputValue = input.value.trim();
+        const inputSplit = inputValue.split(/\s+/);
+        const moduleName = inputSplit[0];
+        const args = inputSplit.slice(1);
+
+        try {
+            const module = await import(`/${moduleName}.js`);
+            if ("main" in module) {
+                const result = await module.main(target, ...args);
+                if (result?.targets?.length) {
+                    for (const targetElement of result.targets) {
+                        main(targetElement);
+                    }
+                    container.remove();
+                }
+            }
+
+        } catch (error) {
+            console.error('Failed to load module:', error.message);
+        }
+    }
+
+    input.addEventListener('keypress', async (e) => {
+        if (e.key === 'Enter') {
+            handleLoad();
+        }
+    });
+
+    container.appendChild(input);
+    target.appendChild(container);
+    input.focus();
+}
+

webui/js/row.js

···
+
+const sheet = new CSSStyleSheet();
+sheet.replaceSync(`
+    .vsplitter {
+        width: 1px;
+        background-color: var(--main-faded);
+        cursor: col-resize;
+        user-select: none;
+        -webkit-user-drag: none;
+        overflow: visible;
+    }
+
+    .vsplitter::before {
+        content: '';
+        position: relative;
+        display: inline-block;
+        top: 0;
+        left: -6px;
+        width: 13px;
+        height: 100%;
+        /*background-color: rgba(255, 0, 0, 0.1);*/
+    }
+
+    .target.left {
+        padding-right: 1rem;
+    }
+
+    .target.right {
+        padding-left: 1rem;
+    }
+
+    .target.middle {
+        padding-left: 1rem;
+        padding-right: 1rem;
+    }
+`);
+document.adoptedStyleSheets = [...document.adoptedStyleSheets, sheet];
+
+export function main(target, n = 2) {
+    n = parseInt(n, 10);
+    if (!Number.isInteger(n) || n < 2) {
+        n = 2;
+    }
+
+    const container = document.createElement('div');
+    container.style.cssText = 'display: flex; height: 100%; width: 100%;';
+
+    const targets = [];
+    const splitters = [];
+
+    function createDragHandler(splitter, i) {
+        splitter.onmousedown = (e) => {
+            e.preventDefault();
+
+            function resizeCallback(e) {
+                const containerRect = container.getBoundingClientRect();
+                const relativeX = e.clientX - containerRect.left;
+                const percent = (relativeX / containerRect.width) * 100;
+
+                // Calculate total width of targets to the left of this splitter
+                let leftTotal = 0;
+                for (let j = 0; j <= i; j++) {
+                    if (j === i) {
+                        const newWidth = Math.max(1, Math.min(99, percent - leftTotal));
+                        targets[j].style.width = newWidth + '%';
+                    } else {
+                        leftTotal += parseFloat(targets[j].style.width);
+                    }
+                }
+            }
+
+            function cleanup() {
+                document.removeEventListener('mousemove', resizeCallback);
+                document.removeEventListener('mouseup', cleanup);
+                document.removeEventListener('mouseleave', cleanup);
+            }
+
+            document.addEventListener('mousemove', resizeCallback);
+            document.addEventListener('mouseup', cleanup);
+            document.addEventListener('mouseleave', cleanup);
+        };
+    }
+
+    for (let i = 0; i < n; i++) {
+        const target = document.createElement('div');
+        target.style.width = `${100/n}%`;
+
+        if (i === 0) {
+            target.className = "target left";
+        } else if (i === n - 1) {
+            target.className = "target right";
+            target.style.flex = '1'; // Last one gets flex to handle rounding
+            target.style.width = 'auto';
+        } else {
+            target.className = "target middle";
+        }
+
+        targets.push(target);
+        container.appendChild(target);
+
+        if (i === n - 1) continue;
+
+        const splitter = document.createElement('div');
+        splitter.className = 'vsplitter';
+        splitters.push(splitter);
+        container.appendChild(splitter);
+
+        createDragHandler(splitter, i);
+    }
+
+    target.appendChild(container);
+
+    return {
+        targets: targets
+    };
+}
+

webui/readme

···
+old_fic_svelte contains the old .svelte files from the original ficscript UI
+
+the styling and scripts are probably still a useful baseline / reference for a lot of what i want to do with this repo
+
+but i don't want to use svelte, so a little bit of work needs to be done to strip it down to html/css/js

webui/style/main.css

···
     -webkit-box-sizing: border-box;
 }
 
+body {
+    height: 100vh;
+    width: 100vw;
+    padding: 1rem;
+}
+
+.target {
+    height: 100%;
+    width: 100%;
+}
+
 pre {
     font-family: var(--code-font);
 }