Support multiple models in StructureElement.Loci.toScriptExpression

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

@@ -309,24 +309,24 @@ namespace StructureElement {
         }
 
         export function toScriptExpression(loci: Loci) {
-            if (loci.structure.models.length > 1) {
-                console.warn('toScriptExpression is only supported for Structure with single model, returning empty expression.');
-                return MS.struct.generator.empty();
-            }
             if (loci.elements.length === 0) return MS.struct.generator.empty();
 
-            const sourceIndexMap = new Map<string, UniqueArray<number, number>>();
+            const models = loci.structure.models;
+            const sourceIndexMap = new Map<string, { modelLabel: string, modelIndex: number, xs: UniqueArray<number, number> }>();
             const el = StructureElement.create(), p = StructureProperties.atom.sourceIndex;
             for (const e of loci.elements) {
                 const { indices } = e;
                 const { elements } = e.unit;
-                const opName = e.unit.conformation.operator.name;
+
+                const key = models.length === 1
+                    ? e.unit.conformation.operator.name
+                    : `${e.unit.conformation.operator.name} ${e.unit.model.label} ${e.unit.model.modelNum}`;
 
                 let sourceIndices: UniqueArray<number, number>;
-                if (sourceIndexMap.has(opName)) sourceIndices = sourceIndexMap.get(opName)!;
+                if (sourceIndexMap.has(key)) sourceIndices = sourceIndexMap.get(key)!.xs;
                 else {
                     sourceIndices = UniqueArray.create<number, number>();
-                    sourceIndexMap.set(opName, sourceIndices);
+                    sourceIndexMap.set(key, { modelLabel: e.unit.model.label, modelIndex: e.unit.model.modelNum, xs: sourceIndices });
                 }
 
                 el.unit = e.unit;
@@ -342,7 +342,8 @@ namespace StructureElement {
             while (true) {
                 const k = keys.next();
                 if (k.done) break;
-                byOpName.push(getOpNameQuery(k.value, sourceIndexMap.get(k.value)!.array));
+                const e = sourceIndexMap.get(k.value)!;
+                byOpName.push(getOpNameQuery(k.value, e.xs.array, models.length > 1, e.modelLabel, e.modelIndex));
             }
 
             return MS.struct.modifier.union([
@@ -350,7 +351,7 @@ namespace StructureElement {
             ]);
         }
 
-        function getOpNameQuery(opName: string, xs: number[]) {
+        function getOpNameQuery(opName: string, xs: number[], multimodel: boolean, modelLabel: string, modelIndex: number) {
             sortArray(xs);
 
             const ranges: number[] = [];
@@ -384,10 +385,19 @@ namespace StructureElement {
                 tests[tests.length] = MS.core.rel.inRange([siProp, ranges[2 * rI], ranges[2 * rI + 1]]);
             }
 
-            return MS.struct.generator.atomGroups({
-                'atom-test': tests.length > 1 ? MS.core.logic.or(tests) : tests[0],
-                'chain-test': MS.core.rel.eq([MS.struct.atomProperty.core.operatorName(), opName])
-            });
+            return multimodel
+                ? MS.struct.generator.atomGroups({
+                    'atom-test': tests.length > 1 ? MS.core.logic.or(tests) : tests[0],
+                    'chain-test': MS.core.rel.eq([MS.struct.atomProperty.core.operatorName(), opName]),
+                    'entity-test': MS.core.logic.and([
+                        MS.core.rel.eq([MS.struct.atomProperty.core.modelLabel(), modelLabel]),
+                        MS.core.rel.eq([MS.struct.atomProperty.core.modelIndex(), modelIndex]),
+                    ])
+                })
+                : MS.struct.generator.atomGroups({
+                    'atom-test': tests.length > 1 ? MS.core.logic.or(tests) : tests[0],
+                    'chain-test': MS.core.rel.eq([MS.struct.atomProperty.core.operatorName(), opName])
+                });
         }
     }
 }

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

@@ -113,6 +113,8 @@ const entity = {
 const unit = {
     id: StructureElement.property(l => l.unit.id),
     operator_name: StructureElement.property(l => l.unit.conformation.operator.name),
+    model_index: StructureElement.property(l => l.unit.model.modelNum),
+    model_label: StructureElement.property(l => l.unit.model.label),
     hkl: StructureElement.property(l => l.unit.conformation.operator.hkl),
     spgrOp: StructureElement.property(l => l.unit.conformation.operator.spgrOp),
 

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

@@ -218,6 +218,10 @@ class Structure {
         return SortedArray.has(this.unitMap.get(e.unit.id).elements, e.element);
     }
 
+    getModelIndex(m: Model) {
+        return this.model
+    }
+
     private initUnits(units: ArrayLike<Unit>) {
         const map = IntMap.Mutable<Unit>();
         let elementCount = 0;

src/mol-script/language/symbol-table/structure-query.ts

@@ -254,6 +254,8 @@ const atomProperty = {
 
         sourceIndex: atomProp(Type.Num, 'Index of the atom/element in the input file.'),
         operatorName: atomProp(Type.Str, 'Name of the symmetry operator applied to this element.'),
+        modelIndex: atomProp(Type.Num, 'Index of the model in the input file.'),
+        modelLabel: atomProp(Type.Str, 'Label/header of the model in the input file.')
     },
 
     topology: {

src/mol-script/runtime/query/table.ts

@@ -247,6 +247,8 @@ const symbols = [
     D(MolScript.structureQuery.atomProperty.core.z, atomProp(StructureProperties.atom.z)),
     D(MolScript.structureQuery.atomProperty.core.sourceIndex, atomProp(StructureProperties.atom.sourceIndex)),
     D(MolScript.structureQuery.atomProperty.core.operatorName, atomProp(StructureProperties.unit.operator_name)),
+    D(MolScript.structureQuery.atomProperty.core.modelIndex, atomProp(StructureProperties.unit.model_index)),
+    D(MolScript.structureQuery.atomProperty.core.modelLabel, atomProp(StructureProperties.unit.model_label)),
     D(MolScript.structureQuery.atomProperty.core.atomKey, (ctx, _) => cantorPairing(ctx.element.unit.id, ctx.element.element)),
 
     // TODO:

src/mol-script/script/mol-script/symbols.ts

@@ -199,6 +199,8 @@ export const SymbolTable = [
             Alias(MolScript.structureQuery.atomProperty.core.z, 'atom.z'),
             Alias(MolScript.structureQuery.atomProperty.core.sourceIndex, 'atom.src-index'),
             Alias(MolScript.structureQuery.atomProperty.core.operatorName, 'atom.op-name'),
+            Alias(MolScript.structureQuery.atomProperty.core.modelIndex, 'atom.model-index'),
+            Alias(MolScript.structureQuery.atomProperty.core.modelLabel, 'atom.model-label'),
             Alias(MolScript.structureQuery.atomProperty.core.atomKey, 'atom.key'),
             Alias(MolScript.structureQuery.atomProperty.core.bondCount, 'atom.bond-count'),