TmMolStar/src/mol-geo/representation/structure/utils.ts

/**
 * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
 *
 * @author Alexander Rose <alexander.rose@weirdbyte.de>
 * @author David Sehnal <david.sehnal@gmail.com>
 */

import { Unit, Element } from 'mol-model/structure';
import { Mat4, Vec3 } from 'mol-math/linear-algebra'

import { createUniformColor, ColorData } from '../../util/color-data';
import { createUniformSize } from '../../util/size-data';
import { elementSizeData } from '../../theme/structure/size/element';
import VertexMap from '../../shape/vertex-map';
import { ColorTheme, SizeTheme } from '../../theme';
import { elementIndexColorData, elementSymbolColorData, instanceIndexColorData, chainIdColorData } from '../../theme/structure/color';
import { ValueCell } from 'mol-util';
import { Mesh } from '../../shape/mesh';
import { Task } from 'mol-task';
import { icosahedronVertexCount } from '../../primitive/icosahedron';
import { MeshBuilder } from '../../shape/mesh-builder';
import { TextureImage } from 'mol-gl/renderable/util';
import { applyFlagAction, FlagAction } from '../../util/flag-data';
import { Loci, isEveryLoci } from 'mol-model/loci';

export function createTransforms({ units }: Unit.SymmetryGroup, transforms?: ValueCell<Float32Array>) {
    const unitCount = units.length
    const n = unitCount * 16
    const array = transforms && transforms.ref.value.length >= n ? transforms.ref.value : new Float32Array(n)
    for (let i = 0; i < unitCount; i++) {
        Mat4.toArray(units[i].conformation.operator.matrix, array, i * 16)
    }
    return transforms ? ValueCell.update(transforms, array) : ValueCell.create(array)
}

export function createColors(group: Unit.SymmetryGroup, vertexMap: VertexMap, props: ColorTheme, colorData?: ColorData) {
    switch (props.name) {
        case 'atom-index':
            return elementIndexColorData({ group, vertexMap }, colorData)
        case 'chain-id':
            return chainIdColorData({ group, vertexMap }, colorData)
        case 'element-symbol':
            return elementSymbolColorData({ group, vertexMap }, colorData)
        case 'instance-index':
            return instanceIndexColorData({ group, vertexMap }, colorData)
        case 'uniform':
            return createUniformColor(props, colorData)
    }
}

export function createSizes(group: Unit.SymmetryGroup, vertexMap: VertexMap, props: SizeTheme) {
    switch (props.name) {
        case 'uniform':
            return createUniformSize(props)
        case 'vdw':
            return elementSizeData({ group, vertexMap })
    }
}

export function createSphereMesh(unit: Unit, radius: Element.Property<number>, detail: number, mesh?: Mesh) {
    return Task.create('Sphere mesh', async ctx => {
        const { elements } = unit;
        const elementCount = elements.length;
        const vertexCount = elementCount * icosahedronVertexCount(detail)
        const meshBuilder = MeshBuilder.create(vertexCount, vertexCount / 2, mesh)

        const v = Vec3.zero()
        const m = Mat4.identity()

        const { x, y, z } = unit.conformation
        const l = Element.Location()
        l.unit = unit

        for (let i = 0; i < elementCount; i++) {
            l.element = elements[i]
            v[0] = x(l.element); v[1] = y(l.element); v[2] = z(l.element)
            Mat4.setTranslation(m, v)

            meshBuilder.setId(i)
            meshBuilder.addIcosahedron(m, { radius: radius(l), detail })

            if (i % 10000 === 0 && ctx.shouldUpdate) {
                await ctx.update({ message: 'Sphere mesh', current: i, max: elementCount });
            }
        }

        return meshBuilder.getMesh()
    })
}


export function applyElementFlags(tFlag: ValueCell<TextureImage>, group: Unit.SymmetryGroup, loci: Loci, action: FlagAction) {
    let changed = false
    const elementCount = group.elements.length
    const instanceCount = group.units.length
    const array = tFlag.ref.value.array
    if (isEveryLoci(loci)) {
        applyFlagAction(array, 0, elementCount * instanceCount, action)
        changed = true
    } else if (Element.isLoci(loci)) {
        for (const e of loci.elements) {
            const unitIdx = Unit.findUnitById(e.unit.id, group.units)
            if (unitIdx !== -1) {
                for (let i = 0, il = e.indices.length; i < il; ++i) {
                    const idx = unitIdx * elementCount + e.indices[i]
                    if (applyFlagAction(array, idx, idx + 1, action) && !changed) {
                        changed = true
                    }
                }
            }
        }
    } else {
        return
    }
    if (changed) {
        ValueCell.update(tFlag, tFlag.ref.value)
    }
}