Skip to content
Snippets Groups Projects
Commit f028ad05 authored by David Sehnal's avatar David Sehnal
Browse files

mol-model: Added Atom and Entity level custom props

parent 828c4dfd
No related branches found
No related tags found
No related merge requests found
...@@ -140,7 +140,7 @@ function createIssueMapFromJson(modelData: Model, data: any): StructureQualityRe ...@@ -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 { 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 ...@@ -153,5 +153,5 @@ function createIssueMapFromCif(modelData: Model, data: Table<typeof StructureQua
ret.set(idx, issues.value(i)); ret.set(idx, issues.value(i));
} }
return IndexedCustomProperty.fromResidueMap(ret, Unit.Kind.Atomic); return IndexedCustomProperty.fromResidueMap(ret);
} }
\ No newline at end of file
...@@ -4,14 +4,85 @@ ...@@ -4,14 +4,85 @@
* @author David Sehnal <david.sehnal@gmail.com> * @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 { Unit, Structure, StructureElement } from '../../../structure';
import { Segmentation } from 'mol-data/int'; import { Segmentation } from 'mol-data/int';
import { UUID } from 'mol-util'; import { UUID } from 'mol-util';
import { CifWriter } from 'mol-io/writer/cif'; import { CifWriter } from 'mol-io/writer/cif';
import { Model } from '../../model'; 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 id: UUID = UUID.create();
readonly kind: Unit.Kind; readonly kind: Unit.Kind;
has(idx: Idx): boolean { return this.map.has(idx); } has(idx: Idx): boolean { return this.map.has(idx); }
...@@ -25,7 +96,7 @@ export class IndexedCustomProperty<Idx extends IndexedCustomProperty.Index, T = ...@@ -25,7 +96,7 @@ export class IndexedCustomProperty<Idx extends IndexedCustomProperty.Index, T =
const unitGroups = structure.unitSymmetryGroups; const unitGroups = structure.unitSymmetryGroups;
const loci: StructureElement[] = []; const loci: StructureElement[] = [];
const segments = this.segmentGetter(models[0]) const segments = this.segmentGetter(models[0]);
for (const unitGroup of unitGroups) { for (const unitGroup of unitGroups) {
const unit = unitGroup.units[0]; const unit = unitGroup.units[0];
...@@ -52,37 +123,101 @@ export class IndexedCustomProperty<Idx extends IndexedCustomProperty.Index, T = ...@@ -52,37 +123,101 @@ export class IndexedCustomProperty<Idx extends IndexedCustomProperty.Index, T =
return { elements, property: i => this.get(index[elements[i].element])! }; 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; this.kind = kind;
} }
} }
export namespace IndexedCustomProperty { class ElementMappedCustomProperty<T = any> implements IndexedCustomProperty<ElementIndex, T> {
export type Index = ResidueIndex | ChainIndex 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> { private getStructureElements(structure: Structure) {
elements: StructureElement[], const models = structure.models;
property(index: number): T 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;
} }
export function getCifDataSource<Idx extends Index, T>(structure: Structure, prop: IndexedCustomProperty<Idx, T> | undefined, cache: any): CifWriter.Category.Instance['source'][0] { const elements = unit.elements;
if (!prop) return { rowCount: 0 }; for (let i = 0, _i = elements.length; i < _i; i++) {
if (cache && cache[prop.id]) return cache[prop.id]; const e = elements[i];
const data = prop.getExportContext(structure); if (!this.has(e) || seenIndices.has(e)) continue;
const ret = { data, rowCount: data.elements.length }; seenIndices.add(elements[i]);
if (cache) cache[prop.id] = ret; loci[loci.length] = StructureElement.create(unit, e);
return ret; }
} }
export type Residue<T> = IndexedCustomProperty<ResidueIndex, T> loci.sort((x, y) => x.element - y.element);
const getResidueSegments = (model: Model) => model.atomicHierarchy.residueAtomSegments; return loci;
export function fromResidueMap<T>(map: Map<ResidueIndex, T>, kind: Unit.Kind): Residue<T> {
return new IndexedCustomProperty(map, getResidueSegments, kind);
} }
export type Chain<T> = IndexedCustomProperty<ChainIndex, T> getExportContext(structure: Structure): IndexedCustomProperty.ExportCtx<T> {
const getChainSegments = (model: Model) => model.atomicHierarchy.chainAtomSegments; const elements = this.getStructureElements(structure);
export function fromChainMap<T>(map: Map<ChainIndex, T>, kind: Unit.Kind): Chain<T> { return { elements, property: i => this.get(elements[i].element)! };
return new IndexedCustomProperty(map, getChainSegments, kind); }
constructor(private map: Map<ElementIndex, T>) {
this.kind = Unit.Kind.Atomic;
}
}
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;
} }
} }
\ No newline at end of file
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment