refactoring CIF exporter

src/apps/cif2bcif/converter.ts

@@ -27,7 +27,7 @@ async function getCIF(ctx: RuntimeContext, path: string) {
 function getCategoryInstanceProvider(cat: CifCategory, fields: CifWriter.Field[]): CifWriter.Category {
     return {
         name: cat.name,
-        instance: () => ({ data: cat, fields, rowCount: cat.rowCount })
+        instance: () => CifWriter.categoryInstance(fields, { data: cat, rowCount: cat.rowCount })
     };
 }
 

src/mol-io/writer/cif.ts

@@ -40,6 +40,10 @@ export namespace CifWriter {
         fixedPoint3: E.by(E.fixedPoint(1000)).and(E.delta).and(E.integerPacking),
     };
 
+    export function categoryInstance<Key, Data>(fields: Field<Key, Data>[], source: Category.DataSource): Category.Instance {
+        return { fields, source: [source] };
+    }
+
     export function createEncodingProviderFromCifFrame(frame: CifFrame): EncodingProvider {
         return {
             get(c, f) {

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

@@ -28,17 +28,21 @@ const _Descriptor = ModelPropertyDescriptor({
         categories: [{
             name: 'pdbe_structure_quality_report',
             instance(ctx) {
-                const issues = StructureQualityReport.get(ctx.model);
-                if (typeof ctx.globalCache.pdbe_structure_quality_report !== 'undefined' && ctx.globalCache.pdbe_structure_quality_report !== ctx.model.modelNum) return CifWriter.Category.Empty;
-                ctx.globalCache.pdbe_structure_quality_report = ctx.model.modelNum;
-                return { fields: _structure_quality_report_fields, rowCount: 1, data: issues ? issues.updated : 'n/a' }
+                const structure = ctx.structures[0];
+                const issues = StructureQualityReport.get(structure.model);
+                return {
+                    fields: _structure_quality_report_fields,
+                    source: [{ data: issues ? issues.updated : 'n/a', rowCount: 1 }]
+                };
             }
         }, {
             name: 'pdbe_structure_quality_report_issues',
             instance(ctx) {
-                const issues = StructureQualityReport.get(ctx.model);
-                if (!issues || !issues.map) return CifWriter.Category.Empty;
-                return ResidueCustomProperty.createCifCategory(ctx, issues.map, _structure_quality_report_issues_fields);
+                return {
+                    fields: _structure_quality_report_issues_fields,
+                    source: ctx.structures.map(
+                        s => ResidueCustomProperty.getCifDataSource(s, StructureQualityReport.getIssueMap(s.model), ctx.cache))
+                };
             }
         }]
     },
@@ -150,6 +154,11 @@ export namespace StructureQualityReport {
         return model._dynamicPropertyData.__StructureQualityReport__;
     }
 
+    export function getIssueMap(model: Model): IssueMap | undefined {
+        const data = get(model);
+        return data && data.map;
+    }
+
     const _emptyArray: string[] = [];
     export function getIssues(e: StructureElement) {
         if (!Unit.isAtomic(e.unit)) return _emptyArray;

src/mol-model-props/rcsb/symmetry.ts

@@ -22,7 +22,7 @@ const { str, int, float, Aliased, Vector, List } = Column.Schema;
 
 function getInstance(name: keyof AssemblySymmetry.Schema): (ctx: CifExportContext) => CifWriter.Category.Instance<any, any> {
     return function(ctx: CifExportContext) {
-        const assemblySymmetry = AssemblySymmetry.get(ctx.model);
+        const assemblySymmetry = AssemblySymmetry.get(ctx.structures[0].model);
         return assemblySymmetry ? Category.ofTable(assemblySymmetry.db[name]) : CifWriter.Category.Empty;
     }
 }

src/mol-model/structure/export/categories/atom_site.ts

@@ -51,12 +51,14 @@ const atom_site_fields = CifWriter.fields<StructureElement, Structure>()
 
 export const _atom_site: CifCategory<CifExportContext> = {
     name: 'atom_site',
-    instance({ structure }: CifExportContext) {
+    instance({ structures }: CifExportContext) {
         return {
             fields: atom_site_fields,
-            data: structure,
-            rowCount: structure.elementCount,
-            keys: () => structure.elementLocations()
+            source: structures.map(s => ({
+                data: s,
+                rowCount: s.elementCount,
+                keys: () => s.elementLocations()
+            }))
         };
     }
 }

src/mol-model/structure/export/categories/modified-residues.ts

@@ -27,7 +27,9 @@ const pdbx_struct_mod_residue_fields: CifField<number, StructureElement[]>[] = [
     CifField.str('details', (i, xs) => xs[i].unit.model.properties.modifiedResidues.details.get(P.residue.label_comp_id(xs[i]))!)
 ];
 
-function getModifiedResidues({ model, structure }: CifExportContext): StructureElement[] {
+function getModifiedResidues({ structures }: CifExportContext): StructureElement[] {
+    // TODO: can different models have differnt modified residues?
+    const structure = structures[0], model = structure.model;
     const map = model.properties.modifiedResidues.parentId;
     if (!map.size) return [];
 
@@ -54,6 +56,9 @@ export const _pdbx_struct_mod_residue: CifCategory<CifExportContext> = {
     name: 'pdbx_struct_mod_residue',
     instance(ctx) {
         const residues = getModifiedResidues(ctx);
-        return { fields: pdbx_struct_mod_residue_fields, data: residues, rowCount: residues.length };
+        return {
+            fields: pdbx_struct_mod_residue_fields,
+            source: [{ data: residues, rowCount: residues.length }]
+        };
     }
 }

src/mol-model/structure/export/categories/secondary-structure.ts

@@ -18,7 +18,10 @@ export const _struct_conf: CifCategory<CifExportContext> = {
     name: 'struct_conf',
     instance(ctx) {
         const elements = findElements(ctx, 'helix');
-        return { fields: struct_conf_fields, data: elements, rowCount: elements.length };
+        return {
+            fields: struct_conf_fields,
+            source: [{ data: elements, rowCount: elements.length }]
+        };
     }
 };
 
@@ -26,7 +29,10 @@ export const _struct_sheet_range: CifCategory<CifExportContext> = {
     name: 'struct_sheet_range',
     instance(ctx) {
         const elements = (findElements(ctx, 'sheet') as SSElement<SecondaryStructure.Sheet>[]).sort(compare_ssr);
-        return { fields: struct_sheet_range_fields, data: elements, rowCount: elements.length };
+        return {
+            fields: struct_sheet_range_fields,
+            source: [{ data: elements, rowCount: elements.length }]
+        };
     }
 };
 
@@ -63,11 +69,12 @@ interface SSElement<T extends SecondaryStructure.Element> {
 }
 
 function findElements<T extends SecondaryStructure.Element>(ctx: CifExportContext, kind: SecondaryStructure.Element['kind']) {
-    const { key, elements } = ctx.model.properties.secondaryStructure;
+    // TODO: encode secondary structure for different models?
+    const { key, elements } = ctx.structures[0].model.properties.secondaryStructure;
 
     const ssElements: SSElement<any>[] = [];
 
-    for (const unit of ctx.structure.units) {
+    for (const unit of ctx.structures[0].units) {
         // currently can only support this for "identity" operators.
         if (!Unit.isAtomic(unit) || !unit.conformation.operator.isIdentity) continue;
 

src/mol-model/structure/export/mmcif.ts

@@ -8,32 +8,30 @@
 import { CifWriter } from 'mol-io/writer/cif'
 import { mmCIF_Schema } from 'mol-io/reader/cif/schema/mmcif'
 import { Structure } from '../structure'
-import { Model } from '../model'
 import { _atom_site } from './categories/atom_site';
 import CifCategory = CifWriter.Category
 import { _struct_conf, _struct_sheet_range } from './categories/secondary-structure';
 import { _pdbx_struct_mod_residue } from './categories/modified-residues';
 
 export interface CifExportContext {
-    structure: Structure,
-    model: Model,
-    localCache: any,
-    /** useful when exporting multiple models at the same time */
-    globalCache: any
+    structures: Structure[],
+    cache: any
 }
 
 export namespace CifExportContext {
-    export function create(structures: Structure | Structure[]): CifExportContext[] {
-        const globalCache = Object.create(null);
-        if (Array.isArray(structures)) return structures.map(structure => ({ structure, model: structure.models[0], localCache: Object.create(null), globalCache }));
-        return [{ structure: structures, model: structures.models[0], localCache: Object.create(null), globalCache }];
+    export function create(structures: Structure | Structure[]): CifExportContext {
+        return {
+            structures: Array.isArray(structures) ? structures : [structures],
+            cache: Object.create(null)
+        };
     }
 }
 
 function copy_mmCif_category(name: keyof mmCIF_Schema): CifCategory<CifExportContext> {
     return {
         name,
-        instance({ model }) {
+        instance({ structures }) {
+            const model = structures[0].model;
             if (model.sourceData.kind !== 'mmCIF') return CifCategory.Empty;
             const table = model.sourceData.data[name];
             if (!table || !table._rowCount) return CifCategory.Empty;
@@ -44,9 +42,9 @@ function copy_mmCif_category(name: keyof mmCIF_Schema): CifCategory<CifExportCon
 
 const _entity: CifCategory<CifExportContext> = {
     name: 'entity',
-    instance({ structure, model}) {
-        const keys = Structure.getEntityKeys(structure);
-        return CifCategory.ofTable(model.entities.data, keys);
+    instance({ structures}) {
+        const keys = Structure.getEntityKeys(structures[0]);
+        return CifCategory.ofTable(structures[0].model.entities.data, keys);
     }
 }
 
@@ -103,13 +101,13 @@ export const mmCIF_Export_Filters = {
 }
 
 /** Doesn't start a data block */
-export function encode_mmCIF_categories(encoder: CifWriter.Encoder, structures: Structure | Structure[], params?: { skipCategoryNames?: Set<string>, exportCtx?: CifExportContext[] }) {
+export function encode_mmCIF_categories(encoder: CifWriter.Encoder, structures: Structure | Structure[], params?: { skipCategoryNames?: Set<string>, exportCtx?: CifExportContext }) {
     const first = Array.isArray(structures) ? structures[0] : (structures as Structure);
     const models = first.models;
     if (models.length !== 1) throw 'Can\'t export stucture composed from multiple models.';
 
     const _params = params || { };
-    const ctx: CifExportContext[] = params && params.exportCtx ? params.exportCtx : CifExportContext.create(structures);
+    const ctx: CifExportContext = params && params.exportCtx ? params.exportCtx : CifExportContext.create(structures);
 
     for (const cat of Categories) {
         if (_params.skipCategoryNames && _params.skipCategoryNames.has(cat.name)) continue;

src/mol-model/structure/model/formats/mmcif/bonds/comp.ts

@@ -26,10 +26,10 @@ export namespace ComponentBond {
             categories: [{
                 name: 'chem_comp_bond',
                 instance(ctx) {
-                    const chem_comp_bond = getChemCompBond(ctx.model);
+                    const chem_comp_bond = getChemCompBond(ctx.structures[0].model);
                     if (!chem_comp_bond) return CifWriter.Category.Empty;
 
-                    const comp_names = getUniqueResidueNames(ctx.structure);
+                    const comp_names = getUniqueResidueNames(ctx.structures[0]);
                     const { comp_id, _rowCount } = chem_comp_bond;
                     const indices: number[] = [];
                     for (let i = 0; i < _rowCount; i++) {

src/mol-model/structure/model/formats/mmcif/bonds/struct_conn.ts

@@ -31,10 +31,11 @@ export namespace StructConn {
             categories: [{
                 name: 'struct_conn',
                 instance(ctx) {
-                    const struct_conn = getStructConn(ctx.model);
+                    const structure = ctx.structures[0], model = structure.model;
+                    const struct_conn = getStructConn(model);
                     if (!struct_conn) return CifWriter.Category.Empty;
 
-                    const strConn = get(ctx.model);
+                    const strConn = get(model);
                     if (!strConn || strConn.entries.length === 0) return CifWriter.Category.Empty;
 
                     const foundAtoms = new Set<ElementIndex>();
@@ -45,7 +46,7 @@ export namespace StructConn {
                         for (let i = 0, _i = partners.length; i < _i; i++) {
                             const atom = partners[i].atomIndex;
                             if (foundAtoms.has(atom)) continue;
-                            if (hasAtom(ctx.structure, atom)) {
+                            if (hasAtom(structure, atom)) {
                                 foundAtoms.add(atom);
                             } else {
                                 hasAll = false;

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

@@ -7,7 +7,6 @@
 import { ResidueIndex } from '../../indexing';
 import { Unit, Structure, StructureElement } from '../../../structure';
 import { Segmentation } from 'mol-data/int';
-import { CifExportContext } from '../../../export/mmcif';
 import { UUID } from 'mol-util';
 import { CifWriter } from 'mol-io/writer/cif';
 
@@ -20,25 +19,23 @@ export interface ResidueCustomProperty<T = any> {
 
 export namespace ResidueCustomProperty {
     export interface ExportCtx<T> {
-        exportCtx: CifExportContext,
         elements: StructureElement[],
         property(index: number): T
     };
 
-    function getExportCtx<T>(exportCtx: CifExportContext, prop: ResidueCustomProperty<T>): ExportCtx<T> {
-        if (exportCtx.localCache[prop.id]) return exportCtx.localCache[prop.id];
-        const residueIndex = exportCtx.model.atomicHierarchy.residueAtomSegments.index;
-        const elements = getStructureElements(exportCtx.structure, prop);
-        return {
-            exportCtx,
-            elements,
-            property: i => prop.get(residueIndex[elements[i].element])!
-        }
+    function getExportCtx<T>(prop: ResidueCustomProperty<T>, structure: Structure): ExportCtx<T> {
+        const residueIndex = structure.model.atomicHierarchy.residueAtomSegments.index;
+        const elements = getStructureElements(structure, prop);
+        return { elements, property: i => prop.get(residueIndex[elements[i].element])! };
     }
 
-    export function createCifCategory<T>(ctx: CifExportContext, prop: ResidueCustomProperty<T>, fields: CifWriter.Field<number, ExportCtx<T>>[]): CifWriter.Category.Instance {
-        const data = getExportCtx(ctx, prop);
-        return { fields, data, rowCount: data.elements.length };
+    export function getCifDataSource<T>(structure: Structure, prop: ResidueCustomProperty<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 = getExportCtx(prop, structure);
+        const ret = { data, rowCount: data.elements.length };
+        if (cache) cache[prop.id] = ret;
+        return ret;
     }
 
     class FromMap<T> implements ResidueCustomProperty<T> {

src/mol-model/structure/structure/structure.ts

@@ -38,6 +38,7 @@ class Structure {
         unitSymmetryGroups?: ReadonlyArray<Unit.SymmetryGroup>,
         carbohydrates?: Carbohydrates,
         models?: ReadonlyArray<Model>,
+        model?: Model,
         hashCode: number,
         elementCount: number,
         polymerResidueCount: number,
@@ -128,6 +129,15 @@ class Structure {
         return this._props.models;
     }
 
+    /** If the structure is based on a single model, return it. Otherwise throw an exception. */
+    get model(): Model {
+        if (this._props.model) return this._props.model;
+        const models = this.models;
+        if (models.length > 1) throw new Error('The structre is based on multiple models.');
+        this._props.model = models[0];
+        return this._props.model;
+    }
+
     hasElement(e: StructureElement) {
         if (!this.unitMap.has(e.unit.id)) return false;
         return SortedArray.has(this.unitMap.get(e.unit.id).elements, e.element);

src/perf-tests/cif-encoder.ts

@@ -23,7 +23,7 @@ function getCat(name: string): CifWriter.Category {
     return {
         name,
         instance(ctx: { fields: CifWriter.Field[], rowCount: number }) {
-            return { data: void 0, fields: ctx.fields, rowCount: ctx.rowCount };
+            return { fields: ctx.fields, source: [{ rowCount: ctx.rowCount }] };
         }
     };
 }

src/servers/model/preprocess/converter.ts

@@ -12,7 +12,7 @@ import { Task } from 'mol-task';
 function getCategoryInstanceProvider(cat: CifCategory, fields: CifWriter.Field[]): CifWriter.Category {
     return {
         name: cat.name,
-        instance: () => ({ data: cat, fields, rowCount: cat.rowCount })
+        instance: () => CifWriter.categoryInstance(fields, { data: cat, rowCount: cat.rowCount })
     };
 }
 

src/servers/model/preprocess/preprocess.ts

@@ -36,7 +36,7 @@ export async function preprocessFile(filename: string, propertyProvider?: ModelP
     }
 }
 
-function encode(structure: Structure, header: string, categories: CifWriter.Category[], encoder: CifWriter.Encoder, exportCtx: CifExportContext[], writer: Writer) {
+function encode(structure: Structure, header: string, categories: CifWriter.Category[], encoder: CifWriter.Encoder, exportCtx: CifExportContext, writer: Writer) {
     const skipCategoryNames = new Set<string>(categories.map(c => c.name));
     encoder.startDataBlock(header);
     for (const cat of categories) {

src/servers/model/server/query.ts

@@ -70,8 +70,8 @@ export async function resolveJob(job: Job): Promise<CifWriter.Encoder<any>> {
 
         perf.start('encode');
         encoder.startDataBlock(sourceStructures[0].models[0].label.toUpperCase());
-        encoder.writeCategory(_model_server_result, [job]);
-        encoder.writeCategory(_model_server_params, [job]);
+        encoder.writeCategory(_model_server_result, job);
+        encoder.writeCategory(_model_server_params, job);
 
         // encoder.setFilter(mmCIF_Export_Filters.onlyPositions);
         encode_mmCIF_categories(encoder, result);
@@ -84,7 +84,7 @@ export async function resolveJob(job: Job): Promise<CifWriter.Encoder<any>> {
             encodeTimeMs: perf.time('encode')
         };
 
-        encoder.writeCategory(_model_server_stats, [stats]);
+        encoder.writeCategory(_model_server_stats, stats);
         encoder.encode();
         ConsoleLogger.logId(job.id, 'Query', 'Encoded.');
         return encoder;
@@ -101,9 +101,9 @@ function getEncodingProvider(structure: StructureWrapper) {
 
 function doError(job: Job, e: any) {
     const encoder = CifWriter.createEncoder({ binary: job.responseFormat.isBinary, encoderName: `ModelServer ${Version}` });
-    encoder.writeCategory(_model_server_result, [job]);
-    encoder.writeCategory(_model_server_params, [job]);
-    encoder.writeCategory(_model_server_error, ['' + e]);
+    encoder.writeCategory(_model_server_result, job);
+    encoder.writeCategory(_model_server_params, job);
+    encoder.writeCategory(_model_server_error, '' + e);
     encoder.encode();
     return encoder;
 }
@@ -155,12 +155,12 @@ const _model_server_stats_fields: CifField<number, Stats>[] = [
 
 const _model_server_result: CifWriter.Category<Job> = {
     name: 'model_server_result',
-    instance: (job) => ({ data: job, fields: _model_server_result_fields, rowCount: 1 })
+    instance: (job) => CifWriter.categoryInstance(_model_server_result_fields,{ data: job, rowCount: 1 })
 };
 
 const _model_server_error: CifWriter.Category<string> = {
     name: 'model_server_error',
-    instance: (message) => ({ data: message, fields: _model_server_error_fields, rowCount: 1 })
+    instance: (message) => CifWriter.categoryInstance(_model_server_error_fields, { data: message, rowCount: 1 })
 };
 
 const _model_server_params: CifWriter.Category<Job> = {
@@ -170,17 +170,12 @@ const _model_server_params: CifWriter.Category<Job> = {
         for (const k of Object.keys(job.normalizedParams)) {
             params.push([k, JSON.stringify(job.normalizedParams[k])]);
         }
-        return {
-            data: params,
-
-            fields: _model_server_params_fields,
-            rowCount: params.length
-        }
+        return CifWriter.categoryInstance(_model_server_params_fields, { data: params, rowCount: params.length });
     }
 };
 
 
 const _model_server_stats: CifWriter.Category<Stats> = {
     name: 'model_server_stats',
-    instance: (stats) => ({ data: stats, fields: _model_server_stats_fields, rowCount: 1 })
+    instance: (stats) => CifWriter.categoryInstance(_model_server_stats_fields, { data: stats, rowCount: 1 })
 }

src/servers/volume/server/query/encode.ts

@@ -102,7 +102,7 @@ const _volume_data_3d_info: CifWriter.Category<ResultContext> = {
             sampledValuesInfo: result.query.data.header.sampling[result.query.samplingInfo.sampling.index].valuesInfo[result.channelIndex]
         };
 
-        return { data: ctx, fields: _volume_data_3d_info_fields, rowCount: 1 }
+        return { fields: _volume_data_3d_info_fields, source: [{ data: ctx, rowCount: 1 }] }
     }
 };
 
@@ -136,7 +136,7 @@ const _volume_data_3d: CifWriter.Category<ResultContext> = {
         }
 
         const fields = [CifWriter.Field.float('values', _volume_data_3d_number, { encoder, typedArray, digitCount: 6 })];
-        return { data, fields, rowCount: data.length };
+        return CifWriter.categoryInstance(fields, { data, rowCount: data.length });
     }
 }
 
@@ -174,12 +174,12 @@ const _density_server_result_fields = [
 
 const _density_server_result: CifWriter.Category<Data.QueryContext> = {
     name: 'density_server_result',
-    instance: ctx => ({ data: ctx, fields: _density_server_result_fields, rowCount: 1 })
+    instance: ctx => CifWriter.categoryInstance(_density_server_result_fields, { data: ctx, rowCount: 1 })
 }
 
 function write(encoder: CifWriter.Encoder, query: Data.QueryContext) {
     encoder.startDataBlock('SERVER');
-    encoder.writeCategory(_density_server_result, [query]);
+    encoder.writeCategory(_density_server_result, query);
 
     switch (query.kind) {
         case 'Data':
@@ -189,7 +189,7 @@ function write(encoder: CifWriter.Encoder, query: Data.QueryContext) {
         const header = query.data.header;
         for (let i = 0; i < header.channels.length; i++) {
             encoder.startDataBlock(header.channels[i]);
-            const ctx: ResultContext[] = [{ query, channelIndex: i }];
+            const ctx: ResultContext = { query, channelIndex: i };
 
             encoder.writeCategory(_volume_data_3d_info, ctx);
             encoder.writeCategory(_volume_data_3d, ctx);