1d:{"title":"GPUs, FFTs, and the 1-D Schrödinger Equation","cells":[{"id":0,"value":"$1f","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":"$20","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":"$21","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":"$22","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":"$23","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":"$24","pinned":true,"hidden":false,"mode":"ojs"},{"id":1437,"value":"$25","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":"$26","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":"$27","pinned":true,"hidden":false,"mode":"ojs"},{"id":1675,"value":"$28","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":"$29","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":"$2a","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":"$2b","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":"$W2c","stdlib":"1"}