TmMolStar/src/extensions/tmdet/transformation.ts

import { MolScriptBuilder as MS } from '../../mol-script/language/builder';
import { PluginUIContext } from '../../mol-plugin-ui/context';
import { Mat4, Vec3 } from '../../mol-math/linear-algebra';
import { PDBTMDescriptor, PDBTMTransformationMatrix, PMS } from './types';
import { createStructureRepresentationParams } from '../../mol-plugin-state/helpers/structure-representation-params';
import { StateTransforms } from '../../mol-plugin-state/transforms';
import { StateObjectRef, StateBuilder } from '../../mol-state';
import { Expression } from '../../mol-script/language/expression';
import { MembraneOrientation } from './prop';
import { TmDetColorThemeProvider } from './tmdet-color-theme';

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

    console.log('BIOMX', annotations.biomatrix);
    if (annotations.biomatrix) {
        annotations.biomatrix.matrix_list.forEach(function(mx) {
            mx.apply_to_chain_list.forEach(function(chainPair) {
                let id  = chainPair.chain_id;
                chainTransformation(plugin, pdbtmDescriptor, mx.transformation_matrix, id);
            });
        });
    }

    // membrane transformation
    const membraneTransformation = transformationForStateTransform(pdbtmDescriptor.additional_entry_annotations.membrane.transformation_matrix);
    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.TransformStructureConformation, {
            "transform": { name: "matrix", params: { data: membraneTransformation, transpose: false } }
        })
        .apply(
            StateTransforms.Representation.StructureRepresentation3D,
            createStructureRepresentationParams(plugin, structure.obj?.data, {
                type: 'cartoon',
                color: TmDetColorThemeProvider.name as any, colorParams: { pdbtmDescriptor }
            })
        );
    update.commit();

}

/**
 * Sets color of specified chain to red.
 *
 * @param plugin UI context
 * @param pdbtmDescriptor TM-specific information
 * @param transformationMatrix 4x4 matrix describes transformation and translation
 * @param chainId Id of chain to be selected for transformation
 */
 export function chainTransformation(plugin: PluginUIContext, pdbtmDescriptor: PDBTMDescriptor, transformationMatrix: PDBTMTransformationMatrix, chainId: string): void {
    const query: Expression = getChainExpression(chainId);

    const transformation = transformationForStateTransform(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: chainId, expression: query }
        )
        .apply(StateTransforms.Model.TransformStructureConformation, {
            "transform": { name: "matrix", params: { data: transformation, transpose: false } }
        })
        .apply(
            StateTransforms.Representation.StructureRepresentation3D,
            createStructureRepresentationParams(plugin, update.selector.data, {
                type: 'cartoon',
                color: TmDetColorThemeProvider.name as any, colorParams: { pdbtmDescriptor }
            })
        );
    update.commit();
}

function vadd(a: Vec3, b: Vec3): Vec3 {
    return Vec3.add(Vec3.zero(), a, b);
}

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

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

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

        // (NOTE: the TMDET extension calculates and sets it during applying preset)
        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: console.log('MAT4\n', Mat4.makeTable(mx));
    return mx;
}

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