Issue #2: Behavior refactoring: update colorization

src/apps/viewer/index.html

@@ -59,7 +59,7 @@
             function loadPdb() {
                 molstar.loadWithUNITMPMembraneRepresentation(viewer.plugin, {
                     structureUrl: `https://cs.litemol.org/${pdbId}/full`,
-                    regionDescriptorUrl: `http://localhost:8080/${pdbId}_regions.php`,
+                    regionDescriptorUrl: `http://localhost:8080/build/data/${pdbId}.json`,
                 });
             }
         </script>

src/extensions/tmdet/behavior.ts

@@ -28,6 +28,13 @@ 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;
+
+type StructureComponentType = StateObjectSelector<
+        PMS,
+        StateTransformer<StateObject<any, StateObject.Type<any>>,
+        StateObject<any, StateObject.Type<any>>, any>
+    > | undefined;
 
 const Tag = MembraneOrientation.Tag;
 
@@ -109,18 +116,18 @@ let membrane: MembraneOrientation;
 
 
 export async function loadWithUNITMPMembraneRepresentation(plugin: PluginUIContext, params: any) {
-    const regionDescriptors: any = await downloadRegionDescriptor(plugin, params);
+    const pdbtmDescriptor: any = await downloadRegionDescriptor(plugin, params);
 
-    membrane = createMembraneOrientation(regionDescriptors);
+    membrane = createMembraneOrientation(pdbtmDescriptor);
 
     console.log('DEBUG-02', membrane.planePoint2);
 
     localStorage.setItem(MEMBRANE_STORAGE_KEY, JSON.stringify(membrane));
 
     // load structure
-    await loadStructure(plugin, params, regionDescriptors);
+    await loadStructure(plugin, params, pdbtmDescriptor);
 
-    await createStructureRepresentation(plugin, regionDescriptors);
+    await createStructureRepresentation(plugin, pdbtmDescriptor);
 
     //
     // reset the camera because the membranes render 1st and the structure might not be fully visible
@@ -128,9 +135,11 @@ export async function loadWithUNITMPMembraneRepresentation(plugin: PluginUIConte
     requestAnimationFrame(() => plugin.canvas3d?.requestCameraReset());
 }
 
-function createMembraneOrientation(regionDescriptors: any): MembraneOrientation {
+type PDBTMDescriptor = any;
+
+function createMembraneOrientation(pdbtmDescriptor: PDBTMDescriptor): MembraneOrientation {
     const membraneNormal: Vec3 = Vec3.fromObj(
-        regionDescriptors['membrane-normal']
+        pdbtmDescriptor.additional_entry_annotations.membrane.normal
     );
     membraneNormal[0] = membraneNormal[2];
     membraneNormal[2] = 0;
@@ -144,32 +153,52 @@ function createMembraneOrientation(regionDescriptors: any): MembraneOrientation
         normalVector: membraneNormal,
 
         // (NOTE: the TMDET extension calculates and sets it during applying preset)
-        radius: regionDescriptors['radius']
+        radius: pdbtmDescriptor['radius']
     };
     membraneOrientation.planePoint2[0] *= -1;
     return membraneOrientation;
 }
 
-async function createStructureRepresentation(plugin: PluginUIContext, regionDescriptors: any) {
+async function createStructureRepresentation(plugin: PluginUIContext, pdbtmDescriptor: any) {
     // get the first structure of the first model
-    const structure: StateObjectRef<PluginStateObject.Molecule.Structure> = plugin.managers.structure.hierarchy.current.models[0].structures[0].cell;
+    const structure: StateObjectRef<PMS> = plugin.managers.structure.hierarchy.current.models[0].structures[0].cell;
     const components = await createStructureComponents(plugin, structure);
 
     await buildStructureRepresentation(plugin, components);
 
-    regionDescriptors.chains.forEach((chain: any) => {
+    const siteColors = [
+        [255, 0, 0], // Side1
+        [255, 0, 0], // Side2
+        [255,255,0], // TM alpha
+        [255,255,0], // TM beta
+        [255,255,0], // TM re-entrant loop
+        [0,0,  255], // Interfacial Helix
+        [0,0,  255], // Unknow localization
+        [0,0,  255], // Membrane Inside
+    ];
+    const siteDefinitions = pdbtmDescriptor.sites.map(
+        (siteDefinition: any) => siteDefinition.label
+    );
+    pdbtmDescriptor.chains.forEach((chain: any) => {
 
-        for (let regionKey in chain.regions) {
-            const regionUpdates = plugin.build();
-            const region = chain.regions[regionKey];
-            createRegionRepresentation(plugin, chain.chain_id, region, regionUpdates.to(structure));
-            regionUpdates.commit();
-        }
+        // console.log('PDBTM descriptor', chain);
 
+
+        const regionUpdates = plugin.build();
+        for (let residueItem of chain.residues) {
+            const siteIndex = residueItem.site_data[0].site_id_ref - 1;
+            const residue = {
+                "auth_ids": [ Number(residueItem.pdb_res_label) ],
+                "color": siteColors[siteIndex],
+                "site": siteDefinitions[siteIndex]
+            };
+            createRegionRepresentation(plugin, chain.chain_label, residue, regionUpdates.to(structure));
+        }
+        regionUpdates.commit();
     });
 }
 
-async function createStructureComponents(plugin: PluginUIContext, structure: StateObjectCell<PluginStateObject.Molecule.Structure, StateTransform<StateTransformer<StateObject<any, StateObject.Type<any>>, StateObject<any, StateObject.Type<any>>, any>>>) {
+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'),
@@ -177,10 +206,12 @@ async function createStructureComponents(plugin: PluginUIContext, structure: Sta
     };
 }
 
-function createRegionRepresentation(plugin: PluginUIContext, chain: string, region: any, update: StateBuilder.To<any, any>) {
-    const regionLabel: string = `${chain} | ${region.name}`;
-    const color: Color = Color.fromArray(region.color, 0);
-    const query: Expression = getQuery(chain, region.auth_ids as number[]);
+function createRegionRepresentation(plugin: PluginUIContext, chain: string, residue: any, update: StateBuilder.To<any, any>) {
+    console.log('RESI', residue);
+    const regionLabel: string = `${chain} | ${residue.site}`;
+    const color: Color = Color.fromArray(residue.color, 0);
+    const query: Expression = getAtomGroupExpression(chain, residue.auth_ids as number[]);
+
 
     // based on https://github.com/molstar/molstar/issues/209
     update
@@ -191,12 +222,6 @@ function createRegionRepresentation(plugin: PluginUIContext, chain: string, regi
         }));
 }
 
-type StructureComponentType = StateObjectSelector<
-        PluginStateObject.Molecule.Structure,
-        StateTransformer<StateObject<any, StateObject.Type<any>>,
-        StateObject<any, StateObject.Type<any>>, any>
-    > | undefined;
-
 async function buildStructureRepresentation(plugin: PluginUIContext, components: { polymer: StructureComponentType; ligand: StructureComponentType; water: StructureComponentType }) {
     const builder = plugin.builders.structure.representation;
     const update = plugin.build();
@@ -209,11 +234,11 @@ async function buildStructureRepresentation(plugin: PluginUIContext, components:
     await update.commit();
 }
 
-async function loadStructure(ctx: PluginUIContext, params: any, regionDescriptors: any): Promise<void> {
+async function loadStructure(ctx: PluginUIContext, params: any, pdbtmDescriptor: any): Promise<void> {
     const builders = ctx.builders;
     const data = await builders.data.download({
         url: params.structureUrl,
-        label: `UniTMP: ${regionDescriptors['pdb-id']}`,
+        label: `UniTMP: ${pdbtmDescriptor['pdb-id']}`,
         isBinary: false
     }); // , { state: { isGhost: true } });
     const trajectory = await builders.structure.parseTrajectory(data, 'mmcif');
@@ -222,7 +247,7 @@ async function loadStructure(ctx: PluginUIContext, params: any, regionDescriptor
         trajectory, 'default', { representationPreset: 'preset-membrane-orientation' as any });
 }
 
-function getQuery(chainId: string, auth_array: number[]): Expression {
+function getAtomGroupExpression(chainId: string, auth_array: number[]): Expression {
     const query: Expression =
         MS.struct.generator.atomGroups({
             'residue-test': MS.core.set.has([MS.set( ...auth_array ), MS.ammp('auth_seq_id')]),
@@ -234,8 +259,8 @@ function getQuery(chainId: string, auth_array: number[]): Expression {
 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 regionDescriptors: any = JSON.parse(downloadResult);
-    return regionDescriptors;
+    const pdbtmDescriptor: any = JSON.parse(downloadResult);
+    return pdbtmDescriptor;
 }
 
 // //////////////////////////// END OF TMDET VIEWER SECTION
@@ -313,7 +338,7 @@ export const MembraneOrientationPreset = StructureRepresentationPresetProvider({
     }
 });
 
-export function tryCreateMembraneOrientation(plugin: PluginContext, structure: StateObjectRef<PluginStateObject.Molecule.Structure>, params?: StateTransformer.Params<MembraneOrientation3D>, initialState?: Partial<StateTransform.State>) {
+export function tryCreateMembraneOrientation(plugin: PluginContext, structure: StateObjectRef<PMS>, params?: StateTransformer.Params<MembraneOrientation3D>, initialState?: Partial<StateTransform.State>) {
     const state = plugin.state.data;
     const membraneOrientation = state.build().to(structure)
         .applyOrUpdateTagged('membrane-orientation-3d', MembraneOrientation3D, params, { state: initialState });