ponder.ooo
🐍🐍🐍
continue overhaul
+186
-27
+61
lib/spaces.py
+61
lib/spaces.py
···
1
+
2
+
import torch
3
+
4
+
type fp_range = tuple[float, float]
5
+
type fp_region2 = tuple[fp_range, fp_range]
6
+
type fp_coords2 = tuple[float, float]
7
+
type hw = tuple[int, int]
8
+
type region_mapping = tuple[fp_region2, hw]
9
+
10
+
def insert_at_coords(coords, values, target, mapping: region_mapping):
11
+
(region, hw) = mapping
12
+
(xrange, yrange) = region
13
+
(h,w) = hw
14
+
(x_min, x_max) = xrange
15
+
(y_min, y_max) = yrange
16
+
17
+
mask = torch.ones([coords.shape[1]])
18
+
mask *= (coords[1] >= x_min) * (coords[1] <= x_max)
19
+
mask *= (coords[0] >= y_min) * (coords[0] <= y_max)
20
+
in_range = mask.nonzero().squeeze()
21
+
22
+
# TODO: combine coord & value tensors so there's only one index_select necessary
23
+
coords_filtered = torch.index_select(coords.permute(1,0), 0, in_range)
24
+
values_filtered = torch.index_select(values, 0, in_range)
25
+
26
+
coords_filtered[:,1] -= x_min
27
+
coords_filtered[:,1] *= (w-1) / (x_max - x_min)
28
+
coords_filtered[:,0] -= y_min
29
+
coords_filtered[:,0] *= (h-1) / (y_max - y_min)
30
+
indices = coords_filtered.long()
31
+
32
+
target.index_put_((indices[:,0],indices[:,1]), values_filtered, accumulate=True)
33
+
34
+
# maps a 2d region of space to a canvas
35
+
def map_space(origin, span, zooms, stretch, scale) -> region_mapping:
36
+
x_min = origin[0] - (span[0] / 2)
37
+
y_min = origin[1] - (span[1] / 2)
38
+
39
+
for ((xa, xb), (ya, yb)) in zooms:
40
+
x_min += span[0] * xa
41
+
y_min += span[1] * ya
42
+
span = span[0] * (xb - xa), span[1] * (yb - ya)
43
+
44
+
x_max = x_min + span[0]
45
+
y_max = y_min + span[1]
46
+
47
+
aspect = span[0] * stretch[0] / (span[1] * stretch[1])
48
+
49
+
if aspect < 1:
50
+
h = scale
51
+
w = int(scale * aspect)
52
+
else:
53
+
w = scale
54
+
h = int(scale / aspect)
55
+
56
+
x_range = (x_min, x_max)
57
+
y_range = (y_min, y_max)
58
+
region = (x_range, y_range)
59
+
return (region, (h,w))
60
+
61
+
-25
lib/util.py
-25
lib/util.py
···
4
4
from safetensors.torch import save_file as sft_save
5
5
import io
6
6
7
-
class Settings(dict):
8
-
__getattr__ = dict.__getitem__
9
-
__setattr__ = dict.__setitem__
10
-
__delattr__ = dict.__delitem__
11
-
12
-
def __init__(self, _vars):
13
-
for key in _vars:
14
-
self[key] = _vars[key]
15
-
16
-
class _Settings():
17
-
def __init__(self, _get, _src):
18
-
self.get = _get
19
-
self.src = _src
20
-
21
-
def settings(f):
22
-
import inspect
23
-
f_src = inspect.getsource(f)
24
-
f_lines = f_src.split('\n')
25
-
f_lines = [l for l in f_lines if f"@settings" not in l]
26
-
f_lines.append(" return Settings(vars())")
27
-
new_src = "\n".join(f_lines)
28
-
evaluation_scope = {}
29
-
exec(new_src, locals=evaluation_scope)
30
-
return _Settings(eval(f.__name__, locals=evaluation_scope),new_src)
31
-
32
7
def first(l, p):
33
8
return next((idx,value) for idx,value in enumerate(l) if p(value))
34
9
+5
-2
main.py
+5
-2
main.py
···
1
1
2
2
import inspect
3
+
import traceback
3
4
from argparse import ArgumentParser as ArgParser
4
5
from importlib import import_module as load
5
6
from time import perf_counter
···
28
29
try:
29
30
settings.update(sandbox[fns[0][0]]()) # run the fn, dump its locals into the inner settings context
30
31
except Exception as e:
31
-
print(f"[Error in {fns[0][0]}: {e}]\n")
32
+
traceback.print_exception(e)
33
+
#print(f"[Error in {fns[0][0]}: {e}]\n")
32
34
return
33
35
if len(fns) == 1:
34
36
print(f"[begin {run_id}]")
···
36
38
try:
37
39
exec(sketch.run.__code__, globals=settings)
38
40
except Exception as e:
39
-
print(f"Error: {e}")
41
+
traceback.print_exception(e)
42
+
#print(f"Error: {e}")
40
43
#print(f"[done in {perf_counter() - t0:.3f}s]\n")
41
44
print()
42
45
else:
+120
sketch/blus.py
+120
sketch/blus.py
···
1
+
2
+
import math
3
+
import time
4
+
from pathlib import Path
5
+
from random import random
6
+
7
+
import torch
8
+
9
+
from lib.util import *
10
+
from lib.spaces import insert_at_coords, map_space
11
+
12
+
seeds = range(50)
13
+
name = "blus"
14
+
15
+
device = "cuda"
16
+
torch.set_default_device(device)
17
+
18
+
t_real = torch.double
19
+
t_complex = torch.cdouble
20
+
21
+
tau = 3.14159265358979323 * 2
22
+
23
+
randomize = False
24
+
random_range = tau #/ 50
25
+
show_gridlines = False
26
+
27
+
particle_count = 300000
28
+
29
+
30
+
# 2 ** 9 = 512; 10 => 1024; 11 => 2048
31
+
scale_power = 11
32
+
scale = 2 ** scale_power
33
+
34
+
origin = 0, 0
35
+
span = 5, 5
36
+
37
+
stretch = 1, 1
38
+
39
+
zooms = [
40
+
]
41
+
42
+
43
+
p_positions = (torch.rand([particle_count], device=device, dtype=t_complex) - (0.5 + 0.5j)) * 4
44
+
p_colors = torch.ones([particle_count,3], device=device, dtype=t_real)
45
+
color_rotation = torch.linspace(0, tau / 4, particle_count)
46
+
47
+
p_colors[:,0] = p_positions.real / 4 + 0.5
48
+
p_colors[:,2] = p_positions.imag / 4 + 0.5
49
+
p_colors[:,1] *= 0
50
+
51
+
ones = torch.ones_like(p_positions)
52
+
53
+
def blus(a, b):
54
+
theta_a = a.angle()
55
+
return torch.polar(a.abs() + b.abs() * torch.cos(b.angle() - theta_a), theta_a)
56
+
57
+
def get_transformation(direction, flow, ebb, rescaling):
58
+
def transformation(p):
59
+
#return blus(p, ones) - ones
60
+
#if randomize:
61
+
# direction = flow * torch.polar(ones.real, (random() * random_range - random_range / 2) * ones.real)
62
+
result = blus(p, direction * flow) - direction * ebb
63
+
res_abs = result.abs()
64
+
if rescaling:
65
+
result = torch.polar(2 * res_abs / res_abs.max(), result.angle())
66
+
return result
67
+
return transformation
68
+
69
+
70
+
def settings_0():
71
+
if seed % 3 == 0:
72
+
ebb = 0
73
+
iterations = 90
74
+
flow = 2 / iterations
75
+
rescaling = False
76
+
elif seed % 3 == 1:
77
+
flow = 0
78
+
iterations = 30
79
+
ebb = 1.8 / iterations
80
+
rescaling = False
81
+
elif seed % 3 == 2:
82
+
flow = 0
83
+
ebb = 0
84
+
iterations = 10
85
+
rescaling = True
86
+
angle = 0 #if seed < 2 else tau / 4
87
+
direction = torch.polar(ones.real, 1 * tau * ones.real)
88
+
next_positions = get_transformation(direction, flow, ebb, rescaling)
89
+
90
+
91
+
92
+
frame_index = [0]
93
+
94
+
run_dir = time.strftime(f"%d.%m.%Y/{name}_t%H.%M.%S")
95
+
mapping = map_space(origin, span, zooms, stretch, scale)
96
+
(_, (h,w)) = mapping
97
+
Path("out/" + run_dir).mkdir(parents=True, exist_ok=True)
98
+
Path("out/" + run_dir + "/aggregate").mkdir(parents=True, exist_ok=True)
99
+
Path("out/" + run_dir + "/frame").mkdir(parents=True, exist_ok=True)
100
+
101
+
def run():
102
+
scratch = torch.zeros([h, w, 3], device=device, dtype=t_real)
103
+
104
+
def project(p):
105
+
return torch.view_as_real(p).permute(1,0)
106
+
107
+
for iteration in range(iterations):
108
+
frame_index[0] += 1
109
+
p_projected = project(p_positions).clone()
110
+
111
+
if iteration % 1 == 0:
112
+
scratch *= 0
113
+
insert_at_coords(p_projected, p_colors, scratch, mapping)
114
+
if show_gridlines:
115
+
scratch[:,:,1] += gridlines
116
+
scratch[:,:,2] += gridlines
117
+
save(scratch.permute(2,0,1).sqrt(), f"{run_dir}/frame/{frame_index[0]:06d}")
118
+
119
+
p_positions.copy_(next_positions(p_positions))
120
+