Issue #2: change to loci label provider

src/extensions/tmdet/behavior.ts

@@ -26,7 +26,7 @@ import { MembraneOrientationProvider, MembraneOrientation, TmDetDescriptorCache
 import { applyTransformations, createMembraneOrientation } from './transformation';
 import { ComponentsType, PDBTMDescriptor, PMS } from './types';
 import { registerTmDetSymmetry } from './symmetry';
-import { LabeledResidues } from './labeling';
+import { TmDetLabelProvider } from './labeling';
 import { TmDetColorThemeProvider } from './tmdet-color-theme';
 import { loadInitialSnapshot, rotateCamera, storeCameraSnapshot } from './camera';
 
@@ -54,9 +54,7 @@ export const TMDETMembraneOrientation = PluginBehavior.create<{ autoAttach: bool
             this.ctx.query.structure.registry.add(isTransmembrane);
 
             this.ctx.representation.structure.themes.colorThemeRegistry.add(TmDetColorThemeProvider);
-            this.ctx.managers.lociLabels.addProvider(LabeledResidues.labelProvider!);
-            this.ctx.customModelProperties.register(LabeledResidues.propertyProvider, true);
-
+            this.ctx.managers.lociLabels.addProvider(TmDetLabelProvider);
 
             this.ctx.genericRepresentationControls.set(Tag.Representation, selection => {
                 const refs: GenericRepresentationRef[] = [];
@@ -86,6 +84,9 @@ export const TMDETMembraneOrientation = PluginBehavior.create<{ autoAttach: bool
 
             this.ctx.genericRepresentationControls.delete(Tag.Representation);
             this.ctx.builders.structure.representation.unregisterPreset(MembraneOrientationPreset);
+
+            this.ctx.representation.structure.themes.colorThemeRegistry.remove(TmDetColorThemeProvider);
+            this.ctx.managers.lociLabels.removeProvider(TmDetLabelProvider);
         }
     },
     params: () => ({

src/extensions/tmdet/labeling.ts

@@ -1,16 +1,9 @@
-/**
- * Copyright (c) 2019-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
- *
- * @author David Sehnal <david.sehnal@gmail.com>
- * @author Alexander Rose <alexander.rose@weirdbyte.de>
- */
-
-import { DebugUtil } from '../../apps/tm-viewer';
-import { CustomElementProperty } from '../../mol-model-props/common/custom-element-property';
-import { Model, ElementIndex } from '../../mol-model/structure';
+import { Loci } from '../../mol-model/loci';
+import { StructureElement } from '../../mol-model/structure';
+import { LociLabel, LociLabelProvider } from '../../mol-plugin-state/manager/loci-label';
 import { TmDetChainListCache, TmDetDescriptorCache } from './prop';
-import { createResidueListsPerChain } from './tmdet-color-theme';
-import { ChainList, getResidue } from './types';
+import { createResidueListsPerChain, getChainAndResidueIds } from './tmdet-color-theme';
+import { getResidue } from './types';
 
 const siteLabels = [
     "Side1",
@@ -25,55 +18,33 @@ const siteLabels = [
 const DefaultResidueLabel = 6; // Unknown localization
 
 
-export const LabeledResidues = CustomElementProperty.create<number>({
-    label: 'TMDet Topology Site Labels',
-    name: 'tmdet-topology-site-labels',
-    getData(model: Model) {
-        const pdbId = model.entryId;
-        DebugUtil.log('start getData', model.label, pdbId);
+export const TmDetLabelProvider: LociLabelProvider = {
+    label: (loci: Loci): LociLabel => {
+        let labelText = siteLabels[DefaultResidueLabel];
 
-        const map = new Map<ElementIndex, number>();
-        const descriptor = TmDetDescriptorCache.get(pdbId);
-        if (!descriptor) {
-            return { value: map };
-        }
+        if (loci.kind == 'element-loci') {
+            const unit = loci.elements[0].unit;
+            const pdbId = unit.model.entryId;
 
-        // TODO: integrate with the coloring version
-        const getSiteId = function(chains: ChainList, chainId: string, residueId: string): number {
-            let siteId = DefaultResidueLabel;
+            const descriptor = TmDetDescriptorCache.get(pdbId);
+            if (!descriptor) {
+                return labelText;
+            }
 
-            const residue = getResidue(chains, chainId, residueId);
-            if (residue) {
-                siteId = residue.siteId;
+            let chainList =  TmDetChainListCache.get(pdbId);
+            if (!chainList) {
+                chainList = createResidueListsPerChain(descriptor.chains);
+                TmDetChainListCache.set(pdbId, chainList);
             }
 
-            return siteId;
-        };
+            const location = StructureElement.Loci.getFirstLocation(loci);
+            const { chainId, residueId } = getChainAndResidueIds(location!);
+            const residue = getResidue(chainList, chainId!, residueId!);
+            if (residue) {
+                labelText = siteLabels[residue?.siteId];
+            }
 
-        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 chainIdx = chainIndex[i];
-            const residueAuthId = model.atomicHierarchy.residues.auth_seq_id.value(residueIdx).toString();
-            const chainAuthId = model.atomicHierarchy.chains.auth_asym_id.value(chainIdx).toString();
-            // const residueName = model.atomicHierarchy.atoms.auth_comp_id.value(i).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, chainAuthId, residueAuthId));
         }
-        DebugUtil.log('end of getData');
-        return { value: map };
-    },
-    getLabel(e) {
-        return siteLabels[e];
+        return labelText;
     }
-});
+}

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

@@ -65,17 +65,30 @@ function getColor(location: Location, chains: ChainList): Color {
     let color = DefaultResidueColor;
 
     // TODO: How to solve cases when chain operation occurs?
+    if (StructureElement.Location.is(location) && Unit.isAtomic(location.unit)) {
+
+        const { chainId, residueId } = getChainAndResidueIds(location);
+        color = residueColor(chains, chainId!, residueId!);
+    }
+    return color;
+}
+
+export function getChainAndResidueIds(location: Location) {
+    let chainId = undefined;
+    let residueId = undefined;
+
     if (StructureElement.Location.is(location) && Unit.isAtomic(location.unit)) {
         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(chains, chainId, authId);
+        residueId = atomicHierarchy.residues.auth_seq_id.value(residueIdx).toString();
+        chainId = atomicHierarchy.chains.auth_asym_id.value(chainIdx).toString();
     }
-    return color;
+    return {
+        chainId: chainId,
+        residueId: residueId
+    };
 }
 
 export function createResidueListsPerChain(chains: PDBTMChain[]) {