refactor AlphaFold coloring to generic pLDDT coloring

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

@@ -5,22 +5,22 @@
  */
 
 import { OrderedSet } from 'molstar/lib/mol-data/int';
-import { AlphaFoldConfidence, AlphaFoldConfidenceProvider } from './prop';
-import { AlphaFoldConfidenceColorThemeProvider } from './color';
+import { PLDDTConfidence, PLDDTConfidenceProvider } from './prop';
+import { PLDDTConfidenceColorThemeProvider } 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 AlphaFoldConfidenceScore = PluginBehavior.create<{ autoAttach: boolean, showTooltip: boolean }>({
-    name: 'af-confidence-prop',
+export const PLLDTConfidenceScore = PluginBehavior.create<{ autoAttach: boolean, showTooltip: boolean }>({
+    name: 'plddt-confidence-prop',
     category: 'custom-props',
     display: {
-        name: 'AlphaFold Confidence Score',
-        description: 'AlphaFold Confidence Score.'
+        name: 'pLDDT Confidence Score',
+        description: 'pLDDT Confidence Score.'
     },
     ctor: class extends PluginBehavior.Handler<{ autoAttach: boolean, showTooltip: boolean }> {
-        private provider = AlphaFoldConfidenceProvider;
+        private provider = PLDDTConfidenceProvider;
 
         private labelProvider = {
             label: (loci: Loci): string | undefined => {
@@ -31,10 +31,10 @@ export const AlphaFoldConfidenceScore = PluginBehavior.create<{ autoAttach: bool
                         if (loci.elements.length === 0) return;
                         const e = loci.elements[0];
                         const u = e.unit;
-                        if (!u.model.customProperties.hasReference(AlphaFoldConfidenceProvider.descriptor)) return;
+                        if (!u.model.customProperties.hasReference(PLDDTConfidenceProvider.descriptor)) return;
 
                         const se = StructureElement.Location.create(loci.structure, u, u.elements[OrderedSet.getAt(e.indices, 0)]);
-                        const confidenceScore = AlphaFoldConfidence.getConfidenceScore(se);
+                        const confidenceScore = PLDDTConfidence.getConfidenceScore(se);
                         return confidenceScore ? `Confidence score: ${confidenceScore[0]} <small>( ${confidenceScore[1]} )</small>` : `No confidence score`;
 
                     default: return;
@@ -46,7 +46,7 @@ export const AlphaFoldConfidenceScore = PluginBehavior.create<{ autoAttach: bool
             this.ctx.customModelProperties.register(this.provider, this.params.autoAttach);
             this.ctx.managers.lociLabels.addProvider(this.labelProvider);
 
-            this.ctx.representation.structure.themes.colorThemeRegistry.add(AlphaFoldConfidenceColorThemeProvider);
+            this.ctx.representation.structure.themes.colorThemeRegistry.add(PLDDTConfidenceColorThemeProvider);
         }
 
         update(p: { autoAttach: boolean, showTooltip: boolean }) {
@@ -58,9 +58,9 @@ export const AlphaFoldConfidenceScore = PluginBehavior.create<{ autoAttach: bool
         }
 
         unregister() {
-            this.ctx.customModelProperties.unregister(AlphaFoldConfidenceProvider.descriptor.name);
+            this.ctx.customModelProperties.unregister(PLDDTConfidenceProvider.descriptor.name);
             this.ctx.managers.lociLabels.removeProvider(this.labelProvider);
-            this.ctx.representation.structure.themes.colorThemeRegistry.remove(AlphaFoldConfidenceColorThemeProvider);
+            this.ctx.representation.structure.themes.colorThemeRegistry.remove(PLDDTConfidenceColorThemeProvider);
         }
     },
     params: () => ({

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

@@ -5,7 +5,7 @@
  * @author Sebastian Bittrich <sebastian.bittrich@rcsb.org>
  */
 
-import { AlphaFoldConfidence, AlphaFoldConfidenceProvider } from './prop';
+import { PLDDTConfidence, PLDDTConfidenceProvider } 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';
@@ -26,8 +26,8 @@ const ConfidenceColors: { [k: string]: Color } = {
 
 const ConfidenceColorLegend = TableLegend(Object.entries(ConfidenceColors));
 
-export function getAlphaFoldConfidenceColorThemeParams(ctx: ThemeDataContext) {
-    const categories = AlphaFoldConfidence.getCategories(ctx.structure);
+export function getPLDDTConfidenceColorThemeParams(ctx: ThemeDataContext) {
+    const categories = PLDDTConfidence.getCategories(ctx.structure);
     if (categories.length === 0) {
         return {
             type: PD.MappedStatic('score', {
@@ -45,14 +45,14 @@ export function getAlphaFoldConfidenceColorThemeParams(ctx: ThemeDataContext) {
         })
     };
 }
-export type AlphaFoldConfidenceColorThemeParams = ReturnType<typeof getAlphaFoldConfidenceColorThemeParams>
+export type PLDDTConfidenceColorThemeParams = ReturnType<typeof getPLDDTConfidenceColorThemeParams>
 
-export function AlphaFoldConfidenceColorTheme(ctx: ThemeDataContext, props: PD.Values<AlphaFoldConfidenceColorThemeParams>): ColorTheme<AlphaFoldConfidenceColorThemeParams> {
+export function PLDDTConfidenceColorTheme(ctx: ThemeDataContext, props: PD.Values<PLDDTConfidenceColorThemeParams>): ColorTheme<PLDDTConfidenceColorThemeParams> {
     let color: LocationColor = () => DefaultColor;
 
-    if (ctx.structure && ctx.structure.models.length > 0 && ctx.structure.models[0].customProperties.has(AlphaFoldConfidenceProvider.descriptor)) {
+    if (ctx.structure && ctx.structure.models.length > 0 && ctx.structure.models[0].customProperties.has(PLDDTConfidenceProvider.descriptor)) {
         const getColor = (location: StructureElement.Location): Color => {
-            const score: string = AlphaFoldConfidence.getConfidenceScore(location)[1];
+            const score: string = PLDDTConfidence.getConfidenceScore(location)[1];
 
             if (props.type.name !== 'score') {
                 const categoryProp = props.type.params.kind;
@@ -71,25 +71,25 @@ export function AlphaFoldConfidenceColorTheme(ctx: ThemeDataContext, props: PD.V
     }
 
     return {
-        factory: AlphaFoldConfidenceColorTheme,
+        factory: PLDDTConfidenceColorTheme,
         granularity: 'group',
         color,
         props,
-        description: 'Assigns residue colors according to the AlphaFold Confidence score.',
+        description: 'Assigns residue colors according to the pLDDT Confidence score.',
         legend: ConfidenceColorLegend
     };
 }
 
-export const AlphaFoldConfidenceColorThemeProvider: ColorTheme.Provider<AlphaFoldConfidenceColorThemeParams, 'af-confidence'> = {
-    name: 'af-confidence',
-    label: 'AlphaFold Confidence',
+export const PLDDTConfidenceColorThemeProvider: ColorTheme.Provider<PLDDTConfidenceColorThemeParams, 'plddt-confidence'> = {
+    name: 'plddt-confidence',
+    label: 'pLDDT Confidence',
     category: ColorTheme.Category.Validation,
-    factory: AlphaFoldConfidenceColorTheme,
-    getParams: getAlphaFoldConfidenceColorThemeParams,
-    defaultValues: PD.getDefaultValues(getAlphaFoldConfidenceColorThemeParams({})),
-    isApplicable: (ctx: ThemeDataContext) => AlphaFoldConfidence.isApplicable(ctx.structure?.models[0]),
+    factory: PLDDTConfidenceColorTheme,
+    getParams: getPLDDTConfidenceColorThemeParams,
+    defaultValues: PD.getDefaultValues(getPLDDTConfidenceColorThemeParams({})),
+    isApplicable: (ctx: ThemeDataContext) => PLDDTConfidence.isApplicable(ctx.structure?.models[0]),
     ensureCustomProperties: {
-        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)
+        attach: (ctx: CustomProperty.Context, data: ThemeDataContext) => data.structure ? PLDDTConfidenceProvider.attach(ctx, data.structure.models[0], void 0, true) : Promise.resolve(),
+        detach: (data) => data.structure && PLDDTConfidenceProvider.ref(data.structure.models[0], false)
     }
 };

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

@@ -20,14 +20,19 @@ import { CustomPropertyDescriptor } from 'molstar/lib/mol-model/custom-property'
 import { arraySetAdd } from 'molstar/lib/mol-util/array';
 import { dateToUtcString } from 'molstar/lib/mol-util/date';
 
-export type AlphaFoldConfidence = PropertyWrapper<{
+export type PLDDTConfidence = PropertyWrapper<{
     score: IndexedCustomProperty.Residue<[number, string]>,
     category: string[]
 }>
 
-export namespace AlphaFoldConfidence {
+export namespace PLDDTConfidence {
+    const AlphaFoldNamespace = 'AF-';
+    const BakerNamespace = 'MA-BAK-';
+
     export function isApplicable(model?: Model): boolean {
         if (!model || !MmcifFormat.is(model.sourceData)) return false;
+        const entryId = model.entryId.toUpperCase();
+        if (!entryId.startsWith(AlphaFoldNamespace) && !entryId.startsWith(BakerNamespace)) return false;
         return model.sourceData.data.frame.categoryNames.includes('ma_qa_metric_local');
     }
 
@@ -60,7 +65,7 @@ export namespace AlphaFoldConfidence {
     }
 
 
-    export function fromCif(ctx: CustomProperty.Context, model: Model): AlphaFoldConfidence | undefined {
+    export function fromCif(ctx: CustomProperty.Context, model: Model): PLDDTConfidence | undefined {
         const info = tryGetInfoFromCif('ma_qa_metric_local', model);
         if (!info) return;
         const data = getCifData(model);
@@ -68,14 +73,14 @@ export namespace AlphaFoldConfidence {
         return { info, data: metricMap };
     }
 
-    export async function obtain(ctx: CustomProperty.Context, model: Model, _props: AlphaFoldConfidenceProps): Promise<CustomProperty.Data<any>> {
+    export async function obtain(ctx: CustomProperty.Context, model: Model, _props: PLDDTConfidenceProps): Promise<CustomProperty.Data<any>> {
         const cif = fromCif(ctx, model);
         return { value: cif };
     }
 
     export function getConfidenceScore(e: StructureElement.Location): [number, string] {
         if (!Unit.isAtomic(e.unit)) return [-1, 'No Score'];
-        const prop = AlphaFoldConfidenceProvider.get(e.unit.model).value;
+        const prop = PLDDTConfidenceProvider.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'];
@@ -84,7 +89,7 @@ export namespace AlphaFoldConfidence {
     const _emptyArray: string[] = [];
     export function getCategories(structure?: Structure) {
         if (!structure) return _emptyArray;
-        const prop = AlphaFoldConfidenceProvider.get(structure.models[0]).value;
+        const prop = PLDDTConfidenceProvider.get(structure.models[0]).value;
         if (!prop || !prop.data) return _emptyArray;
         return prop.data.category;
     }
@@ -97,26 +102,26 @@ export namespace AlphaFoldConfidence {
     }
 }
 
-export const AlphaFoldConfidenceParams = {};
-export type AlphaFoldConfidenceParams = typeof AlphaFoldConfidenceParams
-export type AlphaFoldConfidenceProps = PD.Values<AlphaFoldConfidenceParams>
+export const PLDDTConfidenceParams = {};
+export type PLDDTConfidenceParams = typeof PLDDTConfidenceParams
+export type PLDDTConfidenceProps = PD.Values<PLDDTConfidenceParams>
 
-export const AlphaFoldConfidenceProvider: CustomModelProperty.Provider<AlphaFoldConfidenceParams, AlphaFoldConfidence> = CustomModelProperty.createProvider({
-    label: 'AlphaFold Confidence Score',
+export const PLDDTConfidenceProvider: CustomModelProperty.Provider<PLDDTConfidenceParams, PLDDTConfidence> = CustomModelProperty.createProvider({
+    label: 'pLDDT Confidence Score',
     descriptor: CustomPropertyDescriptor({
-        name: 'af_confidence_score'
+        name: 'plddt_confidence_score'
     }),
     type: 'static',
-    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);
+    defaultParams: PLDDTConfidenceParams,
+    getParams: () => PLDDTConfidenceParams,
+    isApplicable: (data: Model) => PLDDTConfidence.isApplicable(data),
+    obtain: async (ctx: CustomProperty.Context, data: Model, props: Partial<PLDDTConfidenceProps>) => {
+        const p = { ...PD.getDefaultValues(PLDDTConfidenceParams), ...props };
+        return await PLDDTConfidence.obtain(ctx, data, p);
     }
 });
 
-function createScoreMapFromCif(modelData: Model, residueData: Table<typeof AlphaFoldConfidence.Schema.local_metric_values>): AlphaFoldConfidence['data'] {
+function createScoreMapFromCif(modelData: Model, residueData: Table<typeof PLDDTConfidence.Schema.local_metric_values>): PLDDTConfidence['data'] {
     const { label_asym_id, label_seq_id, metric_value, _rowCount } = residueData;
 
     const ret = new Map<ResidueIndex, [number, string]>();

src/viewer/index.ts

@@ -37,7 +37,7 @@ import { DefaultPluginUISpec, PluginUISpec } from 'molstar/lib/mol-plugin-ui/spe
 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 { AlphaFoldConfidenceScore } from './helpers/af-confidence/behavior';
+import { PLLDTConfidenceScore } from './helpers/plddt-confidence/behavior';
 import { PluginContext } from 'molstar/lib/mol-plugin/context';
 import { TrajectoryHierarchyPresetProvider } from 'molstar/lib/mol-plugin-state/builder/structure/hierarchy-preset';
 
@@ -54,7 +54,7 @@ const Extensions = {
     'rcsb-assembly-symmetry': PluginSpec.Behavior(RCSBAssemblySymmetry),
     'rcsb-validation-report': PluginSpec.Behavior(RCSBValidationReport),
     'anvil-membrane-orientation': PluginSpec.Behavior(ANVILMembraneOrientation),
-    'af-confidence': PluginSpec.Behavior(AlphaFoldConfidenceScore)
+    'plddt-confidence': PluginSpec.Behavior(PLLDTConfidenceScore)
 };
 
 const DefaultViewerProps = {

src/viewer/ui/validation.tsx

@@ -11,8 +11,8 @@ import { ValidationReportGeometryQualityPreset } from 'molstar/lib/extensions/rc
 import { ActionMenu } from 'molstar/lib/mol-plugin-ui/controls/action-menu';
 import { Model } from 'molstar/lib/mol-model/structure/model';
 import { MmcifFormat } from 'molstar/lib/mol-model-formats/structure/mmcif';
-import { AlphaFoldConfidence } from '../helpers/af-confidence/prop';
-import { AlphaFoldConfidenceColorThemeProvider } from '../helpers/af-confidence/color';
+import { PLDDTConfidence } from '../helpers/plddt-confidence/prop';
+import { PLDDTConfidenceColorThemeProvider } from '../helpers/plddt-confidence/color';
 
 interface ValidationReportState extends CollapsableState {
     errorStates: Set<string>
@@ -66,7 +66,7 @@ export class ValidationReportControls extends CollapsableControls<{}, Validation
         if (!pivot.obj || pivot.obj.data.models.length !== 1) return false;
         const model = pivot.obj.data.models[0];
         // all supported options must be registered here
-        return ValidationReport.isApplicable(model) || AlphaFoldConfidence.isApplicable(model);
+        return ValidationReport.isApplicable(model) || PLDDTConfidence.isApplicable(model);
     }
 
     get noValidationReport() {
@@ -76,11 +76,11 @@ export class ValidationReportControls extends CollapsableControls<{}, Validation
         return !model || !this.isFromPdbArchive(model);
     }
 
-    get alphaFoldData() {
+    get pLDDTData() {
         const structure = this.pivot.cell.obj?.data;
         if (!structure || structure.models.length !== 1) return false;
         const model = structure.models[0];
-        return AlphaFoldConfidence.isApplicable(model);
+        return PLDDTConfidence.isApplicable(model);
     }
 
     isFromPdbArchive(model: Model) {
@@ -102,8 +102,8 @@ export class ValidationReportControls extends CollapsableControls<{}, Validation
         }
     };
 
-    requestAlphaFoldConfidenceColoring = async () => {
-        await this.plugin.managers.structure.component.updateRepresentationsTheme(this.pivot.components, { color: AlphaFoldConfidenceColorThemeProvider.name as any });
+    requestPLDDTConfidenceColoring = async () => {
+        await this.plugin.managers.structure.component.updateRepresentationsTheme(this.pivot.components, { color: PLDDTConfidenceColorThemeProvider.name as any });
     };
 
     get actions(): ActionMenu.Items {
@@ -118,11 +118,11 @@ export class ValidationReportControls extends CollapsableControls<{}, Validation
             },
         ];
 
-        if (this.alphaFoldData) {
+        if (this.pLDDTData) {
             out.push({
                 kind: 'item',
-                label: 'AlphaFold Confidence Scores',
-                value: this.requestAlphaFoldConfidenceColoring
+                label: 'pLDDT Confidence Scores',
+                value: this.requestPLDDTConfidenceColoring
             });
         }