Issue #2: A working custom theme from basic wrapper

src/apps/tm-viewer/index.html

@@ -46,7 +46,9 @@
                 collapseLeftPanel: false
             });
             // Set PDB Id here
-            var pdbId = '2atk';
+            var pdbId = '1a0s';
+            // pdbId = '2atk';
+            // pdbId = '1afo';
             tm_molstar.loadWithUNITMPMembraneRepresentation(viewer.plugin, {
                 structureUrl: `https://cs.litemol.org/${pdbId}/full`,
                 regionDescriptorUrl: `http://localhost:8000/build/data/${pdbId}.json`,

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

src/extensions/tmdet/behavior.ts

@@ -28,6 +28,8 @@ import { StateObjectSelector } from "../../mol-state/object";
 import { applyTransformations, createMembraneOrientation, getAtomGroupExpression } from './transformation';
 import { PDBTMDescriptor, PDBTMRegion, PMS } from './types';
 import { registerTmDetSymmetry } from './symmetry';
+import { StripedResidues } from './coloring';
+import { CustomColorThemeProvider } from './custom-theme';
 
 type StructureComponentType = StateObjectSelector<
         PMS,
@@ -55,6 +57,12 @@ export const TMDETMembraneOrientation = PluginBehavior.create<{ autoAttach: bool
             this.ctx.representation.structure.registry.add(MembraneOrientationRepresentationProvider);
             this.ctx.query.structure.registry.add(isTransmembrane);
 
+            this.ctx.representation.structure.themes.colorThemeRegistry.add(StripedResidues.colorThemeProvider!);
+            this.ctx.representation.structure.themes.colorThemeRegistry.add(CustomColorThemeProvider);
+            this.ctx.managers.lociLabels.addProvider(StripedResidues.labelProvider!);
+            this.ctx.customModelProperties.register(StripedResidues.propertyProvider, true);
+
+
             this.ctx.genericRepresentationControls.set(Tag.Representation, selection => {
                 const refs: GenericRepresentationRef[] = [];
                 selection.structures.forEach(structure => {
@@ -129,6 +137,18 @@ export async function loadWithUNITMPMembraneRepresentation(plugin: PluginUIConte
     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 });
+        }
+    });
+}
+
+
 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;
@@ -163,7 +183,7 @@ async function createStructureRepresentation(plugin: PluginUIContext, pdbtmDescr
     pdbtmDescriptor.chains.forEach((chain: any) => {
 
         let regionsBySite: PDBTMRegion[] = siteColors.map((color):PDBTMRegion => {
-            // TODO: set site appropriately
+            // TODO: set 'site' appropriately
             return { site: "valami", auth_ids: [], color: color };
         });
 
@@ -172,10 +192,14 @@ async function createStructureRepresentation(plugin: PluginUIContext, pdbtmDescr
             regionsBySite[siteIndex].auth_ids.push(Number(residueItem.pdb_res_label));
         }
 
-        //console.log(`REGIONS-${chain.chain_label}`, regionsBySite);
-        regionsBySite.forEach((region: PDBTMRegion) => {
+        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();
-            createRegionRepresentation(plugin, chain.chain_label, region, regionUpdates.to(structure));
+            await createRegionRepresentation(plugin, chain.chain_label, region, regionUpdates.to(structure));
             regionUpdates.commit();
         });
 
@@ -190,11 +214,12 @@ async function createStructureComponents(plugin: PluginUIContext, structure: Sta
     };
 }
 
-function createRegionRepresentation(plugin: PluginUIContext, chain: string, residue: PDBTMRegion, update: StateBuilder.To<any, any>) {
+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

src/extensions/tmdet/coloring.ts

@@ -0,0 +1,31 @@
+/**
+ * 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 { CustomElementProperty } from '../../mol-model-props/common/custom-element-property';
+import { Model, ElementIndex } from '../../mol-model/structure';
+import { Color } from '../../mol-util/color';
+
+export const StripedResidues = CustomElementProperty.create<number>({
+    label: 'Residue Stripes',
+    name: 'basic-wrapper-residue-striping',
+    getData(model: Model) {
+        console.log('getData');
+        const map = new Map<ElementIndex, number>();
+        const residueIndex = model.atomicHierarchy.residueAtomSegments.index;
+        for (let i = 0, _i = model.atomicHierarchy.atoms._rowCount; i < _i; i++) {
+            map.set(i as ElementIndex, residueIndex[i] % 2);
+        }
+        return { value: map };
+    },
+    coloring: {
+        getColor(e) { return e === 0 ? Color(0xff0000) : Color(0x0000ff); },
+        defaultColor: Color(0x777777)
+    },
+    getLabel(e) {
+        return e === 0 ? 'Odd stripe' : 'Even stripe';
+    }
+});

src/extensions/tmdet/custom-theme.ts

@@ -0,0 +1,51 @@
+import { isPositionLocation } from '../../mol-geo/util/location-iterator';
+import { Vec3 } from '../../mol-math/linear-algebra';
+import { ColorTheme } from '../../mol-theme/color';
+import { ThemeDataContext } from '../../mol-theme/theme';
+import { Color } from '../../mol-util/color';
+import { ColorNames } from '../../mol-util/color/names';
+import { ParamDefinition as PD } from '../../mol-util/param-definition';
+
+export function CustomColorTheme(
+    ctx: ThemeDataContext,
+    props: PD.Values<{}>
+): ColorTheme<{}> {
+    const { radius, center } = ctx.structure?.boundary.sphere!;
+    const radiusSq = Math.max(radius * radius, 0.001);
+    const scale = ColorTheme.PaletteScale;
+
+    return {
+        factory: CustomColorTheme,
+        granularity: 'vertex',
+        color: location => {
+            if (!isPositionLocation(location)) return ColorNames.black;
+            const dist = Vec3.squaredDistance(location.position, center);
+            const t = Math.min(dist / radiusSq, 1);
+            return ((t * scale) | 0) as Color;
+        },
+        palette: {
+            filter: 'nearest',
+            colors: [
+                ColorNames.red,
+                ColorNames.pink,
+                ColorNames.violet,
+                ColorNames.orange,
+                ColorNames.yellow,
+                ColorNames.green,
+                ColorNames.blue
+            ]
+        },
+        props: props,
+        description: '',
+    };
+}
+
+export const CustomColorThemeProvider: ColorTheme.Provider<{}, 'basic-wrapper-custom-color-theme'> = {
+    name: 'basic-wrapper-custom-color-theme',
+    label: 'Custom Color Theme',
+    category: ColorTheme.Category.Misc,
+    factory: CustomColorTheme,
+    getParams: () => ({}),
+    defaultValues: { },
+    isApplicable: (ctx: ThemeDataContext) => true,
+};