add PDBe af-confidence behavior

README.md

@@ -5,6 +5,8 @@
 RCSB PDB implementation of [Mol* (/'mol-star/)](https://github.com/molstar/molstar).
 Try it [here](https://rcsb.org/3d-view/).
 
+PDBe also maintains a flavor of Mol* called [PDBe Molstar](https://github.com/PDBeurope/pdbe-molstar).
+
 ## Functionality
 Provides custom features used in the Mol* viewer on [rcsb.org](https://www.rcsb.org/3d-view):
 - visualization of structure alignment

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

@@ -0,0 +1,71 @@
+/**
+ * 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

@@ -0,0 +1,104 @@
+/**
+ * Copyright (c) 2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Mandar Deshpande <mandar@ebi.ac.uk>
+ */
+
+import { AfConfidence, AfConfidenceProvider } 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';
+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';
+
+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
+};
+
+export const AfConfidenceColorThemeParams = {
+    type: PD.MappedStatic('score', {
+        'score': PD.Group({}),
+        'category': PD.Group({
+            kind: PD.Text()
+        })
+    })
+};
+
+type Params = typeof AfConfidenceColorThemeParams
+
+export function AfConfidenceColorTheme(ctx: ThemeDataContext, props: PD.Values<Params>): ColorTheme<Params> {
+    let color: LocationColor;
+
+    if (ctx.structure && !ctx.structure.isEmpty && ctx.structure.models[0].customProperties.has(AfConfidenceProvider.descriptor)) {
+        const getConfidenceScore = AfConfidence.getConfidenceScore;
+
+        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'];
+            };
+        }
+
+    } else {
+        color = () => ConfidenceColors['No Score'];
+    }
+
+    return {
+        factory: AfConfidenceColorTheme,
+        granularity: 'group',
+        color: color,
+        props: props,
+        description: 'Assigns residue colors according to the AF Confidence score'
+    };
+}
+
+export const AfConfidenceColorThemeProvider: ColorTheme.Provider<Params, 'af-confidence'> =  {
+    name: 'af-confidence',
+    label: 'AF 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]),
+    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)
+    }
+};

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

@@ -0,0 +1,149 @@
+/**
+ * Copyright (c) 2018-2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Mandar Deshpande <mandar@ebi.ac.uk>
+ * @author David Sehnal <david.sehnal@gmail.com>
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+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 { 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';
+
+export { AfConfidence };
+
+type AfConfidence = PropertyWrapper<{
+    score: IndexedCustomProperty.Residue<[number, string]>,
+    category: string[]
+} | undefined>
+
+namespace AfConfidence {
+    export const DefaultServerUrl = '';
+
+    export function isApplicable(model?: Model): boolean {
+        return !!model && Model.isFromPdbArchive(model);
+    }
+
+    export interface Info {
+        timestamp_utc: string
+    }
+
+    export const Schema = {
+        local_metric_values: {
+            label_asym_id: Column.Schema.str,
+            label_comp_id: Column.Schema.str,
+            label_seq_id: Column.Schema.int,
+            metric_id: Column.Schema.int,
+            metric_value: Column.Schema.float,
+            model_id: Column.Schema.int,
+            ordinal_id: Column.Schema.int
+        }
+    };
+    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)) {
+            return;
+        }
+
+        const timestampField = model.sourceData.data.frame.categories[categoryName].getField('metric_value');
+        if (!timestampField || timestampField.rowCount === 0) return;
+
+        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);
+        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> {
+        const cif = fromCif(ctx, model);
+        return { value: cif };
+    }
+
+    export function getConfidenceScore(e: StructureElement.Location) {
+        if (!Unit.isAtomic(e.unit)) return [-1, 'No Score'];
+        const prop = AfConfidenceProvider.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'];
+    }
+
+    const _emptyArray: string[] = [];
+    export function getCategories(structure?: Structure) {
+        if (!structure) return _emptyArray;
+        const prop = AfConfidenceProvider.get(structure.models[0]).value;
+        if (!prop || !prop.data) return _emptyArray;
+        return prop.data.category;
+    }
+
+    function getCifData(model: Model) {
+        if (!MmcifFormat.is(model.sourceData)) throw new Error('Data format must be mmCIF.');
+        return {
+            residues: toTable(Schema.local_metric_values, model.sourceData.data.frame.categories.ma_qa_metric_local),
+        };
+    }
+}
+
+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 AfConfidenceProvider: CustomModelProperty.Provider<AfConfidenceParams, AfConfidence> = CustomModelProperty.createProvider({
+    label: 'AF 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);
+    }
+});
+
+function createScoreMapFromCif(modelData: Model, residueData: Table<typeof AfConfidence.Schema.local_metric_values>): AfConfidence['data'] | undefined {
+    const ret = new Map<ResidueIndex, [number, string]>();
+    const { label_asym_id, label_seq_id, metric_value, _rowCount } = residueData;
+
+    const categories: string[] = [];
+
+    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';
+        }
+
+        ret.set(idx, [confidenceScore, confidencyCategory]);
+        arraySetAdd(categories, confidencyCategory);
+    }
+
+    return {
+        score: IndexedCustomProperty.fromResidueMap(ret),
+        category: categories
+    };
+}