compute.ts

/**
 * 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 { DefineSpec, 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 { ValueCell } from '../../../mol-util';
import { arrayMin } from '../../../mol-util/array';
import { isLittleEndian } from '../../../mol-util/is-little-endian';
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', 'rgba', 'float', 'nearest'),
    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'),
    uMaxCoeffs: UniformSpec('i'),
    uLittleEndian: UniformSpec('b')
};
const AlphaOrbitalsShaderCode = ShaderCode('postprocessing', quad_vert, shader_frag);
type AlphaOrbitalsRenderable = ComputeRenderable<Values<typeof AlphaOrbitalsSchema>>

function createTextureData({
    basis,
    sphericalOrder,
    alphaOrbitals,
    cutoffThreshold
}: 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(4 * 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 maxCoeffs = 0;