TmMolStar/src/extensions/tmdet/behavior.ts

/**
 * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
 *
 * @author Alexander Rose <alexander.rose@weirdbyte.de>
 * @author Sebastian Bittrich <sebastian.bittrich@rcsb.org>
 */

import { ParamDefinition as PD } from '../../mol-util/param-definition';
import { StructureRepresentationPresetProvider, PresetStructureRepresentations } from '../../mol-plugin-state/builder/structure/representation-preset';
import { StateObject, StateObjectRef, StateObjectCell, StateTransformer, StateTransform } from '../../mol-state';
import { Task } from '../../mol-task';
import { PluginBehavior } from '../../mol-plugin/behavior';
import { PluginStateObject, PluginStateTransform } from '../../mol-plugin-state/objects';
import { PluginContext } from '../../mol-plugin/context';
import { DefaultQueryRuntimeTable } from '../../mol-script/runtime/query/compiler';
import { StructureSelectionQuery, StructureSelectionCategory } from '../../mol-plugin-state/helpers/structure-selection-query';
import { MolScriptBuilder as MS } from '../../mol-script/language/builder';
import { GenericRepresentationRef } from '../../mol-plugin-state/manager/structure/hierarchy-state';
import { PluginUIContext } from '../../mol-plugin-ui/context';



// TMDET imports
import { MembraneOrientationRepresentationProvider, MembraneOrientationParams, MembraneOrientationRepresentation } from './representation';
import { MembraneOrientationProvider, MembraneOrientation, TmDetDescriptorCache } from './prop';
import { applyTransformations, createMembraneOrientation } from './transformation';
import { ComponentsType, PDBTMDescriptor, PMS } from './types';
import { registerTmDetSymmetry } from './symmetry';
import { LabeledResidues } from './labeling';
import { TmDetColorThemeProvider } from './tmdet-color-theme';
import { loadInitialSnapshot, rotateCamera, storeCameraSnapshot } from './camera';



const Tag = MembraneOrientation.Tag;
const TMDET_MEMBRANE_ORIENTATION = 'Membrane Orientation';

export const TMDETMembraneOrientation = PluginBehavior.create<{ autoAttach: boolean }>({
    name: 'tmdet-membrane-orientation-prop',
    category: 'custom-props',
    display: {
        name: TMDET_MEMBRANE_ORIENTATION,
        description: 'Data calculated with TMDET algorithm.'
    },
    ctor: class extends PluginBehavior.Handler<{ autoAttach: boolean }> {
        private provider = MembraneOrientationProvider

        register(): void {
            DefaultQueryRuntimeTable.addCustomProp(this.provider.descriptor);

            this.ctx.customStructureProperties.register(this.provider, this.params.autoAttach);

            this.ctx.representation.structure.registry.add(MembraneOrientationRepresentationProvider);
            this.ctx.query.structure.registry.add(isTransmembrane);

            this.ctx.representation.structure.themes.colorThemeRegistry.add(TmDetColorThemeProvider);
            this.ctx.managers.lociLabels.addProvider(LabeledResidues.labelProvider!);
            this.ctx.customModelProperties.register(LabeledResidues.propertyProvider, true);


            this.ctx.genericRepresentationControls.set(Tag.Representation, selection => {
                const refs: GenericRepresentationRef[] = [];
                selection.structures.forEach(structure => {
                    const memRepr = structure.genericRepresentations?.filter(r => r.cell.transform.transformer.id === MembraneOrientation3D.id)[0];
                    if (memRepr) refs.push(memRepr);
                });
                return [refs, 'Membrane Orientation'];
            });
            this.ctx.builders.structure.representation.registerPreset(MembraneOrientationPreset);
        }

        update(p: { autoAttach: boolean }) {
            let updated = this.params.autoAttach !== p.autoAttach;
            this.params.autoAttach = p.autoAttach;
            this.ctx.customStructureProperties.setDefaultAutoAttach(this.provider.descriptor.name, this.params.autoAttach);
            return updated;
        }

        unregister() {
            DefaultQueryRuntimeTable.removeCustomProp(this.provider.descriptor);

            this.ctx.customStructureProperties.unregister(this.provider.descriptor.name);

            this.ctx.representation.structure.registry.remove(MembraneOrientationRepresentationProvider);
            this.ctx.query.structure.registry.remove(isTransmembrane);

            this.ctx.genericRepresentationControls.delete(Tag.Representation);
            this.ctx.builders.structure.representation.unregisterPreset(MembraneOrientationPreset);
        }
    },
    params: () => ({
        autoAttach: PD.Boolean(false)
    })
});

//

export const isTransmembrane = StructureSelectionQuery('Residues Embedded in Membrane', MS.struct.modifier.union([
    MS.struct.modifier.wholeResidues([
        MS.struct.modifier.union([
            MS.struct.generator.atomGroups({
                'chain-test': MS.core.rel.eq([MS.ammp('objectPrimitive'), 'atomistic']),
                'atom-test': MembraneOrientation.symbols.isTransmembrane.symbol(),
            })
        ])
    ])
]), {
    description: 'Select residues that are embedded between the membrane layers.',
    category: StructureSelectionCategory.Residue,
    ensureCustomProperties: (ctx, structure) => {
        return MembraneOrientationProvider.attach(ctx, structure);
    }
});




//
// //////////////////////////// TMDET VIEWER FUNCTIONS
//




export let membraneOrientation: MembraneOrientation;

export async function loadWithUNITMPMembraneRepresentation(plugin: PluginUIContext, params: any) {
    storeCameraSnapshot(plugin); // store if it is not stored yet

    loadInitialSnapshot(plugin); // load if there is a stored one
    setTimeout(() => plugin.clear(), 100); // clear scene after some delay


    setTimeout(async () => {
        const pdbtmDescriptor: PDBTMDescriptor = await downloadRegionDescriptor(plugin, params);
        TmDetDescriptorCache.add(pdbtmDescriptor);

        membraneOrientation = createMembraneOrientation(pdbtmDescriptor);

        // load structure
        await loadStructure(plugin, params, pdbtmDescriptor);
        // cartoon, colors etc.
        await createStructureRepresentation(plugin, pdbtmDescriptor);

        //
        // It also resets the camera because the membranes render 1st and the structure might not be fully visible
        //
        rotateCamera(plugin);
    }, 500);
}

async function downloadRegionDescriptor(plugin: PluginUIContext, params: any): Promise<any> {
    // run a fetch task
    const downloadResult: string = await plugin.runTask(plugin.fetch({ url: params.regionDescriptorUrl })) as string;
    const pdbtmDescriptor: any = JSON.parse(downloadResult);
    return pdbtmDescriptor;
}

async function createStructureRepresentation(plugin: PluginUIContext, pdbtmDescriptor: any) {
    // get the first structure of the first model
    const structure: StateObjectRef<PMS> = plugin.managers.structure.hierarchy.current.models[0].structures[0].cell;
    const components = await createStructureComponents(plugin, structure);

    await applyTransformations(plugin, pdbtmDescriptor);

    await buildStructureRepresentation(plugin, pdbtmDescriptor, components);
}

async function createStructureComponents(plugin: PluginUIContext, structure: StateObjectCell<PMS, StateTransform<StateTransformer<StateObject<any, StateObject.Type<any>>, StateObject<any, StateObject.Type<any>>, any>>>) {
    return {
        polymer: await plugin.builders.structure.tryCreateComponentStatic(structure, 'polymer'),
        ligand: await plugin.builders.structure.tryCreateComponentStatic(structure, 'ligand'),
        water: await plugin.builders.structure.tryCreateComponentStatic(structure, 'water'),
    };
}

async function buildStructureRepresentation(plugin: PluginUIContext, pdbtmDescriptor: PDBTMDescriptor, components: ComponentsType) {
    const builder = plugin.builders.structure.representation;
    const update = plugin.build();
    if (components.polymer)
        builder.buildRepresentation(update, components.polymer, {
                type: 'cartoon',
                color: TmDetColorThemeProvider.name as any, colorParams: { pdbtmDescriptor }
            },
            { 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();
}

async function loadStructure(ctx: PluginUIContext, params: any, pdbtmDescriptor: PDBTMDescriptor): Promise<void> {

    // replace original symmetry format function
    registerTmDetSymmetry(pdbtmDescriptor);

    const builders = ctx.builders;
    const data = await builders.data.download({
        url: params.structureUrl,
        label: `${pdbtmDescriptor.pdb_id}`,
        isBinary: false
    }); // , { state: { isGhost: true } });
    const trajectory = await builders.structure.parseTrajectory(data, 'mmcif');

    // create membrane representation
    await builders.structure.hierarchy.applyPreset(
        trajectory, 'default', { representationPreset: 'preset-membrane-orientation' as any });
}




//
// //////////////////////////// END OF TMDET VIEWER SECTION
//




type MembraneOrientation3DType = typeof MembraneOrientation3D
const MembraneOrientation3D = PluginStateTransform.BuiltIn({
    name: Tag.Representation,
    display: {
        name: TMDET_MEMBRANE_ORIENTATION,
        description: 'Membrane Orientation planes and rims. Data calculated with TMDET algorithm.'
    },
    from: PluginStateObject.Molecule.Structure,
    to: PluginStateObject.Shape.Representation3D,
    params: () => {
        return {
            ...MembraneOrientationParams,
        };
    }
})({
    canAutoUpdate({ oldParams, newParams }) {
        return true;
    },
    apply({ a, params }, plugin: PluginContext) {
        return Task.create(TMDET_MEMBRANE_ORIENTATION, async ctx => {
            await MembraneOrientationProvider.attach({ runtime: ctx, assetManager: plugin.managers.asset }, a.data);
            const repr = MembraneOrientationRepresentation({ webgl: plugin.canvas3d?.webgl, ...plugin.representation.structure.themes }, () => MembraneOrientationParams);
            await repr.createOrUpdate(params, a.data).runInContext(ctx);
            return new PluginStateObject.Shape.Representation3D({ repr, sourceData: a.data }, { label: 'Membrane Orientation' });
        });
    },
    update({ a, b, newParams }, plugin: PluginContext) {
        return Task.create(TMDET_MEMBRANE_ORIENTATION, async ctx => {
            await MembraneOrientationProvider.attach({ runtime: ctx, assetManager: plugin.managers.asset }, a.data);
            const props = { ...b.data.repr.props, ...newParams };
            await b.data.repr.createOrUpdate(props, a.data).runInContext(ctx);
            b.data.sourceData = a.data;
            return StateTransformer.UpdateResult.Updated;
        });
    },
    isApplicable(a) {
        return MembraneOrientationProvider.isApplicable(a.data);
    }
});

export const MembraneOrientationPreset = StructureRepresentationPresetProvider({
    id: 'preset-membrane-orientation',
    display: {
        name: TMDET_MEMBRANE_ORIENTATION, group: 'Annotation',
        description: 'Shows orientation of membrane layers. Data calculated with TMDET algorithm.' // TODO add ' or obtained via RCSB PDB'
    },
    isApplicable(a) {
        return MembraneOrientationProvider.isApplicable(a.data);
    },
    params: () => StructureRepresentationPresetProvider.CommonParams,
    async apply(ref, params, plugin) {
        const structureCell = StateObjectRef.resolveAndCheck(plugin.state.data, ref);
        const structure  = structureCell?.obj?.data;
        if (!structureCell || !structure) return {};

        if (!MembraneOrientationProvider.get(structure).value) {
            await plugin.runTask(Task.create(TMDET_MEMBRANE_ORIENTATION, async runtime => {
                await MembraneOrientationProvider.attach({ runtime, assetManager: plugin.managers.asset }, structure);
            }));
        }

        const membraneOrientation = await tryCreateMembraneOrientation(plugin, structureCell);
        const colorTheme =  TmDetColorThemeProvider.name as any;
        const preset = await PresetStructureRepresentations.auto.apply(ref, { ...params, theme: { globalName: colorTheme, focus: { name: colorTheme } } }, plugin);

        return { components: preset.components, representations: { ...preset.representations, membraneOrientation } };
    }
});

export function tryCreateMembraneOrientation(plugin: PluginContext, structure: StateObjectRef<PMS>, params?: StateTransformer.Params<MembraneOrientation3DType>, initialState?: Partial<StateTransform.State>) {
    const state = plugin.state.data;
    const membraneOrientation = state.build().to(structure)
        .applyOrUpdateTagged(Tag.Representation, MembraneOrientation3D, params, { state: initialState });
    return membraneOrientation.commit({ revertOnError: true });
}