Issue #2: getData of labeling, ChainList cache

src/extensions/tmdet/behavior.ts

@@ -54,7 +54,7 @@ export const TMDETMembraneOrientation = PluginBehavior.create<{ autoAttach: bool
             this.ctx.query.structure.registry.add(isTransmembrane);
 
             this.ctx.representation.structure.themes.colorThemeRegistry.add(TmDetColorThemeProvider);
-            LabeledResidues.labelProvider!.label = TmLabelProvider.label;
+            //LabeledResidues.labelProvider!.label = TmLabelProvider.label;
             this.ctx.managers.lociLabels.addProvider(LabeledResidues.labelProvider!);
             this.ctx.customModelProperties.register(LabeledResidues.propertyProvider, true);
 

src/extensions/tmdet/labeling.ts

@@ -10,9 +10,9 @@ import { CustomElementProperty } from '../../mol-model-props/common/custom-eleme
 import { Model, ElementIndex } from '../../mol-model/structure';
 import { Loci } from '../../mol-model/loci';
 import { LociLabel, LociLabelProvider } from '../../mol-plugin-state/manager/loci-label';
-import { TmDetDescriptorCache } from './prop';
+import { TmDetChainListCache, TmDetDescriptorCache } from './prop';
 import { createResidueListsPerChain } from './tmdet-color-theme';
-import { ChainList } from './types';
+import { ChainList, getResidue } from './types';
 import { OrderedSet } from '../../mol-data/int';
 
 const siteLabels = [
@@ -32,45 +32,46 @@ export const LabeledResidues = CustomElementProperty.create<number>({
     label: 'TMDet Topology Site Labels',
     name: 'tmdet-topology-site-labels',
     getData(model: Model) {
-        DebugUtil.log('start getData', model.label, model.entryId);
+        const pdbId = model.entryId;
+        DebugUtil.log('start getData', model.label, pdbId);
 
         const map = new Map<ElementIndex, number>();
-        const descriptor = TmDetDescriptorCache.get(model.entryId);
+        const descriptor = TmDetDescriptorCache.get(pdbId);
         if (!descriptor) {
             return { value: map };
         }
 
-        const getSiteId = function(chains: ChainList, residueId: string): number {
-            // TODO: integrate with the coloring version
-            let label = DefaultResidueLabel;
+        // TODO: integrate with the coloring version
+        const getSiteId = function(chains: ChainList, chainId: string, residueId: string): number {
+            let siteId = DefaultResidueLabel;
 
-            for (let chain of chains) {
-                const residue = chain.residues.filter((res) => res.authId === residueId)[0];
-                if (residue) {
-                    label = residue.siteId;
-                    break;
-                }
+            const residue = getResidue(chains, chainId, residueId);
+            if (residue) {
+                siteId = residue.siteId;
             }
-            //const chain = chains.filter((chain) => chain.chainId === chainId)[0];
-            // if (chain) {
-            //     const residue = chain.residues.filter((res) => res.authId === residueId)[0];
-            //     if (residue) {
-            //         label = siteLabels[residue.siteId];
-            //     }
-            // }
 
-            return label;
+            return siteId;
         };
 
-        const chainList = createResidueListsPerChain(descriptor.chains);
+        let chainList =  TmDetChainListCache.get(pdbId);
+        if (!chainList) {
+            chainList = createResidueListsPerChain(descriptor.chains);
+            TmDetChainListCache.set(pdbId, chainList);
+        }
         DebugUtil.log('getData', model.atomicHierarchy.chainAtomSegments);
         // DebugUtil.log('getData: auth_comp_id', model.atomicHierarchy.atoms.auth_comp_id.toArray());
         const residueIndex = model.atomicHierarchy.residueAtomSegments.index;
+        const chainIndex = model.atomicHierarchy.chainAtomSegments.index;
         for (let i = 0, _i = model.atomicHierarchy.atoms._rowCount; i < _i; i++) {
             const residueIdx = residueIndex[i];
-            const authId = model.atomicHierarchy.residues.auth_seq_id.value(residueIdx).toString();
+            const chainIdx = chainIndex[i];
+            const residueAuthId = model.atomicHierarchy.residues.auth_seq_id.value(residueIdx).toString();
+            const residueName = model.atomicHierarchy.atoms.auth_comp_id.value(i).toString();
+            const chainAuthId = model.atomicHierarchy.chains.auth_asym_id.value(chainIdx).toString();
+            const chainLabel = model.atomicHierarchy.chains.label_asym_id.value(chainIdx).toString();
+            //DebugUtil.log(`${i} ${residueIdx} ${residueName} ${residueAuthId} ${chainIdx} ${chainAuthId} ${chainLabel}`);
     
-            map.set(i as ElementIndex, getSiteId(chainList, authId));
+            map.set(i as ElementIndex, getSiteId(chainList, chainAuthId, residueAuthId));
         }
         DebugUtil.log('end of getData');
         return { value: map };
@@ -80,6 +81,17 @@ export const LabeledResidues = CustomElementProperty.create<number>({
     }
 });
 
+// function residueLabel(chains: ChainList, chainId: string, residueId: string): string {
+//     let label = DefaultResidueLabel;
+
+//     const residue = getResidue(chains, chainId, residueId);
+//     if (residue) {
+//         label = siteLabels[residue.siteId];
+//     }
+
+//     return color;
+// }
+
 
 export const TmLabelProvider: LociLabelProvider = {
     label: (loci: Loci): LociLabel => {

src/extensions/tmdet/prop.ts

@@ -16,7 +16,7 @@ import { QuerySymbolRuntime } from '../../mol-script/runtime/query/base';
 import { CustomPropSymbol } from '../../mol-script/language/symbol';
 import { Type } from '../../mol-script/language/type';
 import { membraneOrientation } from './behavior';
-import { PDBTMDescriptor } from './types';
+import { ChainList, PDBTMDescriptor } from './types';
 
 export const MembraneOrientationParams = {
     ...TMDETParams
@@ -53,6 +53,27 @@ class DescriptorCache {
 }
 export const TmDetDescriptorCache = new DescriptorCache();
 
+class ChainListCache {
+    private map: Map<string, ChainList>;
+
+    constructor() {
+        this.map = new Map<string, ChainList>();
+    }
+
+    set(key: string, chains: ChainList) {
+        key = key.toLowerCase();
+        if (!this.map.has(key)) {
+            this.map.set(key, chains);
+        }
+    }
+
+    get(key: string) {
+        key = key.toLowerCase();
+        return this.map.get(key);
+    }
+}
+export const TmDetChainListCache = new ChainListCache();
+
 interface MembraneOrientation {
     // point in membrane boundary
     readonly planePoint1: Vec3,

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

@@ -5,8 +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 { ChainList, PDBTMChain } from './types';
-import { TmDetDescriptorCache } from './prop';
+import { ChainList, getResidue, PDBTMChain } from './types';
+import { TmDetChainListCache, TmDetDescriptorCache } from './prop';
 
 
 export type TmDetColorThemeParams = {
@@ -19,10 +19,11 @@ export function TmDetColorTheme(
 ): ColorTheme<TmDetColorThemeParams> {
     let descriptorChains: PDBTMChain[] = [];
     let pdbtmDescriptor = props.pdbtmDescriptor;
+    const pdbId = ctx.structure?.model.entryId;
+
     // If it is not given as parameter,
     // but there is one in the cache for this structure.
     if (!pdbtmDescriptor) {
-        const pdbId = ctx.structure?.model.entryId;
         if (pdbId) {
             pdbtmDescriptor = TmDetDescriptorCache.get(pdbId);
         }
@@ -31,16 +32,21 @@ export function TmDetColorTheme(
         descriptorChains = pdbtmDescriptor.chains;
     }
 
+    const chainList = createResidueListsPerChain(descriptorChains);
+    if (pdbId) {
+        TmDetChainListCache.set(pdbId, chainList);
+    }
+
     return {
         factory: TmDetColorTheme,
         granularity: 'group', //'residue' as any,
-        color: (location: Location) => getColor(location, descriptorChains),
+        color: (location: Location) => getColor(location, chainList),
         props: props,
         description: 'TMDet...',
     };
 }
 
-const DefaultResidueColor = ColorNames.black;
+const DefaultResidueColor = ColorNames.lightgrey;
 
 //TODO: colors of not curated sites
 const siteColors = [
@@ -55,11 +61,10 @@ const siteColors = [
 ];
 
 
-function getColor(location: Location, chains: PDBTMChain[]): Color {
+function getColor(location: Location, chains: ChainList): Color {
     let color = DefaultResidueColor;
 
     // TODO: How to solve cases when chain operation occurs?
-    const chainList = createResidueListsPerChain(chains);
     if (StructureElement.Location.is(location) && Unit.isAtomic(location.unit)) {
         const atomicHierarchy = location.unit.model.atomicHierarchy;
         const residueIdx = StructureElement.residueIndex(location);
@@ -68,7 +73,7 @@ function getColor(location: Location, chains: PDBTMChain[]): Color {
         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);
+        color = residueColor(chains, chainId, authId);
     }
     return color;
 }
@@ -92,12 +97,9 @@ export function createResidueListsPerChain(chains: PDBTMChain[]) {
 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];
-        }
+    const residue = getResidue(chains, chainId, residueId);
+    if (residue) {
+        color = siteColors[residue.siteId];
     }
 
     return color;

src/extensions/tmdet/types.ts

@@ -58,10 +58,20 @@ type StructureComponentType = StateObjectSelector<
 export type ComponentsType = { polymer: StructureComponentType; ligand: StructureComponentType; water: StructureComponentType };
 
 // for coloring and labeling
+type ResidueItem = {
+    authId: string|undefined,
+    siteId: number
+};
 export type ChainList = {
     chainId: string,
-    residues: {
-        authId: string|undefined,
-        siteId: number
-    }[]
+    residues: ResidueItem[]
 }[];
+export function getResidue(chains: ChainList, chainId: string, residueId: string): ResidueItem|undefined {
+    let result = undefined;
+    const chain = chains.filter((chain) => chain.chainId === chainId)[0];
+    if (chain) {
+        result = chain.residues.filter((res) => res.authId === residueId)[0];
+    }
+
+    return result;
+}