Issue #2: refactoring: transformation and type modules

src/apps/viewer/index.ts

@@ -22,7 +22,7 @@ require('mol-plugin-ui/skin/light.scss');
 
 export { PLUGIN_VERSION as version } from '../../mol-plugin/version';
 export { setDebugMode, setProductionMode } from '../../mol-util/debug';
-export { loadWithUNITMPMembraneRepresentation, red } from '../../extensions/tmdet/behavior';
+export { loadWithUNITMPMembraneRepresentation } from '../../extensions/tmdet/behavior';
 
 const Extensions = {
     'tmdet-membrane-orientation': PluginSpec.Behavior(TMDETMembraneOrientation)

src/extensions/tmdet/behavior.ts

@@ -27,8 +27,8 @@ import { PluginUIContext } from '../../mol-plugin-ui/context';
 import { StateObjectSelector } from "../../mol-state/object";
 import { MEMBRANE_STORAGE_KEY } from '../../extensions/tmdet/algorithm';
 import { Vec3 } from '../../mol-math/linear-algebra';
-
-type PMS = PluginStateObject.Molecule.Structure;
+import { applyTransformations, getAtomGroupExpression } from './transformation';
+import { PDBTMDescriptor, PDBTMRegion, PMS } from './types';
 
 type StructureComponentType = StateObjectSelector<
         PMS,
@@ -133,86 +133,6 @@ export async function loadWithUNITMPMembraneRepresentation(plugin: PluginUIConte
     requestAnimationFrame(() => plugin.canvas3d?.requestCameraReset());
 }
 
-/**
- * Util/test function. Preparation for chain-level operations.
- * Sets color of specified chain to red.
- *
- * @param plugin UI context
- * @param chainId Id of chain to be selected for transformation
- */
-export function red(plugin: PluginUIContext, chainId: string): void {
-    const color: Color = Color.fromArray([255, 0, 0], 0);
-    const query: Expression = getChainExpression(chainId);
-
-    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: 'RED', expression: query }
-        )
-        .apply(
-            StateTransforms.Representation.StructureRepresentation3D,
-            createStructureRepresentationParams(plugin, update.selector.data, {
-                type: 'gaussian-surface',
-                color: 'uniform',
-                colorParams: { value: color }
-            })
-        );
-    update.commit();
-}
-
-
-type PDBTMChainLabel = string;
-type PDBTMChain = {
-    chain_label: PDBTMChainLabel,
-    additional_chain_annotations: {
-        type: string,
-        num_tm: number
-    },
-    residues: {
-        pdb_res_label?: string,
-        aa_type?: string,
-        site_data?: { site_id_ref: number, confidence_classification: string }[]
-    }[]
-};
-type PDBTMRegion = { site: string, auth_ids: number[], color: number[] };
-type PDBTMVec3 = { x: number, y: number, z: number };
-type PDBTMTransformationMatrixRow = { x: number, y: number, z: number, t: number };
-type PDBTMTransformationMatrix = {
-    rowx: PDBTMTransformationMatrixRow,
-    rowy: PDBTMTransformationMatrixRow,
-    rowz: PDBTMTransformationMatrixRow
-};
-type PDBTMDescriptor = {
-    pdb_id: string,
-    chains: PDBTMChain[],
-    site: { site_id: number, label: string }[],
-    additional_entry_annotations: {
-        membrane: {
-            normal: PDBTMVec3,
-            transformation_matrix: PDBTMTransformationMatrix,
-            radius: number
-        },
-        biomatrix: {
-            chain_deletes: string[],
-            matrix_list: {
-                matrix_id: string,
-                apply_to_chain_list: {
-                    chain_id: PDBTMChainLabel,
-                    new_chain_id: PDBTMChainLabel
-                }[],
-                transformation_matrix: PDBTMTransformationMatrix
-            }[]
-        }
-    }
-};
-type TrajectoryType = StateObjectSelector<
-    PluginStateObject.Molecule.Trajectory,
-    StateTransformer<StateObject<any, StateObject.Type<any>>,
-    StateObject<any, StateObject.Type<any>>, any>
->;
-
 function createMembraneOrientation(pdbtmDescriptor: PDBTMDescriptor): MembraneOrientation {
     const normal = pdbtmDescriptor.additional_entry_annotations.membrane.normal;
     const membraneNormal: Vec3 = Vec3.fromObj(normal);
@@ -238,6 +158,9 @@ async function createStructureRepresentation(plugin: PluginUIContext, pdbtmDescr
     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, components);
 
     // TODO: colors of not curated sites
@@ -258,6 +181,7 @@ async function createStructureRepresentation(plugin: PluginUIContext, pdbtmDescr
     pdbtmDescriptor.chains.forEach((chain: any) => {
 
         let regionsBySite: PDBTMRegion[] = siteColors.map((color):PDBTMRegion => {
+            // TODO: set site appropriately
             return { site: "valami", auth_ids: [], color: color };
         });
 
@@ -285,7 +209,8 @@ async function createStructureComponents(plugin: PluginUIContext, structure: Sta
 }
 
 function createRegionRepresentation(plugin: PluginUIContext, chain: string, residue: PDBTMRegion, update: StateBuilder.To<any, any>) {
-    const regionLabel: string = `${chain}: ${residue.auth_ids[0]} | ${residue.site}`;
+    const lastIdIndex: number = residue.auth_ids.length - 1;
+    const regionLabel: string = `${chain}: ${residue.auth_ids[0]}-${residue.auth_ids[lastIdIndex]} | ${residue.site}`;
     const color: Color = Color.fromArray(residue.color, 0);
     const query: Expression = getAtomGroupExpression(chain, residue.auth_ids);
 
@@ -320,88 +245,18 @@ async function loadStructure(ctx: PluginUIContext, params: any, pdbtmDescriptor:
     const trajectory = await builders.structure.parseTrajectory(data, 'mmcif');
 
     // TODO: DEBUG logs
-    console.log('ATOMIC CONFORMATION', trajectory.data!.representative.atomicConformation);
-    console.log('ATOMIC RANGES', trajectory.data!.representative.atomicRanges);
-    console.log('ATOMIC HIER', trajectory.data!.representative.atomicHierarchy);
-    console.log('ATOMIC ENTITIES', trajectory.data!.representative.entities);
-
-
-    //       ctx.managers.structure.hierarchy.getStructuresWithSelection();
-
+    // console.log('ATOMIC CONFORMATION', trajectory.data!.representative.atomicConformation);
+    // console.log('ATOMIC RANGES', trajectory.data!.representative.atomicRanges);
+    // console.log('ATOMIC HIER', trajectory.data!.representative.atomicHierarchy);
+    // console.log('ATOMIC ENTITIES', trajectory.data!.representative.entities);
 
-
-    await applyTransformations(pdbtmDescriptor, trajectory);
+    //ctx.managers.structure.hierarchy.getStructuresWithSelection();
 
     // create membrane representation
     await builders.structure.hierarchy.applyPreset(
         trajectory, 'default', { representationPreset: 'preset-membrane-orientation' as any });
 }
 
-async function applyTransformations(pdbtmDescriptor: PDBTMDescriptor, trajectory: TrajectoryType) {
-    console.log('BIOMX', pdbtmDescriptor.additional_entry_annotations.biomatrix);
-    if (trajectory.data != null) {
-        let atomicConformation = trajectory.data.representative.atomicConformation;
-        let tmatrix = pdbtmDescriptor.additional_entry_annotations.membrane.transformation_matrix;
-        let newConformation = {
-            atomId: [] as number[],
-            x: [] as number[],
-            y: [] as number[],
-            z: [] as number[],
-            B_iso_or_equiv: [] as number[],
-            occupancy: [] as number[]
-        };
-
-        for (let i = 0; i < atomicConformation.x.length; i++) {
-            let coords = {
-                x: atomicConformation.x[i],
-                y: atomicConformation.y[i],
-                z: atomicConformation.z[i]
-            };
-            coords = applyTransformationMatrix(coords, tmatrix);
-            newConformation.x.push(coords.x);
-            newConformation.y.push(coords.y);
-            newConformation.z.push(coords.z);
-
-            // copy old values
-        }
-        atomicConformation.x = newConformation.x;
-        atomicConformation.y = newConformation.y;
-        atomicConformation.z = newConformation.z;
-    }
-}
-
-function applyTransformationMatrix(coords: PDBTMVec3, tmatrix: PDBTMTransformationMatrix): PDBTMVec3 {
-    return {
-        x: tmatrix.rowx.t + vectorMultiply(coords, tmatrix.rowx),
-        y: tmatrix.rowy.t + vectorMultiply(coords, tmatrix.rowy),
-        z: tmatrix.rowz.t + vectorMultiply(coords, tmatrix.rowz)
-    };
-}
-
-function vectorMultiply(v1: PDBTMVec3, v2: PDBTMVec3): number {
-    return v1.x*v2.x + v1.y*v2.y + v1.z*v2.z
-}
-
-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 {
-    // TODO console.log('auth_array:', auth_array);
-    const query: Expression =
-        MS.struct.generator.atomGroups({
-            'chain-test': MS.core.rel.eq([chainId, MS.ammp('label_asym_id')])
-        });
-    return query;
-}
-
-
 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;

src/extensions/tmdet/transformation.ts

@@ -0,0 +1,97 @@
+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 { Color } from '../../mol-util/color';
+import { Expression } from '../../mol-script/language/expression';
+
+export async function applyTransformations(ctx: PluginUIContext, pdbtmDescriptor: PDBTMDescriptor) {
+
+    console.log('BIOMX', pdbtmDescriptor.additional_entry_annotations.biomatrix);
+    pdbtmDescriptor.additional_entry_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;
+            console.log(`before RED ${id} -> ${newId}`);
+            red(ctx, mx.transformation_matrix, id);
+            console.log("after RED");
+        });
+    });
+}
+
+/**
+ * Util/test function. Preparation for chain-level operations.
+ * Sets color of specified chain to red.
+ *
+ * @param plugin UI context
+ * @param transformationMatrix 4x4 matrix describes transformation and translation
+ * @param chainId Id of chain to be selected for transformation
+ */
+ export function red(plugin: PluginUIContext, transformationMatrix: PDBTMTransformationMatrix, chainId: string): void {
+    const color: Color = Color.fromArray([255, 0, 0], 0);
+    const query: Expression = getChainExpression(chainId);
+
+    let 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: 'uniform',
+                colorParams: { value: color }
+            })
+        );
+    update.commit();
+}
+
+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
+    ));
+    console.log('is rot and trans?\n', Mat4.isRotationAndTranslation(mx));
+    console.log('with translation\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 {
+    // TODO console.log('auth_array:', auth_array);
+    const query: Expression =
+        MS.struct.generator.atomGroups({
+            'chain-test': MS.core.rel.eq([chainId, MS.ammp('label_asym_id')])
+        });
+    return query;
+}

src/extensions/tmdet/types.ts

@@ -0,0 +1,54 @@
+import { PluginStateObject } from "../../mol-plugin-state/objects";
+
+export type PMS = PluginStateObject.Molecule.Structure;
+
+type PDBTMChainLabel = string;
+type PDBTMChain = {
+    chain_label: PDBTMChainLabel,
+    additional_chain_annotations: {
+        type: string,
+        num_tm: number
+    },
+    residues: {
+        pdb_res_label?: string,
+        aa_type?: string,
+        site_data?: { site_id_ref: number, confidence_classification: string }[]
+    }[]
+};
+export type PDBTMRegion = { site: string, auth_ids: number[], color: number[] };
+export type PDBTMVec3 = { x: number, y: number, z: number };
+type PDBTMTransformationMatrixRow = { x: number, y: number, z: number, t: number };
+export type PDBTMTransformationMatrix = {
+    rowx: PDBTMTransformationMatrixRow,
+    rowy: PDBTMTransformationMatrixRow,
+    rowz: PDBTMTransformationMatrixRow
+};
+export type PDBTMDescriptor = {
+    pdb_id: string,
+    chains: PDBTMChain[],
+    site: { site_id: number, label: string }[],
+    additional_entry_annotations: {
+        membrane: {
+            normal: PDBTMVec3,
+            transformation_matrix: PDBTMTransformationMatrix,
+            radius: number
+        },
+        biomatrix: {
+            chain_deletes: string[],
+            matrix_list: {
+                matrix_id: string,
+                apply_to_chain_list: {
+                    chain_id: PDBTMChainLabel,
+                    new_chain_id: PDBTMChainLabel
+                }[],
+                transformation_matrix: PDBTMTransformationMatrix
+            }[]
+        }
+    }
+};
+
+// type TrajectoryType = StateObjectSelector<
+//     PluginStateObject.Molecule.Trajectory,
+//     StateTransformer<StateObject<any, StateObject.Type<any>>,
+//     StateObject<any, StateObject.Type<any>>, any>
+// >;