TmMolStar/src/examples/assembly-tm/tmdet-extension/transformation.ts

/**
 * 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 { MolScriptBuilder as MS } from 'molstar/lib/mol-script/language/builder';
import { Mat4, Quat, Vec3 } from 'molstar/lib/mol-math/linear-algebra';
import { PDBTMDescriptor, PDBTMTransformationMatrix, PMS } from './types';
import { createStructureRepresentationParams } from 'molstar/lib/mol-plugin-state/helpers/structure-representation-params';
import { StateTransforms } from 'molstar/lib/mol-plugin-state/transforms';
import { StateObjectRef, StateBuilder } from 'molstar/lib/mol-state';
import { Expression } from 'molstar/lib/mol-script/language/expression';
import { MembraneOrientation } from './prop';
import { TmDetColorThemeProvider } from './tmdet-color-theme';
import { DebugUtil } from './debug-utils';
import { PluginContext } from 'molstar/lib/mol-plugin/context';

export function applyTransformations(plugin: PluginContext, pdbtmDescriptor: PDBTMDescriptor) {
    const annotations = pdbtmDescriptor.additional_entry_annotations;

    const membraneTransformation = transformationForStateTransform(pdbtmDescriptor.additional_entry_annotations.membrane.transformation_matrix);

    if (annotations?.biomatrix?.matrix_list) {
        annotations.biomatrix.matrix_list.forEach(function(mx) {
            mx.apply_to_chain_list.forEach(function(chainPair) {
                let id  = chainPair.chain_id;
                let newId = chainPair.new_chain_id;
                if (annotations.biomatrix.chain_deletes?.includes(newId)) {
                    DebugUtil.log(`${id} -> ${newId} transformation skipped due to delete rule`);
                    return;
                }
                const mtx = transformationForStateTransform(mx.transformation_matrix);
                const composedTransformation = Mat4.mul(Mat4(), membraneTransformation, mtx);
                chainTransformation(plugin, composedTransformation, id, newId);
                // await plugin.runTask(Task.create(`TMDET: Transform '${id}' into '${newId}'`, async () => {
                //     chainTransformation(plugin, mx.transformation_matrix, id, newId);
                // }));
            });
        });
    }

    // WARNING: the components cannot be accessed here (created by chainTransformations in the above loop)
    //          Maybe due to some kind of synchronization behavior.
    //          plugin.runTask with "await" also cannot synchronize here.
    //
    // So this function call transforms only already existing components due to a side effect of another issue.
    transformWholeModel(plugin, pdbtmDescriptor.additional_entry_annotations.membrane.transformation_matrix);
    // await plugin.runTask(Task.create(`TMDET: Apply membrane transformation`, async () => {
    //     transformWholeModel(plugin, pdbtmDescriptor.additional_entry_annotations.membrane.transformation_matrix);
    // }));

}

export function transformWholeModel(plugin: PluginContext, membraneMatrix: PDBTMTransformationMatrix) {
    //
    // membrane transformation
    //
    const membraneTransformation = transformationForStateTransform(membraneMatrix);

    // transform each component
    DebugUtil.log("STRUCTURES", getCurrentHierarchy(plugin).structures);
    getCurrentHierarchy(plugin).structures[0].components.forEach(component => {
        const structure: StateObjectRef<PMS> = component.cell;
        const update: StateBuilder.To<any, any> = plugin.build().to(structure);

        const label = component.cell.obj!.label.toUpperCase();
        DebugUtil.log("memb.transform of", label, component.cell.obj?.id);
        DebugUtil.log(`${label} component.cell.obj:`, component.cell.obj);
        // DebugUtil.log(`${label} component:`, component);

        update
            .apply(StateTransforms.Model.TransformStructureConformation, {
                "transform": { name: "matrix", params: { data: membraneTransformation, transpose: false } }
            });
        update.commit();
    });
}

/**
 * Perform transformation on a chain.
 *
 * @param plugin UI context
 * @param transformationMatrix 4x4 matrix describes transformation and translation
 * @param chainId Id of chain to be selected for transformation
 */
export function chainTransformation(plugin: PluginContext, transformationMatrix: Mat4, chainId: string, newId: string): void {
    const query: Expression = getChainExpression(chainId);

//    const transformation = transformationForStateTransform(transformationMatrix);
    const transformation = transformationMatrix;
    const structure: StateObjectRef<PMS> = plugin.managers.structure.hierarchy.current.models[0].structures[0].cell;
    const update: StateBuilder.To<any, any> = plugin.build().to(structure);

    update
        .apply(
            StateTransforms.Model.StructureSelectionFromExpression,
            { label: newId, expression: query }
        )
        .apply(StateTransforms.Model.TransformStructureConformation, {
            "transform": { name: "matrix", params: { data: transformation, transpose: false } }
        })
        .apply(
            StateTransforms.Representation.StructureRepresentation3D,
            createStructureRepresentationParams(plugin, structure.obj?.data, {
                type: 'cartoon',
                color: TmDetColorThemeProvider.name as any //, colorParams: { pdbtmDescriptor }
            })
        );
    update.commit();
    DebugUtil.log(`${chainId}->${newId} DONE`);
}

function vneg(u: Vec3): Vec3 {
    return Vec3.negate(Vec3.zero(), u);
}

export function createMembraneOrientation(pdbtmDescriptor: PDBTMDescriptor): MembraneOrientation {
    const membrane = pdbtmDescriptor.additional_entry_annotations.membrane;

    let membraneNormal: Vec3 = Vec3.fromObj(membrane.normal);
    const result: MembraneOrientation = {
        planePoint1: membraneNormal,
        planePoint2: vneg(membraneNormal),
        centroid: Vec3.zero(),
        normalVector: membraneNormal,
        radius: membrane.radius
    };

    return result;
}

export function transformationForStateTransform(tmatrix: PDBTMTransformationMatrix): Mat4 {
    // matrix expected in column-major order
    const mx: Mat4 = Mat4.fromArray(Mat4.zero(),
        [
            tmatrix.rowx.x, tmatrix.rowy.x, tmatrix.rowz.x, 0,
            tmatrix.rowx.y, tmatrix.rowy.y, tmatrix.rowz.y, 0,
            tmatrix.rowx.z, tmatrix.rowy.z, tmatrix.rowz.z, 0,
                         0,              0,              0, 1
        ], 0
    );
    Mat4.setTranslation(mx, Vec3.create(
        tmatrix.rowx.t,
        tmatrix.rowy.t,
        tmatrix.rowz.t
    ));
    // TODO: DebugUtil.log('MAT4\n', Mat4.makeTable(mx));
    return mx;
}

export function getAtomGroupExpression(chainId: string, auth_array: number[]): Expression {
    // TODO DebugUtil.log('auth_array:', auth_array);
    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;
}

export function getChainExpression(chainId: string): Expression {
    const query: Expression =
        MS.struct.generator.atomGroups({
            'chain-test': MS.core.rel.eq([chainId, MS.ammp('label_asym_id')])
        });
    return query;
}

export function getCurrentHierarchy(plugin: PluginContext) {
    return plugin.managers.structure.hierarchy.current;
}