alpha-orbitals: use square texture for GPU comp

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

@@ -27,6 +27,7 @@ const AlphaOrbitalsSchema = {
     tInfo: TextureSpec('image-float32', 'rgba', 'float', 'nearest'),
     tCoeff: TextureSpec('image-float32', 'rgb', 'float', 'nearest'),
     tAlpha: TextureSpec('image-float32', 'alpha', 'float', 'nearest'),
+    uWidth: UniformSpec('f'),
     uNCenters: UniformSpec('i'),
     uNAlpha: UniformSpec('i'),
     uNCoeff: UniformSpec('i'),
@@ -114,11 +115,15 @@ function createTextureData({
 function getPostprocessingRenderable(ctx: WebGLContext, params: CollocationParams): AlphaOrbitalsRenderable {
     const data = createTextureData(params);
 
+    const [nx, ny, nz] = params.grid.dimensions;
+    const width = Math.ceil(Math.sqrt(nx * ny * nz));
+
     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),
+        uWidth: ValueCell.create(width),
         uNCenters: ValueCell.create(data.nCenters),
         uNAlpha: ValueCell.create(data.nAlpha),
         uNCoeff: ValueCell.create(data.nCoeff),
@@ -146,34 +151,32 @@ function normalizeParams(webgl: WebGLContext) {
 
 export function gpuComputeAlphaOrbitalsGridValues(webgl: WebGLContext, params: CollocationParams) {
     const [nx, ny, nz] = params.grid.dimensions;
-
     normalizeParams(webgl);
 
+    const renderable = getPostprocessingRenderable(webgl, params);
+    const width = renderable.values.uWidth.ref.value;
+
     if (!webgl.computeTargets['alpha-oribtals']) {
-        webgl.computeTargets['alpha-oribtals'] = webgl.createRenderTarget(nx, ny * nz, false, 'uint8', 'nearest');
+        webgl.computeTargets['alpha-oribtals'] = webgl.createRenderTarget(width, width, false, 'uint8', 'nearest');
     } else {
-        webgl.computeTargets['alpha-oribtals'].setSize(nx, ny * nz);
+        webgl.computeTargets['alpha-oribtals'].setSize(width, width);
     }
 
     const target = webgl.computeTargets['alpha-oribtals'];
-    const renderable = getPostprocessingRenderable(webgl, params);
-
-    const width = nx;
-    const height = ny * nz;
 
     const { gl, state } = webgl;
     target.bind();
-    gl.viewport(0, 0, width, height);
-    gl.scissor(0, 0, width, height);
+    gl.viewport(0, 0, width, width);
+    gl.scissor(0, 0, width, width);
     state.disable(gl.SCISSOR_TEST);
     state.disable(gl.BLEND);
     state.disable(gl.DEPTH_TEST);
     state.depthMask(false);
     renderable.render();
 
-    const array = new Uint8Array(width * height * 4);
-    webgl.readPixels(0, 0, width, height, array);
-    const floats = new Float32Array(array.buffer, array.byteOffset, width * height);
+    const array = new Uint8Array(width * width * 4);
+    webgl.readPixels(0, 0, width, width, array);
+    const floats = new Float32Array(array.buffer, array.byteOffset, nx * ny * nz);
     renderable.dispose();
 
     return floats;

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

@@ -19,6 +19,8 @@ uniform sampler2D tInfo;
 uniform sampler2D tCoeff;
 uniform sampler2D tAlpha;
 
+uniform float uWidth;
+
 #ifndef uNCenters
     uniform int uNCenters;
 #endif
@@ -65,11 +67,11 @@ vec4 floatToRgba(float texelFloat) {
     );
 }
 
-float L1(const in vec3 p, const in float a0, const in float a1, const in float a2) {
+float L1(vec3 p, float a0, float a1, float a2) {
     return a0 * p.z + a1 * p.x + a2 * p.y;
 }
 
-float L2(const in vec3 p, const in float a0, const in float a1, const in float a2, const in float a3, const in float a4) {
+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 (
@@ -81,7 +83,7 @@ float L2(const in vec3 p, const in float a0, const in float a1, const in float a
     );
 }
 
-float L3(const in vec3 p, const in float a0, const in float a1, const in float a2, const in float a3, const in float a4, const in float a5, const in float a6) {
+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;
@@ -96,7 +98,7 @@ float L3(const in vec3 p, const in float a0, const in float a1, const in float a
     );
 }
 
-float L4(const in vec3 p, const in float a0, const in float a1, const in float a2, const in float a3, const in float a4, const in float a5, const in float a6, const in float a7, const in float a8) {
+float L4(vec3 p, float a0, float a1, float a2, float a3, float a4, float a5, float a6, float a7, float a8) {
     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;
@@ -114,11 +116,11 @@ float L4(const in vec3 p, const in float a0, const in float a1, const in float a
     );
 }
 
-float alpha(const in float offset, const in float f) {
+float alpha(float offset, float f) {
     return texture2D(tAlpha, vec2(offset * f, 0.5)).x;
 }
 
-float Y(const in int L, const in vec3 X, const in float aO, const in float fA) {
+float Y(int L, vec3 X, float aO, float fA) {
     if (L == 0) {
         return alpha(aO, fA);
     } else if (L == 1) {
@@ -145,7 +147,7 @@ float Y(const in int L, const in vec3 X, const in float aO, const in float fA) {
 }
 
 #ifndef uMaxCoeffs
-    float R(const in float R2, const in int start, const in int end, const in float fCoeff) {
+    float R(float R2, int start, int end, float fCoeff) {
         float gauss = 0.0;
         for (int i = start; i < end; i++) {
             vec2 c = texture2D(tCoeff, vec2(float(i) * fCoeff, 0.5)).xy;
@@ -156,7 +158,7 @@ float Y(const in int L, const in vec3 X, const in float aO, const in float fA) {
 #endif
 
 #ifdef uMaxCoeffs
-    float R(const in float R2, const in int start, const in int end, const in float fCoeff) {
+    float R(float R2, int start, int end, float fCoeff) {
         float gauss = 0.0;
         int o = start;
         for (int i = 0; i < uMaxCoeffs; i++) {
@@ -170,11 +172,11 @@ float Y(const in int L, const in vec3 X, const in float aO, const in float fA) {
     }
 #endif
 
-float intDiv(const in float a, const in float b) { return float(int(a) / int(b)); }
-float intMod(const in float a, const in float b) { return a - b * float(int(a) / int(b)); }
+float intDiv(float a, float b) { return float(int(a) / int(b)); }
+float intMod(float a, float b) { return a - b * float(int(a) / int(b)); }
 
 void main(void) { 
-    float offset = round(floor(gl_FragCoord.x) + floor(gl_FragCoord.y) * uDimensions.x);
+    float offset = round(floor(gl_FragCoord.x) + floor(gl_FragCoord.y) * uWidth);
     
     // axis order fast to slow Z, Y, X
     // TODO: support arbitrary axis orders?