TmMolStar/src/apps/tm-viewer/index.ts

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

import { TMDETMembraneOrientation } from '../../extensions/tmdet/behavior';
import { StateTransforms } from '../../mol-plugin-state/transforms';
import { createPlugin } from '../../mol-plugin-ui';
import { PluginUIContext } from '../../mol-plugin-ui/context';
import { PluginLayoutControlsDisplay } from '../../mol-plugin/layout';
import { DefaultPluginUISpec, PluginUISpec } from '../../mol-plugin-ui/spec';
import { PluginConfig } from '../../mol-plugin/config';
import { PluginSpec } from '../../mol-plugin/spec';
import { StateBuilder, StateObjectRef } from '../../mol-state';
import { Color } from '../../mol-util/color';
import '../../mol-util/polyfill';
import { ObjectKeys } from '../../mol-util/type-helpers';
import './embedded.html';
import './favicon.ico';
import './index.html';
import { MolScriptBuilder as MS } from '../../mol-script/language/builder';
import { createStructureRepresentationParams } from '../../mol-plugin-state/helpers/structure-representation-params';
import { Expression } from '../../mol-script/language/expression';
import { PluginStateObject } from '../../mol-plugin-state/objects';
import { MembraneOrientation } from '../../extensions/tmdet/prop';
import { MEMBRANE_STORAGE_KEY } from '../../extensions/tmdet/algorithm';
import { Quat, Vec3 } from '../../mol-math/linear-algebra';
import { PluginCommands } from '../../mol-plugin/commands';

require('mol-plugin-ui/skin/light.scss');

export { PLUGIN_VERSION as version } from '../../mol-plugin/version';
export { setDebugMode, setProductionMode } from '../../mol-util/debug';

const Extensions = {
    'tmdet-membrane-orientation': PluginSpec.Behavior(TMDETMembraneOrientation)
};

const DefaultViewerOptions = {
    extensions: ObjectKeys(Extensions),
    layoutIsExpanded: true,
    layoutShowControls: true,
    layoutShowRemoteState: true,
    layoutControlsDisplay: 'reactive' as PluginLayoutControlsDisplay,
    layoutShowSequence: true,
    layoutShowLog: true,
    layoutShowLeftPanel: true,
    collapseLeftPanel: false,
    disableAntialiasing: false,
    pixelScale: 1,
    enableWboit: true,

    viewportShowExpand: PluginConfig.Viewport.ShowExpand.defaultValue,
    viewportShowControls: PluginConfig.Viewport.ShowControls.defaultValue,
    viewportShowSettings: PluginConfig.Viewport.ShowSettings.defaultValue,
    viewportShowSelectionMode: PluginConfig.Viewport.ShowSelectionMode.defaultValue,
    viewportShowAnimation: PluginConfig.Viewport.ShowAnimation.defaultValue,
    pluginStateServer: PluginConfig.State.DefaultServer.defaultValue,
    volumeStreamingServer: PluginConfig.VolumeStreaming.DefaultServer.defaultValue,
    volumeStreamingDisabled: !PluginConfig.VolumeStreaming.Enabled.defaultValue,
    pdbProvider: PluginConfig.Download.DefaultPdbProvider.defaultValue,
    emdbProvider: PluginConfig.Download.DefaultEmdbProvider.defaultValue,
};
type ViewerOptions = typeof DefaultViewerOptions;


export var membrane: MembraneOrientation;

export class Viewer {
    plugin: PluginUIContext





    ////////////////////////////// UNITMP VIEWER PROTOTYPING SECTION

    async loadWithUNITMPMembraneRepresentation(url: string, regionDescriptors: any) {
        const membraneNormal: Vec3 = Vec3.fromObj(
            regionDescriptors['membrane-normal']
        );
        const membrane: MembraneOrientation = {
            planePoint1: Vec3.fromArray(Vec3.zero(), membraneNormal, 0),
            planePoint2: Vec3.fromArray(Vec3.zero(), membraneNormal, 0),
            // NOTE: centroid is not 0,0,0. It is x,y,0. Right?
            centroid: Vec3.fromArray(
                Vec3.zero(), [ membraneNormal[0], membraneNormal[1], 0 ], 0
            ),
            normalVector: membraneNormal,

            // TODO: radius is still just a dummy value now.
            //       Can we send a precalculated value by our backend?
            //
            // (NOTE: the TMDET extension calculates and sets it during applying preset)
            radius: regionDescriptors[ 'radius' ]
        };
        membrane.planePoint2[2] *= -1;

        localStorage.setItem(MEMBRANE_STORAGE_KEY, JSON.stringify(membrane));

        const isBinary = false;

        const data = await this.plugin.builders.data.download({
            url, label: regionDescriptors['pdb-id'], isBinary
        }); //, { state: { isGhost: true } });
        const trajectory = await this.plugin.builders.structure.parseTrajectory(data, regionDescriptors.format);
        // create membrane representation
        await this.plugin.builders.structure.hierarchy.applyPreset(
            trajectory, 'default', { representationPreset: 'preset-membrane-orientation' as any });

        const structure: StateObjectRef<PluginStateObject.Molecule.Structure> =
            this.plugin.managers.structure.hierarchy.current.models[0].structures[0].cell;
        const components = {
            polymer: await this.plugin.builders.structure.tryCreateComponentStatic(structure, 'polymer'),
            ligand: await this.plugin.builders.structure.tryCreateComponentStatic(structure, 'ligand'),
            water: await this.plugin.builders.structure.tryCreateComponentStatic(structure, 'water'),
        };

        const builder = this.plugin.builders.structure.representation;
        const update = this.plugin.build();
        if (components.polymer) builder.buildRepresentation(update, components.polymer, { type: 'cartoon', typeParams: { alpha: 0.5 } }, { tag: 'polymer' });
        if (components.ligand) builder.buildRepresentation(update, components.ligand, { type: 'ball-and-stick' }, { tag: 'ligand' });
        if (components.water) builder.buildRepresentation(update, components.water, { type: 'ball-and-stick', typeParams: { alpha: 0.6 } }, { tag: 'water' });
        await update.commit();

        regionDescriptors.chains.forEach((chain: any) => {

            for(let regionKey in chain.regions) {
              const update = this.plugin.build();
              const region = chain.regions[regionKey];
              this.createRegionRepresentation(chain.chain_id, region, update.to(structure));
              update.commit();
            }

        });

        //
        // reset the camera because the membranes render 1st and the structure might not be fully visible
        //
        this.rotateCamera();

        //requestAnimationFrame(() => this.plugin.canvas3d?.requestCameraReset());
    }

    private createRegionRepresentation(chain: string, region: any, update: StateBuilder.To<any, any>) {
        const regionLabel: string = `${chain} | ${region.name}`;
        const color: Color = Color.fromArray(region.color, 0);
        const query: Expression = this.getQuery(chain, region.auth_ids as number[]);

        // based on https://github.com/molstar/molstar/issues/209
        update
            .apply(StateTransforms.Model.StructureSelectionFromExpression, { label: regionLabel, expression: query })
            .apply(StateTransforms.Representation.StructureRepresentation3D, createStructureRepresentationParams(this.plugin, update.selector.data, {
                color: 'uniform',
                colorParams: { value: color }
            }));
    }

    private getQuery(chainId: string, auth_array: number[]): Expression {
        const query: Expression =
            MS.struct.generator.atomGroups({
                'residue-test': MS.core.set.has([MS.set( ...auth_array ), MS.ammp('auth_seq_id')]),
                'chain-test': MS.core.rel.eq([chainId, MS.ammp('label_asym_id')])
            });
        return query;
    }

    private rotateCamera(): void {
        function rot90q(v: Vec3, axis: Vec3 = Vec3.create(1, 0, 0)): Vec3 {
            const q = Quat.setAxisAngle(Quat(), axis, -Math.PI/2);
            return Vec3.transformQuat(Vec3(), v, q);
        }
        function sub(v: Vec3, u: Vec3): Vec3 {
            return Vec3.sub(Vec3(), v, u);
        }
        function add(v: Vec3, u: Vec3): Vec3 {
            return Vec3.add(Vec3(), v, u);
        }


        const cam = this.plugin.canvas3d!.camera;
        const snapshot = cam.getSnapshot();
        const newSnapshot = {
            ...snapshot,
            // target + rotateBy90(postition - target)
            position: add(snapshot.target, rot90q(sub(snapshot.position, snapshot.target))),
            target: snapshot.target,
            up: Vec3.negUnitZ
        };

        PluginCommands.Camera.Reset(this.plugin, { snapshot: newSnapshot, durationMs: 1000 });
        //requestAnimationFrame(() => this.plugin.canvas3d?.requestCameraReset({ snapshot: newSnapshot }));
    }

    ////////////////////////////// END OF PROTOTYPING SECTION






    constructor(elementOrId: string | HTMLElement, options: Partial<ViewerOptions> = {}) {
        const o = { ...DefaultViewerOptions, ...options };
        const defaultSpec = DefaultPluginUISpec();

        const spec: PluginUISpec = {
            actions: defaultSpec.actions,
            behaviors: [
                ...defaultSpec.behaviors,
                ...o.extensions.map(e => Extensions[e]),
            ],
            animations: [...defaultSpec.animations || []],
            customParamEditors: defaultSpec.customParamEditors,
            layout: {
                initial: {
                    isExpanded: o.layoutIsExpanded,
                    showControls: o.layoutShowControls,
                    controlsDisplay: o.layoutControlsDisplay,
                    regionState: {
                        bottom: 'full',
                        left: o.collapseLeftPanel ? 'collapsed' : 'full',
                        right: 'full',
                        top: 'full',
                    }
                },
            },
            components: {
                ...defaultSpec.components,
                controls: {
                    ...defaultSpec.components?.controls,
                    top: o.layoutShowSequence ? undefined : 'none',
                    bottom: o.layoutShowLog ? undefined : 'none',
                    left: o.layoutShowLeftPanel ? undefined : 'none',
                },
                remoteState: o.layoutShowRemoteState ? 'default' : 'none',
            },
            config: [
                [PluginConfig.General.DisableAntialiasing, o.disableAntialiasing],
                [PluginConfig.General.PixelScale, o.pixelScale],
                [PluginConfig.General.EnableWboit, o.enableWboit],
                [PluginConfig.Viewport.ShowExpand, o.viewportShowExpand],
                [PluginConfig.Viewport.ShowControls, o.viewportShowControls],
                [PluginConfig.Viewport.ShowSettings, o.viewportShowSettings],
                [PluginConfig.Viewport.ShowSelectionMode, o.viewportShowSelectionMode],
                [PluginConfig.Viewport.ShowAnimation, o.viewportShowAnimation],
                [PluginConfig.State.DefaultServer, o.pluginStateServer],
                [PluginConfig.State.CurrentServer, o.pluginStateServer],
                [PluginConfig.VolumeStreaming.DefaultServer, o.volumeStreamingServer],
                [PluginConfig.VolumeStreaming.Enabled, !o.volumeStreamingDisabled],
                [PluginConfig.Download.DefaultPdbProvider, o.pdbProvider],
                [PluginConfig.Download.DefaultEmdbProvider, o.emdbProvider]
            ]
        };

        const element = typeof elementOrId === 'string'
            ? document.getElementById(elementOrId)
            : elementOrId;
        if (!element) throw new Error(`Could not get element with id '${elementOrId}'`);
        this.plugin = createPlugin(element, spec);
    }

    handleResize() {
        this.plugin.layout.events.updated.next();
    }
}