Issue #2: TmDetColorThemeProvider with applied pdbtm descriptor

src/apps/tm-viewer/index.html

@@ -47,7 +47,7 @@
             });
             // Set PDB Id here
             var pdbId = '1a0s';
-            // pdbId = '2atk';
+            pdbId = '2atk';
             // pdbId = '1afo';
             tm_molstar.loadWithUNITMPMembraneRepresentation(viewer.plugin, {
                 structureUrl: `https://cs.litemol.org/${pdbId}/full`,

src/apps/tm-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, customColor, customTheme } from '../../extensions/tmdet/behavior';
+export { loadWithUNITMPMembraneRepresentation } from '../../extensions/tmdet/behavior';
 
 const Extensions = {
     'tmdet-membrane-orientation': PluginSpec.Behavior(TMDETMembraneOrientation)

src/extensions/tmdet/behavior.ts

@@ -8,7 +8,7 @@
 import { ParamDefinition as PD } from '../../mol-util/param-definition';
 import { StructureRepresentationPresetProvider, PresetStructureRepresentations } from '../../mol-plugin-state/builder/structure/representation-preset';
 import { MembraneOrientationProvider, MembraneOrientation } from './prop';
-import { StateObject, StateObjectRef, StateObjectCell, StateTransformer, StateTransform, StateBuilder } from '../../mol-state';
+import { StateObject, StateObjectRef, StateObjectCell, StateTransformer, StateTransform } from '../../mol-state';
 import { Task } from '../../mol-task';
 import { PluginBehavior } from '../../mol-plugin/behavior';
 import { MembraneOrientationRepresentationProvider, MembraneOrientationParams, MembraneOrientationRepresentation } from './representation';
@@ -19,24 +19,13 @@ import { DefaultQueryRuntimeTable } from '../../mol-script/runtime/query/compile
 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 { Expression } from '../../mol-script/language/expression';
-import { StateTransforms } from '../../mol-plugin-state/transforms';
-import { Color } from '../../mol-util/color';
-import { createStructureRepresentationParams } from '../../mol-plugin-state/helpers/structure-representation-params';
 import { PluginUIContext } from '../../mol-plugin-ui/context';
-import { StateObjectSelector } from "../../mol-state/object";
-import { applyTransformations, createMembraneOrientation, getAtomGroupExpression } from './transformation';
-import { PDBTMDescriptor, PDBTMRegion, PMS } from './types';
+import { applyTransformations, createMembraneOrientation } from './transformation';
+import { ComponentsType, PDBTMDescriptor, PMS } from './types';
 import { registerTmDetSymmetry } from './symmetry';
 import { StripedResidues } from './coloring';
 import { TmDetColorThemeProvider } from './tmdet-color-theme';
 
-type StructureComponentType = StateObjectSelector<
-        PMS,
-        StateTransformer<StateObject<any, StateObject.Type<any>>,
-        StateObject<any, StateObject.Type<any>>, any>
-    > | undefined;
-
 const Tag = MembraneOrientation.Tag;
 
 export const TMDETMembraneOrientation = PluginBehavior.create<{ autoAttach: boolean }>({
@@ -128,46 +117,15 @@ export async function loadWithUNITMPMembraneRepresentation(plugin: PluginUIConte
 
     // load structure
     await loadStructure(plugin, params, pdbtmDescriptor);
-
+    // cartoon, colors etc.
     await createStructureRepresentation(plugin, pdbtmDescriptor);
 
-    await customTheme(plugin);
-
     //
     // reset the camera because the membranes render 1st and the structure might not be fully visible
     //
     requestAnimationFrame(() => plugin.canvas3d?.requestCameraReset());
 }
 
-export async function customColor(plugin: PluginUIContext) {
-    // color theme code:
-    plugin.dataTransaction(async () => {
-        console.log("data transaction...");
-        for (const s of plugin.managers.structure.hierarchy.current.structures) {
-            console.log("structure.components:", s.components);
-            await plugin.managers.structure.component.updateRepresentationsTheme(s.components, { color: StripedResidues.propertyProvider.descriptor.name as any });
-        }
-    });
-}
-
-export async function customTheme(plugin: PluginUIContext) {
-    plugin.dataTransaction(async () => {
-        console.log("data transaction...");
-        for (const s of plugin.managers.structure.hierarchy.current.structures) {
-            console.log("structure.components:", s.components);
-            await plugin.managers.structure.component.updateRepresentationsTheme(
-                s.components,
-                {
-                    color: TmDetColorThemeProvider.name as any,
-                    colorParams: { pdbtmDescriptor: 3 }
-                }
-            );
-        }
-    });
-}
-
-
-
 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;
@@ -182,47 +140,7 @@ async function createStructureRepresentation(plugin: PluginUIContext, pdbtmDescr
 
     await applyTransformations(plugin, pdbtmDescriptor);
 
-    await buildStructureRepresentation(plugin, components);
-
-    // TODO: colors of not curated sites
-    const siteColors = [
-        [255,100,100], // Side1
-        [100,100,255], // Side2
-        [255,255,0],   // TM alpha
-        [255,255,0],   // TM beta
-        [255,127,0],   // TM re-entrant loop
-        [0,255,  0],   // Interfacial Helix
-        [196,196,196], // Unknow localization
-        [0,255,  0],   // Membrane Inside
-    ];
-    // const siteDefinitions = pdbtmDescriptor.sites.map(
-    //     (siteDefinition: any) => siteDefinition.label
-    // );
-
-    pdbtmDescriptor.chains.forEach((chain: any) => {
-
-        let regionsBySite: PDBTMRegion[] = siteColors.map((color):PDBTMRegion => {
-            // TODO: set 'site' appropriately
-            return { site: "valami", auth_ids: [], color: color };
-        });
-
-        for (let residueItem of chain.residues) {
-            const siteIndex = residueItem.site_data[0].site_id_ref - 1;
-            regionsBySite[siteIndex].auth_ids.push(Number(residueItem.pdb_res_label));
-        }
-
-        console.log(`REGIONS-${chain.chain_label}`, regionsBySite);
-        regionsBySite.forEach(async (region: PDBTMRegion) => {
-            if (region.auth_ids.length == 0) { return; }
-
-            console.log("auth_ids: ", region.auth_ids);
-            console.log("color: ", region.color);
-            const regionUpdates = plugin.build();
-            await createRegionRepresentation(plugin, chain.chain_label, region, regionUpdates.to(structure));
-            regionUpdates.commit();
-        });
-
-    });
+    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>>>) {
@@ -233,27 +151,16 @@ async function createStructureComponents(plugin: PluginUIContext, structure: Sta
     };
 }
 
-async function createRegionRepresentation(plugin: PluginUIContext, chain: string, residue: PDBTMRegion, update: StateBuilder.To<any, any>) {
-    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);
-    console.log("regionLabel:", regionLabel);
-
-    // based on https://github.com/molstar/molstar/issues/209
-    update
-        .apply(StateTransforms.Model.StructureSelectionFromExpression, { label: regionLabel, expression: query })
-        .apply(StateTransforms.Representation.StructureRepresentation3D, createStructureRepresentationParams(plugin, update.selector.data, {
-            color: 'uniform',
-            colorParams: { value: color }
-        }));
-}
-
-async function buildStructureRepresentation(plugin: PluginUIContext, components: { polymer: StructureComponentType; ligand: StructureComponentType; water: StructureComponentType }) {
+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', typeParams: { alpha: 0.5 } }, { tag: '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)
@@ -274,14 +181,6 @@ async function loadStructure(ctx: PluginUIContext, params: any, pdbtmDescriptor:
     }); // , { state: { isGhost: true } });
     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();
-
     // create membrane representation
     await builders.structure.hierarchy.applyPreset(
         trajectory, 'default', { representationPreset: 'preset-membrane-orientation' as any });

src/extensions/tmdet/tmdet-color-theme.ts

@@ -5,7 +5,8 @@ import { Color } from '../../mol-util/color';
 import { ColorNames } from '../../mol-util/color/names';
 import { ParamDefinition as PD } from '../../mol-util/param-definition';
 import { Location } from '../../mol-model/location';
-import { PDBTMDescriptor } from './types';
+import { PDBTMChain } from './types';
+
 
 export type TmDetColorThemeParams = {
     pdbtmDescriptor: any
@@ -16,35 +17,92 @@ export function TmDetColorTheme(
     props: PD.Values<TmDetColorThemeParams>
 ): ColorTheme<TmDetColorThemeParams> {
     console.log('ColorTheme CTX', ctx);
-    console.log('ColorTheme PROPS', props);
+    const descriptorChains = props.pdbtmDescriptor.chains as PDBTMChain[];
+    console.log('ColorTheme PROPS', descriptorChains);
 
     return {
         factory: TmDetColorTheme,
         granularity: 'group', //'residue' as any,
-        color: (location: Location) => getColor(location, props),
+        color: (location: Location) => getColor(location, descriptorChains),
         props: props,
         description: 'TMDet...',
     };
 }
 
 const DefaultResidueColor = ColorNames.black;
-const ColorMap = [
-    Color(0xff0000),
-    Color(0x0000ff),
-    Color(0xffff00)
+
+//TODO: colors of not curated sites
+const siteColors = [
+    Color.fromArray([255,100,100], 0), // Side1
+    Color.fromArray([100,100,255], 0), // Side2
+    Color.fromArray([255,255,  0], 0), // TM alpha
+    Color.fromArray([255,255,  0], 0), // TM beta
+    Color.fromArray([255,127,  0], 0), // TM re-entrant loop
+    Color.fromArray([0,255,    0], 0), // Interfacial Helix
+    Color.fromArray([196,196,196], 0), // Unknow localization
+    Color.fromArray([0,255,    0], 0)  // Membrane Inside
 ];
 
-function getColor(location: Location, props: any): Color {
+
+function getColor(location: Location, chains: PDBTMChain[]): Color {
     let color = DefaultResidueColor;
 
-    console.log('getColor params:', props);
+    // TODO: How to solve cases when chain operation occurs?
+    const chainList = createResidueListsPerChain(chains);
     if (StructureElement.Location.is(location) && Unit.isAtomic(location.unit)) {
-        const compId = StructureElement.residueIndex(location) % 3;
-        color = ColorMap[compId]
+        const atomicHierarchy = location.unit.model.atomicHierarchy;
+        const residueIdx = StructureElement.residueIndex(location);
+        const chainIdx = StructureElement.Location.chainIndex(location);
+
+        const authId = atomicHierarchy.residues.auth_seq_id.value(residueIdx).toString();
+        const chainId = atomicHierarchy.chains.auth_asym_id.value(chainIdx).toString();
+
+        color = residueColor(chainList, chainId, authId);
+        console.log(`chain: ${chainId} res: ${authId}`);
+    }
+    return color;
+}
+
+type ChainList = {
+    chainId: string,
+    residues: {
+        authId: string|undefined,
+        siteId: number
+    }[]
+}[];
+
+function createResidueListsPerChain(chains: PDBTMChain[]) {
+    const chainList: ChainList = [];
+
+    chains.forEach((chain: PDBTMChain) => {
+        chainList.push({ chainId: chain.chain_label, residues: [] });
+        chain.residues.forEach((residue) => {
+            chainList[chainList.length - 1].residues.push({
+                authId: residue.pdb_res_label,
+                siteId: residue.site_data![0].site_id_ref - 1
+            });
+        });
+    });
+
+    return chainList;
+}
+
+function residueColor(chains: ChainList, chainId: string, residueId: string): Color {
+    let color = DefaultResidueColor;
+
+    const chain = chains.filter((chain) => chain.chainId === chainId)[0];
+    if (chain) {
+        const residue = chain.residues.filter((res) => res.authId === residueId)[0];
+        if (residue) {
+            color = siteColors[residue.siteId];
+        }
     }
+
     return color;
 }
 
+// Provider
+
 export const TmDetColorThemeProvider: ColorTheme.Provider<TmDetColorThemeParams, 'tmdet-custom-color-theme'> = {
     name: 'tmdet-custom-color-theme',
     label: 'TMDet Topology Theme',

src/extensions/tmdet/transformation.ts

@@ -7,6 +7,7 @@ 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(ctx: PluginUIContext, pdbtmDescriptor: PDBTMDescriptor) {
     const annotations = pdbtmDescriptor.additional_entry_annotations;
@@ -21,21 +22,21 @@ export async function applyTransformations(ctx: PluginUIContext, pdbtmDescriptor
             let id  = chainPair.chain_id;
             let newId = chainPair.new_chain_id;
             console.log(`before RED ${id} -> ${newId}`);
-            red(ctx, mx.transformation_matrix, id);
+            chainTransformation(ctx, pdbtmDescriptor, 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 pdbtmDescriptor TM-specific information
  * @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 {
+ export function chainTransformation(plugin: PluginUIContext, pdbtmDescriptor: PDBTMDescriptor, transformationMatrix: PDBTMTransformationMatrix, chainId: string): void {
     // TODO: const color: Color = Color.fromArray([255, 0, 0], 0);
     const query: Expression = getChainExpression(chainId);
 
@@ -55,9 +56,7 @@ export async function applyTransformations(ctx: PluginUIContext, pdbtmDescriptor
             StateTransforms.Representation.StructureRepresentation3D,
             createStructureRepresentationParams(plugin, update.selector.data, {
                 type: 'cartoon',
-                color: 'chain-id',
-                // TODO: Can we copy colors of parent asymmetry?
-                // colorParams: { value: color }
+                color: TmDetColorThemeProvider.name as any, colorParams: { pdbtmDescriptor }
             })
         );
     update.commit();

src/extensions/tmdet/types.ts

@@ -1,9 +1,11 @@
 import { PluginStateObject } from "../../mol-plugin-state/objects";
+import { StateObject, StateTransformer } from '../../mol-state';
+import { StateObjectSelector } from "../../mol-state/object";
 
 export type PMS = PluginStateObject.Molecule.Structure;
 
 type PDBTMChainLabel = string;
-type PDBTMChain = {
+export type PDBTMChain = {
     chain_label: PDBTMChainLabel,
     additional_chain_annotations: {
         type: string,
@@ -47,8 +49,10 @@ export type PDBTMDescriptor = {
     }
 };
 
-// type TrajectoryType = StateObjectSelector<
-//     PluginStateObject.Molecule.Trajectory,
-//     StateTransformer<StateObject<any, StateObject.Type<any>>,
-//     StateObject<any, StateObject.Type<any>>, any>
-// >;
+type StructureComponentType = StateObjectSelector<
+        PMS,
+        StateTransformer<StateObject<any, StateObject.Type<any>>,
+        StateObject<any, StateObject.Type<any>>, any>
+    > | undefined;
+
+export type ComponentsType = { polymer: StructureComponentType; ligand: StructureComponentType; water: StructureComponentType };