TmMolStar/src/mol-repr/structure/visual/label-text.ts

/**
 * Copyright (c) 2019-2020 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 { ParamDefinition as PD } from '../../../mol-util/param-definition';
import { VisualUpdateState } from '../../../mol-repr/util';
import { VisualContext } from '../../../mol-repr/visual';
import { Structure, StructureElement, StructureProperties } from '../../../mol-model/structure';
import { Theme } from '../../../mol-theme/theme';
import { Text } from '../../../mol-geo/geometry/text/text';
import { TextBuilder } from '../../../mol-geo/geometry/text/text-builder';
import { ComplexTextVisual, ComplexTextParams, ComplexVisual } from '../complex-visual';
import { ElementIterator, getSerialElementLoci, eachSerialElement } from './util/element';
import { ColorNames } from '../../../mol-util/color/names';
import { Vec3 } from '../../../mol-math/linear-algebra';
import { PhysicalSizeTheme } from '../../../mol-theme/size/physical';
import { BoundaryHelper } from '../../../mol-math/geometry/boundary-helper';

export const LabelTextParams = {
    ...ComplexTextParams,
    background: PD.Boolean(true),
    backgroundMargin: PD.Numeric(0, { min: 0, max: 1, step: 0.01 }),
    backgroundColor: PD.Color(ColorNames.black),
    backgroundOpacity: PD.Numeric(0.5, { min: 0, max: 1, step: 0.01 }),
    level: PD.Select('residue', [['chain', 'Chain'], ['residue', 'Residue'], ['element', 'Element']]),
    chainScale: PD.Numeric(10, { min: 0, max: 20, step: 0.1 }),
    residueScale: PD.Numeric(1, { min: 0, max: 20, step: 0.1 }),
    elementScale: PD.Numeric(0.5, { min: 0, max: 20, step: 0.1 }),
}
export type LabelTextParams = typeof LabelTextParams
export type LabelTextProps = PD.Values<LabelTextParams>
export type LabelLevels = LabelTextProps['level']

export function LabelTextVisual(materialId: number): ComplexVisual<LabelTextParams> {
    return ComplexTextVisual<LabelTextParams>({
        defaultProps: PD.getDefaultValues(LabelTextParams),
        createGeometry: createLabelText,
        createLocationIterator: ElementIterator.fromStructure,
        getLoci: getSerialElementLoci,
        eachLocation: eachSerialElement,
        setUpdateState: (state: VisualUpdateState, newProps: PD.Values<LabelTextParams>, currentProps: PD.Values<LabelTextParams>) => {
            state.createGeometry = (
                newProps.level !== currentProps.level ||
                (newProps.level === 'chain' && newProps.chainScale !== currentProps.chainScale) ||
                (newProps.level === 'residue' && newProps.residueScale !== currentProps.residueScale) ||
                (newProps.level === 'element' && newProps.elementScale !== currentProps.elementScale)
            )
        }
    }, materialId)
}

function createLabelText(ctx: VisualContext, structure: Structure, theme: Theme, props: LabelTextProps, text?: Text): Text {

    switch (props.level) {
        case 'chain': return createChainText(ctx, structure, theme, props, text)
        case 'residue': return createResidueText(ctx, structure, theme, props, text)
        case 'element': return createElementText(ctx, structure, theme, props, text)
    }
}

//

const tmpVec = Vec3();
const boundaryHelper = new BoundaryHelper();

function createChainText(ctx: VisualContext, structure: Structure, theme: Theme, props: LabelTextProps, text?: Text): Text {
    const l = StructureElement.Location.create(structure);
    const { units, serialMapping } = structure;
    const { auth_asym_id, label_asym_id } = StructureProperties.chain;
    const { cumulativeUnitElementCount } = serialMapping

    const count = units.length
    const { chainScale } = props
    const builder = TextBuilder.create(props, count, count / 2, text)

    for (let i = 0, il = units.length; i < il; ++i) {
        const unit = units[i]
        l.unit = unit
        l.element = unit.elements[0]
        const { center, radius } = unit.lookup3d.boundary.sphere
        Vec3.transformMat4(tmpVec, center, unit.conformation.operator.matrix)
        const authId = auth_asym_id(l)
        const labelId = label_asym_id(l)
        const text = authId === labelId ? labelId : `${labelId} [${authId}]`
        builder.add(text, tmpVec[0], tmpVec[1], tmpVec[2], radius, chainScale, cumulativeUnitElementCount[i])
    }

    return builder.getText()
}

function createResidueText(ctx: VisualContext, structure: Structure, theme: Theme, props: LabelTextProps, text?: Text): Text {
    const l = StructureElement.Location.create(structure);
    const { units, serialMapping } = structure;
    const { auth_seq_id, label_comp_id } = StructureProperties.residue;
    const { cumulativeUnitElementCount } = serialMapping

    const count = structure.polymerResidueCount * 2
    const { residueScale } = props
    const builder = TextBuilder.create(props, count, count / 2, text)

    for (let i = 0, il = units.length; i < il; ++i) {
        const unit = units[i]
        const pos = unit.conformation.position;
        const { elements } = unit
        l.unit = unit
        l.element = unit.elements[0]

        const residueIndex = unit.model.atomicHierarchy.residueAtomSegments.index;
        const groupOffset = cumulativeUnitElementCount[i]

        let j = 0, jl = elements.length;
        while (j < jl) {
            const start = j, rI = residueIndex[elements[j]];
            j++;
            while (j < jl && residueIndex[elements[j]] === rI) j++;

            boundaryHelper.reset(0);
            for (let eI = start; eI < j; eI++) {
                pos(elements[eI], tmpVec);
                boundaryHelper.boundaryStep(tmpVec, 0);
            }
            boundaryHelper.finishBoundaryStep();
            for (let eI = start; eI < j; eI++) {
                pos(elements[eI], tmpVec);
                boundaryHelper.extendStep(tmpVec, 0);
            }

            l.element = elements[start];

            const { center, radius } = boundaryHelper
            const authSeqId = auth_seq_id(l)
            const compId = label_comp_id(l)

            const text = `${compId} ${authSeqId}`
            builder.add(text, center[0], center[1], center[2], radius, residueScale, groupOffset + start)
        }
    }

    return builder.getText()
}

function createElementText(ctx: VisualContext, structure: Structure, theme: Theme, props: LabelTextProps, text?: Text): Text {
    const l = StructureElement.Location.create(structure);
    const { units, serialMapping } = structure;
    const { label_atom_id, label_alt_id } = StructureProperties.atom;
    const { cumulativeUnitElementCount } = serialMapping

    const sizeTheme = PhysicalSizeTheme({}, {})

    const count = structure.elementCount
    const { elementScale } = props
    const builder = TextBuilder.create(props, count, count / 2, text)

    for (let i = 0, il = units.length; i < il; ++i) {
        const unit = units[i]
        const pos = unit.conformation.position;
        const { elements } = unit
        l.unit = unit

        const groupOffset = cumulativeUnitElementCount[i]

        for (let j = 0, _j = elements.length; j < _j; j++) {
            l.element = elements[j];
            pos(l.element, tmpVec);
            const atomId = label_atom_id(l)
            const altId = label_alt_id(l)
            const text = altId ? `${atomId}%${altId}` : atomId
            builder.add(text, tmpVec[0], tmpVec[1], tmpVec[2], sizeTheme.size(l), elementScale, groupOffset + j)
        }
    }

    return builder.getText()
}