alpha-orbitals: gpu wip

src/examples/alpha-orbitals/index.html

@@ -14,8 +14,8 @@
                 position: absolute;
                 left: 160px;
                 top: 100px;
-                width: 600px;
-                height: 600px;
+                width: 800px;
+                height: 800px;
                 border: 1px solid #ccc;
             }
         </style>

src/examples/alpha-orbitals/index.ts

@@ -4,8 +4,11 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { Basis, computeIsocontourValues } from '../../extensions/alpha-orbitals/cubes';
+import { Basis, computeIsocontourValues, CubeGridInfo } from '../../extensions/alpha-orbitals/cubes';
+import { AlphaOrbitalsPass } from '../../extensions/alpha-orbitals/gpu/pass';
 import { SphericalBasisOrder } from '../../extensions/alpha-orbitals/orbitals';
+import { Box3D } from '../../mol-math/geometry';
+import { Vec3 } from '../../mol-math/linear-algebra';
 import { createPluginAsync, DefaultPluginSpec } from '../../mol-plugin';
 import { createVolumeRepresentationParams } from '../../mol-plugin-state/helpers/volume-representation-params';
 import { StateTransforms } from '../../mol-plugin-state/transforms';
@@ -13,6 +16,7 @@ import { PluginContext } from '../../mol-plugin/context';
 import { ColorNames } from '../../mol-util/color/names';
 import { DemoMoleculeSDF, DemoOrbitals } from './example-data';
 import './index.html';
+import { TestWaterParams } from './test-water';
 import { CreateOrbitalVolume, StaticBasisAndOrbitals } from './transforms';
 require('mol-plugin-ui/skin/light.scss');
 
@@ -73,7 +77,6 @@ class AlphaOrbitalsExample {
 
         const repr = this.plugin.build().to(volumeRef);
 
-
         if (positive !== void 0) {
             repr.apply(StateTransforms.Representation.VolumeRepresentation3D, createVolumeRepresentationParams(this.plugin, volumeRef.data!, {
                 type: 'isosurface',
@@ -93,7 +96,29 @@ class AlphaOrbitalsExample {
         }
 
         await repr.commit();
+
+        this.gpu();
+    }
+
+    gpu() {
+        // const pass = new AlphaOrbitalsPass(this.plugin.canvas3d!.webgl, TestWaterParams);
+        // pass.getData();
+        // pass.getData();
     }
 }
 
+function createCubeGrid(): CubeGridInfo {
+    const box = Box3D.create(Vec3.create(-1, -1, -1), Vec3.create(1, 1, 1));
+    const dimensions = Vec3.create(50, 50, 50);
+    const size = Box3D.size(Vec3(), box);
+
+    return {
+        box,
+        dimensions,
+        npoints: dimensions[0] * dimensions[1] * dimensions[2],
+        size,
+        delta: Vec3.div(Vec3(), size, Vec3.subScalar(Vec3(), dimensions, 1))
+    };
+}
+
 (window as any).AlphaOrbitalsExample = new AlphaOrbitalsExample();

src/examples/alpha-orbitals/test-water.ts

@@ -0,0 +1,190 @@
+import { CollocationParams } from '../../extensions/alpha-orbitals/collocation';
+import { Basis } from '../../extensions/alpha-orbitals/cubes';
+import { Box3D } from '../../mol-math/geometry';
+import { Vec3 } from '../../mol-math/linear-algebra';
+
+const _testBasis: Basis = {
+    'atoms': [
+        {
+            'center': [
+                0.025886090588624934,
+                0.019164790004065606,
+                -0.013539970104105408
+            ] as Vec3,
+            'shells': [
+                {
+                    'angularMomentum': [0],
+                    'coefficients': [
+                        [
+                            -0.004151277818987536,
+                            -0.02067024147993795,
+                            -0.05150303336984537,
+                            0.33462711739899537,
+                            0.5621061300983125,
+                            0.17129946969948573
+                        ]
+                    ],
+                    'exponents': [
+                        152.28769660788095,
+                        27.928015215973073,
+                        7.848374792384515,
+                        1.1223350202705642,
+                        0.5093846587907856,
+                        0.24292266532510307
+                    ]
+                },
+                {
+                    'angularMomentum': [1],
+                    'coefficients': [
+                        [
+                            0.007924233646294425,
+                            0.051441048251911314,
+                            0.18984000600705359,
+                            0.4049863191150474,
+                            0.40123628611490797,
+                            0.1051855189039082
+                        ]
+                    ],
+                    'exponents': [
+                        27.203421487167727,
+                        7.09409912597673,
+                        2.5383362605345954,
+                        1.0610730767843852,
+                        0.4851948916410433,
+                        0.22938302550642545
+                    ]
+                }
+            ]
+        },
+        {
+            'center': [
+                0.5082729578468134,
+                1.6880351220025265,
+                0.4963443067810461
+            ] as Vec3,
+            'shells': [
+                {
+                    'angularMomentum': [0],
+                    'coefficients': [
+                        [
+                            0.009163596280542963,
+                            0.04936149294292479,
+                            0.16853830490998634,
+                            0.37056279972195677,
+                            0.4164915298246781,
+                            0.13033408410772263
+                        ]
+                    ],
+                    'exponents': [
+                        33.710073211949485,
+                        6.180705022740464,
+                        1.7291385346152253,
+                        0.5940057549921978,
+                        0.2306698170449518,
+                        0.09500256906284119
+                    ]
+                },
+                {
+                    'angularMomentum': [0],
+                    'coefficients': [
+                        [
+                            -0.32279868167000036,
+                            3.209629817295221,
+                            2.4672629224617935,
+                            -0.048487066612842224,
+                            -0.2611850111200143,
+                            -0.8917817597810863,
+                            -1.9607480081275706,
+                            -2.203769342520311,
+                            -0.6896328935259993
+                        ]
+                    ],
+                    'exponents': [
+                        10.256286070314905,
+                        0.6227965325875392,
+                        0.2391007667853915,
+                        33.710073211949485,
+                        6.180705022740464,
+                        1.7291385346152253,
+                        0.5940057549921978,
+                        0.2306698170449518,
+                        0.09500256906284119
+                    ]
+                }
+            ]
+        },
+        {
+            'center': [
+                1.1367367844436005,
+                -0.47018519422670163,
+                -1.356802622574504
+            ] as Vec3,
+            'shells': [
+                {
+                    'angularMomentum': [0],
+                    'coefficients': [
+                        [
+                            0.009163596280542963,
+                            0.04936149294292479,
+                            0.16853830490998634,
+                            0.37056279972195677,
+                            0.4164915298246781,
+                            0.13033408410772263
+                        ]
+                    ],
+                    'exponents': [
+                        33.710073211949485,
+                        6.180705022740464,
+                        1.7291385346152253,
+                        0.5940057549921978,
+                        0.2306698170449518,
+                        0.09500256906284119
+                    ]
+                },
+                {
+                    'angularMomentum': [0],
+                    'coefficients': [
+                        [
+                            -0.32279868167000036,
+                            3.209629817295221,
+                            2.4672629224617935,
+                            -0.048487066612842224,
+                            -0.2611850111200143,
+                            -0.8917817597810863,
+                            -1.9607480081275706,
+                            -2.203769342520311,
+                            -0.6896328935259993
+                        ]
+                    ],
+                    'exponents': [
+                        10.256286070314905,
+                        0.6227965325875392,
+                        0.2391007667853915,
+                        33.710073211949485,
+                        6.180705022740464,
+                        1.7291385346152253,
+                        0.5940057549921978,
+                        0.2306698170449518,
+                        0.09500256906284119
+                    ]
+                }
+            ]
+        }
+    ]
+};
+
+const grid = {
+    box: Box3D.create(Vec3.create(-1, -1, -1), Vec3.create(1, 1, 1)),
+    delta: Vec3.create(2, 2, 2),
+    dimensions: Vec3.create(2, 2, 2),
+    npoints: 8,
+    size: Vec3.create(2, 2, 2)
+};
+
+export const TestWaterParams: CollocationParams = {
+    grid,
+    alphaOrbitals: [-2.2623991420609075e-16, 0.6360205395000592, 0.6672122399886391, -0.3876927909355508, -1.6780131293332933e-16, 2.844782862661151e-16, 4.977960694176068e-19, -2.3945919908996803e-16],
+    basis: _testBasis,
+    cutoffThreshold: 0,
+    sphericalOrder: 'cca-reverse'
+};

src/examples/alpha-orbitals/transforms.ts

@@ -11,6 +11,7 @@ import { Task } from '../../mol-task';
 import { CustomProperties } from '../../mol-model/custom-property';
 import { SphericalBasisOrder } from '../../extensions/alpha-orbitals/orbitals';
 import { Volume } from '../../mol-model/volume';
+import { PluginContext } from '../../mol-plugin/context';
 
 export class BasisAndOrbitals extends PluginStateObject.Create<{ basis: Basis, order: SphericalBasisOrder, orbitals: { energy: number, alpha: number[] }[] }>({ name: 'Basis', typeClass: 'Object' }) { }
 
@@ -46,7 +47,7 @@ export const CreateOrbitalVolume = PluginStateTransform.BuiltIn({
         };
     }
 })({
-    apply({ a, params }) {
+    apply({ a, params }, plugin: PluginContext) {
         return Task.create('Orbital Volume', async ctx => {
             const data = await createSphericalCollocationGrid({
                 basis: a.data.basis,
@@ -60,7 +61,7 @@ export const CreateOrbitalVolume = PluginStateTransform.BuiltIn({
                     [25, 0.4],
                     [0, 0.33],
                 ],
-            }).runInContext(ctx);
+            }, plugin.canvas3d?.webgl).runInContext(ctx);
             const volume: Volume = {
                 grid: data.grid,
                 sourceData: { name: 'custom grid', kind: 'custom', data },

src/extensions/alpha-orbitals/collocation.ts

@@ -12,8 +12,8 @@ import { arrayMin } from '../../mol-util/array';
 import { Basis, CubeGridInfo } from './cubes';
 import {
     normalizeBasicOrder,
-    SphericalFunctions,
-    SphericalBasisOrder,
+
+    SphericalBasisOrder, SphericalFunctions
 } from './orbitals';
 
 export interface CollocationParams {
@@ -51,6 +51,8 @@ export async function sphericalCollocation(
 
     const matrix = new Float32Array(grid.npoints);
 
+    let ii = 0;
+
     let baseIndex = 0;
     for (const atom of basis.atoms) {
         for (const shell of atom.shells) {
@@ -135,7 +137,6 @@ function collocationBasis(
         for (let j = jMin; j <= jMax; j++) {
             const y = sy + gdy * j - cy;
             const oY = oX + j * nz;
-
             for (let k = kMin; k <= kMax; k++) {
                 const z = sz + gdz * k - cz;
                 const R2 = x * x + y * y + z * z;
@@ -150,8 +151,9 @@ function collocationBasis(
                         coefficients[c] * Math.exp(-exponents[c] * R2);
                 }
 
-                const sphericalSum =
-                    L === 0 ? alpha[0] : sphericalFunc(alpha, x, y, z);
+                const sphericalSum = L === 0 ? alpha[0] : sphericalFunc(alpha, x, y, z);
+
+
                 matrix[k + oY] += gaussianSum * sphericalSum;
             }
         }

src/extensions/alpha-orbitals/cubes.ts

@@ -7,12 +7,14 @@
  */
 
 import { sortArray } from '../../mol-data/util';
+import { WebGLContext } from '../../mol-gl/webgl/context';
 import { Box3D } from '../../mol-math/geometry';
 import { Mat4, Tensor, Vec3 } from '../../mol-math/linear-algebra';
 import { Grid } from '../../mol-model/volume';
 import { Task } from '../../mol-task';
 import { arrayMax, arrayMin, arrayRms } from '../../mol-util/array';
-import { sphericalCollocation } from './collocation';
+import { CollocationParams, sphericalCollocation } from './collocation';
+import { AlphaOrbitalsPass } from './gpu/pass';
 import { SphericalBasisOrder } from './orbitals';
 
 export interface CubeGridInfo {
@@ -60,7 +62,7 @@ export interface SphericalCollocationParams {
 }
 
 export function createSphericalCollocationGrid(
-    params: SphericalCollocationParams
+    params: SphericalCollocationParams, webgl?: WebGLContext
 ): Task<CubeGrid> {
     return Task.create('Spherical Collocation Grid', async (ctx) => {
         const grid = initBox(
@@ -69,24 +71,52 @@ export function createSphericalCollocationGrid(
             params.boxExpand
         );
 
-        const matrix = await sphericalCollocation(
-            {
-                grid,
-                basis: params.basis,
-                alphaOrbitals: params.alphaOrbitals,
-                cutoffThreshold: params.cutoffThreshold,
-                sphericalOrder: 'cca-reverse', // renamed entos ordering
-            },
-            ctx
-        );
+        let matrix: Float32Array;
+
+        const cParams: CollocationParams = {
+            grid,
+            basis: params.basis,
+            alphaOrbitals: params.alphaOrbitals,
+            cutoffThreshold: params.cutoffThreshold,
+            sphericalOrder: params.sphericalOrder
+        };
+
+
+        console.time('gpu');
+        const pass = new AlphaOrbitalsPass(webgl!, cParams);
+        const matrixGL = pass.getData();
+        console.timeEnd('gpu');
+
+        console.time('gpu');
+        const pass0 = new AlphaOrbitalsPass(webgl!, cParams);
+        pass0.getData();
+        console.timeEnd('gpu');
 
-        return createCubeGrid(grid, matrix);
+        if (false && webgl) {
+        } else {
+            console.time('cpu');
+            matrix = await sphericalCollocation(cParams, ctx);
+            console.timeEnd('cpu');
+        }
+
+        // console.log(matrixGL);
+        // console.log(matrix);
+
+        // for (let i = 0; i < matrixGL.length; i++) {
+        //     if (Math.abs(matrixGL[i] - matrix[i]) > 1e-4) {
+        //         console.log('err', i, matrixGL[i], matrix[i]);
+        //         // console.log()
+        //         break;
+        //     }
+        // }
+
+        return createCubeGrid(grid, matrixGL, [0, 1, 2]);
     });
 }
 
 const BohrToAngstromFactor = 0.529177210859;
 
-function createCubeGrid(gridInfo: CubeGridInfo, values: Float32Array) {
+function createCubeGrid(gridInfo: CubeGridInfo, values: Float32Array, axisOrder: number[]) {
     const boxSize = Box3D.size(Vec3(), gridInfo.box);
     const boxOrigin = Vec3.clone(gridInfo.box.min);
 
@@ -107,7 +137,7 @@ function createCubeGrid(gridInfo: CubeGridInfo, values: Float32Array) {
     const grid: Grid = {
         transform: { kind: 'matrix', matrix },
         cells: Tensor.create(
-            Tensor.Space(gridInfo.dimensions, [0, 1, 2], Float32Array),
+            Tensor.Space(gridInfo.dimensions, axisOrder, Float32Array),
             (values as any) as Tensor.Data
         ),
         stats: {
@@ -118,8 +148,12 @@ function createCubeGrid(gridInfo: CubeGridInfo, values: Float32Array) {
         },
     };
 
+    // TODO: when using GPU rendering, the cumulative sum can be computed
+    // along the ray on the fly
     const isovalues = computeIsocontourValues(values, 0.85);
 
+    console.log(isovalues);
+
     return { grid, isovalues };
 }
 
@@ -146,7 +180,9 @@ function initBox(
         if (spacingThresholds[i][0] <= count) break;
     }
 
-    const dimensions = Vec3.ceil(Vec3(), Vec3.scale(Vec3(), size, 1 / s));
+    // dimensions need to be powers of 2 otherwise it leads to roudning error
+    // TODO: possible to avoid?
+    const dimensions = Vec3.create(64, 64, 64); //   Vec3.ceil(Vec3(), Vec3.scale(Vec3(), size, 1 / s));
     return {
         box,
         dimensions,

src/extensions/alpha-orbitals/gpu/pass.ts

@@ -0,0 +1,192 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { QuadSchema, QuadValues } from '../../../mol-gl/compute/util';
+import { ComputeRenderable, createComputeRenderable } from '../../../mol-gl/renderable';
+import { TextureSpec, UniformSpec, Values } from '../../../mol-gl/renderable/schema';
+import { ShaderCode } from '../../../mol-gl/shader-code';
+import quad_vert from '../../../mol-gl/shader/quad.vert';
+import { WebGLContext } from '../../../mol-gl/webgl/context';
+import { createComputeRenderItem } from '../../../mol-gl/webgl/render-item';
+import { RenderTarget } from '../../../mol-gl/webgl/render-target';
+import { ValueCell } from '../../../mol-util';
+import { CollocationParams } from '../collocation';
+import { normalizeBasicOrder } from '../orbitals';
+import shader_frag from './shader.frag';
+
+const AlphaOrbitalsSchema = {
+    ...QuadSchema,
+    uDimensions: UniformSpec('v3'),
+    uMin: UniformSpec('v3'),
+    uDelta: UniformSpec('v3'),
+    tCenters: TextureSpec('image-float32', 'rgb', 'float', 'nearest'),
+    tInfo: TextureSpec('image-float32', 'rgba', 'float', 'nearest'),
+    tCoeff: TextureSpec('image-float32', 'rgb', 'float', 'nearest'),
+    tAlpha: TextureSpec('image-float32', 'alpha', 'float', 'nearest'),
+    uNCenters: UniformSpec('i'),
+    uNAlpha: UniformSpec('i'),
+    uNCoeff: UniformSpec('i'),
+    uLittleEndian: UniformSpec('i') // TODO: boolean uniforms
+};
+const AlphaOrbitalsShaderCode = ShaderCode('postprocessing', quad_vert, shader_frag);
+type AlphaOrbitalsRenderable = ComputeRenderable<Values<typeof AlphaOrbitalsSchema>>
+
+function createTextureData({
+    basis,
+    sphericalOrder,
+    alphaOrbitals,
+}: CollocationParams) {
+    let centerCount = 0;
+    let baseCount = 0;
+    let coeffCount = 0;
+    for (const atom of basis.atoms) {
+        for (const shell of atom.shells) {
+            for (const L of shell.angularMomentum) {
+                if (L > 4) {
+                    // TODO: will L > 4 be required? Would need to precompute more functions in that case.
+                    throw new Error('Angular momentum L > 4 not supported.');
+                }
+
+                centerCount++;
+                baseCount += 2 * L + 1;
+                coeffCount += shell.exponents.length;
+            }
+        }
+    }
+
+    const centers = new Float32Array(3 * centerCount);
+    // L, alpha_offset, coeff_offset_start, coeff_offset_end
+    const info = new Float32Array(4 * centerCount);
+    const alpha = new Float32Array(baseCount);
+    const coeff = new Float32Array(3 * coeffCount);
+
+    let cO = 0, aO = 0, coeffO = 0;
+    for (const atom of basis.atoms) {
+        for (const shell of atom.shells) {
+
+            let amIndex = 0;
+            for (const L of shell.angularMomentum) {
+                const a0 = normalizeBasicOrder(L, alphaOrbitals.slice(aO, aO + 2 * L + 1), sphericalOrder);
+
+                if (cO === 1) {
+                    console.log('y', atom.center[1]);
+                }
+
+                centers[3 * cO + 0] = atom.center[0];
+                centers[3 * cO + 1] = atom.center[1];
+                centers[3 * cO + 2] = atom.center[2];
+
+                info[4 * cO + 0] = L;
+                info[4 * cO + 1] = aO;
+                info[4 * cO + 2] = coeffO;
+                info[4 * cO + 3] = coeffO + shell.exponents.length;
+
+                for (let i = 0; i < a0.length; i++) alpha[aO + i] = a0[i];
+
+                const c0 = shell.coefficients[amIndex++];
+                for (let i = 0; i < shell.exponents.length; i++) {
+                    coeff[3 * (coeffO + i) + 0] = c0[i];
+                    coeff[3 * (coeffO + i) + 1] = shell.exponents[i];
+                }
+
+                cO++;
+                aO += 2 * L + 1;
+                coeffO += shell.exponents.length;
+            }
+        }
+    }
+
+    return { nCenters: centerCount, nAlpha: baseCount, nCoeff: coeffCount, centers, info, alpha, coeff };
+}
+
+function getPostprocessingRenderable(ctx: WebGLContext, params: CollocationParams): AlphaOrbitalsRenderable {
+    const data = createTextureData(params);
+
+    console.log(data);
+
+    const values: Values<typeof AlphaOrbitalsSchema> = {
+        ...QuadValues,
+        uDimensions: ValueCell.create(params.grid.dimensions),
+        uMin: ValueCell.create(params.grid.box.min),
+        uDelta: ValueCell.create(params.grid.delta),
+        uNCenters: ValueCell.create(data.nCenters),
+        uNAlpha: ValueCell.create(data.nAlpha),
+        uNCoeff: ValueCell.create(data.nCoeff),
+        tCenters: ValueCell.create({ width: data.nCenters, height: 1, array: data.centers }),
+        tInfo: ValueCell.create({ width: data.nCenters, height: 1, array: data.info }),
+        tCoeff: ValueCell.create({ width: data.nCoeff, height: 1, array: data.coeff }),
+        tAlpha: ValueCell.create({ width: data.nAlpha, height: 1, array: data.alpha }),
+        uLittleEndian: ValueCell.create(isLittleEndian()),
+    };
+
+    const schema = { ...AlphaOrbitalsSchema };
+    const renderItem = createComputeRenderItem(ctx, 'triangles', AlphaOrbitalsShaderCode, schema, values);
+
+    return createComputeRenderable(renderItem, values);
+}
+
+export class AlphaOrbitalsPass {
+    target: RenderTarget
+    renderable: AlphaOrbitalsRenderable
+
+    constructor(private webgl: WebGLContext, private params: CollocationParams) {
+        const [nx, ny, nz] = params.grid.dimensions;
+
+        // TODO: add single component float32 render target option for WebGL2?
+        // TODO: figure out the ordering so that it does not have to be remapped in the shader
+        this.target = webgl.createRenderTarget(nx, ny * nz, false, 'uint8', 'nearest');
+        this.renderable = getPostprocessingRenderable(webgl, params);
+    }
+
+    render() {
+        const [nx, ny, nz] = this.params.grid.dimensions;
+        const width = nx;
+        const height = ny * nz;
+
+        // const { x, y, width, height } = this.camera.viewport;
+
+        const { gl, state } = this.webgl;
+        this.target.bind();
+        // this.setQuadShift(0, 0);
+        gl.viewport(0, 0, width, height);
+        gl.scissor(0, 0, width, height);
+        state.disable(gl.SCISSOR_TEST);
+        state.disable(gl.BLEND);
+        state.disable(gl.DEPTH_TEST);
+        state.depthMask(false);
+        this.renderable.render();
+    }
+
+    getData() {
+        const [nx, ny, nz] = this.params.grid.dimensions;
+        const width = nx;
+        const height = ny * nz;
+
+        this.render();
+        this.target.bind();
+        const array = new Uint8Array(width * height * 4);
+        this.webgl.readPixels(0, 0, width, height, array);
+        // PixelData.flipY({ array, width, height });
+        const floats = new Float32Array(array.buffer, array.byteOffset, width * height);
+
+        // console.log(array);
+        // console.log(floats);
+
+        return floats;
+
+        // return new ImageData(new Uint8ClampedArray(array), width, height);
+    }
+}
+
+function isLittleEndian() {
+    const arrayBuffer = new ArrayBuffer(2);
+    const uint8Array = new Uint8Array(arrayBuffer);
+    const uint16array = new Uint16Array(arrayBuffer);
+    uint8Array[0] = 0xAA; // set first byte
+    uint8Array[1] = 0xBB; // set second byte
+    if(uint16array[0] === 0xBBAA) return 1;
+    return 0;
+}

src/extensions/alpha-orbitals/gpu/shader.frag.ts

@@ -0,0 +1,172 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+export default `
+precision highp float;
+precision highp int;
+precision highp sampler2D;
+
+uniform vec2 uQuadShift;
+uniform vec3 uDimensions;
+uniform vec3 uMin;
+uniform vec3 uDelta;
+
+uniform sampler2D tCenters;
+uniform sampler2D tInfo;
+uniform sampler2D tCoeff;
+uniform sampler2D tAlpha;
+
+uniform int uNCenters;
+uniform int uNCoeff;
+uniform int uNAlpha;
+
+uniform int uLittleEndian;
+
+float shiftRight (float v, float amt) {
+  v = floor(v) + 0.5;
+  return floor(v / exp2(amt));
+}
+float shiftLeft (float v, float amt) {
+    return floor(v * exp2(amt) + 0.5);
+}
+float maskLast (float v, float bits) {
+    return mod(v, shiftLeft(1.0, bits));
+}
+float extractBits (float num, float from, float to) {
+    from = floor(from + 0.5); to = floor(to + 0.5);
+    return maskLast(shiftRight(num, from), to - from);
+}
+// Adapted from https://github.com/equinor/glsl-float-to-rgba
+// MIT License, Copyright (c) 2020 Equinor
+vec4 floatToRgba(float texelFloat) {
+    if (texelFloat == 0.0) return vec4(0, 0, 0, 0);
+    float sign = texelFloat > 0.0 ? 0.0 : 1.0;
+    texelFloat = abs(texelFloat);
+    float exponent = floor(log2(texelFloat));
+    float biased_exponent = exponent + 127.0;
+    float fraction = ((texelFloat / exp2(exponent)) - 1.0) * 8388608.0;
+    float t = biased_exponent / 2.0;
+    float last_bit_of_biased_exponent = fract(t) * 2.0;
+    float remaining_bits_of_biased_exponent = floor(t);
+    float byte4 = extractBits(fraction, 0.0, 8.0) / 255.0;
+    float byte3 = extractBits(fraction, 8.0, 16.0) / 255.0;
+    float byte2 = (last_bit_of_biased_exponent * 128.0 + extractBits(fraction, 16.0, 23.0)) / 255.0;
+    float byte1 = (sign * 128.0 + remaining_bits_of_biased_exponent) / 255.0;
+    return (
+      uLittleEndian > 0
+      ? vec4(byte4, byte3, byte2, byte1)
+      : vec4(byte1, byte2, byte3, byte4)
+    );
+}
+
+float L1(vec3 p, float a0, float a1, float a2) {
+    return a0 * p.z + a1 * p.x + a2 * p.y;
+}
+
+float L2(vec3 p, float a0, float a1, float a2, float a3, float a4) {
+    float x = p.x, y = p.y, z = p.z;
+    float xx = x * x, yy = y * y, zz = z * z;
+    return (
+        a0 * (-0.5 * xx - 0.5 * yy + zz) +
+        a1 * (1.7320508075688772 * x * z) +
+        a2 * (1.7320508075688772 * y * z) +
+        a3 * (0.8660254037844386 * xx - 0.8660254037844386 * yy) +
+        a4 * (1.7320508075688772 * x * y)
+    );
+}
+
+float L3(vec3 p, float a0, float a1, float a2, float a3, float a4, float a5, float a6) {
+    float x = p.x, y = p.y, z = p.z;
+    float xx = x * x, yy = y * y, zz = z * z;
+    float xxx = xx * x, yyy = yy * y, zzz = zz * z;
+    return (
+        a0 * (-1.5 * xx * z - 1.5 * yy * z + zzz) +
+        a1 * (-0.6123724356957945 * xxx - 0.6123724356957945 * x * yy + 2.449489742783178 * x * zz) +
+        a2 * (-0.6123724356957945 * xx * y - 0.6123724356957945 * yyy + 2.449489742783178 * y * zz) +
+        a3 * (1.9364916731037085 * xx * z - 1.9364916731037085 * yy * z) +
+        a4 * (3.872983346207417 * x * y * z) +
+        a5 * (0.7905694150420949 * xxx - 2.3717082451262845 * x * yy) +
+        a6 * (2.3717082451262845 * xx * y - 0.7905694150420949 * yyy)
+    );
+}
+
+void main(void) {
+    // TODO: use "integer division" to avoid rounding errors for non power of 2 dimensions
+    vec3 idx = vec3(gl_FragCoord.x - 0.5, mod(gl_FragCoord.y - 0.5, uDimensions.z), floor((gl_FragCoord.y - 0.5) / uDimensions.y));
+    float offset = idx.x + idx.y * uDimensions.x + idx.z * uDimensions.x * uDimensions.y;
+
+    // TODO: is there render target ordering that does not require the remapping of the coords?
+    float k = mod(offset, uDimensions.z), kk = floor(offset / uDimensions.z);
+    float j = mod(kk, uDimensions.y);
+    float i = floor(kk / uDimensions.y);
+    
+    vec3 xyz = uMin + uDelta * vec3(i, j, k);
+
+    float fCenter = 1.0 / float(uNCenters - 1);
+    float fCoeff = 1.0 / float(uNCoeff - 1);
+    float fAlpha = 1.0 / float(uNAlpha - 1);
+
+    float v = 0.0;
+
+    for (int i = 0; i < uNCenters; i++) {
+        vec2 cCoord = vec2(float(i) * fCenter, 0.5);
+
+        vec3 X = xyz - texture2D(tCenters, cCoord).xyz;
+        vec4 info = texture2D(tInfo, cCoord);
+
+        int L = int(info.x);
+        float alphaOffset = info.y;
+        int coeffStart = int(info.z);
+        int coeffEnd = int(info.w);
+
+        float R2 = dot(X, X);
+
+        float gauss = 0.0;
+        for (int j = coeffStart; j < coeffEnd; j++) {
+            vec2 c = texture2D(tCoeff, vec2(float(j) * fCoeff, 0.5)).xy;
+            gauss += c.x * exp(-c.y * R2);
+        }
+
+        float spherical = 0.0;
+        if (L == 0) {
+            spherical = texture2D(tAlpha, vec2(alphaOffset * fAlpha, 0.5)).x;
+        } else if (L == 1) {
+            spherical = L1(
+                X,
+                texture2D(tAlpha, vec2(alphaOffset * fAlpha, 0.5)).x,
+                texture2D(tAlpha, vec2((alphaOffset + 1.0) * fAlpha, 0.5)).x,
+                texture2D(tAlpha, vec2((alphaOffset + 2.0) * fAlpha, 0.5)).x
+            );
+        } else if (L == 2) {
+            spherical = L2(
+                X,
+                texture2D(tAlpha, vec2(alphaOffset * fAlpha, 0.5)).x,
+                texture2D(tAlpha, vec2((alphaOffset + 1.0) * fAlpha, 0.5)).x,
+                texture2D(tAlpha, vec2((alphaOffset + 2.0) * fAlpha, 0.5)).x,
+                texture2D(tAlpha, vec2((alphaOffset + 3.0) * fAlpha, 0.5)).x,
+                texture2D(tAlpha, vec2((alphaOffset + 4.0) * fAlpha, 0.5)).x
+            );
+        } else if (L == 3) {
+            spherical = L3(
+                X,
+                texture2D(tAlpha, vec2(alphaOffset * fAlpha, 0.5)).x,
+                texture2D(tAlpha, vec2((alphaOffset + 1.0) * fAlpha, 0.5)).x,
+                texture2D(tAlpha, vec2((alphaOffset + 2.0) * fAlpha, 0.5)).x,
+                texture2D(tAlpha, vec2((alphaOffset + 3.0) * fAlpha, 0.5)).x,
+                texture2D(tAlpha, vec2((alphaOffset + 4.0) * fAlpha, 0.5)).x,
+                texture2D(tAlpha, vec2((alphaOffset + 5.0) * fAlpha, 0.5)).x,
+                texture2D(tAlpha, vec2((alphaOffset + 6.0) * fAlpha, 0.5)).x
+            );
+        } 
+        // TODO: do we need L > 3?
+
+        v += gauss * spherical;
+    }
+
+    // TODO: render to single component float32 texture in WebGL2
+    gl_FragColor = floatToRgba(v);
+}
+`;