31:T508,md`# GPUs, FFTs, and the 1-D Schrödinger Equation This notebook simulates the one-dimensional [Schrödinger Equation](https://en.wikipedia.org/wiki/Schr%C3%B6dinger_equation) in zero potential, ${tex`\displaystyle{i \hbar \frac{\partial \psi}{\partial t} = - \frac{\hbar^2}{2m}\frac{\partial^2 \psi}{\partial x^2}}.`} Without diving into Quantum Mechanics right here, the short version is that it's a complex-valued wave equation similar to the more familiar [second order wave equation](https://en.wikipedia.org/wiki/Wave_equation), except instead of waves all traveling at the same speed (think: sound), the wave components travel at a speed which depends linearly on frequency. Although there's a straightforward analytical Greens function solution and also a straightforward method for stepping the solution in the temporal frequency domain, this exercise uses fourth order explicit Runge-Kutta time integration and iterates the solution in the spatial frequency domain. The time integration is a bit overkill, but taking derivatives in the spatial frequency domain instead of via simple finite differences effectively eliminates spurious dispersion. Additionally, it initializes, computes, and visualizes the solution entirely within WebGL, without data ever touching the CPU.`32:T415,md`Since there's an analytical solution, we could step arbitrarily far into the future just by shifting phases, but since we're bothering to integrating in time, we'll need to be very careful about our time step. If you've ever solved a wave equation on a computer, you may have noticed that as the spatial resolution increases, you start having to use *very* small time steps to prevent the solution from blowing up. It leads to big headaches. Physically speaking, in a given amount of time, information is traveling farther and farther relative to the grid spacing. For problems solved on a discretized grid, the [Courant–Friedrichs–Lewy (CFL) condition](https://en.wikipedia.org/wiki/Courant%E2%80%93Friedrichs%E2%80%93Lewy_condition) dictates that information should not travel further than a single grid cell per timestep. If it did, information would have to *skip over* a grid cell in a single timestep. Not all methods require this condition, but many do. We'll select a CFL number slightly less than the maximum allowable (1.0).`33:T486,timeIntegrationPass = regl({ frag: ` precision highp float; varying vec2 uv; uniform sampler2D uSrc; uniform vec2 uResolution; uniform float uDt, uDx; const float TWOPI = 3.14159265358979 * 2.0; const float hbar = 1.0; const float m = 1.0; vec2 dydt (vec2 y, float k2) { return k2 * (hbar * 0.5 / m) * vec2(y.y, -y.x); } float wavenumber (float resolution, float dx) { float x = (gl_FragCoord.x - 0.5) * resolution; return ((x < 0.5) ? x : x - 1.0) * (TWOPI / dx); } void main () { float kx = wavenumber(uResolution.x, uDx); float kx2 = kx * kx; // Get the current state vec2 y0 = texture2D(uSrc, uv).xy; // Perform RK-4 time integration vec2 k1 = uDt * dydt(y0, kx2); vec2 k2 = uDt * dydt(y0 + k1 * 0.5, kx2); vec2 k3 = uDt * dydt(y0 + k2 * 0.5, kx2); vec2 k4 = uDt * dydt(y0 + k3, kx2); vec2 y1 = y0 + (k1 + k4 + 2.0 * (k2 + k3)) / 6.0; // Return the stepped solution gl_FragColor = vec4(y1, vec2(0)); } `, uniforms: { uSrc: regl.prop('input'), }, framebuffer: regl.prop('output'), })34:T7d1,performFFTPasses = regl({ frag: ` precision highp float; uniform sampler2D uSrc; uniform vec2 uResolution; uniform float uSubtransformSize, uNormalization; uniform bool uHorizontal, uForward; const float TWOPI = 6.283185307179586; vec4 fft ( sampler2D src, vec2 resolution, float subtransformSize, bool horizontal, bool forward, float normalization ) { vec2 evenPos, oddPos, twiddle, outputA, outputB; vec4 even, odd; float index, evenIndex, twiddleArgument; index = (horizontal ? gl_FragCoord.x : gl_FragCoord.y) - 0.5; evenIndex = floor(index / subtransformSize) * (subtransformSize * 0.5) + mod(index, subtransformSize * 0.5) + 0.5; if (horizontal) { evenPos = vec2(evenIndex, gl_FragCoord.y); oddPos = vec2(evenIndex, gl_FragCoord.y); } else { evenPos = vec2(gl_FragCoord.x, evenIndex); oddPos = vec2(gl_FragCoord.x, evenIndex); } evenPos *= resolution; oddPos *= resolution; if (horizontal) { oddPos.x += 0.5; } else { oddPos.y += 0.5; } even = texture2D(src, evenPos); odd = texture2D(src, oddPos); twiddleArgument = (forward ? TWOPI : -TWOPI) * (index / subtransformSize); twiddle = vec2(cos(twiddleArgument), sin(twiddleArgument)); return (even.rgba + vec4( twiddle.x * odd.xz - twiddle.y * odd.yw, twiddle.y * odd.xz + twiddle.x * odd.yw ).xzyw) * normalization; } void main () { gl_FragColor = fft(uSrc, uResolution, uSubtransformSize, uHorizontal, uForward, uNormalization); }`, uniforms: { uSrc: regl.prop('input'), uResolution: regl.prop('resolution'), uSubtransformSize: regl.prop('subtransformSize'), uHorizontal: regl.prop('horizontal'), uForward: regl.prop('forward'), uNormalization: regl.prop('normalization'), }, framebuffer: regl.prop('output'), })35:Taab,function planFFT (opts) { function isPowerOfTwo(n) { return n !== 0 && (n & (n - 1)) === 0 } function checkPOT (label, value) { if (!isPowerOfTwo(value)) { throw new Error(label + ' must be a power of two. got ' + label + ' = ' + value); } } var i, ping, pong, uniforms, tmp, width, height; opts = opts || {}; opts.forward = opts.forward === undefined ? true : opts.forward; opts.splitNormalization = opts.splitNormalization === undefined ? true : opts.splitNormalization; function swap () { tmp = ping; ping = pong; pong = tmp; } if (opts.size !== undefined) { width = height = opts.size; checkPOT('size', width); } else if (opts.width !== undefined && opts.height !== undefined) { width = opts.width; height = opts.height; checkPOT('width', width); checkPOT('height', width); } else { throw new Error('either size or both width and height must provided.'); } // Swap to avoid collisions with the input: ping = opts.ping; if (opts.input === opts.pong) { ping = opts.pong; } pong = ping === opts.ping ? opts.pong : opts.ping; var passes = []; var xIterations = Math.round(Math.log(width) / Math.log(2)); var yIterations = Math.round(Math.log(height) / Math.log(2)); var iterations = xIterations + yIterations; // Swap to avoid collisions with output: if (opts.output === ((iterations % 2 === 0) ? pong : ping)) { swap(); } // If we've avoiding collision with output creates an input collision, // then you'll just have to rework your framebuffers and try again. if (opts.input === pong) { throw new Error([ 'not enough framebuffers to compute without copying data. You may perform', 'the computation with only two framebuffers, but the output must equal', 'the input when an even number of iterations are required.' ].join(' ')); } for (i = 0; i < iterations; i++) { uniforms = { input: ping, output: pong, horizontal: i < xIterations, forward: !!opts.forward, resolution: [1.0 / width, 1.0 / height] }; if (i === 0) { uniforms.input = opts.input; } else if (i === iterations - 1) { uniforms.output = opts.output; } if (i === 0) { if (!!opts.splitNormalization) { uniforms.normalization = 1.0 / Math.sqrt(width * height); } else if (!opts.forward) { uniforms.normalization = 1.0 / width / height; } else { uniforms.normalization = 1; } } else { uniforms.normalization = 1; } uniforms.subtransformSize = Math.pow(2, (uniforms.horizontal ? i : (i - xIterations)) + 1); passes.push(uniforms); swap(); } return passes; }36:T7ed,function drawToScreen (opts) { regl.clear({color: [1, 1, 1, 1]}); configureViewport(() => { // Draw the axes drawLineSegments({ positions: axisLinesBuffers.positions, colors: axisLinesBuffers.colors, count: axisLines.positions.length / 2, lineWidth: 1 }); // Draw lines for the real and imaginary components of the solution drawLinesFromTexture([{ x: x, positionsData: y[0], lookupTable: textureLookupTableBuffer, count: N, color: [211 / 255, 45 / 255, 161 / 255, opts.envelope ? 0.5 : 0.5], mask: [1, 0, 0, 0], lineWidth: 1.5 }, { x: x, positionsData: y[0], lookupTable: textureLookupTableBuffer, count: N, color: [45 / 255, 161 / 255, 211 / 255, opts.envelope ? 0.5 : 0.5], mask: [0, 1, 0, 0], lineWidth: 1.5 }]); //if (!envelope) { // Draw points for the real and imaginary components of the solution drawPointsFromTexture([{ x: x, positionsData: y[0], lookupTable: textureLookupTableBuffer, count: N, color: [211 / 255, 45 / 255, 161 / 255, opts.envelope ? 0.3 : 1], mask: [1, 0, 0, 0], pointSize: 4, }, { x: x, positionsData: y[0], lookupTable: textureLookupTableBuffer, count: N, color: [45 / 255, 161 / 255, 211 / 255, opts.envelope ? 0.3 : 1], mask: [0, 1, 0, 0], pointSize: 4, }]); //} // Draw lines for the real and imaginary components of the solution if (opts.envelope) { drawEnvelopeFromTexture([{ x: x, positionsData: y[0], lookupTable: textureLookupTableBuffer, count: N, color: [0.3, 0.1, 0.3, 0.7], sign: 1, lineWidth: 2 }, { x: x, positionsData: y[0], lookupTable: textureLookupTableBuffer, count: N, color: [0.3, 0.1, 0.3, 0.7], sign: -1, lineWidth: 2 }]); } }); }37:T5b4,md`Drawing is a little complicated. *We can't just draw texture data as lines in WebGL.* Draw calls must be backed by attributes. Textures are not attributes, so we'll have to perform a bit of indirection. Two main steps are involved. First, we create an attribute buffer with coordinates representing the location of values in the texture. For each solution point, we start with this vertex attribute and obtain the solution value with a texture lookup. This first step is adequate for points with a single vertex and a nice \`GL_POINTS\` primitive, but [drawing lines is hard](https://mattdesl.svbtle.com/drawing-lines-is-hard) (and of course we simply can't deal with horribly inadequate \`GL_LINES\`). Lines require a second step. We render each line segment as two triangles and [adjust their offset in screen space](https://github.com/substack/screen-projected-lines). (I don't address miter joins here.) This second step *would* involve a lot of shuffling of data to create triangles between adjacent vertices, but the \`ANGLE_instanced_arrays\` extension gives us some tricks so that we don't actually need to shuffle any data to draw lines either between adjacent data points (or strided points to achieve disconnected line segments, for that matter). Instead we pass the same buffer to two different attributes with different offsets so that each vertex receives two adjacent points. This is hands down my favorite sort of GPU rendering trick.`38:T444,drawPointsFromTexture = regl({ vert: ` precision highp float; uniform mat4 uProjectionView; uniform vec4 uMask; uniform float uPointSize; attribute vec2 aPoint; attribute float aX; uniform sampler2D uSrc; void main () { float y = dot(texture2D(uSrc, aPoint.xy), uMask); vec2 xy = vec2(aX, y); gl_Position = uProjectionView * vec4(xy, 0, 1); gl_PointSize = uPointSize; }`, frag: ` precision highp float; uniform vec4 uColor; uniform float uPointSize; void main () { float r = length(gl_PointCoord.xy * 2.0 - 1.0) * uPointSize; float alpha = smoothstep(uPointSize, uPointSize - 2.0, r); gl_FragColor = vec4(uColor.rgb, uColor.a * alpha); }`, attributes: { aPoint: regl.prop('lookupTable'), aX: regl.prop('x'), }, uniforms: { uSrc: regl.prop('positionsData'), uMask: regl.prop('mask'), uColor: regl.prop('color'), uPointSize: (ctx, props) => ctx.pixelRatio * props.pointSize, }, depth: {enable: false}, primitive: 'points', count: regl.prop('count'), });39:T6de,drawLinesFromTexture = regl({ vert: ` precision highp float; uniform mat4 uProjectionView; uniform vec4 uMask; uniform float uLineWidth, uAspect; attribute vec2 aLookup, aNextLookup; attribute float aX, aNextX; attribute vec2 aLinePosition; uniform sampler2D uSrc; vec2 lineNormal (vec4 p, vec4 n, float aspect) { return normalize((p.yx / p.w - n.yx / n.w) * vec2(1, aspect)) * vec2(-1.0 / aspect, 1); } void main () { float y = dot(texture2D(uSrc, aLookup), uMask); float yNext = dot(texture2D(uSrc, aNextLookup), uMask); vec2 xy = vec2(aX, y); vec2 xyNext = vec2(aNextX, yNext); vec4 p = uProjectionView * vec4(xy, 0, 1); vec4 n = uProjectionView * vec4(xyNext, 0, 1); // Select one or the other positions gl_Position = mix(p, n, aLinePosition.y); // Apply the screen-space offset gl_Position.xy += lineNormal(p, n, uAspect) * aLinePosition.x * uLineWidth * gl_Position.w; }`, frag: ` precision mediump float; uniform vec4 uColor; void main () { gl_FragColor = uColor; }`, attributes: { aLookup: {buffer: regl.prop('lookupTable'), divisor: 1, stride: 8}, aNextLookup: {buffer: regl.prop('lookupTable'), divisor: 1, offset: 8, stride: 8}, aX: {buffer: regl.prop('x'), divisor: 1}, aNextX: {buffer: regl.prop('x'), offset: 4, divisor: 1}, aLinePosition: [-1, 0, 1, 0, -1, 1, 1, 1], }, uniforms: { uSrc: regl.prop('positionsData'), uLineWidth: (ctx, props) => props.lineWidth / ctx.framebufferHeight * ctx.pixelRatio, uMask: regl.prop('mask'), uColor: regl.prop('color'), }, primitive: 'triangle strip', instances: (ctx, props) => props.count - 1, count: 4, });3a:T6f0,drawEnvelopeFromTexture = regl({ vert: ` precision highp float; uniform mat4 uProjectionView; uniform float uSign; uniform float uLineWidth, uAspect; attribute vec2 aLookup, aNextLookup; attribute float aX, aNextX; attribute vec2 aLinePosition; uniform sampler2D uSrc; vec2 lineNormal (vec4 p, vec4 n, float aspect) { return normalize((p.yx / p.w - n.yx / n.w) * vec2(1, aspect)) * vec2(-1.0 / aspect, 1); } void main () { float y = length(texture2D(uSrc, aLookup).xy) * uSign; float yNext = length(texture2D(uSrc, aNextLookup).xy) * uSign; vec2 xy = vec2(aX, y); vec2 xyNext = vec2(aNextX, yNext); vec4 p = uProjectionView * vec4(xy, 0, 1); vec4 n = uProjectionView * vec4(xyNext, 0, 1); // Select one or the other positions gl_Position = mix(p, n, aLinePosition.y); // Apply the screen-space offset gl_Position.xy += lineNormal(p, n, uAspect) * aLinePosition.x * uLineWidth * gl_Position.w; }`, frag: ` precision mediump float; uniform vec4 uColor; void main () { gl_FragColor = uColor; }`, attributes: { aLookup: {buffer: regl.prop('lookupTable'), divisor: 1, stride: 8}, aNextLookup: {buffer: regl.prop('lookupTable'), divisor: 1, offset: 8, stride: 8}, aX: {buffer: regl.prop('x'), divisor: 1}, aNextX: {buffer: regl.prop('x'), offset: 4, divisor: 1}, aLinePosition: [-1, 0, 1, 0, -1, 1, 1, 1], }, uniforms: { uSrc: regl.prop('positionsData'), uLineWidth: (ctx, props) => props.lineWidth / ctx.framebufferHeight * ctx.pixelRatio, uSign: regl.prop('sign'), uColor: regl.prop('color'), }, primitive: 'triangle strip', instances: (ctx, props) => props.count - 1, count: 4, });3b:T537,axisLines = { let dark = [0.5, 0.5, 0.5, 1]; let light = [0.92, 0.92, 0.92, 1]; let positions = []; let colors = []; let labelPosition = []; let labelText = []; let labelAlign = []; for (let i = -10; i <= 10; i += (width > 600 ? 1 : 2)) { let x = i * 0.1; labelPosition.push([x, -1]); labelText.push(x.toFixed(1).replace(/([^\.])0$/,'$1')); labelAlign.push('center'); if (~[-10, 0, 10].indexOf(i)) continue; positions.push([x, -1], [x, 1]); colors.push(light, light); positions.push([x, -1], [x, -0.975]); colors.push(dark, dark); } for (let i = -4; i <= 4; i += (width > 600 ? 1 : 2)) { let y = i * 0.25; labelPosition.push([-1, y]); labelText.push(y.toFixed(2).replace(/([^\.])0$/,'$1')); labelAlign.push('end') if (~[-4, 0, 4].indexOf(i)) continue; positions.push([-1, y], [1, y]); colors.push(light, light); positions.push([-1, y], [-0.99, y]); colors.push(dark, dark); } positions.push([-1, 0], [1, 0], [0, -1], [0, 1], [-1, -1], [-1, 1], [-1, 1], [1, 1], [1, 1], [1, -1], [1, -1], [-1, -1]); colors.push(dark, dark, dark, dark, dark, dark, dark, dark, dark, dark, dark, dark); return {positions, colors, labelPosition, labelText, labelAlign}; }3c:T651,drawLineSegments = regl({ vert: ` precision highp float; uniform mat4 uProjectionView; uniform float uLineWidth, uAspect; attribute vec2 aPosition, aNextPosition; attribute vec4 aColor, aNextColor; attribute vec2 aLinePosition; varying vec4 vColor; vec2 lineNormal (vec4 p, vec4 n, float aspect) { return normalize((p.yx / p.w - n.yx / n.w) * vec2(1, aspect)) * vec2(-1.0 / aspect, 1); } void main () { vec4 p = uProjectionView * vec4(aPosition, 0, 1); vec4 n = uProjectionView * vec4(aNextPosition, 0, 1); // Select one or the other positions gl_Position = mix(p, n, aLinePosition.y); vColor = mix(aColor, aNextColor, aLinePosition.y); // Apply the screen-space offset gl_Position.xy += lineNormal(p, n, uAspect) * aLinePosition.x * uLineWidth * gl_Position.w; }`, frag: ` precision mediump float; varying vec4 vColor; void main () { gl_FragColor = vColor; }`, attributes: { aPosition: {buffer: regl.prop('positions'), divisor: 1, offset: 0, stride: 16}, aNextPosition: {buffer: regl.prop('positions'), divisor: 1, offset: 8, stride: 16}, aColor: {buffer: regl.prop('colors'), divisor: 1, offset: 0, stride: 32}, aNextColor: {buffer: regl.prop('colors'), divisor: 1, offset: 16, stride: 32}, aLinePosition: [-1, 0, 1, 0, -1, 1, 1, 1], }, uniforms: { uLineWidth: (ctx, props) => props.lineWidth / ctx.framebufferHeight * ctx.pixelRatio, uColor: regl.prop('color'), }, primitive: 'triangle strip', instances: (ctx, props) => props.count, count: 4, });3d:T5ac,function createREGLInstance () { const container = document.createElement('div'); container.style.position = "relative"; const w = width; const h = height; const pixelRatio = window.devicePixelRatio; container.style.width = w + 'px'; container.style.height = h + 'px'; const reglCanvas = document.createElement('canvas'); reglCanvas.style.position = "absolute"; reglCanvas.style.top = 0 reglCanvas.style.left = 0; reglCanvas.style.width = '100%'; reglCanvas.style.height = '100%'; reglCanvas.style.zIndex = 1; reglCanvas.width = Math.floor(w * pixelRatio); reglCanvas.height = Math.floor(h * pixelRatio); const labelCanvas = document.createElement('canvas'); labelCanvas.style.position = "absolute"; labelCanvas.style.top = 0 labelCanvas.style.left = 0; labelCanvas.style.width = '100%'; labelCanvas.style.height = '100%'; labelCanvas.style.zIndex = 2; labelCanvas.width = Math.floor(w * pixelRatio); labelCanvas.height = Math.floor(h * pixelRatio); container.appendChild(reglCanvas); container.appendChild(labelCanvas); const regl = createREGL({ canvas: reglCanvas, attributes: { antialias: true, }, extensions: [ 'OES_texture_float', 'ANGLE_instanced_arrays' ], optionalExtensions: [ 'OES_texture_half_float', ] }); regl.container = container; regl.reglCanvas = reglCanvas; regl.labelCanvas = labelCanvas; return regl; }3e:[] 2f:{"title":"GPUs, FFTs, and the 1-D Schrödinger Equation","cells":[{"id":0,"value":"$31","pinned":false,"hidden":false,"mode":"ojs"},{"id":1136,"value":"md`## The results\n\nTrue to stateless, reactive form, we'll start with the result. The plot below displays the real (red) and imaginary (blue) components of the solution with ${tex`\\hbar = 1`} and ${tex`m = 1`}. You may adjust the initial conditions using the controls below.\n\nAs noted above, waves with higher frequencies (momenta) travel at higher speeds. Since a waveform with finite spatial extent isn't made up of a single pure frequency component, finite wave packets contain a range of frequency components and spread out—\"disperse\"—as they move.\n\nIn fact, the more localized a wave is in space, the less localized it is in frequency (momentum) space. Try narrowing the envelope width and observe that the wave disperses more quickly.\n\n${preferHalfFloat ? html`BEWARE! Your device is using half floating point precision which may result in an unnaturally large amount of dissipation! For more accurate results, try a desktop computer with floating point texture support.`: ''}`","pinned":false,"hidden":false,"mode":"ojs"},{"id":718,"value":"viewof N = {\n const form = html``;\n form.i.oninput = () => form.o.value = `N = ${(form.value = Math.pow(2, form.i.valueAsNumber)).toFixed(0)}`;\n form.i.oninput();\n return form;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":624,"value":"viewof envelopeWidth = {\n const form = html``;\n form.i.oninput = () => form.o.value = `envelope width = ${(form.value = form.i.valueAsNumber).toFixed(3)}`;\n form.i.oninput();\n return form;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":639,"value":"viewof wavenumber = {\n // Allow up to *half* of the Nyquist frequency for the wavenumber. More than\n // that and the solution starts to look like nonsense.\n const form = html``;\n form.i.oninput = () => form.o.value = `wavenumber = ${(form.value = N * form.i.valueAsNumber * dx * Math.PI).toFixed(2)}`;\n form.i.oninput();\n return form;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":1684,"value":"viewof envelope = {\n const form = html``;\n form.i.oninput = () => {\n form.value = form.i.checked\n }\n form.i.oninput();\n return form;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":131,"value":"{\n N;\n regl.poll();\n configureComputePasses({dt, dx, N}, () => {\n initializeSolutionPass({\n output: y[0],\n wavenumber,\n envelopeWidth,\n });\n });\n \n while (true) {\n regl.poll();\n \n configureComputePasses({dt, dx, N}, () => {\n performFFTPasses(forwardFFTPasses({input: y[0], output: y[0]}));\n \n for (let i = 0; i < 10; i++) {\n timeIntegrationPass([\n {input: y[0], output: y[1]},\n {input: y[1], output: y[0]}\n ]);\n }\n\n performFFTPasses(inverseFFTPasses({input: y[0], output: y[0]}));\n })\n\n drawToScreen({envelope});\n \n yield regl.container\n }\n}","pinned":true,"hidden":false,"mode":"ojs"},{"id":1336,"value":"md`Above is our main iteration loop. Since there's a nice and linear solution, we start each iteration by computing the FFT and end each iteration by inverting the FFT. The astute reader will roll their eyes and note that if we didn't want to visualize the solution between steps, we could cut out the Fourier transforms entirely and step the solution with nothing more than a bit of multiplication in the temporal frequency domain.\n\nNonetheless, this approach gets a lot of nice machinery flowing which will be good to have locked down before we continue on to bigger things.`","pinned":false,"hidden":false,"mode":"ojs"},{"id":1343,"value":"md`## Solution setup\n\nWe start by creating a \\`regl\\` buffer containing x-values. Note in particular that we have no use for a JavaScript array of x-values! Once we've initialized a WebGL buffer with x-values, we never again touch them in JavaScript.`","pinned":false,"hidden":false,"mode":"ojs"},{"id":1272,"value":"x = regl.buffer(new Array(N).fill(0).map((d, i) => (i / N) * 2.0 - 1.0))","pinned":true,"hidden":false,"mode":"ojs"},{"id":1365,"value":"md`The y-values are a bit different, though we still don't need JavaScript storage. We create framebuffers with texture attachments so that we'll be able to write values to them via WebGL shaders. We create three. The first is our primary solution vector while the second and third are strictly for temporary storage while we iterate. \n\nSomewhere in the future we find ourselves able to use the GPU to perform computation on *data* in a browser, but as far as WebGL 1.0 is concerned, GPUs are for graphics. That means we'll be running our computation on a texture with four floating point color channels.\n\nOur usage of the four components of the texture is:\n- red: real component of first value\n- green: imaginary component of first value\n- blue: real component of second value *(unused)*\n- alpha: imaginary component of second value *(unused)*\n`","pinned":false,"hidden":false,"mode":"ojs"},{"id":116,"value":"y = new Array(3).fill(0).map(() => regl.framebuffer({\n color: regl.texture({\n width: N,\n height: 1,\n\n // (r,b,g,a) => (real0, imag0, real1, imag1)\n format: 'rgba',\n type: preferHalfFloat ? 'half float' : 'float',\n\n // We have discrete solution components; we need no interpolation.\n min: 'nearest',\n mag: 'nearest',\n })\n}))","pinned":true,"hidden":false,"mode":"ojs"},{"id":1375,"value":"$32","pinned":false,"hidden":false,"mode":"ojs"},{"id":892,"value":"CFL = 0.95","pinned":true,"hidden":false,"mode":"ojs"},{"id":864,"value":"dx = 1.0 / N","pinned":true,"hidden":false,"mode":"ojs"},{"id":754,"value":"dt = Math.min(1e-6, CFL * Math.pow(dx * 0.5, 2))","pinned":true,"hidden":false,"mode":"ojs"},{"id":1396,"value":"md`## Computation\n\nAt last, the shaders! We start with initialization. Apart from the particular syntax, there's not a whole lot to say about it since it just sets up real and imaginary sine waves with a 90 degree phase shift and an envelope.\n\nThe only important thing to reiterate is that an infinite sine wave has a single pure frequency, but as soon as we restrict its spatial domain, its frequencies start to spread out.`","pinned":false,"hidden":false,"mode":"ojs"},{"id":1166,"value":"initializeSolutionPass = regl({\n frag: `\n precision highp float;\n varying vec2 uv;\n uniform vec2 uResolution;\n uniform float uDx, uK, uEnvelopeWidth;\n \n void main () { \n float x = ((gl_FragCoord.x - 0.5) * uResolution.x - 0.5) / uResolution.x * uDx;\n\n gl_FragColor = vec4(\n vec2(\n cos(uK * x),\n sin(uK * x)\n ) * exp(-0.5 * pow(x / uEnvelopeWidth, 2.0)),\n vec2(0)\n );\n }\n `,\n uniforms: {\n uSrc: regl.prop('input'),\n uK: regl.prop('wavenumber'),\n uEnvelopeWidth: regl.prop('envelopeWidth'),\n },\n framebuffer: regl.prop('output'),\n})","pinned":true,"hidden":false,"mode":"ojs"},{"id":1404,"value":"md`Below is the time integration pass. It implements fourth order Runge-Kutta stepping and the derivative from the equation at the top.`","pinned":false,"hidden":false,"mode":"ojs"},{"id":1162,"value":"$33","pinned":true,"hidden":false,"mode":"ojs"},{"id":1409,"value":"md`This method performs some boilerplate shader setup. Recall that to fake \"compute\" shaders in WebGL 1.0, we must render [a big triangle](https://www.npmjs.com/package/a-big-triangle) which covers the whole screen and perform all of the work in the fragment shader.`","pinned":false,"hidden":false,"mode":"ojs"},{"id":298,"value":"configureComputePasses = regl({\n vert: `\n precision highp float;\n attribute vec2 xy;\n varying vec2 uv;\n void main () {\n uv = xy * 0.5 + 0.5;\n gl_Position = vec4(xy, 0, 1);\n }`,\n attributes: {\n xy: [-4, -4, 4, -4, 0, 4],\n },\n uniforms: {\n uDt: regl.prop('dt'),\n uDx: regl.prop('dx'),\n uResolution: (ctx, props) => [1 / props.N, 1],\n },\n depth: {enable: false},\n count: 3\n})","pinned":true,"hidden":false,"mode":"ojs"},{"id":1414,"value":"md`## Fast Fourier Transform\n\nThe final computation shader is the Fourier transform. The heavy lifting is lifted from [glsl-fft](https://github.com/rreusser/glsl-fft). It performs an FFT on two real-complex pairs stored in the four color channels.`","pinned":false,"hidden":false,"mode":"ojs"},{"id":1233,"value":"$34","pinned":true,"hidden":false,"mode":"ojs"},{"id":1418,"value":"md`\\`regl\\` has a really nice feature in which it renders an *array* of passes with a single regl call. We utilize this heavily to create a really slick interface for performing the FFT. The only caveat is that we defer deciding the input and output destinations until we actually call it.\n\nNote that for both the forward and inverse passes, we ping-pong back and forth between the second and third framebuffers but always start and end with the first.`","pinned":false,"hidden":false,"mode":"ojs"},{"id":1229,"value":"inverseFFTPasses = {\n let inversePasses = planFFT({\n width: N,\n height: 1,\n input: y[0],\n ping: y[1],\n pong: y[2],\n output: y[0],\n forward: false\n })\n return function (opts) {\n inversePasses[0].input = opts.input;\n inversePasses[inversePasses.length - 1].output = opts.output;\n return inversePasses;\n }\n}","pinned":true,"hidden":false,"mode":"ojs"},{"id":1223,"value":"forwardFFTPasses = {\n let forwardPasses = planFFT({\n width: N,\n height: 1,\n input: y[0],\n ping: y[1],\n pong: y[2],\n output: y[0],\n forward: true\n })\n return function (opts) {\n forwardPasses[0].input = opts.input;\n forwardPasses[forwardPasses.length - 1].output = opts.output;\n return forwardPasses;\n }\n}","pinned":true,"hidden":false,"mode":"ojs"},{"id":1421,"value":"md`The final function just does annoying bookkeeping to set up the proper metadata for the FFT shader passes.`","pinned":false,"hidden":false,"mode":"ojs"},{"id":1226,"value":"$35","pinned":true,"hidden":false,"mode":"ojs"},{"id":1424,"value":"md`## Rendering\n\nWow! That's a lot of setup, but we're not done. If we're going to create nice plots, we need to write some rendering routines.\n\nWe start with the final draw call definition. `","pinned":false,"hidden":false,"mode":"ojs"},{"id":312,"value":"$36","pinned":true,"hidden":false,"mode":"ojs"},{"id":1437,"value":"$37","pinned":false,"hidden":false,"mode":"ojs"},{"id":185,"value":"textureLookupTable = createTextureLookupTable(N, 1)","pinned":true,"hidden":false,"mode":"ojs"},{"id":187,"value":"textureLookupTableBuffer = regl.buffer(textureLookupTable)","pinned":true,"hidden":false,"mode":"ojs"},{"id":175,"value":"function createTextureLookupTable (width, height, stride) {\n stride = stride || 2;\n let n = width * height * stride;\n\n let out = new Float32Array(n);\n\n for (let i = 0, iStride = 0; iStride < n; i++, iStride += stride) {\n out[iStride] = ((i % width) + 0.5) / width;\n out[iStride + 1] = (((i / width) | 0) + 0.5) / height;\n }\n\n return out;\n}","pinned":true,"hidden":false,"mode":"ojs"},{"id":1451,"value":"md`The command below draws points from the texture containing solution data.`","pinned":false,"hidden":false,"mode":"ojs"},{"id":727,"value":"$38","pinned":true,"hidden":false,"mode":"ojs"},{"id":1453,"value":"md`The command below draws lines between adjacent solution points.`","pinned":false,"hidden":false,"mode":"ojs"},{"id":167,"value":"$39","pinned":true,"hidden":false,"mode":"ojs"},{"id":1675,"value":"$3a","pinned":true,"hidden":false,"mode":"ojs"},{"id":1432,"value":"md`We now create arrays containing the positions of the axis and grid lines.`","pinned":false,"hidden":false,"mode":"ojs"},{"id":1043,"value":"$3b","pinned":true,"hidden":false,"mode":"ojs"},{"id":1047,"value":"axisLinesBuffers = {\n return {\n positions: regl.buffer(axisLines.positions),\n colors: regl.buffer(axisLines.colors)\n };\n}","pinned":true,"hidden":false,"mode":"ojs"},{"id":1850,"value":"md`Well this part is just a regrettable mess. 🤫`","pinned":false,"hidden":false,"mode":"ojs"},{"id":1758,"value":"{\n let canvas = regl.labelCanvas;\n let ctx = regl.labelCanvas.getContext('2d');\n let pixelRatio = window.devicePixelRatio;\n let w = canvas.width;\n let h = canvas.height;\n ctx.clearRect(0, 0, w, h);\n ctx.fillColor = '#000';\n let font = `sans-serif`\n let fontSizePx = 10 * pixelRatio;\n ctx.font = `${fontSizePx}px ${font}`;\n ctx.textAlign = \"center\";\n let labelPosition = axisLines.labelPosition;\n let labelText = axisLines.labelText;\n for (let i = 0; i < labelPosition.length; i++) {\n let position = labelPosition[i];\n let text = labelText[i];\n let align = axisLines.labelAlign[i];\n ctx.textAlign = align;\n let xShift = align === 'center' ? 0 : -4 * pixelRatio;\n let yShift = align === 'center' ? 2 * pixelRatio + fontSizePx : fontSizePx * 0.4;\n let xy = vec3transformMat4([], [position[0], position[1], 0], projectionView);\n let ii = w * (0.5 + 0.5 * xy[0]);\n let jj = h * (0.5 - 0.5 * xy[1]);\n ctx.fillText(text, ii + xShift, jj + yShift);\n \n }\n}","pinned":true,"hidden":false,"mode":"ojs"},{"id":1780,"value":"function vec3transformMat4(out, a, m) {\n var x = a[0], y = a[1], z = a[2],\n w = m[3] * x + m[7] * y + m[11] * z + m[15]\n w = w || 1.0\n out[0] = (m[0] * x + m[4] * y + m[8] * z + m[12]) / w\n out[1] = (m[1] * x + m[5] * y + m[9] * z + m[13]) / w\n out[2] = (m[2] * x + m[6] * y + m[10] * z + m[14]) / w\n return out\n}","pinned":true,"hidden":false,"mode":"ojs"},{"id":1469,"value":"md`Just one more line drawing command. This one is a little simpler. We just draw point attributes directly to the screen, without any texture lookups.`","pinned":false,"hidden":false,"mode":"ojs"},{"id":924,"value":"$3c","pinned":true,"hidden":false,"mode":"ojs"},{"id":1472,"value":"md`Finally we set up some common draw properties.`","pinned":false,"hidden":false,"mode":"ojs"},{"id":1782,"value":"projectionView = [\n (width - 35) / width, 0, 0, 0,\n 0, (height - 35) / height, 0, 0,\n 0, 0, 1, 0,\n 35 * 0.5 / width, 35 / height / 2, 0, 1\n]","pinned":true,"hidden":false,"mode":"ojs"},{"id":222,"value":"configureViewport = regl({\n uniforms: {\n uProjectionView: projectionView,\n uAspect: ctx => ctx.framebufferWidth / ctx.framebufferHeight,\n },\n blend: {\n enable: true,\n func: {\n srcRGB: 'src alpha',\n srcAlpha: 1,\n dstRGB: 'one minus src alpha',\n dstAlpha: 1,\n },\n equation: {\n rgb: 'add',\n alpha: 'add',\n }\n },\n depth: {enable: false},\n})","pinned":true,"hidden":false,"mode":"ojs"},{"id":1474,"value":"md`And we require and initialize [\\`regl\\`](https://github.com/regl-project/regl).`","pinned":false,"hidden":false,"mode":"ojs"},{"id":1139,"value":"createREGL = require(\"regl\")","pinned":true,"hidden":false,"mode":"ojs"},{"id":1142,"value":"regl = createREGLInstance()","pinned":true,"hidden":false,"mode":"ojs"},{"id":1880,"value":"height = Math.max(250, Math.floor(width * 0.5));","pinned":true,"hidden":false,"mode":"ojs"},{"id":681,"value":"$3d","pinned":true,"hidden":false,"mode":"ojs"},{"id":1504,"value":"function isTouchDevice() {\n return (\n !!(typeof window !== 'undefined' &&\n ('ontouchstart' in window ||\n (window.DocumentTouch &&\n typeof document !== 'undefined' &&\n document instanceof window.DocumentTouch))) ||\n !!(typeof navigator !== 'undefined' &&\n (navigator.maxTouchPoints || navigator.msMaxTouchPoints))\n );\n}","pinned":true,"hidden":false,"mode":"ojs"},{"id":1506,"value":"preferHalfFloat = regl.hasExtension('OES_texture_half_float') && isTouchDevice()","pinned":true,"hidden":false,"mode":"ojs"}],"files":"$W3e","stdlib":"1","titling":"auto"}