822 lines
28 KiB
JavaScript
822 lines
28 KiB
JavaScript
/*
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
https://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
*/
|
|
|
|
/*
|
|
Usage:
|
|
const gui = new dat.GUI();
|
|
const ca = new CA(gl, models_json, [W, H], gui); // gui is optional
|
|
ca.step();
|
|
|
|
ca.paint(x, y, radius, modelIndex);
|
|
ca.clearCircle(x, y, radius;
|
|
|
|
const stats = ca.benchmark();
|
|
ca.draw();
|
|
ca.draw(zoom);
|
|
*/
|
|
|
|
const vs_code = `
|
|
attribute vec4 position;
|
|
varying vec2 uv;
|
|
void main() {
|
|
uv = position.xy*0.5 + 0.5;
|
|
gl_Position = position;
|
|
}
|
|
`
|
|
|
|
function defInput(name) {
|
|
return `
|
|
uniform Tensor ${name};
|
|
uniform sampler2D ${name}_tex;
|
|
|
|
vec4 ${name}_read(vec2 pos, float ch) {return _read(${name}, ${name}_tex, pos, ch);}
|
|
vec4 ${name}_read01(vec2 pos, float ch) {return _read01(${name}, ${name}_tex, pos, ch);}
|
|
vec4 ${name}_readUV(vec2 uv) {return _readUV(${name}, ${name}_tex, uv);}
|
|
`
|
|
}
|
|
|
|
const PREFIX = `
|
|
#extension GL_OES_standard_derivatives : enable
|
|
precision highp float;
|
|
|
|
const float PI = 3.14159265359;
|
|
|
|
// "Hash without Sine" by David Hoskins (https://www.shadertoy.com/view/4djSRW)
|
|
float hash13(vec3 p3) {
|
|
p3 = fract(p3 * .1031);
|
|
p3 += dot(p3, p3.yzx + 33.33);
|
|
return fract((p3.x + p3.y) * p3.z);
|
|
}
|
|
vec2 hash23(vec3 p3)
|
|
{
|
|
p3 = fract(p3 * vec3(.1031, .1030, .0973));
|
|
p3 += dot(p3, p3.yzx+33.33);
|
|
return fract((p3.xx+p3.yz)*p3.zy);
|
|
}
|
|
|
|
struct Tensor {
|
|
vec2 size;
|
|
vec2 gridSize;
|
|
float depth, depth4;
|
|
vec2 packScaleZero;
|
|
};
|
|
uniform Tensor u_output;
|
|
|
|
vec4 _readUV(Tensor tensor, sampler2D tex, vec2 uv) {
|
|
vec4 v = texture2D(tex, uv);
|
|
vec2 p = tensor.packScaleZero;
|
|
v = (v-p.y)*p.x;
|
|
return v;
|
|
}
|
|
vec2 _getUV(Tensor tensor, vec2 pos, float ch) {
|
|
ch += 0.5;
|
|
float tx = floor(mod(ch, tensor.gridSize.x));
|
|
float ty = floor(ch / tensor.gridSize.x);
|
|
vec2 p = fract(pos/tensor.size) + vec2(tx, ty);
|
|
p /= tensor.gridSize;
|
|
return p;
|
|
}
|
|
vec4 _read01(Tensor tensor, sampler2D tex, vec2 pos, float ch) {
|
|
return texture2D(tex, _getUV(tensor, pos, ch));
|
|
}
|
|
vec4 _read(Tensor tensor, sampler2D tex, vec2 pos, float ch) {
|
|
vec2 p = _getUV(tensor, pos, ch);
|
|
return _readUV(tensor, tex, p);
|
|
}
|
|
vec2 getOutputXY() {
|
|
return mod(gl_FragCoord.xy, u_output.size);
|
|
}
|
|
float getOutputChannel() {
|
|
vec2 xy = floor(gl_FragCoord.xy/u_output.size);
|
|
return xy.y*u_output.gridSize.x+xy.x;
|
|
}
|
|
|
|
void setOutput(vec4 v) {
|
|
vec2 p = u_output.packScaleZero;
|
|
v = v/p.x + p.y;
|
|
gl_FragColor = v;
|
|
}
|
|
|
|
#ifdef SPARSE_UPDATE
|
|
uniform sampler2D u_shuffleTex, u_unshuffleTex;
|
|
uniform vec2 u_shuffleOfs;
|
|
#endif
|
|
|
|
${defInput('u_input')}
|
|
|
|
uniform float u_angle, u_alignment;
|
|
const float u_hexGrid = 1.0;
|
|
|
|
mat2 rotate(float ang) {
|
|
float s = sin(ang), c = cos(ang);
|
|
return mat2(c, s, -s, c);
|
|
}
|
|
|
|
vec2 ang2vec(float a) {
|
|
return vec2(cos(a), sin(a));
|
|
}
|
|
|
|
${defInput('u_alignTex')}
|
|
vec2 getCellDirection(vec2 xy) {
|
|
return u_alignTex_read(xy, 0.0).xy;
|
|
}
|
|
|
|
vec4 conv3x3(vec2 xy, float inputCh, mat3 filter) {
|
|
vec4 a = vec4(0.0);
|
|
for (int y=0; y<3; ++y)
|
|
for (int x=0; x<3; ++x) {
|
|
vec2 p = xy+vec2(float(x-1), float(y-1));
|
|
a += filter[y][x] * u_input_read(p, inputCh);
|
|
}
|
|
return a;
|
|
}
|
|
|
|
// https://www.shadertoy.com/view/Xljczw
|
|
// https://www.shadertoy.com/view/MlXyDl
|
|
// returns xy - in cell pos, zw - skewed cell id
|
|
vec4 getHex(vec2 u) {
|
|
vec2 s = vec2(1., mix(2.0, 1.732, u_hexGrid));
|
|
vec2 p = vec2(0.5*u_hexGrid, 0.5);
|
|
vec2 a = mod(u ,s)*2.-s;
|
|
vec2 b = mod(u+s*p,s)*2.-s;
|
|
vec2 ai = floor(u/s);
|
|
vec2 bi = floor(u/s+p);
|
|
// skewed coords
|
|
ai = vec2(ai.x-ai.y*u_hexGrid, ai.y*2.0+1.0);
|
|
bi = vec2(bi.x-bi.y*u_hexGrid, bi.y*2.0);
|
|
return dot(a,a)<dot(b,b) ? vec4(a, ai) : vec4(b, bi);
|
|
}
|
|
|
|
float hex(in vec2 p){
|
|
vec2 s = vec2(1., 1.732);
|
|
p = abs(p);
|
|
return max(dot(p, s*.5), p.x); // Hexagon.
|
|
}
|
|
|
|
// https://www.shadertoy.com/view/XtXcWs
|
|
vec2 cmul(vec2 a, vec2 b) {
|
|
return vec2(a.x*b.x - a.y*b.y, a.x*b.y + a.y*b.x);
|
|
}
|
|
|
|
vec2 cdiv(vec2 a, vec2 b) {
|
|
return cmul(a, vec2(b.x, -b.y)) / dot(b, b);
|
|
}
|
|
|
|
uniform vec2 u_viewSize;
|
|
|
|
struct Hexel {
|
|
vec2 cellXY;
|
|
vec2 p;
|
|
float zoom;
|
|
};
|
|
|
|
Hexel screen2hex(vec2 xy) {
|
|
xy /= u_viewSize;
|
|
xy.y = 1.0-xy.y;
|
|
vec2 normViewSize = u_viewSize/length(u_viewSize);
|
|
xy = (xy*0.85+0.25) * normViewSize;
|
|
|
|
Hexel h;
|
|
float nxy = length(xy);
|
|
h.zoom = 4.0/(nxy*nxy);
|
|
xy = cmul(xy, xy);
|
|
xy = cmul(xy, xy);
|
|
xy *= 160.0;
|
|
vec4 r = getHex(xy);
|
|
h.cellXY = r.zw;
|
|
h.p = r.xy;
|
|
return h;
|
|
}
|
|
|
|
float calcMouseDist(vec2 mousePosScr) {
|
|
Hexel h = screen2hex(mousePosScr);
|
|
h.cellXY = mod(h.cellXY, u_output.size);
|
|
vec2 diff = abs(getOutputXY()-h.cellXY-0.5);
|
|
return length(min(diff, u_output.size-diff))*h.zoom;
|
|
}
|
|
`;
|
|
|
|
const PROGRAMS = {
|
|
paint: `
|
|
uniform vec2 u_pos;
|
|
uniform float u_r;
|
|
uniform vec4 u_brush;
|
|
|
|
void main() {
|
|
if (u_r>0.0 && calcMouseDist(u_pos)>=80.0)
|
|
discard;
|
|
setOutput(u_brush);
|
|
}`,
|
|
peek: `
|
|
uniform vec2 u_pos;
|
|
|
|
vec2 getPeekPos(float i) {
|
|
float a = i*0.61803398875*2.0*PI;
|
|
float r = (u_viewSize.x+u_viewSize.y)/1000.0;
|
|
return vec2(cos(a), sin(a)) * sqrt(i) * r;
|
|
}
|
|
|
|
void main() {
|
|
float out_i = getOutputXY().x;
|
|
float i = floor(out_i / u_input.depth4);
|
|
float channel = floor(mod(out_i, u_input.depth4));
|
|
Hexel h = screen2hex(u_pos + getPeekPos(i));
|
|
setOutput(u_input_read(h.cellXY, channel));
|
|
}`,
|
|
align: `
|
|
uniform vec2 u_pos;
|
|
uniform float u_r;
|
|
uniform float u_init;
|
|
|
|
const mat3 blur = mat3(1.0/9.0);
|
|
const mat3 blurHex = mat3(0.0, 1.0, 1.0,
|
|
1.0, 1.0, 1.0,
|
|
1.0, 1.0, 0.0)/7.0;
|
|
|
|
void main() {
|
|
vec2 xy = getOutputXY();
|
|
vec4 v = conv3x3(xy, 0.0, blur*(1.0-u_hexGrid) + blurHex*u_hexGrid);
|
|
v.xy = normalize(mix(u_input_read(xy, 0.0).xy, v.xy, 1.0));
|
|
setOutput(v);
|
|
|
|
if (u_init > 0.0) {
|
|
if (u_r>0.0 && calcMouseDist(u_pos)>=80.0)
|
|
return;
|
|
float a = hash13(vec3(xy+vec2(34299.0, -56593.0), u_init)) * 2.0 * PI;
|
|
vec2 v = normalize(ang2vec(a)+0.2*ang2vec(u_init));
|
|
setOutput(vec4(v, 0.0, 0.0));
|
|
}
|
|
}`,
|
|
perception: `
|
|
const mat3 sobelX = mat3(-1.0, 0.0, 1.0, -2.0, 0.0, 2.0, -1.0, 0.0, 1.0)/8.0;
|
|
const mat3 sobelY = mat3(-1.0,-2.0,-1.0, 0.0, 0.0, 0.0, 1.0, 2.0, 1.0)/8.0;
|
|
const mat3 gauss = mat3(1.0, 2.0, 1.0, 2.0, 4.0-16.0, 2.0, 1.0, 2.0, 1.0)/8.0;
|
|
const mat3 sobelXhex = mat3( 0.0, -1.0, 1.0,
|
|
-2.0, 0.0, 2.0,
|
|
-1.0, 1.0, 0.0)/8.0;
|
|
|
|
const mat3 sobelYhex = mat3( 0.0, -2.0,-2.0,
|
|
0.0, 0.0, 0.0,
|
|
2.0, 2.0, 0.0)/8.0;
|
|
|
|
const mat3 gaussHex = mat3(0.0, 2.0, 2.0,
|
|
2.0, 4.0-16.0, 2.0,
|
|
2.0, 2.0, 0.0)/8.0;
|
|
|
|
void main() {
|
|
vec2 xy = getOutputXY();
|
|
#ifdef SPARSE_UPDATE
|
|
xy = texture2D(u_shuffleTex, xy/u_output.size).xy*255.0+0.5 + u_shuffleOfs;
|
|
xy = mod(xy, u_input.size);
|
|
#endif
|
|
float ch = getOutputChannel();
|
|
if (ch >= u_output.depth4)
|
|
return;
|
|
|
|
float filterBand = floor((ch+0.5)/u_input.depth4);
|
|
float inputCh = ch-filterBand*u_input.depth4;
|
|
if (filterBand < 0.5) {
|
|
setOutput(u_input_read(xy, inputCh));
|
|
} else if (filterBand < 2.5) {
|
|
vec4 dx = conv3x3(xy, inputCh, sobelX*(1.0-u_hexGrid) + sobelXhex*u_hexGrid);
|
|
vec4 dy = conv3x3(xy, inputCh, sobelY*(1.0-u_hexGrid) + sobelYhex*u_hexGrid);
|
|
vec2 dir = getCellDirection(xy);
|
|
float s = dir.x, c = dir.y;
|
|
setOutput(filterBand < 1.5 ? dx*c-dy*s : dx*s+dy*c);
|
|
} else {
|
|
setOutput(conv3x3(xy, inputCh, gauss*(1.0-u_hexGrid) + gaussHex*u_hexGrid));
|
|
}
|
|
}`,
|
|
dense: `
|
|
${defInput('u_control')}
|
|
uniform sampler2D u_weightTex;
|
|
uniform float u_seed, u_fuzz;
|
|
uniform vec2 u_weightCoefs; // scale, center
|
|
uniform vec2 u_layout;
|
|
|
|
const float MAX_PACKED_DEPTH = 32.0;
|
|
|
|
vec4 readWeightUnscaled(vec2 p) {
|
|
vec4 w = texture2D(u_weightTex, p);
|
|
return w-u_weightCoefs.y;
|
|
}
|
|
|
|
void main() {
|
|
vec2 xy = getOutputXY();
|
|
float ch = getOutputChannel();
|
|
if (ch >= u_output.depth4)
|
|
return;
|
|
|
|
float dy = 1.0/(u_input.depth+1.0)/u_layout.y;
|
|
vec2 p = vec2((ch+0.5)/u_output.depth4, dy*0.5);
|
|
vec2 fuzz = (hash23(vec3(xy, u_seed+ch))-0.5)*u_fuzz;
|
|
|
|
vec2 realXY = xy;
|
|
#ifdef SPARSE_UPDATE
|
|
realXY = texture2D(u_shuffleTex, xy/u_output.size).xy*255.0+0.5 + u_shuffleOfs;
|
|
#endif
|
|
float modelIdx = u_control_read(realXY+fuzz, 0.0).x+0.5;
|
|
p.x += floor(mod(modelIdx, u_layout.x));
|
|
p.y += floor(modelIdx/u_layout.x);
|
|
p /= u_layout;
|
|
vec4 result = vec4(0.0);
|
|
for (float i=0.0; i < MAX_PACKED_DEPTH; i+=1.0) {
|
|
vec4 inVec = u_input_read(xy, i);
|
|
result += inVec.x * readWeightUnscaled(p); p.y += dy;
|
|
result += inVec.y * readWeightUnscaled(p); p.y += dy;
|
|
result += inVec.z * readWeightUnscaled(p); p.y += dy;
|
|
result += inVec.w * readWeightUnscaled(p); p.y += dy;
|
|
if (i+1.5>u_input.depth4) {
|
|
break;
|
|
}
|
|
}
|
|
result += readWeightUnscaled(p); // bias
|
|
setOutput(result*u_weightCoefs.x);
|
|
}`,
|
|
update: `
|
|
${defInput('u_update')}
|
|
uniform float u_seed, u_updateProbability;
|
|
|
|
varying vec2 uv;
|
|
|
|
void main() {
|
|
vec2 xy = getOutputXY();
|
|
vec4 state = u_input_readUV(uv);
|
|
vec4 update = vec4(0.0);
|
|
#ifdef SPARSE_UPDATE
|
|
vec4 shuffleInfo = texture2D(u_unshuffleTex, fract((xy-u_shuffleOfs)/u_output.size));
|
|
if (shuffleInfo.z > 0.5) {
|
|
update = u_update_read(shuffleInfo.xy*255.0+0.5, getOutputChannel());
|
|
}
|
|
#else
|
|
if (hash13(vec3(xy, u_seed)) <= u_updateProbability) {
|
|
update = u_update_readUV(uv);
|
|
}
|
|
#endif
|
|
setOutput(state + update);
|
|
}`,
|
|
vis: `
|
|
uniform float u_raw;
|
|
uniform float u_zoom;
|
|
uniform float u_perceptionCircle, u_arrows;
|
|
uniform float u_devicePixelRatio;
|
|
|
|
varying vec2 uv;
|
|
|
|
float clip01(float x) {
|
|
return min(max(x, 0.0), 1.0);
|
|
}
|
|
|
|
float peak(float x, float r) {
|
|
float y = x/r;
|
|
return exp(-y*y);
|
|
}
|
|
|
|
float getElement(vec4 v, float i) {
|
|
if (i<1.0) return v.x;
|
|
if (i<2.0) return v.y;
|
|
if (i<3.0) return v.z;
|
|
return v.w;
|
|
}
|
|
|
|
vec3 onehot3(float i) {
|
|
if (i<1.0) return vec3(1.0, 0.0, 0.0);
|
|
if (i<2.0) return vec3(0.0, 1.0, 0.0);
|
|
return vec3(0.0, 0.0, 1.0);
|
|
}
|
|
|
|
float sdTriangleIsosceles( in vec2 p, in vec2 q ) {
|
|
p.x = abs(p.x);
|
|
vec2 a = p - q*clamp( dot(p,q)/dot(q,q), 0.0, 1.0 );
|
|
vec2 b = p - q*vec2( clamp( p.x/q.x, 0.0, 1.0 ), 1.0 );
|
|
float s = -sign( q.y );
|
|
vec2 d = min( vec2( dot(a,a), s*(p.x*q.y-p.y*q.x) ),
|
|
vec2( dot(b,b), s*(p.y-q.y) ));
|
|
return -sqrt(d.x)*sign(d.y);
|
|
}
|
|
|
|
float aastep(float v) {
|
|
return clip01(v/fwidth(v)/u_devicePixelRatio);
|
|
}
|
|
|
|
float smoothstep(float t) {
|
|
t = clip01(t);
|
|
return t * t * (3.0 - 2.0 * t);
|
|
}
|
|
|
|
void spot(vec2 pos, float v, vec2 xy, inout vec3 rgb) {
|
|
v = sqrt(abs(v))*sign(v);
|
|
pos *= v*0.6;
|
|
float r = abs(v)*0.30;
|
|
rgb += clip01((r-length(xy-pos))/r)*0.2;
|
|
}
|
|
|
|
float sdBox( in vec2 p, in vec2 b )
|
|
{
|
|
vec2 d = abs(p)-b;
|
|
return length(max(d,0.0)) + min(max(d.x,d.y),0.0);
|
|
}
|
|
|
|
void main() {
|
|
vec2 xy = vec2(uv.x, 1.0-uv.y);
|
|
if (u_raw > 0.5) {
|
|
gl_FragColor = texture2D(u_input_tex, xy);
|
|
gl_FragColor.a = 1.0;
|
|
} else {
|
|
vec2 screenPos = xy*u_viewSize;
|
|
Hexel h = screen2hex(screenPos);
|
|
vec2 p = h.p;
|
|
h.cellXY += 0.5;
|
|
|
|
vec3 rgb = u_input_read(h.cellXY, 0.0).rgb/2.0+0.5;
|
|
if (4.0<h.zoom) {
|
|
vec2 dir = getCellDirection(floor(h.cellXY)+0.5);
|
|
float s = dir.x, c = dir.y;
|
|
float fade = clip01((h.zoom-4.0)/4.0);
|
|
float r = clip01((1.0-hex(p))*8.0);
|
|
r = pow(r, 0.2);
|
|
rgb *= mix(1.0, r, fade);
|
|
|
|
p = mat2(c, s, -s, c) * p;
|
|
|
|
if (12.0 < h.zoom) {
|
|
float da = PI/12.0;
|
|
float a = -da;
|
|
vec4 v4;
|
|
vec3 spots;
|
|
for (float ch=0.0; ch<2.5; ++ch) {
|
|
v4 = (u_input_read01(h.cellXY, ch)-127.0/255.0)*2.0;
|
|
spot(ang2vec(a+=da), v4.x, p, spots);
|
|
spot(ang2vec(a+=da), v4.y, p, spots);
|
|
spot(ang2vec(a+=da), v4.z, p, spots);
|
|
spot(ang2vec(a+=da), v4.w, p, spots);
|
|
}
|
|
spots *= clip01((h.zoom-12.0)/3.0);
|
|
rgb += spots;
|
|
}
|
|
}
|
|
gl_FragColor = vec4(rgb, 1.0);
|
|
}
|
|
}`
|
|
}
|
|
|
|
function createPrograms(gl, defines) {
|
|
defines = defines || '';
|
|
const res = {};
|
|
for (const name in PROGRAMS) {
|
|
const fs_code = defines + PREFIX + PROGRAMS[name];
|
|
const progInfo = twgl.createProgramInfo(gl, [vs_code, fs_code]);
|
|
progInfo.name = name;
|
|
res[name] = progInfo;
|
|
}
|
|
return res;
|
|
}
|
|
|
|
function createTensor(gl, w, h, depth, packScaleZero) {
|
|
const depth4 = Math.ceil(depth / 4);
|
|
const gridW = Math.ceil(Math.sqrt(depth4));
|
|
const gridH = Math.floor((depth4 + gridW - 1) / gridW);
|
|
const texW = w * gridW, texH = h * gridH;
|
|
|
|
const attachments = [{ minMag: gl.NEAREST }];
|
|
const fbi = twgl.createFramebufferInfo(gl, attachments, texW, texH);
|
|
const tex = fbi.attachments[0];
|
|
return {
|
|
_type: 'tensor',
|
|
fbi, w, h, depth, gridW, gridH, depth4, tex, packScaleZero
|
|
};
|
|
}
|
|
|
|
function setTensorUniforms(uniforms, name, tensor) {
|
|
uniforms[name + '.size'] = [tensor.w, tensor.h];
|
|
uniforms[name + '.gridSize'] = [tensor.gridW, tensor.gridH];
|
|
uniforms[name + '.depth'] = tensor.depth;
|
|
uniforms[name + '.depth4'] = tensor.depth4;
|
|
uniforms[name + '.packScaleZero'] = tensor.packScaleZero;
|
|
if (name != 'u_output') {
|
|
uniforms[name + '_tex'] = tensor.tex;
|
|
}
|
|
}
|
|
|
|
function decodeBase64(b64) {
|
|
const bin = atob(b64);
|
|
const bytes = new Uint8Array(bin.length);
|
|
for (let i = 0; i < bin.length; i++) {
|
|
bytes[i] = bin.charCodeAt(i);
|
|
}
|
|
return bytes;
|
|
}
|
|
|
|
function createDenseInfo(gl, params, onready) {
|
|
const coefs = [params.scale, 127.0 / 255.0];
|
|
const [in_n, out_n] = params.shape;
|
|
const info = { coefs, layout: params.layout, in_n: in_n - 1, out_n,
|
|
quantScaleZero: params.quant_scale_zero, ready: false };
|
|
// workaround against iOS WebKit bug (https://bugs.webkit.org/show_bug.cgi?id=138477)
|
|
// non-premultiplied PNG were decoded incorrectly
|
|
const img = UPNG.decode(decodeBase64(params.data.split(',')[1]));
|
|
const data = new Uint8Array(UPNG.toRGBA8(img)[0]);
|
|
info.tex = twgl.createTexture(gl, {
|
|
width: img.width, height: img.height,
|
|
minMag: gl.NEAREST, src: data,//, flipY: false, premultiplyAlpha: false,
|
|
}, ()=>{
|
|
//info.ready = true;
|
|
//onready();
|
|
});
|
|
setTimeout(()=>{
|
|
info.ready = true;
|
|
onready();
|
|
}, 0);
|
|
return info;
|
|
}
|
|
|
|
class CA {
|
|
constructor(gl, models, gridSize, onready) {
|
|
this.onready = onready || (()=>{});
|
|
this.gl = gl;
|
|
this.gridSize = gridSize || [96, 96];
|
|
gl.getExtension('OES_standard_derivatives');
|
|
|
|
this.updateProbability = 0.5;
|
|
this.shuffledMode = true;
|
|
|
|
this.rotationAngle = 0.0;
|
|
this.alignment = 1;
|
|
this.fuzz = 8.0;
|
|
this.perceptionCircle = 0.0;
|
|
this.arrowsCoef = 0.0;
|
|
this.visMode = 'color';
|
|
this.hexGrid = 1.0;
|
|
this.devicePixelRatio = globalThis.devicePixelRatio || 1;
|
|
|
|
this.layers = [];
|
|
this.setWeights(models);
|
|
|
|
this.progs = createPrograms(gl, this.shuffledMode ? '#define SPARSE_UPDATE\n' : '');
|
|
this.quad = twgl.createBufferInfoFromArrays(gl, {
|
|
position: [-1, -1, 0, 1, -1, 0, -1, 1, 0, -1, 1, 0, 1, -1, 0, 1, 1, 0],
|
|
});
|
|
|
|
this.setupBuffers();
|
|
const visNames = Object.getOwnPropertyNames(this.buf);
|
|
visNames.push('color');
|
|
|
|
this.clearCircle(0, 0, -1);
|
|
this.disturb();
|
|
}
|
|
|
|
disturb() {
|
|
this.runLayer(this.progs.align, this.buf.align, {
|
|
u_input: this.buf.newAlign, u_hexGrid: this.hexGrid, u_init: Math.random()*1000+1,
|
|
u_r: -1,
|
|
});
|
|
}
|
|
|
|
disturbCircle(x, y, r, viewSize) {
|
|
viewSize = viewSize || [128, 128];
|
|
this.runLayer(this.progs.align, this.buf.align, {
|
|
u_input: this.buf.newAlign, u_hexGrid: this.hexGrid, u_init: Math.random()*1000+1,
|
|
u_pos: [x, y], u_r: r, u_viewSize: viewSize,
|
|
});
|
|
}
|
|
|
|
|
|
setupBuffers() {
|
|
const gl = this.gl;
|
|
const [gridW, gridH] = this.gridSize;
|
|
const shuffleH = Math.ceil(gridH * this.updateProbability);
|
|
const shuffleCellN = shuffleH * gridW;
|
|
const totalCellN = gridW * gridH;
|
|
const shuffleBuf = new Uint8Array(shuffleCellN * 4);
|
|
const unshuffleBuf = new Uint8Array(totalCellN * 4);
|
|
let k = 0;
|
|
for (let i = 0; i < totalCellN; ++i) {
|
|
if (Math.random() < (shuffleCellN - k) / (totalCellN - i)) {
|
|
shuffleBuf[k * 4 + 0] = i % gridW;
|
|
shuffleBuf[k * 4 + 1] = Math.floor(i / gridW);
|
|
unshuffleBuf[i * 4 + 0] = k % gridW;
|
|
unshuffleBuf[i * 4 + 1] = Math.floor(k / gridW);
|
|
unshuffleBuf[i * 4 + 2] = 255;
|
|
k += 1;
|
|
}
|
|
}
|
|
this.shuffleTex = twgl.createTexture(gl, { minMag: gl.NEAREST, width: gridW, height: shuffleH, src: shuffleBuf});
|
|
this.unshuffleTex = twgl.createTexture(gl, { minMag: gl.NEAREST, width: gridW, height: gridH, src: unshuffleBuf});
|
|
this.shuffleOfs = [0, 0];
|
|
|
|
const updateH = this.shuffledMode ? shuffleH : gridH;
|
|
const perception_n = this.layers[0].in_n;
|
|
const lastLayer = this.layers[this.layers.length-1];
|
|
const channel_n = lastLayer.out_n;
|
|
this.channel_n = channel_n;
|
|
const stateQuantization = lastLayer.quantScaleZero;
|
|
const sonicN = 16;
|
|
this.buf = {
|
|
control: createTensor(gl, gridW, gridH, 4, [255.0, 0.0]),
|
|
align: createTensor(gl, gridW, gridH, 4, [2.0, 127.0 / 255.0]),
|
|
newAlign: createTensor(gl, gridW, gridH, 4, [2.0, 127.0 / 255.0]),
|
|
state: createTensor(gl, gridW, gridH, channel_n, stateQuantization),
|
|
newState: createTensor(gl, gridW, gridH, channel_n, stateQuantization),
|
|
perception: createTensor(gl, gridW, updateH, perception_n, stateQuantization),
|
|
sonic: createTensor(gl, sonicN*channel_n/4, 1, 4, stateQuantization),
|
|
};
|
|
{
|
|
const {width, height} = this.buf.sonic.fbi;
|
|
this.sonicBuf = new Uint8Array(height*width*4);
|
|
}
|
|
|
|
for (let i=0; i<this.layers.length; ++i) {
|
|
const layer = this.layers[i];
|
|
this.buf[`layer${i}`] = createTensor(gl, gridW, updateH, layer.out_n, layer.quantScaleZero);
|
|
}
|
|
}
|
|
|
|
step(stage) {
|
|
stage = stage || 'all';
|
|
if (!this.layers.every(l=>l.ready))
|
|
return;
|
|
|
|
if (stage == 'all') {
|
|
const [gridW, gridH] = this.gridSize;
|
|
this.shuffleOfs = [Math.floor(Math.random() * gridW), Math.floor(Math.random() * gridH)];
|
|
}
|
|
|
|
if (stage == 'all' || stage == 'align') {
|
|
this.runLayer(this.progs.align, this.buf.newAlign, {
|
|
u_input: this.buf.align, u_hexGrid: this.hexGrid, u_init: 0.0
|
|
});
|
|
}
|
|
if (stage == 'all' || stage == 'perception') {
|
|
this.runLayer(this.progs.perception, this.buf.perception, {
|
|
u_input: this.buf.state, u_angle: this.rotationAngle / 180.0 * Math.PI,
|
|
u_alignTex: this.buf.newAlign,
|
|
u_alignment: this.alignment, u_hexGrid: this.hexGrid
|
|
});
|
|
}
|
|
let inputBuf = this.buf.perception;
|
|
for (let i=0; i<this.layers.length; ++i) {
|
|
if (stage == 'all' || stage == `layer${i}`)
|
|
this.runDense(this.buf[`layer${i}`], inputBuf, this.layers[i]);
|
|
inputBuf = this.buf[`layer${i}`];
|
|
}
|
|
if (stage == 'all' || stage == 'newState') {
|
|
this.runLayer(this.progs.update, this.buf.newState, {
|
|
u_input: this.buf.state, u_update: inputBuf,
|
|
u_unshuffleTex: this.unshuffleTex,
|
|
u_seed: Math.random() * 1000, u_updateProbability: this.updateProbability
|
|
});
|
|
}
|
|
|
|
if (stage == 'all') {
|
|
[this.buf.state, this.buf.newState] = [this.buf.newState, this.buf.state];
|
|
[this.buf.align, this.buf.newAlign] = [this.buf.newAlign, this.buf.align];
|
|
}
|
|
}
|
|
|
|
benchmark() {
|
|
const gl = this.gl;
|
|
const flushBuf = new Uint8Array(4);
|
|
const flush = buf=>{
|
|
buf = buf || this.buf.state;
|
|
// gl.flush/finish don't seem to do anything, so reading a single
|
|
// pixel from the state buffer to flush the GPU command pipeline
|
|
twgl.bindFramebufferInfo(gl, buf.fbi);
|
|
gl.readPixels(0, 0, 1, 1, gl.RGBA, gl.UNSIGNED_BYTE, flushBuf);
|
|
}
|
|
|
|
flush();
|
|
const stepN = 100;
|
|
const start = Date.now();
|
|
for (let i = 0; i < stepN; ++i)
|
|
this.step();
|
|
flush();
|
|
const total = (Date.now() - start) / stepN;
|
|
|
|
const ops = ['align', 'perception'];
|
|
for (let i=0; i<this.layers.length; ++i)
|
|
ops.push(`layer${i}`);
|
|
ops.push('newState');
|
|
let perOpTotal = 0.0;
|
|
const perOp = [];
|
|
for (const op of ops) {
|
|
const start = Date.now();
|
|
for (let i = 0; i < stepN; ++i) {
|
|
this.step(op);
|
|
}
|
|
flush(this.buf[op]);
|
|
const dt = (Date.now() - start) / stepN;
|
|
perOpTotal += dt
|
|
perOp.push([op, dt]);
|
|
}
|
|
const perOpStr = perOp.map((p) => {
|
|
const [programName, dt] = p;
|
|
const percent = 100.0 * dt / perOpTotal;
|
|
return `${programName}: ${percent.toFixed(1)}%`;
|
|
}).join(', ');
|
|
return `${(total).toFixed(2)} ms/step, ${(1000.0 / total).toFixed(2)} step/sec\n` + perOpStr + '\n\n';
|
|
}
|
|
|
|
paint(x, y, r, brush, viewSize) {
|
|
viewSize = viewSize || [128, 128];
|
|
this.runLayer(this.progs.paint, this.buf.control, {
|
|
u_pos: [x, y], u_r: r, u_brush: [brush, 0, 0, 0], u_viewSize: viewSize,
|
|
});
|
|
}
|
|
|
|
peek(x, y, viewSize) {
|
|
this.runLayer(this.progs.peek, this.buf.sonic, {
|
|
u_pos: [x, y], u_viewSize: viewSize, u_input: this.buf.state
|
|
});
|
|
const {width, height} = this.buf.sonic.fbi;
|
|
const gl = this.gl;
|
|
twgl.bindFramebufferInfo(gl, this.buf.sonic.fbi);
|
|
gl.readPixels(0, 0, width, height, gl.RGBA, gl.UNSIGNED_BYTE, this.sonicBuf);
|
|
return {buf: this.sonicBuf, tex: this.buf.sonic.tex, pos: [x, y]};
|
|
}
|
|
|
|
clearCircle(x, y, r, viewSize) {
|
|
viewSize = viewSize || [128, 128];
|
|
this.runLayer(this.progs.paint, this.buf.state, {
|
|
u_pos: [x, y], u_r: r, u_brush: [0, 0, 0, 0], u_viewSize: viewSize,
|
|
});
|
|
}
|
|
|
|
setWeights(models) {
|
|
const gl = this.gl;
|
|
this.layers.forEach(layer=>gl.deleteTexture(layer));
|
|
const onready = ()=>{
|
|
if (this.layers.every(l=>l.ready))
|
|
this.onready();
|
|
}
|
|
this.layers = models.layers.map(layer=>createDenseInfo(gl, layer, onready));
|
|
}
|
|
|
|
runLayer(program, output, inputs) {
|
|
const gl = this.gl;
|
|
inputs = inputs || {};
|
|
const uniforms = {};
|
|
for (const name in inputs) {
|
|
const val = inputs[name];
|
|
if (val._type == 'tensor') {
|
|
setTensorUniforms(uniforms, name, val);
|
|
} else {
|
|
uniforms[name] = val;
|
|
}
|
|
}
|
|
uniforms['u_shuffleTex'] = this.shuffleTex;
|
|
uniforms['u_shuffleOfs'] = this.shuffleOfs;
|
|
setTensorUniforms(uniforms, 'u_output', output);
|
|
|
|
twgl.bindFramebufferInfo(gl, output.fbi);
|
|
gl.useProgram(program.program);
|
|
twgl.setBuffersAndAttributes(gl, program, this.quad);
|
|
twgl.setUniforms(program, uniforms);
|
|
twgl.drawBufferInfo(gl, this.quad);
|
|
return { programName: program.name, output }
|
|
}
|
|
|
|
runDense(output, input, layer) {
|
|
return this.runLayer(this.progs.dense, output, {
|
|
u_input: input, u_control: this.buf.control,
|
|
u_weightTex: layer.tex, u_weightCoefs: layer.coefs, u_layout: layer.layout,
|
|
u_seed: Math.random() * 1000, u_fuzz: this.fuzz
|
|
});
|
|
}
|
|
|
|
draw(viewSize, visMode) {
|
|
visMode = visMode || this.visMode;
|
|
const gl = this.gl;
|
|
|
|
gl.useProgram(this.progs.vis.program);
|
|
twgl.setBuffersAndAttributes(gl, this.progs.vis, this.quad);
|
|
const uniforms = { u_raw: 0.0,
|
|
u_angle: this.rotationAngle / 180.0 * Math.PI,
|
|
u_alignment: this.alignment,
|
|
u_perceptionCircle: this.perceptionCircle,
|
|
u_arrows: this.arrowsCoef,
|
|
u_devicePixelRatio: this.devicePixelRatio,
|
|
u_viewSize: viewSize,
|
|
};
|
|
let inputBuf = this.buf.state;
|
|
if (visMode != 'color') {
|
|
inputBuf = this.buf[visMode];
|
|
uniforms.u_raw = 1.0;
|
|
}
|
|
setTensorUniforms(uniforms, 'u_input', inputBuf);
|
|
setTensorUniforms(uniforms, 'u_alignTex', this.buf.align);
|
|
twgl.setUniforms(this.progs.vis, uniforms);
|
|
twgl.drawBufferInfo(gl, this.quad);
|
|
}
|
|
}
|