TmMolStar/src/examples/assembly-tm/tmdet-extension/behavior.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>
 */
/**
 * Copyright (C) 2022, Protein Bioinformatics Research Group, RCNS
 *
 * Licensed under CC BY-NC 4.0, see LICENSE file for more info.
 *
 * @author Gabor Tusnady <tusnady.gabor@ttk.hu>
 * @author Csongor Gerdan <gerdan.csongor@ttk.hu>
 */

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


// TMDET imports
import { MembraneOrientationRepresentationProvider, MembraneOrientationParams, MembraneOrientationRepresentation } from './representation';
import { MembraneOrientationProvider, TmDetDescriptorCache, isTransmembrane as tmSymbol, setMembraneOrientation } from './prop';
import { applyTransformations, createMembraneOrientation } from './transformation';
import { ComponentsType, PDBTMDescriptor, PMS } from './types';
import { registerTmDetSymmetry } from './symmetry';
import { TmDetLabelProvider } from './labeling';
import { TmDetColorThemeProvider, updateSiteColors } from './tmdet-color-theme';
//import { loadInitialSnapshot, rotateCamera, storeCameraSnapshot } from './camera';
import { DebugUtil } from './debug-utils';

const TMDET_MEMB_ORI_REPRESENTATION_TAG = 'tmdet-membrane-orientation-3d';
const TMDET_MEMBRANE_ORIENTATION = 'TMDET Membrane Orientation';
export const TMDET_STRUCTURE_PRESET_ID = 'tmdet-preset-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 {
            console.log('TMDET plugin behavior REGISTER');
            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(TmDetLabelProvider);

            this.ctx.genericRepresentationControls.set(TMDET_MEMB_ORI_REPRESENTATION_TAG, 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 }) {
            console.log('TMDET plugin behavior UPDATE');
            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() {
            console.log('TMDET plugin behavior 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(TMDET_MEMB_ORI_REPRESENTATION_TAG);
            this.ctx.builders.structure.representation.unregisterPreset(MembraneOrientationPreset);

            this.ctx.representation.structure.themes.colorThemeRegistry.remove(TmDetColorThemeProvider);
            this.ctx.managers.lociLabels.removeProvider(TmDetLabelProvider);
        }
    },
    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': tmSymbol.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 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

    updateSiteColors(params.side1);

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

        const membraneOrientation = createMembraneOrientation(pdbtmDescriptor);
        setMembraneOrientation(membraneOrientation);

        // 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);

    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);

    let format: "mmcif"|"pdb" = "mmcif";
    // check url to determine file format and compression
    const match: string[] = params.structureUrl.match(new RegExp('/(pdb|cif)(.gz)?$/g/'));
    DebugUtil.log(`format: ${params.format}`);
    if ((params.format != 'mmcif' || params.format != 'pdb') && match != null) {
        if (match[0].startsWith('cif')) {
            format = 'mmcif';
        }
        if (match[0].startsWith('pdb')) {
            format = 'pdb';
        }
    } else if (params.format) {
        format = params.format;
    }

    const builders = ctx.builders;
    const data = await downloadData(ctx, params, pdbtmDescriptor);

    const trajectory = await builders.structure.parseTrajectory(data, format);


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


async function downloadData(ctx: PluginUIContext, params: any, pdbtmDescriptor: PDBTMDescriptor) {

    const builders = ctx.builders;
    let downloadResult = await ctx.runTask(ctx.fetch({ url: params.structureUrl, type: "string" }));
    DebugUtil.log('First 50 chars of input data', downloadResult.slice(0, 50));

    return await builders.data.rawData({
        data: downloadResult,
        label: `${pdbtmDescriptor.pdb_id}`,
    }); // , { state: { isGhost: true } });
}




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




type MembraneOrientation3DType = typeof MembraneOrientation3D
const MembraneOrientation3D = PluginStateTransform.BuiltIn({
    name: TMDET_MEMB_ORI_REPRESENTATION_TAG,
    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: TMDET_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: TMDET_STRUCTURE_PRESET_ID,
    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;
        console.log('MembOriPreset apply params:', params);
        console.log('MembOriPreset ref:', ref);
        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(TMDET_MEMB_ORI_REPRESENTATION_TAG, MembraneOrientation3D, params, { state: initialState });
    return membraneOrientation.commit({ revertOnError: true });
}