2a:null 2b:["$","title",null,{"children":"WebGL Tiny Sinkhorn | Zihou Ng | Observable"}] 30:T4c1,pointcloud = webgl.compileModel(gl, { attribs: { index: { data: new Float32Array(linspace(0,1,model.X[0].tex.size[0])), size: 1 } }, num: model.X[0].tex.size[0], mode: gl.POINTS, uniforms: {}, shader: webgl.defShader(gl, { vs: ` void main() { vec3 dat = texture2D(texPos, vec2(index, 0.5)).xyz; vec2 pos = dat.xy; pos.x /= aspect; pos.y += 0.1; gl_Position = vec4(pos, 0.0, 1.0); gl_PointSize = dat.z * 2.0; } `, fs: ` void main() { vec2 pc = gl_PointCoord * 2.0 - 1.0; vec2 fw = fwidth(pc); float delta = min(fw.x, fw.y); float r2 = dot(pc, pc); float t = 1.0 - smoothstep(1.0 - delta, 1.0 + delta, r2); if (t < delta) discard; t = 1.0 - smoothstep(1.0 - delta * 3.0, 1.0 + delta * 3.0, r2); gl_FragColor = vec4(vec3(t), 1.0); }`, declPrefixes: { v: "v_" }, attribs: { index: "float" }, varying: { }, uniforms: { aspect: ["float", gl.drawingBufferWidth / gl.drawingBufferHeight], texPos: ["sampler2D", 0], }, state: { depth: true, cull: true, blend: false, blendFn: webgl.BLEND_NORMAL }, ext: { OES_standard_derivatives: true } }), textures: [] })31:T51d,class Model { constructor(X, Y, n) { this.front = 0; this.X = [ new Buffer(n, 1, X), new Buffer(n, 1) ]; this.Y = new Buffer(n, 1, Y); const V = new Float32Array(n * 2); this.V = [ new Buffer(n, 1, dilate(V, 2, 2)), new Buffer(n, 1) ]; this.C = new Buffer(this.Y.tex.size[0], this.X[0].tex.size[0]); this.P = new Buffer(this.Y.tex.size[0], this.X[0].tex.size[0]); this.costMatrix = new CostMatrix(); this.sinkhorn = new Sinkhorn(this.X[0].tex.size[0], this.Y.tex.size[0]); this.couplingMatrix = new CouplingMatrix(); this.computeVelocity = new ComputeVelocity(this.X[0].tex.size[0]); this.accumulatePos = new AccumulatePos(); } forward(niter = 1) { const front = this.front, back = front ^ 1; this.costMatrix.forward(this.C, this.X[front], this.Y); for (let i = 0; i < niter; i++) this.sinkhorn.forward(this.C, this.sinkhorn.f); this.couplingMatrix.forward(this.P, this.sinkhorn.f, this.sinkhorn.g, this.C); this.computeVelocity.forward(this.V[back], this.P, this.Y, this.X[front], this.V[front]); this.accumulatePos.forward(this.X[back], this.X[front], this.V[back]); this.front = back; } get positions() { return this.X[this.front]; } }32:[{"id":"f47fc3121ede1155f04aaa4313b2c9908dd3f0cbd96eaefa0d40f96f81d4d1d1c2cd24e8dda2cfd4262b41cc8d63688c74419ae8afe46b448edcb6e13d68ec59","href":"https://static.observableusercontent.com/files/f47fc3121ede1155f04aaa4313b2c9908dd3f0cbd96eaefa0d40f96f81d4d1d1c2cd24e8dda2cfd4262b41cc8d63688c74419ae8afe46b448edcb6e13d68ec59","name":"char_circle_packings@2.json","type":"application/json","lastModified":1623634202042.473,"size":111885}] 2e:{"id":"a0861fc9302d1f50","slug":"webgl-tiny-sinkhorn","title":"WebGL Tiny Sinkhorn","thumbnail":"b108c404c3535e6e3d4ea77dec65e3832a023584de55d8991709541b6c7302b5","access_level":"public","create_time":"$D2021-06-13T14:27:43.565Z","update_time":"$D2021-06-19T07:36:26.357Z","view_time":null,"first_public_time":"$D2021-06-14T03:15:13.297Z","fork":true,"stars":6,"starred":false,"owner":{"type":"team","id":"7fad80bf128915e9","login":"stwind","name":"Zihou Ng","avatar_url":"https://avatars-next.observableusercontent.com/ff8a8f73ad220dcaec180c1e813951d692777837355004f062692b49042d2458","tier":"starter_2024"},"version":1439,"publish_version":1439,"role":null,"authored":false,"authors":[{"type":"individual","id":"dbb85afd665c9e5e","login":"stwind","name":"Zihou Ng","avatar_url":"https://avatars-next.observableusercontent.com/ff8a8f73ad220dcaec180c1e813951d692777837355004f062692b49042d2458","tier":null}],"members":[{"type":"individual","id":"dbb85afd665c9e5e","login":"stwind","name":"Zihou Ng","avatar_url":"https://avatars-next.observableusercontent.com/ff8a8f73ad220dcaec180c1e813951d692777837355004f062692b49042d2458","tier":null,"role":"editor"}],"comments":[],"custom_thumbnail":false,"trashed":false,"subscribed":null,"forks":2,"collection":null,"type":"notebook","body":{"title":"WebGL Tiny Sinkhorn","cells":[{"id":0,"value":"md`# WebGL Tiny Sinkhorn`","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":1357,"value":"md`Simple implementation of the [Sinkhorn algorithm](https://nbviewer.jupyter.org/github/gpeyre/numerical-tours/blob/master/python/optimaltransp_6_entropic_adv.ipynb) for [Optimal Transport](https://optimaltransport.github.io/) in PyTorch and WebGL, based on [Alex Mordvintsev](https://twitter.com/zzznah)'s [mini Sinkhorn](https://twitter.com/zzznah/status/1151153125778235393) notebook.\n\n*PyTorch and circle packings generation code are available in this [colab notebook](https://colab.research.google.com/github/stwind/notebooks/blob/master/tiny_sinkhorn.ipynb).*\n\n*_{Not works for mobile safari._{*`","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":1094,"value":"viewof target = Inputs.radio(new Map(d3.map(\"Haru\", (x, i) => [x, i])), {value: 0, label: \"Character\"})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":1328,"value":"viewof paused = Inputs.toggle({label: \"Pause\", value: false})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":17,"value":"viewof gl = {\n const canvas = DOM.canvas(width, width * 0.6);\n const gl = canvas.value = canvas.getContext('webgl');\n gl.getExtension('OES_texture_float');\n // gl.getExtension('WEBGL_draw_buffers');\n // gl.getExtension('WEBGL_color_buffer_float');\n \n return canvas;\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":1349,"value":"md`---\n## Code`","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":1128,"value":"data = {\n const data = await FileAttachment(\"char_circle_packings@2.json\").json();\n return data.map(dat => {\n const d = dat.map(c => [c[0] * 2 - 1, (1 - c[1]) * 2 - 1, c[2]]);\n return dilate(new Float32Array(d.flat(1)), 3, 1);\n });\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":42,"value":"render = {\n update;\n \n gl.clearColor(0.1, 0.1, 0.1, 1);\n gl.viewport(0, 0, gl.drawingBufferWidth, gl.drawingBufferHeight);\n gl.clear(gl.COLOR_BUFFER_BIT);\n \n pointcloud.textures[0] = model.positions.tex;\n webgl.draw(pointcloud);\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":1332,"value":"tick = {\n if (paused) return yield;\n while (true) yield;\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":954,"value":"update = {\n tick;\n model.forward();\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":1123,"value":"model.Y.update(data[target])","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":41,"value":"$30","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":800,"value":"model = {\n const X = dilate(new Float32Array(\n makeCircle(900, 0.5).map(p => [...p, 0.5 + (9 - 0.5) * Math.random()]).flat(1)\n ), 3, 1);\n return new Model(X, data[1], 900);\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":62,"value":"md`### WebGL`","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":752,"value":"$31","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":776,"value":"class AccumulatePos {\n constructor() {\n this.program = defProgram(`\nvoid main() {\n vec4 x = texture2D(texX, v_uv);\n vec4 v = texture2D(texV, v_uv);\n\n gl_FragColor = x + v * 0.1;\n}`, {\n texX: [\"sampler2D\", 0],\n texV: [\"sampler2D\", 1],\n });\n }\n \n forward(out, X, V) {\n this.program.textures = [X.tex, V.tex];\n out.draw(this.program);\n }\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":657,"value":"class ComputeVelocity {\n constructor(n) {\n this.program = defProgram(`\nfloat coord(float i, float n) {\n return (i + 0.5) / n;\n}\n\nvoid main() {\n vec4 g = vec4(0.0);\n for (int i = 0; i < ${n}; i++) {\n float ii = coord(float(i), ${n}.);\n float p = texture2D(texP, vec2(ii, v_uv.x)).x;\n vec4 y = texture2D(texY, vec2(ii, 0.5));\n g += p * y;\n }\n\n vec4 x = texture2D(texX, vec2(v_uv.x, 0.5));\n g -= x;\n\n vec4 v = texture2D(texV, vec2(v_uv.x, 0.5));\n\n v = v * 0.8 + g * 0.2;\n\n gl_FragColor = v;\n}`, {\n texP: [\"sampler2D\", 0],\n texY: [\"sampler2D\", 1],\n texX: [\"sampler2D\", 2],\n texV: [\"sampler2D\", 3]\n });\n }\n \n forward(out, P, Y, X, V) {\n this.program.textures = [P.tex, Y.tex, X.tex, V.tex];\n out.draw(this.program);\n }\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":337,"value":"class CouplingMatrix {\n constructor() {\n this.program = defProgram(`\nvoid main() {\n vec4 f = texture2D(texF, vec2(v_uv.y, 0.5));\n vec4 g = texture2D(texG, vec2(v_uv.x, 0.5));\n vec4 c = texture2D(texC, v_uv);\n\n gl_FragColor = exp(f + g - c);\n}`, {\n texF: [\"sampler2D\", 0],\n texG: [\"sampler2D\", 1],\n texC: [\"sampler2D\", 2]\n });\n }\n \n forward(out, f, g, C) {\n this.program.textures = [f.tex, g.tex, C.tex];\n out.draw(this.program);\n }\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":292,"value":"class Sinkhorn {\n constructor(n, m) {\n this.subC = new SubC();\n this.nlsec = new NegLogSumExpC(n);\n this.subR = new SubR();\n this.nlser = new NegLogSumExpR(m);\n \n this.f = new Buffer(n, 1);\n this.g = new Buffer(m, 1);\n \n Object.defineProperties(this, {\n _outS: { value: new Buffer(m, n) },\n });\n }\n \n forward(C, f) {\n this.subC.forward(this._outS, f, C);\n this.nlsec.forward(this.g, this._outS);\n this.subR.forward(this._outS, this.g, C);\n this.nlser.forward(this.f, this._outS);\n }\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":273,"value":"class NegLogSumExpR {\n constructor(n) {\n this.program = defProgram(`\nfloat coord(float i, float n) {\n return (i + 0.5) / n;\n}\n\nvoid main() {\n float amax = -1e10;\n for (int i = 0; i < ${n}; i++) {\n float v = texture2D(tex, vec2(coord(float(i), ${n}.), v_uv.x)).x;\n amax = max(amax, v);\n }\n\n float val = 0.0;\n for (int i = 0; i < ${n}; i++) {\n float v = texture2D(tex, vec2(coord(float(i), ${n}.), v_uv.x)).x;\n val += exp(v - amax);\n }\n\n val = amax + log(val);\n\n gl_FragColor = vec4(-val, 0.0, 0.0, 0.0);\n}`, {\n tex: [\"sampler2D\", 0]\n });\n }\n \n forward(out, x) {\n this.program.textures = [x.tex];\n out.draw(this.program);\n }\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":235,"value":"class NegLogSumExpC {\n constructor(n) {\n this.program = defProgram(`\nfloat coord(float i, float n) {\n return (i + 0.5) / n;\n}\n\nvoid main() {\n float amax = -1e10;\n for (int i = 0; i < ${n}; i++) {\n float v = texture2D(tex, vec2(v_uv.x, coord(float(i), ${n}.))).x;\n amax = max(amax, v);\n }\n\n float val = 0.0;\n for (int i = 0; i < ${n}; i++) {\n float v = texture2D(tex, vec2(v_uv.x, coord(float(i), ${n}.))).x;\n val += exp(v - amax);\n }\n\n val = amax + log(val);\n\n gl_FragColor = vec4(-val, 0.0, 0.0, 0.0);\n}`, {\n tex: [\"sampler2D\", 0]\n });\n }\n \n forward(out, x) {\n this.program.textures = [x.tex];\n out.draw(this.program);\n }\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":250,"value":"class SubR {\n constructor() {\n this.program = defProgram(`\nvoid main() {\n vec4 g = texture2D(texG, vec2(v_uv.x, 0.5));\n vec4 c = texture2D(texC, v_uv);\n\n gl_FragColor = g - c;\n}`, {\n texG: [\"sampler2D\", 0],\n texC: [\"sampler2D\", 1]\n });\n }\n \n forward(output, g, C) {\n this.program.textures = [g.tex, C.tex];\n output.draw(this.program);\n }\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":176,"value":"class SubC {\n constructor() {\n this.program = defProgram(`\nvoid main() {\n vec4 f = texture2D(texF, vec2(v_uv.y, 0.5));\n vec4 c = texture2D(texC, v_uv);\n\n gl_FragColor = f - c;\n}`, {\n texF: [\"sampler2D\", 0],\n texC: [\"sampler2D\", 1]\n });\n }\n \n forward(output, f, C) {\n this.program.textures = [f.tex, C.tex];\n output.draw(this.program);\n }\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":131,"value":"class CostMatrix {\n constructor() {\n this.program = defProgram(`\nvoid main() {\n vec4 x = texture2D(texX, vec2(v_uv.y, 0.5));\n vec4 y = texture2D(texY, vec2(v_uv.x, 0.5));\n\n vec4 v = pow(abs(x - y), vec4(2.0));\n float val = v.x + v.y + v.z + v.w;\n\n val /= eps;\n\n gl_FragColor = vec4(val, 0.0, 0.0, 0.0);\n}`, {\n texX: [\"sampler2D\", 0],\n texY: [\"sampler2D\", 1],\n eps: [\"float\", 1e-4]\n });\n }\n \n forward(out, X, Y, eps = 1e-4) {\n this.program.uniforms.eps = eps;\n this.program.textures = [X.tex, Y.tex];\n out.draw(this.program);\n }\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":120,"value":"class Buffer {\n constructor(width, height, data = null) {\n this.tex = webgl.defTextureFloat(gl, data, width, height);\n this.fbo = webgl.defFBO(gl, { tex: [this.tex] });\n Object.defineProperties(this, {\n _out: { value: new Float32Array(width * height * 4) }\n });\n }\n \n read(size) {\n const { type, size: [w, h] } = this.tex;\n this.fbo.bind();\n gl.readPixels(0, 0, w, h, gl.RGBA, type, this._out);\n this.fbo.unbind();\n return this._out;\n }\n \n update(data) {\n const { target, type, size: [w, h] } = this.tex;\n this.tex.bind();\n gl.texSubImage2D(target, 0, 0, 0, w, h, gl.RGBA, type, data, 0);\n this.tex.unbind();\n }\n \n draw(program) {\n const { size: [w, h] } = this.tex;\n this.fbo.bind();\n gl.viewport(0, 0, w, h);\n webgl.draw(program);\n this.fbo.unbind();\n }\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":65,"value":"defProgram = (fs, uniforms = {}, shaderOpts = null, spec = null) => webgl.compileModel(gl, {\n ...webgl.defQuadModel(),\n shader: webgl.defShader(gl, {\n fs, uniforms,\n vs: `\nvoid main() {\n v_uv = uv;\n gl_Position = vec4(position, 0.0, 1.0);\n}`,\n declPrefixes: { v: \"v_\" },\n attribs: {\n position: \"vec2\",\n uv: \"vec2\",\n },\n varying: { uv: \"vec2\" },\n state: { blend: false },\n ...shaderOpts\n }),\n ...spec\n})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":29,"value":"md`### Utils`","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":367,"value":"makeCircle = (n, r) => linspace(-math.TAU, math.TAU, n).map(a => [Math.sin(a) * r, Math.cos(a) * r])","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":394,"value":"makeSquare = (n, r) => linspace(-r, r, n).map(y => linspace(-r, r, n).map(x => [x, y])).flat(1)","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":80,"value":"dilate = (arr, stride = 1, padding = 3, val = 0) => {\n if (padding == 0) return arr;\n \n const n = arr.length;\n const p = Math.floor(n / stride);\n const m = (stride + padding) * p + (n % stride);\n const res = new Float32Array(m);\n let ii = 0;\n for (let i = 0; i < n; i++) {\n res[ii++] = arr[i];\n if ((i + 1) % stride == 0) {\n for (let j = 0; j < padding; j++) res[ii++] = val;\n }\n }\n return res;\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":83,"value":"undilate = (arr, stride = 1, padding = 3) => {\n if (padding == 0) return arr;\n \n const n = arr.length, step = (stride + padding);\n const m = Math.floor(n / step) * stride + (n % step);\n \n const res = new Float32Array(m);\n let ii = 0;\n for (let i = 0; i < n; i++) {\n if (i % step >= stride) i += padding - 1;\n else res[ii++] = arr[i];\n }\n return res;\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":365,"value":"linspace = (a, b, n) => d3.range(a, b + 1e-6, (b - a) / (n - 1))","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":683,"value":"allclose = (a, b, eps = 1e-6) => a.length == b.length && a.every((x, i) => math.eqDelta(x, b[i], eps))","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":674,"value":"arrView = (d, n) => [...d3.group(d, (_, i) => i/n|0).values()].map(x => x.map(y => y))","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":908,"value":"fps = {\n let t = performance.now(), delta = 0;\n const last = [t];\n return (cell, size = 10) => {\n last.length >= size && last.shift();\n t = performance.now();\n delta = (t - last[0]) / last.length;\n last.push(t);\n return Math.round(1000 / delta);\n }\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":18,"value":"md`### Appendix`","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":5,"value":"webgl = require(\"https://bundle.run/@thi.ng/webgl@4.0.16\")","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"},{"id":13,"value":"math = require(\"https://bundle.run/@thi.ng/math@4.0.0\")","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined"}],"files":"$W32","stdlib":"1","titling":"auto"}}
2f:"$2e:body"}}