mol-model: Added Atom and Entity level custom props

src/mol-model-props/pdbe/structure-quality-report.ts

@@ -140,7 +140,7 @@ function createIssueMapFromJson(modelData: Model, data: any): StructureQualityRe
         }
     }
 
-    return IndexedCustomProperty.fromResidueMap(ret, Unit.Kind.Atomic);
+    return IndexedCustomProperty.fromResidueMap(ret);
 }
 
 function createIssueMapFromCif(modelData: Model, data: Table<typeof StructureQualityReport.Schema.pdbe_structure_quality_report_issues>): StructureQualityReport.IssueMap | undefined {
@@ -153,5 +153,5 @@ function createIssueMapFromCif(modelData: Model, data: Table<typeof StructureQua
         ret.set(idx, issues.value(i));
     }
 
-    return IndexedCustomProperty.fromResidueMap(ret, Unit.Kind.Atomic);
+    return IndexedCustomProperty.fromResidueMap(ret);
 }

src/mol-model/structure/model/properties/custom/indexed.ts

@@ -4,14 +4,85 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { ResidueIndex, ChainIndex, ElementIndex } from '../../indexing';
+import { ResidueIndex, ChainIndex, ElementIndex, EntityIndex } from '../../indexing';
 import { Unit, Structure, StructureElement } from '../../../structure';
 import { Segmentation } from 'mol-data/int';
 import { UUID } from 'mol-util';
 import { CifWriter } from 'mol-io/writer/cif';
 import { Model } from '../../model';
 
-export class IndexedCustomProperty<Idx extends IndexedCustomProperty.Index, T = any> {
+export interface IndexedCustomProperty<Idx extends IndexedCustomProperty.Index, T = any> {
+    readonly id: UUID,
+    readonly kind: Unit.Kind,
+    readonly level: IndexedCustomProperty.Level,
+    has(idx: Idx): boolean,
+    get(idx: Idx): T | undefined,
+    getExportContext(structure: Structure): IndexedCustomProperty.ExportCtx<T>
+}
+
+export namespace IndexedCustomProperty {
+    export type Index = ElementIndex | ResidueIndex | ChainIndex | EntityIndex
+    export type Level = 'atom' | 'residue' | 'chain' | 'entity'
+
+    export interface ExportCtx<T> {
+        elements: StructureElement[],
+        property(index: number): T
+    }
+
+    export function getCifDataSource<Idx extends Index, T>(structure: Structure, prop: IndexedCustomProperty<Idx, T> | undefined, cache: any): CifWriter.Category.Instance['source'][0] {
+        if (!prop) return { rowCount: 0 };
+        if (cache && cache[prop.id]) return cache[prop.id];
+        const data = prop.getExportContext(structure);
+        const ret = { data, rowCount: data.elements.length };
+        if (cache) cache[prop.id] = ret;
+        return ret;
+    }
+
+    export type Atom<T> = IndexedCustomProperty<ElementIndex, T>
+    export function fromAtomMap<T>(map: Map<ElementIndex, T>): Atom<T> {
+        return new ElementMappedCustomProperty(map);
+    }
+
+    export function fromAtomArray<T>(array: ArrayLike<T>): Atom<T> {
+        // TODO: create "array based custom property" as optimization
+        return new ElementMappedCustomProperty(arrayToMap(array));
+    }
+
+    export type Residue<T> = IndexedCustomProperty<ResidueIndex, T>
+    const getResidueSegments = (model: Model) => model.atomicHierarchy.residueAtomSegments;
+    export function fromResidueMap<T>(map: Map<ResidueIndex, T>): Residue<T> {
+        return new SegmentedMappedIndexedCustomProperty('residue', map, getResidueSegments, Unit.Kind.Atomic);
+    }
+
+    export function fromResidueArray<T>(array: ArrayLike<T>): Residue<T> {
+        // TODO: create "array based custom property" as optimization
+        return new SegmentedMappedIndexedCustomProperty('residue', arrayToMap(array), getResidueSegments, Unit.Kind.Atomic);
+    }
+
+    export type Chain<T> = IndexedCustomProperty<ChainIndex, T>
+    const getChainSegments = (model: Model) => model.atomicHierarchy.chainAtomSegments;
+    export function fromChainMap<T>(map: Map<ChainIndex, T>): Chain<T> {
+        return new SegmentedMappedIndexedCustomProperty('chain', map, getChainSegments, Unit.Kind.Atomic);
+    }
+
+    export function fromChainArray<T>(array: ArrayLike<T>): Chain<T> {
+        // TODO: create "array based custom property" as optimization
+        return new SegmentedMappedIndexedCustomProperty('chain', arrayToMap(array), getChainSegments, Unit.Kind.Atomic);
+    }
+
+    export type Entity<T> = IndexedCustomProperty<EntityIndex, T>
+    export function fromEntityMap<T>(map: Map<EntityIndex, T>): Entity<T> {
+        return new EntityMappedCustomProperty(map);
+    }
+}
+
+function arrayToMap<Idx extends IndexedCustomProperty.Index, T>(array: ArrayLike<T>): Map<Idx, T> {
+    const ret = new Map<Idx, T>();
+    for (let i = 0 as Idx, _i = array.length; i < _i; i++) ret.set(i, array[i]);
+    return ret;
+}
+
+class SegmentedMappedIndexedCustomProperty<Idx extends IndexedCustomProperty.Index, T = any> implements IndexedCustomProperty<Idx, T> {
     readonly id: UUID = UUID.create();
     readonly kind: Unit.Kind;
     has(idx: Idx): boolean { return this.map.has(idx); }
@@ -25,7 +96,7 @@ export class IndexedCustomProperty<Idx extends IndexedCustomProperty.Index, T =
         const unitGroups = structure.unitSymmetryGroups;
         const loci: StructureElement[] = [];
 
-        const segments = this.segmentGetter(models[0])
+        const segments = this.segmentGetter(models[0]);
 
         for (const unitGroup of unitGroups) {
             const unit = unitGroup.units[0];
@@ -52,37 +123,101 @@ export class IndexedCustomProperty<Idx extends IndexedCustomProperty.Index, T =
         return { elements, property: i => this.get(index[elements[i].element])! };
     }
 
-    constructor(private map: Map<Idx, T>, private segmentGetter: (model: Model) => Segmentation<ElementIndex, Idx>, kind: Unit.Kind) {
+    constructor(public level: 'residue' | 'chain', private map: Map<Idx, T>, private segmentGetter: (model: Model) => Segmentation<ElementIndex, Idx>, kind: Unit.Kind) {
         this.kind = kind;
     }
 }
 
-export namespace IndexedCustomProperty {
-    export type Index = ResidueIndex | ChainIndex
+class ElementMappedCustomProperty<T = any> implements IndexedCustomProperty<ElementIndex, T> {
+    readonly id: UUID = UUID.create();
+    readonly kind: Unit.Kind;
+    readonly level = 'atom';
+    has(idx: ElementIndex): boolean { return this.map.has(idx); }
+    get(idx: ElementIndex) { return this.map.get(idx); }
 
-    export interface ExportCtx<T> {
-        elements: StructureElement[],
-        property(index: number): T
+    private getStructureElements(structure: Structure) {
+        const models = structure.models;
+        if (models.length !== 1) throw new Error(`Only works on structures with a single model.`);
+
+        const seenIndices = new Set<ElementIndex>();
+        const unitGroups = structure.unitSymmetryGroups;
+        const loci: StructureElement[] = [];
+
+        for (const unitGroup of unitGroups) {
+            const unit = unitGroup.units[0];
+            if (unit.kind !== this.kind) {
+                continue;
+            }
+
+            const elements = unit.elements;
+            for (let i = 0, _i = elements.length; i < _i; i++) {
+                const e = elements[i];
+                if (!this.has(e) || seenIndices.has(e)) continue;
+                seenIndices.add(elements[i]);
+                loci[loci.length] = StructureElement.create(unit, e);
+            }
+        }
+
+        loci.sort((x, y) => x.element - y.element);
+        return loci;
     }
 
-    export function getCifDataSource<Idx extends Index, T>(structure: Structure, prop: IndexedCustomProperty<Idx, T> | undefined, cache: any): CifWriter.Category.Instance['source'][0] {
-        if (!prop) return { rowCount: 0 };
-        if (cache && cache[prop.id]) return cache[prop.id];
-        const data = prop.getExportContext(structure);
-        const ret = { data, rowCount: data.elements.length };
-        if (cache) cache[prop.id] = ret;
-        return ret;
+    getExportContext(structure: Structure): IndexedCustomProperty.ExportCtx<T> {
+        const elements = this.getStructureElements(structure);
+        return { elements, property: i => this.get(elements[i].element)! };
     }
 
-    export type Residue<T> = IndexedCustomProperty<ResidueIndex, T>
-    const getResidueSegments = (model: Model) => model.atomicHierarchy.residueAtomSegments;
-    export function fromResidueMap<T>(map: Map<ResidueIndex, T>, kind: Unit.Kind): Residue<T> {
-        return new IndexedCustomProperty(map, getResidueSegments, kind);
+    constructor(private map: Map<ElementIndex, T>) {
+        this.kind = Unit.Kind.Atomic;
     }
+}
 
-    export type Chain<T> = IndexedCustomProperty<ChainIndex, T>
-    const getChainSegments = (model: Model) => model.atomicHierarchy.chainAtomSegments;
-    export function fromChainMap<T>(map: Map<ChainIndex, T>, kind: Unit.Kind): Chain<T> {
-        return new IndexedCustomProperty(map, getChainSegments, kind);
+class EntityMappedCustomProperty<T = any> implements IndexedCustomProperty<EntityIndex, T> {
+    readonly id: UUID = UUID.create();
+    readonly kind: Unit.Kind;
+    readonly level = 'entity';
+    has(idx: EntityIndex): boolean { return this.map.has(idx); }
+    get(idx: EntityIndex) { return this.map.get(idx); }
+
+    private getStructureElements(structure: Structure) {
+        const models = structure.models;
+        if (models.length !== 1) throw new Error(`Only works on structures with a single model.`);
+
+        const index = models[0].atomicHierarchy.index;
+        const seenIndices = new Set<EntityIndex>();
+        const unitGroups = structure.unitSymmetryGroups;
+        const loci: StructureElement[] = [];
+
+        const segments = models[0].atomicHierarchy.chainAtomSegments;
+
+        for (const unitGroup of unitGroups) {
+            const unit = unitGroup.units[0];
+            if (unit.kind !== this.kind) {
+                continue;
+            }
+
+            const chains = Segmentation.transientSegments(segments, unit.elements);
+            while (chains.hasNext) {
+                const seg = chains.move();
+                const eI = index.getEntityFromChain(seg.index);
+                if (!this.has(eI) || seenIndices.has(eI)) continue;
+                seenIndices.add(eI);
+                loci[loci.length] = StructureElement.create(unit, unit.elements[seg.start]);
+            }
+        }
+
+        loci.sort((x, y) => x.element - y.element);
+        return loci;
+    }
+
+    getExportContext(structure: Structure): IndexedCustomProperty.ExportCtx<T> {
+        const elements = this.getStructureElements(structure);
+        const chainIndex = structure.model.atomicHierarchy.chainAtomSegments.index;
+        const index = structure.model.atomicHierarchy.index;
+        return { elements, property: i => this.get(index.getEntityFromChain(chainIndex[elements[i].element]))! };
+    }
+
+    constructor(private map: Map<EntityIndex, T>) {
+        this.kind = Unit.Kind.Atomic;
     }
 }