TmMolStar/src/apps/basic-wrapper/helpers.ts

/**
 * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
 *
 * @author David Sehnal <david.sehnal@gmail.com>
 */

import { Mat4, Vec3 } from '../../mol-math/linear-algebra';
import { PluginContext } from '../../mol-plugin/context';
import { PluginStateObject as PSO } from '../../mol-plugin-state/objects';
import { StateTransforms } from '../../mol-plugin-state/transforms';
import { MolScriptBuilder as MS } from '../../mol-script/language/builder';
import { StateBuilder } from '../../mol-state';
import Expression from '../../mol-script/language/expression';
import { ColorTheme } from '../../mol-theme/color';
import { createStructureRepresentationParams } from '../../mol-plugin-state/helpers/structure-representation-params';
type SupportedFormats = 'cif' | 'pdb'

export namespace StateHelper {
    export function download(b: StateBuilder.To<PSO.Root>, url: string, ref?: string) {
        return b.apply(StateTransforms.Data.Download, { url, isBinary: false }, { ref });
    }

    export function getModel(b: StateBuilder.To<PSO.Data.Binary | PSO.Data.String>, format: SupportedFormats, modelIndex = 0) {
        const parsed = format === 'cif'
            ? b.apply(StateTransforms.Data.ParseCif).apply(StateTransforms.Model.TrajectoryFromMmCif)
            : b.apply(StateTransforms.Model.TrajectoryFromPDB);

        return parsed.apply(StateTransforms.Model.ModelFromTrajectory, { modelIndex });
    }

    export function structure(b: StateBuilder.To<PSO.Molecule.Model>) {
        return b.apply(StateTransforms.Model.StructureFromModel, void 0, { tags: 'structure' })
    };

    export function selectChain(b: StateBuilder.To<PSO.Molecule.Structure>, auth_asym_id: string) {
        const expression = MS.struct.generator.atomGroups({
            'chain-test': MS.core.rel.eq([MS.struct.atomProperty.macromolecular.auth_asym_id(), auth_asym_id])
        })
        return b.apply(StateTransforms.Model.StructureSelectionFromExpression, { expression, label: `Chain ${auth_asym_id}` });
    }

    export function select(b: StateBuilder.To<PSO.Molecule.Structure>, expression: Expression) {
        return b.apply(StateTransforms.Model.StructureSelectionFromExpression, { expression });
    }

    export function selectSurroundingsOfFirstResidue(b: StateBuilder.To<PSO.Molecule.Structure>, comp_id: string, radius: number) {
        const expression = MS.struct.modifier.includeSurroundings({
            0: MS.struct.filter.first([
                MS.struct.generator.atomGroups({
                    'residue-test': MS.core.rel.eq([MS.struct.atomProperty.macromolecular.label_comp_id(), comp_id]),
                    'group-by': MS.struct.atomProperty.macromolecular.residueKey()
                })
            ]),
            radius
        })
        return b.apply(StateTransforms.Model.StructureSelectionFromExpression, { expression, label: `Surr. ${comp_id} (${radius} ang)` });
    }

    export function identityTransform(b: StateBuilder.To<PSO.Molecule.Structure>, m: Mat4) {
        return b.apply(StateTransforms.Model.TransformStructureConformation,
            { transform: { name: 'components', params: { axis: Vec3.create(1, 0, 0), angle: 0, translation: Vec3.zero() } } },
            { tags: 'transform' });
    }

    export function transform(b: StateBuilder.To<PSO.Molecule.Structure>, matrix: Mat4) {
        return b.apply(StateTransforms.Model.TransformStructureConformation, { 
            transform: { name: 'matrix', params: matrix }
        }, { tags: 'transform' });
    }

    export function assemble(b: StateBuilder.To<PSO.Molecule.Model>, id?: string) {
        const props = {
            type: {
                name: 'assembly' as const,
                params: { id: id || 'deposited' }
            }
        }
        return b.apply(StateTransforms.Model.StructureFromModel, props, { tags: 'asm' })
    }

    export function visual(ctx: PluginContext, visualRoot: StateBuilder.To<PSO.Molecule.Structure>) {
        visualRoot.apply(StateTransforms.Model.StructureComplexElement, { type: 'atomic-sequence' })
            .apply(StateTransforms.Representation.StructureRepresentation3D,
                createStructureRepresentationParams(ctx, void 0, { type: 'cartoon' }), { tags: 'seq-visual' });
        visualRoot.apply(StateTransforms.Model.StructureComplexElement, { type: 'atomic-het' })
            .apply(StateTransforms.Representation.StructureRepresentation3D,
                createStructureRepresentationParams(ctx, void 0, { type: 'ball-and-stick' }), { tags: 'het-visual' });
        return visualRoot;
    }

    export function ballsAndSticks(ctx: PluginContext, visualRoot: StateBuilder.To<PSO.Molecule.Structure>, expression: Expression, color?: ColorTheme.BuiltIn) {
        visualRoot
            .apply(StateTransforms.Model.StructureSelectionFromExpression, { expression })
            .apply(StateTransforms.Representation.StructureRepresentation3D,
                createStructureRepresentationParams(ctx, void 0, { type: 'ball-and-stick', color }), { tags: 'het-visual' });
        return visualRoot;
    }

}