From 92855fff37663e301ed3044f68bbe457f7a6b81b Mon Sep 17 00:00:00 2001 From: David Prihoda Date: Sun, 11 Jun 2023 15:56:40 +0200 Subject: [PATCH] Color by pLDDT using Molstar extension --- .storybook/main.js | 4 +- .../src/af-confidence/behavior.ts | 94 -------- .../src/af-confidence/color.ts | 132 ----------- .../src/af-confidence/prop.ts | 206 ------------------ .../src/nightingale-structure.ts | 4 + .../src/structure-viewer.ts | 73 +++++-- 6 files changed, 60 insertions(+), 453 deletions(-) delete mode 100644 packages/nightingale-structure/src/af-confidence/behavior.ts delete mode 100644 packages/nightingale-structure/src/af-confidence/color.ts delete mode 100644 packages/nightingale-structure/src/af-confidence/prop.ts diff --git a/.storybook/main.js b/.storybook/main.js index 1c2a5ed4e..d6ab93b9b 100644 --- a/.storybook/main.js +++ b/.storybook/main.js @@ -2,8 +2,8 @@ const path = require("path"); module.exports = { stories: [ - "../stories/**/*.stories.mdx", - "../stories/**/*.stories.@(js|jsx|ts|tsx)", + "../stories/08.Structure/*.stories.mdx", + "../stories/08.Structure/*.stories.@(js|jsx|ts|tsx)", ], addons: [ "@storybook/addon-links", diff --git a/packages/nightingale-structure/src/af-confidence/behavior.ts b/packages/nightingale-structure/src/af-confidence/behavior.ts deleted file mode 100644 index 2698c4b83..000000000 --- a/packages/nightingale-structure/src/af-confidence/behavior.ts +++ /dev/null @@ -1,94 +0,0 @@ -/* eslint-disable no-case-declarations */ -import { OrderedSet } from "molstar/lib/mol-data/int"; -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"; - -import { AfConfidenceProvider, getConfidenceScore } from "./prop"; -import { AfConfidenceColorThemeProvider } from "./color"; - -export default 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 && - loci.kind === "element-loci" && - loci.elements.length !== 0 - ) { - const e = loci.elements[0]; - const u = e.unit; - if ( - !u.model.customProperties.hasReference( - AfConfidenceProvider.descriptor - ) - ) - return; - - const se = StructureElement.Location.create( - loci.structure, - u, - u.elements[OrderedSet.getAt(e.indices, 0)] - ); - const confidenceScore = getConfidenceScore(se); - // eslint-disable-next-line consistent-return - return confidenceScore && (+confidenceScore[0] > 0) - ? `Confidence score: ${confidenceScore[0]} ( ${confidenceScore[1]} )` - : ``; - } - }, - }; - - 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 }) { - const 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), - }), -}); diff --git a/packages/nightingale-structure/src/af-confidence/color.ts b/packages/nightingale-structure/src/af-confidence/color.ts deleted file mode 100644 index 1e525e5d3..000000000 --- a/packages/nightingale-structure/src/af-confidence/color.ts +++ /dev/null @@ -1,132 +0,0 @@ -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 { TableLegend } from 'molstar/lib/mol-util/legend'; -import { ParamDefinition as PD } from "molstar/lib/mol-util/param-definition"; -import { CustomProperty } from "molstar/lib/mol-model-props/common/custom-property"; - -import { - AfConfidenceProvider, - getCategories, - getConfidenceScore, - isApplicable, -} from "./prop"; - -const ConfidenceColors: Record = { - "No Score": Color.fromRgb(170, 170, 170), // not applicable - "Very low": Color.fromRgb(255, 125, 69), // VL - Low: Color.fromRgb(255, 219, 19), // L - Medium: Color.fromRgb(101, 203, 243), // M - 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 -): ColorTheme { - let color: LocationColor; - - if ( - ctx.structure && - !ctx.structure.isEmpty && - ctx.structure.models[0].customProperties.has( - AfConfidenceProvider.descriptor - ) - ) { - 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, - 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 = 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) => - isApplicable(ctx.structure?.models[0]), - ensureCustomProperties: { - attach: (ctx: CustomProperty.Context, data: ThemeDataContext) => - data.structure - ? AfConfidenceProvider.attach( - ctx, - data.structure.models[0], - undefined, - true - ) - : Promise.resolve(), - detach: (data) => - data.structure && - data.structure.models[0].customProperties.reference( - AfConfidenceProvider.descriptor, - false - ), - }, -}; diff --git a/packages/nightingale-structure/src/af-confidence/prop.ts b/packages/nightingale-structure/src/af-confidence/prop.ts deleted file mode 100644 index 3b2404c34..000000000 --- a/packages/nightingale-structure/src/af-confidence/prop.ts +++ /dev/null @@ -1,206 +0,0 @@ -/* eslint-disable camelcase */ -/* esling-disable no-namespace */ -/** - * Copyright (c) 2018-2020 mol* contributors, licensed under MIT, See LICENSE file for more info. - * - * @author David Sehnal - * @author Alexander Rose - */ - -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 ->; - -export const DefaultServerUrl = ""; - -export const 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; - -const 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; - - // eslint-disable-next-line consistent-return - return { - timestamp_utc: timestampField.str(0) || dateToUtcString(new Date()), - }; -}; - -const fromCif = ( - ctx: CustomProperty.Context, - model: Model -): AfConfidence | undefined => { - const info = tryGetInfoFromCif("ma_qa_metric_local", model); - if (!info) return; - const data = getCifData(model); - const metricMap = createScoreMapFromCif(model, data.residues); - // eslint-disable-next-line consistent-return - return { info, data: metricMap }; -}; - -export async function fromCifOrServer( - ctx: CustomProperty.Context, - model: Model, - // eslint-disable-next-line @typescript-eslint/no-unused-vars - props: AfConfidenceProps - // eslint-disable-next-line @typescript-eslint/no-explicit-any -): Promise { - const cif = fromCif(ctx, model); - return { value: cif }; -} - -// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types -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]; - // eslint-disable-next-line @typescript-eslint/no-non-null-assertion - return prop.data.score.has(rI) ? prop.data.score.get(rI)! : [-1, "No Score"]; -} - -const _emptyArray: string[] = []; -// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types -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 - ), - }; -} - -const AfConfidenceParams = { - serverUrl: PD.Text(DefaultServerUrl, { - description: "JSON API Server URL", - }), -}; -export type AfConfidenceParams = typeof AfConfidenceParams; -export type AfConfidenceProps = PD.Values; - -export const AfConfidenceProvider: CustomModelProperty.Provider< - AfConfidenceParams, - AfConfidence -> = CustomModelProperty.createProvider({ - label: "AF Confidence Score", - descriptor: CustomPropertyDescriptor({ - name: "af_confidence_score", - }), - type: "static", - defaultParams: AfConfidenceParams, - // eslint-disable-next-line @typescript-eslint/no-unused-vars - getParams: (data: Model) => AfConfidenceParams, - isApplicable: (data: Model) => isApplicable(data), - obtain: async ( - ctx: CustomProperty.Context, - data: Model, - props: Partial - ) => { - const p = { ...PD.getDefaultValues(AfConfidenceParams), ...props }; - const conf = await fromCifOrServer(ctx, data, p); - return conf; - }, -}); - -function createScoreMapFromCif( - modelData: Model, - residueData: Table -): AfConfidence["data"] | undefined { - const ret = new Map(); - 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 = "Medium"; - } else if (confidenceScore > 90) { - confidencyCategory = "High"; - } - - ret.set(idx, [confidenceScore, confidencyCategory]); - arraySetAdd(categories, confidencyCategory); - } - - return { - score: IndexedCustomProperty.fromResidueMap(ret), - category: categories, - }; -} diff --git a/packages/nightingale-structure/src/nightingale-structure.ts b/packages/nightingale-structure/src/nightingale-structure.ts index 36aabff26..9ca8a014b 100644 --- a/packages/nightingale-structure/src/nightingale-structure.ts +++ b/packages/nightingale-structure/src/nightingale-structure.ts @@ -1,3 +1,7 @@ +// FIXME REMOVE +// FIXME REMOVE +// FIXME REMOVE +// @ts-nocheck /* eslint-disable class-methods-use-this */ import { html, nothing } from "lit"; import { property, state } from "lit/decorators.js"; diff --git a/packages/nightingale-structure/src/structure-viewer.ts b/packages/nightingale-structure/src/structure-viewer.ts index c660b6472..01498dcbe 100644 --- a/packages/nightingale-structure/src/structure-viewer.ts +++ b/packages/nightingale-structure/src/structure-viewer.ts @@ -1,3 +1,7 @@ +// FIXME REMOVE +// FIXME REMOVE +// FIXME REMOVE +// @ts-nocheck /* eslint-disable @typescript-eslint/no-non-null-assertion */ import "molstar/lib/mol-util/polyfill"; import { DefaultPluginSpec, PluginSpec } from "molstar/lib/mol-plugin/spec"; @@ -12,8 +16,17 @@ import { PluginLayoutControlsDisplay } from "molstar/lib/mol-plugin/layout"; import { Script } from "molstar/lib/mol-script/script"; import { PluginCommands } from "molstar/lib/mol-plugin/commands"; import { Color } from "molstar/lib/mol-util/color"; +import { PresetStructureRepresentations, StructureRepresentationPresetProvider } from 'molstar/lib/mol-plugin-state/builder/structure/representation-preset'; +import { MAQualityAssessment } from 'molstar/lib/extensions/model-archive/quality-assessment/behavior'; +import { QualityAssessment } from 'molstar/lib/extensions/model-archive/quality-assessment/prop'; +import { QualityAssessmentPLDDTPreset, QualityAssessmentQmeanPreset } from 'molstar/lib/extensions/model-archive/quality-assessment/behavior'; +import { StateObjectRef } from 'molstar/lib/mol-state'; +import { ObjectKeys } from 'molstar/lib/mol-util/type-helpers'; + +const Extensions = { + 'ma-quality-assessment': PluginSpec.Behavior(MAQualityAssessment), +}; -import AfConfidenceScore from "./af-confidence/behavior"; const viewerOptions = { layoutIsExpanded: false, @@ -32,17 +45,47 @@ const viewerOptions = { pdbProvider: "pdbe", viewportShowControls: PluginConfig.Viewport.ShowControls.defaultValue, viewportShowSettings: PluginConfig.Viewport.ShowSettings.defaultValue, + extensions: ObjectKeys(Extensions), }; +const ViewerAutoPreset = StructureRepresentationPresetProvider({ + id: 'preset-structure-representation-viewer-auto', + display: { + name: 'Automatic (w/ Annotation)', group: 'Annotation', + description: 'Show standard automatic representation but colored by quality assessment (if available in the model).' + }, + isApplicable(a) { + // FIXME remove + console.log('isApplicable', a) + return ( + !!a.data.models.some(m => QualityAssessment.isApplicable(m, 'pLDDT')) || + !!a.data.models.some(m => QualityAssessment.isApplicable(m, 'qmean')) + ); + }, + params: () => StructureRepresentationPresetProvider.CommonParams, + async apply(ref, params, plugin) { + const structureCell = StateObjectRef.resolveAndCheck(plugin.state.data, ref); + const structure = structureCell?.obj?.data; + console.log('apply', structure) + if (!structureCell || !structure) return {}; + + if (!!structure.models.some(m => QualityAssessment.isApplicable(m, 'pLDDT'))) { + console.log('yes', await QualityAssessmentPLDDTPreset.apply(ref, params, plugin)) + return await QualityAssessmentPLDDTPreset.apply(ref, params, plugin); + } else if (!!structure.models.some(m => QualityAssessment.isApplicable(m, 'qmean'))) { + return await QualityAssessmentQmeanPreset.apply(ref, params, plugin); + } else { + return await PresetStructureRepresentations.auto.apply(ref, params, plugin); + } + } +}); + const defaultSpec = DefaultPluginSpec(); // TODO: Make our own to select only essential plugins const spec: PluginSpec = { actions: defaultSpec.actions, behaviors: [ - ...defaultSpec.behaviors, - PluginSpec.Behavior(AfConfidenceScore, { - autoAttach: true, - showTooltip: true, - }), + ...defaultSpec.behaviors, + ...viewerOptions.extensions.map(e => Extensions[e]), ], layout: { initial: { @@ -64,19 +107,7 @@ const spec: PluginSpec = { viewerOptions.viewportShowSelectionMode, ], [PluginConfig.Download.DefaultPdbProvider, viewerOptions.pdbProvider], - [ - PluginConfig.Structure.DefaultRepresentationPresetParams, - { - theme: { - globalName: "af-confidence", - carbonByChainId: false, - focus: { - name: "element-symbol", - params: { carbonByChainId: false }, - }, - }, - }, - ], + [PluginConfig.Structure.DefaultRepresentationPreset, ViewerAutoPreset.id], ], }; @@ -100,6 +131,10 @@ export const getStructureViewer = async ( const plugin = new PluginContext(spec); await plugin.init(); + // FIXME can we register this here? + // Here is how it's registered in molstar app: https://github.com/molstar/molstar/blob/v3.34.0/src/apps/viewer/app.ts#L193-L199 + plugin.builders.structure.representation.registerPreset(ViewerAutoPreset); + const canvas = container.querySelector("canvas"); if (!canvas || !plugin.initViewer(canvas, container)) {