wip ModelServer multiple model support

src/servers/model/preprocess/preprocess.ts

@@ -4,7 +4,7 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { readStructure } from '../server/structure-wrapper';
+import { readStructureWrapper, resolveStructure } from '../server/structure-wrapper';
 import { classifyCif } from './converter';
 // import { ConsoleLogger } from 'mol-util/console-logger';
 import { Structure } from 'mol-model/structure';
@@ -24,18 +24,19 @@ export async function preprocessFile(filename: string, outputCif?: string, outpu
     //const started = now();
     //ConsoleLogger.log(`${linearId}`, `Reading '${filename}'...`);
     // TODO: support the custom prop provider list here.
-    const input = await readStructure('entry', '_local_', filename, void 0);
+    const input = await readStructureWrapper('entry', '_local_', filename, void 0);
+    const inputStructure = (await resolveStructure(input))!;
     //ConsoleLogger.log(`${linearId}`, `Classifying CIF categories...`);
     const categories = await classifyCif(input.cifFrame);
     //clearLine();
 
-    const exportCtx = CifExportContext.create(input.structure, input.structure.models[0]);
+    const exportCtx = CifExportContext.create(inputStructure, inputStructure.models[0]);
 
     if (outputCif) {
         //ConsoleLogger.log(`${linearId}`, `Encoding CIF...`);
         const writer = wrapFileToWriter(outputCif);
         const encoder = CifWriter.createEncoder({ binary: false });
-        await encode(input.structure, input.cifFrame.header, categories, encoder, exportCtx, writer);
+        await encode(inputStructure, input.cifFrame.header, categories, encoder, exportCtx, writer);
         // clearLine();
         writer.end();
     }
@@ -44,7 +45,7 @@ export async function preprocessFile(filename: string, outputCif?: string, outpu
         // ConsoleLogger.log(`${linearId}`, `Encoding BinaryCIF...`);
         const writer = wrapFileToWriter(outputBcif);
         const encoder = CifWriter.createEncoder({ binary: true, binaryAutoClassifyEncoding: true });
-        await encode(input.structure, input.cifFrame.header, categories, encoder, exportCtx, writer);
+        await encode(inputStructure, input.cifFrame.header, categories, encoder, exportCtx, writer);
         //clearLine();
         writer.end();
     }

src/servers/model/server/query.ts

@@ -14,7 +14,7 @@ import { PerformanceMonitor } from 'mol-util/performance-monitor';
 import Config from '../config';
 import Version from '../version';
 import { Job } from './jobs';
-import { getStructure, StructureWrapper } from './structure-wrapper';
+import { createStructureWrapperFromJob, StructureWrapper, resolveStructure } from './structure-wrapper';
 import CifField = CifWriter.Field
 import { createModelPropertiesProviderFromConfig } from '../provider';
 
@@ -31,13 +31,15 @@ const propertyProvider = createModelPropertiesProviderFromConfig();
 export async function resolveJob(job: Job): Promise<CifWriter.Encoder<any>> {
     ConsoleLogger.logId(job.id, 'Query', 'Starting.');
 
-    const wrappedStructure = await getStructure(job, propertyProvider);
+    const wrappedStructure = await createStructureWrapperFromJob(job, propertyProvider);
 
     try {
         perf.start('query');
+        const sourceStructure = await resolveStructure(wrappedStructure);
+        if (!sourceStructure) throw new Error('Model not available');
         const structure = job.queryDefinition.structureTransform
-            ? await job.queryDefinition.structureTransform(job.normalizedParams, wrappedStructure.structure)
-            : wrappedStructure.structure;
+            ? await job.queryDefinition.structureTransform(job.normalizedParams, sourceStructure)
+            : sourceStructure;
         const query = job.queryDefinition.query(job.normalizedParams, structure);
         const result = await StructureSelection.unionStructure(StructureQuery.run(query, structure, Config.maxQueryTimeInMs));
         perf.end('query');

src/servers/model/server/structure-wrapper.ts

@@ -36,20 +36,22 @@ export interface StructureInfo {
 
 export interface StructureWrapper {
     info: StructureInfo,
-
     isBinary: boolean,
     key: string,
     approximateSize: number,
-    structure: Structure,
+    models: ArrayLike<Model>,
+    modelMap: Map<number, Model>,
+    structureModelMap: Map<number, Structure>,
+    propertyProvider: ModelPropertiesProvider | undefined,
     cifFrame: CifFrame
 }
 
-export async function getStructure(job: Job, propertyProvider: ModelPropertiesProvider | undefined, allowCache = true): Promise<StructureWrapper> {
+export async function createStructureWrapperFromJob(job: Job, propertyProvider: ModelPropertiesProvider | undefined, allowCache = true): Promise<StructureWrapper> {
     if (allowCache && Config.cacheParams.useCache) {
         const ret = StructureCache.get(job.key);
         if (ret) return ret;
     }
-    const ret = await readStructure(job.key, job.sourceId, job.entryId, propertyProvider);
+    const ret = await readStructureWrapper(job.key, job.sourceId, job.entryId, propertyProvider);
     if (allowCache && Config.cacheParams.useCache) {
         StructureCache.add(ret);
     }
@@ -86,7 +88,7 @@ async function parseCif(data: string|Uint8Array) {
     return parsed.result;
 }
 
-export async function readStructure(key: string, sourceId: string | '_local_', entryId: string, propertyProvider: ModelPropertiesProvider | undefined) {
+export async function readStructureWrapper(key: string, sourceId: string | '_local_', entryId: string, propertyProvider: ModelPropertiesProvider | undefined) {
     const filename = sourceId === '_local_' ? entryId : Config.mapFile(sourceId, entryId);
     if (!filename) throw new Error(`Cound not map '${key}' to a valid filename.`);
     if (!fs.existsSync(filename)) throw new Error(`Could not find source file for '${key}'.`);
@@ -108,6 +110,11 @@ export async function readStructure(key: string, sourceId: string | '_local_', e
     const models = await Model.create(Format.mmCIF(frame)).run();
     perf.end('createModel');
 
+    const modelMap = new Map<number, Model>();
+    for (const m of models) {
+        modelMap.set(m.modelNum, m);
+    }
+
     perf.start('attachProps');
     if (propertyProvider) {
         const modelProps = propertyProvider(models[0]);
@@ -117,28 +124,47 @@ export async function readStructure(key: string, sourceId: string | '_local_', e
     }
     perf.end('attachProps');
 
-    const structure = Structure.ofModel(models[0]);
-
     const ret: StructureWrapper = {
         info: {
             sourceType: StructureSourceType.File,
             readTime: perf.time('read'),
             parseTime: perf.time('parse'),
             createModelTime: perf.time('createModel'),
-            attachPropsTime: perf.time('attachProps'),
+            attachPropsTime: 0, // perf.time('attachProps'),
             sourceId,
             entryId
         },
         isBinary: /\.bcif/.test(filename),
         key,
         approximateSize: typeof data === 'string' ? 2 * data.length : data.length,
-        structure,
-        cifFrame: frame
+        models,
+        modelMap,
+        structureModelMap: new Map(),
+        cifFrame: frame,
+        propertyProvider
     };
 
     return ret;
 }
 
+export async function resolveStructure(wrapper: StructureWrapper, modelNum?: number) {
+    if (typeof modelNum === 'undefined') modelNum = wrapper.models[0].modelNum;
+    if (wrapper.structureModelMap.has(modelNum)) return wrapper.structureModelMap.get(modelNum)!;
+    if (!wrapper.modelMap.has(modelNum)) {
+        return void 0;
+    }
+
+    const model = wrapper.modelMap.get(modelNum)!;
+    const structure = Structure.ofModel(model);
+    if (wrapper.propertyProvider) {
+        const modelProps = wrapper.propertyProvider(model);
+        for (const p of modelProps) {
+            await tryAttach(wrapper.key, p);
+        }
+    }
+    return structure;
+}
+
 async function tryAttach(key: string, promise: Promise<any>) {
     try {
         await promise;