Issue #2: refactored CustomColorTheme

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

src/extensions/tmdet/behavior.ts

@@ -29,7 +29,7 @@ import { applyTransformations, createMembraneOrientation, getAtomGroupExpression
 import { PDBTMDescriptor, PDBTMRegion, PMS } from './types';
 import { registerTmDetSymmetry } from './symmetry';
 import { StripedResidues } from './coloring';
-import { CustomColorThemeProvider } from './custom-theme';
+import { CustomColorThemeProvider } from './tmdet-color-theme';
 
 type StructureComponentType = StateObjectSelector<
         PMS,
@@ -131,6 +131,8 @@ export async function loadWithUNITMPMembraneRepresentation(plugin: PluginUIConte
 
     await createStructureRepresentation(plugin, pdbtmDescriptor);
 
+    await customTheme(plugin);
+
     //
     // reset the camera because the membranes render 1st and the structure might not be fully visible
     //
@@ -148,6 +150,22 @@ export async function customColor(plugin: PluginUIContext) {
     });
 }
 
+export async function customTheme(plugin: PluginUIContext) {
+    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: CustomColorThemeProvider.name as any
+                }
+            );
+        }
+    });
+}
+
+
 
 async function downloadRegionDescriptor(plugin: PluginUIContext, params: any): Promise<any> {
     // run a fetch task

src/extensions/tmdet/coloring.ts

@@ -9,20 +9,30 @@ import { CustomElementProperty } from '../../mol-model-props/common/custom-eleme
 import { Model, ElementIndex } from '../../mol-model/structure';
 import { Color } from '../../mol-util/color';
 
+const ColorMap = [
+    Color(0xff0000),
+    Color(0x0000ff),
+    Color(0xffff00)
+];
+
 export const StripedResidues = CustomElementProperty.create<number>({
-    label: 'Residue Stripes',
-    name: 'basic-wrapper-residue-striping',
+    label: 'TMDet Topology Colors',
+    name: 'tmdet-topology-colors',
     getData(model: Model) {
-        console.log('getData');
+        console.log('getData', model.atomicHierarchy.residueAtomSegments);
+        console.log('getData: auth_comp_id', model.atomicHierarchy.atoms.auth_comp_id.toArray());
         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);
+            const residueId = residueIndex[i];
+            const value = (residueId === 310) ? 2 : residueId % 2;
+            map.set(i as ElementIndex, value);
         }
+        console.log('map', map);
         return { value: map };
     },
     coloring: {
-        getColor(e) { return e === 0 ? Color(0xff0000) : Color(0x0000ff); },
+        getColor(e) { return ColorMap[e]; },
         defaultColor: Color(0x777777)
     },
     getLabel(e) {

src/extensions/tmdet/custom-theme.ts

@@ -1,51 +0,0 @@
-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,
-};

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

@@ -0,0 +1,55 @@
+import { StructureElement, Unit } from '../../mol-model/structure';
+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';
+import { Location } from '../../mol-model/location';
+import { PDBTMDescriptor } from './types';
+
+export type ResidueNameColorThemeParams = {
+    pdbtmDescriptor: any
+};
+
+export function CustomColorTheme(
+    ctx: ThemeDataContext,
+    props: PD.Values<ResidueNameColorThemeParams>
+): ColorTheme<ResidueNameColorThemeParams> {
+    console.log('ColorTheme CTX', ctx);
+    console.log('ColorTheme PROPS', props);
+
+    return {
+        factory: CustomColorTheme,
+        granularity: 'group', //'residue' as any,
+        color: getColor,
+        props: props,
+        description: 'TMDet...',
+    };
+}
+
+const DefaultResidueColor = ColorNames.black;
+const ColorMap = [
+    Color(0xff0000),
+    Color(0x0000ff),
+    Color(0xffff00)
+];
+
+function getColor(location: Location): Color {
+    let color = DefaultResidueColor;
+
+    if (StructureElement.Location.is(location) && Unit.isAtomic(location.unit)) {
+        const compId = StructureElement.residueIndex(location) % 3;
+        color = ColorMap[compId]
+    }
+    return color;
+}
+
+export const CustomColorThemeProvider: ColorTheme.Provider<ResidueNameColorThemeParams, 'tmdet-custom-color-theme'> = {
+    name: 'tmdet-custom-color-theme',
+    label: 'TMDet Topology Theme',
+    category: 'TMDet',
+    factory: CustomColorTheme,
+    getParams: () => ({ pdbtmDescriptor: { isOptional: true } }),
+    defaultValues: { pdbtmDescriptor: undefined },
+    isApplicable: () => true,
+};