src/Tm3DStandaloneApp/tmdet-viewer/TmViewerStandalone.ts

@@ -25,6 +25,19 @@ import { Mat4 } from 'molstar/lib/commonjs/mol-math/linear-algebra';
 import { TrajectoryHierarchyPresetProvider } from 'molstar/lib/mol-plugin-state/builder/structure/hierarchy-preset';
 import { registerRegionDescriptorData } from '../../TmFv3DApp/UniTmpHelper';
 import { TmDetRcsbPreset } from '../TmTrajectoryHierarchyPreset';
+import { Expression } from 'molstar/lib/mol-script/language/expression';
+import { MolScriptBuilder } from 'molstar/lib/mol-script/language/builder';
+import { StateTransforms } from 'molstar/lib/mol-plugin-state/transforms';
+import { Color } from 'molstar/lib/mol-util/color';
+
+import { StructureElement, Structure, StructureSelection, QueryContext } from 'molstar/lib/mol-model/structure';
+import { OrderedSet } from 'molstar/lib/mol-data/int';
+import { setStructureOverpaint } from 'molstar/lib/mol-plugin-state/helpers/structure-overpaint';
+import { Script } from 'molstar/lib/mol-script/script';
+import { compile } from 'molstar/lib/mol-script/runtime/query/base';
+import { Overpaint } from 'molstar/lib/mol-theme/overpaint';
+import { createStructureRepresentationParams } from 'molstar/lib/mol-plugin-state/helpers/structure-representation-params';
+import { ColorNames } from 'molstar/lib/mol-util/color/names';
 
 export class TmViewerStandalone extends Viewer {
 
@@ -59,8 +72,67 @@ export class TmViewerStandalone extends Viewer {
         return super.loadStructureFromUrl(url, format, isBinary, config);
     }
 
-    unselectChains(chains: string[]) {
-        console.log(TmDetColorThemeProvider);
+    async lociGetter(structure: Structure) {
+        // const structure = this.plugin.managers.structure.hierarchy.current.models[0].structures[0].cell;
+        // const unit = structure.obj!.data!.unitMap.get(0)
+        // structure.obj!.data.entityIndices
+        // StructureElement.Loci(structure.obj!.data, [{ unit: unit, indices: Array.from(unit.elements) }]);
+        // console.log(Structure.Loci());
+        let indices = OrderedSet.ofRange(10 as StructureElement.UnitIndex, 100  as StructureElement.UnitIndex);
+        // let indices = OrderedSet.ofSortedArray([ 15 as StructureElement.UnitIndex ]);
+        return StructureElement.Loci(structure, [
+            { unit: structure.unitMap.get(0), indices: indices },
+            { unit: structure.unitMap.get(1), indices: indices },
+            { unit: structure.unitMap.get(2), indices: indices }
+        ]);
+    }
+
+    chainSelection(auth_asym_id: string) {
+        return MolScriptBuilder.struct.generator.atomGroups({
+            'chain-test': MolScriptBuilder.core.rel.eq([MolScriptBuilder.struct.atomProperty.macromolecular.auth_asym_id(), auth_asym_id])
+        });
+    }
+
+    setChainColor(labelAsymIds: string[], hexRgb: string|false = false) {
+        this.overPaint(labelAsymIds, hexRgb, false);
+    }
+
+    resetChainColor(labelAsymIds: string[]) {
+        this.overPaint(labelAsymIds, false, true);
+    }
+
+    /**
+     * Util function to handle overpaint logic.
+     * @param labelAsymId 
+     * @param hexRgb default color grey, when it is false
+     * @param clear if true, it clears the overpaint color
+     */
+    protected overPaint(labelAsymIds: string[], hexRgb: string|false = false, clear: boolean = false) {
+        const state = this.plugin.state.data;
+
+        const selected = state.selectQ(q => q.ofTransformer(StateTransforms.Representation.StructureRepresentation3D));
+        if (selected.length == 0) {
+            console.log('Representation not found');
+            return;
+        }
+
+        const chainList = labelAsymIds.join(' ');
+        selected.forEach(async (representation) => {
+            const update = this.plugin.build().to(representation);
+            update.apply(StateTransforms.Representation.OverpaintStructureRepresentation3DFromScript, {
+                layers: [
+                    {
+                        script:  Script(
+                            `(sel.atom.atom-groups :chain-test (set.has (set ${chainList}) atom.label_asym_id))`,
+                            'mol-script'
+                        ),
+                        color:  hexRgb !== false ? Color.fromHexString(hexRgb) : ColorNames.grey,
+                        clear: clear
+                    }
+                ]
+            });
+            await update.commit();
+        });
     }
 
     async loadWithUNITMPMembraneRepresentation(params: any) {