wip AlphaFold coloring

src/viewer/helpers/af-confidence/behavior.ts

@@ -1,71 +0,0 @@
-/**
- * Copyright (c) 2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
- *
- * @author Mandar Deshpande <mandar@ebi.ac.uk>
- */
-
-import { OrderedSet } from 'molstar/lib/mol-data/int';
-import { AfConfidence, AfConfidenceProvider } from './prop';
-import { AfConfidenceColorThemeProvider } from './color';
-import { Loci } from 'molstar/lib/mol-model/loci';
-import { StructureElement } from 'molstar/lib/mol-model/structure';
-import { ParamDefinition as PD } from 'molstar/lib/mol-util/param-definition';
-import { PluginBehavior } from 'molstar/lib/mol-plugin/behavior/behavior';
-
-export const AfConfidenceScore = PluginBehavior.create<{ autoAttach: boolean, showTooltip: boolean }>({
-    name: 'af-confidence-prop',
-    category: 'custom-props',
-    display: {
-        name: 'AlphaFold Confidence Score',
-        description: 'AlphaFold Confidence Score.'
-    },
-    ctor: class extends PluginBehavior.Handler<{ autoAttach: boolean, showTooltip: boolean }> {
-
-        private provider = AfConfidenceProvider
-
-        private labelAfConfScore = {
-            label: (loci: Loci): string | undefined => {
-                if (!this.params.showTooltip) return void 0;
-
-                switch (loci.kind) {
-                    case 'element-loci':
-                        if (loci.elements.length === 0) return void 0;
-                        const e = loci.elements[0];
-                        const u = e.unit;
-                        if (!u.model.customProperties.hasReference(AfConfidenceProvider.descriptor)) return void 0;
-
-                        const se = StructureElement.Location.create(loci.structure, u, u.elements[OrderedSet.getAt(e.indices, 0)]);
-                        const confidenceScore = AfConfidence.getConfidenceScore(se);
-                        return confidenceScore ? `Confidence score: ${confidenceScore[0]} <small>( ${confidenceScore[1]} )</small>` : `No confidence score`;
-
-                    default: return void 0;
-                }
-            }
-        }
-
-        register(): void {
-            this.ctx.customModelProperties.register(this.provider, this.params.autoAttach);
-            this.ctx.managers.lociLabels.addProvider(this.labelAfConfScore);
-
-            this.ctx.representation.structure.themes.colorThemeRegistry.add(AfConfidenceColorThemeProvider);
-        }
-
-        update(p: { autoAttach: boolean, showTooltip: boolean }) {
-            let updated = this.params.autoAttach !== p.autoAttach;
-            this.params.autoAttach = p.autoAttach;
-            this.params.showTooltip = p.showTooltip;
-            this.ctx.customModelProperties.setDefaultAutoAttach(this.provider.descriptor.name, this.params.autoAttach);
-            return updated;
-        }
-
-        unregister() {
-            this.ctx.customModelProperties.unregister(AfConfidenceProvider.descriptor.name);
-            this.ctx.managers.lociLabels.removeProvider(this.labelAfConfScore);
-            this.ctx.representation.structure.themes.colorThemeRegistry.remove(AfConfidenceColorThemeProvider);
-        }
-    },
-    params: () => ({
-        autoAttach: PD.Boolean(false),
-        showTooltip: PD.Boolean(true)
-    })
-});

src/viewer/helpers/af-confidence/color.ts

@@ -2,9 +2,10 @@
  * Copyright (c) 2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Mandar Deshpande <mandar@ebi.ac.uk>
+ * @author Sebastian Bittrich <sebastian.bittrich@rcsb.org>
  */
 
-import { AfConfidence, AfConfidenceProvider } from './prop';
+import { AlphaFoldConfidence, AlphaFoldConfidenceProvider } from './prop';
 import { Location } from 'molstar/lib/mol-model/location';
 import { StructureElement } from 'molstar/lib/mol-model/structure';
 import { ColorTheme, LocationColor } from 'molstar/lib/mol-theme/color';
@@ -12,93 +13,83 @@ import { ThemeDataContext } from 'molstar/lib/mol-theme/theme';
 import { Color } from 'molstar/lib/mol-util/color';
 import { ParamDefinition as PD } from 'molstar/lib/mol-util/param-definition';
 import { CustomProperty } from 'molstar/lib/mol-model-props/common/custom-property';
+import { TableLegend } from 'molstar/lib/mol-util/legend';
 
-const ConfidenceColors: any = {
-    'No Score': Color.fromRgb(170, 170, 170), // not applicable
-    'Very low': Color.fromRgb(255, 125, 69), // VL
-    'Low': Color.fromRgb(255, 219, 19), // L
-    'Confident': Color.fromRgb(101, 203, 243), // M
-    'Very high': Color.fromRgb(0, 83, 214), // H
+const DefaultColor = Color(0xaaaaaa);
+const ConfidenceColors: { [k: string]: Color } = {
+    'No Score': DefaultColor,
+    'Very low': Color(0xff7d45),
+    'Low': Color(0xffdb13),
+    'Confident': Color(0x65cbf3),
+    'Very high': Color(0x0053d6)
 };
 
-export const AfConfidenceColorThemeParams = {
-    type: PD.MappedStatic('score', {
-        'score': PD.Group({}),
-        'category': PD.Group({
-            kind: PD.Text()
+const ConfidenceColorLegend = TableLegend(Object.entries(ConfidenceColors));
+
+export function getAlphaFoldConfidenceColorThemeParams(ctx: ThemeDataContext) {
+    const categories = AlphaFoldConfidence.getCategories(ctx.structure);
+    if (categories.length === 0) {
+        return {
+            type: PD.MappedStatic('score', {
+                'score': PD.Group({})
+            })
+        };
+    }
+
+    return {
+        type: PD.MappedStatic('score', {
+            'score': PD.Group({}),
+            'category': PD.Group({
+                kind: PD.Select(categories[0], PD.arrayToOptions(categories))
+            }, { isFlat: true })
         })
-    })
-};
+    };
+}
+export type AlphaFoldConfidenceColorThemeParams = ReturnType<typeof getAlphaFoldConfidenceColorThemeParams>
 
-type Params = typeof AfConfidenceColorThemeParams
+export function AlphaFoldConfidenceColorTheme(ctx: ThemeDataContext, props: PD.Values<AlphaFoldConfidenceColorThemeParams>): ColorTheme<AlphaFoldConfidenceColorThemeParams> {
+    let color: LocationColor = () => DefaultColor;
 
-export function AfConfidenceColorTheme(ctx: ThemeDataContext, props: PD.Values<Params>): ColorTheme<Params> {
-    let color: LocationColor;
+    if (ctx.structure && ctx.structure.models.length > 0 && ctx.structure.models[0].customProperties.has(AlphaFoldConfidenceProvider.descriptor)) {
+        const getColor = (location: StructureElement.Location): Color => {
+            const score: string = AlphaFoldConfidence.getConfidenceScore(location)[1];
 
-    if (ctx.structure && !ctx.structure.isEmpty && ctx.structure.models[0].customProperties.has(AfConfidenceProvider.descriptor)) {
-        const getConfidenceScore = AfConfidence.getConfidenceScore;
+            if (props.type.name !== 'score') {
+                const categoryProp = props.type.params.kind;
+                if (score === categoryProp) return ConfidenceColors[score];
+            }
 
-        if (props.type.name === 'score') {
-            color = (location: Location) => {
-                if (StructureElement.Location.is(location)) {
-                    const confidenceScore = getConfidenceScore(location);
-                    return ConfidenceColors[confidenceScore[1]];
-                }
-                return ConfidenceColors['No Score'];
-            };
-        } else {
-            const categoryProp = props.type.params.kind;
-            color = (location: Location) => {
-                if (StructureElement.Location.is(location)) {
-                    const confidenceScore = getConfidenceScore(location);
-                    if(confidenceScore[1] === categoryProp) return ConfidenceColors[confidenceScore[1]];
-                    return ConfidenceColors['No Score'];
-                }
-                return ConfidenceColors['No Score'];
-            };
-        }
+            return ConfidenceColors[score];
+        };
 
-    } else {
-        color = () => ConfidenceColors['No Score'];
+        color = (location: Location) => {
+            if (StructureElement.Location.is(location)) {
+                return getColor(location);
+            }
+            return DefaultColor;
+        };
     }
 
     return {
-        factory: AfConfidenceColorTheme,
+        factory: AlphaFoldConfidenceColorTheme,
         granularity: 'group',
-        color: color,
-        props: props,
-        description: 'Assigns residue colors according to the AF Confidence score'
+        color,
+        props,
+        description: 'Assigns residue colors according to the AlphaFold Confidence score.',
+        legend: ConfidenceColorLegend
     };
 }
 
-export const AfConfidenceColorThemeProvider: ColorTheme.Provider<Params, 'af-confidence'> =  {
+export const AlphaFoldConfidenceColorThemeProvider: ColorTheme.Provider<AlphaFoldConfidenceColorThemeParams, 'af-confidence'> =  {
     name: 'af-confidence',
-    label: 'AF Confidence',
+    label: 'AlphaFold Confidence',
     category: ColorTheme.Category.Validation,
-    factory: AfConfidenceColorTheme,
-    getParams: ctx => {
-        const categories = AfConfidence.getCategories(ctx.structure);
-        if (categories.length === 0) {
-            return {
-                type: PD.MappedStatic('score', {
-                    'score': PD.Group({})
-                })
-            };
-        }
-
-        return {
-            type: PD.MappedStatic('score', {
-                'score': PD.Group({}),
-                'category': PD.Group({
-                    kind: PD.Select(categories[0], PD.arrayToOptions(categories))
-                }, { isFlat: true })
-            })
-        };
-    },
-    defaultValues: PD.getDefaultValues(AfConfidenceColorThemeParams),
-    isApplicable: (ctx: ThemeDataContext) => AfConfidence.isApplicable(ctx.structure?.models[0]),
+    factory: AlphaFoldConfidenceColorTheme,
+    getParams: getAlphaFoldConfidenceColorThemeParams,
+    defaultValues: PD.getDefaultValues(getAlphaFoldConfidenceColorThemeParams({})),
+    isApplicable: (ctx: ThemeDataContext) => AlphaFoldConfidence.isApplicable(ctx.structure?.models[0]),
     ensureCustomProperties: {
-        attach: (ctx: CustomProperty.Context, data: ThemeDataContext) => data.structure ? AfConfidenceProvider.attach(ctx, data.structure.models[0], void 0, true) : Promise.resolve(),
-        detach: (data) => data.structure && data.structure.models[0].customProperties.reference(AfConfidenceProvider.descriptor, false)
+        attach: (ctx: CustomProperty.Context, data: ThemeDataContext) => data.structure ? AlphaFoldConfidenceProvider.attach(ctx, data.structure.models[0], void 0, true) : Promise.resolve(),
+        detach: (data) => data.structure && AlphaFoldConfidenceProvider.ref(data.structure.models[0], false)
     }
 };

src/viewer/helpers/af-confidence/prop.ts

@@ -4,33 +4,31 @@
  * @author Mandar Deshpande <mandar@ebi.ac.uk>
  * @author David Sehnal <david.sehnal@gmail.com>
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ * @author Sebastian Bittrich <sebastian.bittrich@rcsb.org>
  */
 
 import { Column, Table } from 'molstar/lib/mol-data/db';
 import { toTable } from 'molstar/lib/mol-io/reader/cif/schema';
-import { Model, ResidueIndex, Unit, IndexedCustomProperty } from 'molstar/lib/mol-model/structure';
-import { StructureElement, Structure } from 'molstar/lib/mol-model/structure/structure';
+import { IndexedCustomProperty, Model, ResidueIndex, Unit } from 'molstar/lib/mol-model/structure';
+import { Structure, StructureElement } from 'molstar/lib/mol-model/structure/structure';
 import { ParamDefinition as PD } from 'molstar/lib/mol-util/param-definition';
 import { MmcifFormat } from 'molstar/lib/mol-model-formats/structure/mmcif';
 import { PropertyWrapper } from 'molstar/lib/mol-model-props/common/wrapper';
 import { CustomProperty } from 'molstar/lib/mol-model-props/common/custom-property';
 import { CustomModelProperty } from 'molstar/lib/mol-model-props/common/custom-model-property';
 import { CustomPropertyDescriptor } from 'molstar/lib/mol-model/custom-property';
-import { dateToUtcString } from 'molstar/lib/mol-util/date';
 import { arraySetAdd } from 'molstar/lib/mol-util/array';
+import { dateToUtcString } from 'molstar/lib/mol-util/date';
 
-export { AfConfidence };
-
-type AfConfidence = PropertyWrapper<{
+export type AlphaFoldConfidence = PropertyWrapper<{
     score: IndexedCustomProperty.Residue<[number, string]>,
     category: string[]
-} | undefined>
-
-namespace AfConfidence {
-    export const DefaultServerUrl = '';
+}>
 
+export namespace AlphaFoldConfidence {
     export function isApplicable(model?: Model): boolean {
-        return !!model && Model.isFromPdbArchive(model);
+        if (!model || !MmcifFormat.is(model.sourceData)) return false;
+        return model.sourceData.data.frame.categoryNames.includes('ma_qa_metric_local');
     }
 
     export interface Info {
@@ -48,7 +46,7 @@ namespace AfConfidence {
             ordinal_id: Column.Schema.int
         }
     };
-    export type Schema = typeof Schema
+    export type Schema = typeof Schema;
 
     function tryGetInfoFromCif(categoryName: string, model: Model): undefined | Info {
         if (!MmcifFormat.is(model.sourceData) || !model.sourceData.data.frame.categoryNames.includes(categoryName)) {
@@ -61,22 +59,23 @@ namespace AfConfidence {
         return { timestamp_utc: timestampField.str(0) || dateToUtcString(new Date()) };
     }
 
-    export function fromCif(ctx: CustomProperty.Context, model: Model): AfConfidence | undefined {
-        let info = tryGetInfoFromCif('ma_qa_metric_local', model);
+
+    export function fromCif(ctx: CustomProperty.Context, model: Model): AlphaFoldConfidence | undefined {
+        const info = tryGetInfoFromCif('ma_qa_metric_local', model);
         if (!info) return;
         const data = getCifData(model);
         const metricMap = createScoreMapFromCif(model, data.residues);
         return { info, data: metricMap };
     }
 
-    export async function fromCifOrServer(ctx: CustomProperty.Context, model: Model, props: AfConfidenceProps): Promise<any> {
+    export async function obtain(ctx: CustomProperty.Context, model: Model, _props: AlphaFoldConfidenceProps): Promise<CustomProperty.Data<any>> {
         const cif = fromCif(ctx, model);
         return { value: cif };
     }
 
-    export function getConfidenceScore(e: StructureElement.Location) {
+    export function getConfidenceScore(e: StructureElement.Location): [number, string] {
         if (!Unit.isAtomic(e.unit)) return [-1, 'No Score'];
-        const prop = AfConfidenceProvider.get(e.unit.model).value;
+        const prop = AlphaFoldConfidenceProvider.get(e.unit.model).value;
         if (!prop || !prop.data) return [-1, 'No Score'];
         const rI = e.unit.residueIndex[e.element];
         return prop.data.score.has(rI) ? prop.data.score.get(rI)! : [-1, 'No Score'];
@@ -85,7 +84,7 @@ namespace AfConfidence {
     const _emptyArray: string[] = [];
     export function getCategories(structure?: Structure) {
         if (!structure) return _emptyArray;
-        const prop = AfConfidenceProvider.get(structure.models[0]).value;
+        const prop = AlphaFoldConfidenceProvider.get(structure.models[0]).value;
         if (!prop || !prop.data) return _emptyArray;
         return prop.data.category;
     }
@@ -98,48 +97,45 @@ namespace AfConfidence {
     }
 }
 
-export const AfConfidenceParams = {
-    serverUrl: PD.Text(AfConfidence.DefaultServerUrl, { description: 'JSON API Server URL' })
-};
-export type AfConfidenceParams = typeof AfConfidenceParams
-export type AfConfidenceProps = PD.Values<AfConfidenceParams>
+export const AlphaFoldConfidenceParams = {};
+export type AlphaFoldConfidenceParams = typeof AlphaFoldConfidenceParams
+export type AlphaFoldConfidenceProps = PD.Values<AlphaFoldConfidenceParams>
 
-export const AfConfidenceProvider: CustomModelProperty.Provider<AfConfidenceParams, AfConfidence> = CustomModelProperty.createProvider({
-    label: 'AF Confidence Score',
+export const AlphaFoldConfidenceProvider: CustomModelProperty.Provider<AlphaFoldConfidenceParams, AlphaFoldConfidence> = CustomModelProperty.createProvider({
+    label: 'AlphaFold Confidence Score',
     descriptor: CustomPropertyDescriptor({
         name: 'af_confidence_score'
     }),
     type: 'static',
-    defaultParams: AfConfidenceParams,
-    getParams: () => AfConfidenceParams,
-    isApplicable: (data: Model) => AfConfidence.isApplicable(data),
-    obtain: async (ctx: CustomProperty.Context, data: Model, props: Partial<AfConfidenceProps>) => {
-        const p = { ...PD.getDefaultValues(AfConfidenceParams), ...props };
-        return await AfConfidence.fromCifOrServer(ctx, data, p);
+    defaultParams: AlphaFoldConfidenceParams,
+    getParams: () => AlphaFoldConfidenceParams,
+    isApplicable: (data: Model) => AlphaFoldConfidence.isApplicable(data),
+    obtain: async (ctx: CustomProperty.Context, data: Model, props: Partial<AlphaFoldConfidenceProps>) => {
+        const p = { ...PD.getDefaultValues(AlphaFoldConfidenceParams), ...props };
+        return await AlphaFoldConfidence.obtain(ctx, data, p);
     }
 });
 
-function createScoreMapFromCif(modelData: Model, residueData: Table<typeof AfConfidence.Schema.local_metric_values>): AfConfidence['data'] | undefined {
-    const ret = new Map<ResidueIndex, [number, string]>();
+function createScoreMapFromCif(modelData: Model, residueData: Table<typeof AlphaFoldConfidence.Schema.local_metric_values>): AlphaFoldConfidence['data'] {
     const { label_asym_id, label_seq_id, metric_value, _rowCount } = residueData;
 
+    const ret = new Map<ResidueIndex, [number, string]>();
     const categories: string[] = [];
 
+    const toCategory = (v: number): 'Very low' | 'Low' | 'Confident' | 'Very high' => {
+        if (v > 50 && v <= 70)  return 'Low';
+        if (v > 70 && v <= 90) return 'Confident';
+        if (v > 90) return 'Very high';
+        return 'Very low';
+    };
+
     for (let i = 0; i < _rowCount; i++) {
         const confidenceScore = metric_value.value(i);
-        const idx = modelData.atomicHierarchy.index.findResidue('1', label_asym_id.value(i), label_seq_id.value(i), '');
-
-        let confidencyCategory = 'Very low';
-        if(confidenceScore > 50 && confidenceScore <= 70) {
-            confidencyCategory = 'Low';
-        } else if(confidenceScore > 70 && confidenceScore <= 90) {
-            confidencyCategory = 'Confident';
-        } else if(confidenceScore > 90) {
-            confidencyCategory = 'Very high';
-        }
+        const idx = modelData.atomicHierarchy.index.findResidue('1', label_asym_id.value(i), label_seq_id.value(i));
+        const confidenceCategory = toCategory(confidenceScore);
 
-        ret.set(idx, [confidenceScore, confidencyCategory]);
-        arraySetAdd(categories, confidencyCategory);
+        ret.set(idx, [confidenceScore, confidenceCategory]);
+        arraySetAdd(categories, confidenceCategory);
     }
 
     return {

src/viewer/index.ts

@@ -38,6 +38,7 @@ import { PluginUIContext } from 'molstar/lib/mol-plugin-ui/context';
 import { ANVILMembraneOrientation, MembraneOrientationPreset } from 'molstar/lib/extensions/anvil/behavior';
 import { MembraneOrientationRepresentationProvider } from 'molstar/lib/extensions/anvil/representation';
 import { MotifAlignmentRequest, alignMotifs } from './helpers/superpose/pecos-integration';
+import { AlphaFoldConfidenceColorThemeProvider } from './helpers/af-confidence/color';
 
 /** package version, filled in at bundle build time */
 declare const __RCSB_MOLSTAR_VERSION__: string;
@@ -62,6 +63,7 @@ const DefaultViewerProps = {
     showStructureSourceControls: true,
     showSuperpositionControls: true,
     showMembraneOrientationPreset: false,
+    showAlphaFoldConfidenceColorTheme: false,
     /**
      * Needed when running outside of sierra. If set to true, the strucmotif UI will use an absolute URL to sierra-prod.
      * Otherwise, the link will be relative on the current host.
@@ -173,6 +175,7 @@ export class Viewer {
         this.plugin.init()
             .then(async () => {
                 // hide 'Membrane Orientation' preset from UI - has to happen 'before' react render, apparently
+                // the corresponding behavior must be registered either way, because the 3d-view uses it (even without appearing in the UI)
                 if (!o.showMembraneOrientationPreset) {
                     this.plugin.builders.structure.representation.unregisterPreset(MembraneOrientationPreset);
                     this.plugin.representation.structure.registry.remove(MembraneOrientationRepresentationProvider);
@@ -183,6 +186,10 @@ export class Viewer {
                 const renderer = this.plugin.canvas3d!.props.renderer;
                 await PluginCommands.Canvas3D.SetSettings(this.plugin, { settings: { renderer: { ...renderer, backgroundColor: o.backgroundColor } } });
                 this.plugin.representation.structure.themes.colorThemeRegistry.add(SuperposeColorThemeProvider);
+                // conditionally load PDBe's AlphaFold preset
+                if (o.showAlphaFoldConfidenceColorTheme) {
+                    this.plugin.representation.structure.themes.colorThemeRegistry.add(AlphaFoldConfidenceColorThemeProvider);
+                }
 
                 if (o.showWelcomeToast) {
                     await PluginCommands.Toast.Show(this.plugin, {