mol-model: make StructureElement.Loci.toScriptExpression work on assemblies

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

@@ -254,11 +254,19 @@ namespace StructureElement {
             }
             if (loci.elements.length === 0) return MS.struct.generator.empty();
 
-            const sourceIndices = UniqueArray.create<number, number>();
+            const sourceIndexMap = new Map<string, 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;
+
+                let sourceIndices: UniqueArray<number, number>;
+                if (sourceIndexMap.has(opName)) sourceIndices = sourceIndexMap.get(opName)!;
+                else {
+                    sourceIndices = UniqueArray.create<number, number>();
+                    sourceIndexMap.set(opName, sourceIndices);
+                }
 
                 el.unit = e.unit;
                 for (let i = 0, _i = OrderedSet.size(indices); i < _i; i++) {
@@ -268,7 +276,20 @@ namespace StructureElement {
                 }
             }
 
-            const xs = sourceIndices.array;
+            const byOpName: Expression[] = [];
+            const keys = sourceIndexMap.keys();
+            while (true) {
+                const k = keys.next();
+                if (k.done) break;
+                byOpName.push(getOpNameQuery(k.value, sourceIndexMap.get(k.value)!.array));
+            }
+
+            return MS.struct.modifier.union([
+                byOpName.length === 1 ? byOpName[0] : MS.struct.combinator.merge.apply(null, byOpName)
+            ]);
+        }
+
+        function getOpNameQuery(opName: string, xs: number[]) {
             sortArray(xs);
 
             const ranges: number[] = [];
@@ -304,7 +325,7 @@ namespace StructureElement {
 
             return MS.struct.generator.atomGroups({
                 'atom-test': tests.length > 1 ? MS.core.logic.or.apply(null, tests) : tests[0],
-                'group-by': 0
+                'chain-test': MS.core.rel.eq([MS.struct.atomProperty.core.operatorName(), opName])
             });
         }
     }

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

@@ -249,6 +249,7 @@ const atomProperty = {
         }), Type.Num, 'Number of bonds (by default only covalent bonds are counted).'),
 
         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.'),
     },
 
     topology: {

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

@@ -206,9 +206,14 @@ const symbols = [
         elementRadius: xs['atom-radius']
     })(ctx)),
     D(MolScript.structureQuery.modifier.wholeResidues, (ctx, xs) => Queries.modifiers.wholeResidues(xs[0] as any)(ctx)),
+    D(MolScript.structureQuery.modifier.union, (ctx, xs) => Queries.modifiers.union(xs[0] as any)(ctx)),
     D(MolScript.structureQuery.modifier.expandProperty, (ctx, xs) => Queries.modifiers.expandProperty(xs[0] as any, xs['property'])(ctx)),
     D(MolScript.structureQuery.modifier.exceptBy, (ctx, xs) => Queries.modifiers.exceptBy(xs[0] as any, xs['by'] as any)(ctx)),
 
+    // ============= COMBINATORS ================
+
+    D(MolScript.structureQuery.combinator.merge, (ctx, xs) => Queries.combinators.merge(xs as any)(ctx)),
+
     // ============= ATOM PROPERTIES ================
 
     // ~~~ CORE ~~~
@@ -220,6 +225,7 @@ const symbols = [
     D(MolScript.structureQuery.atomProperty.core.y, atomProp(StructureProperties.atom.y)),
     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.atomKey, (ctx, _) => cantorPairing(ctx.element.unit.id, ctx.element.element)),
 
     // TODO:

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

@@ -7,7 +7,6 @@
 import { UniqueArray } from 'mol-data/generic';
 import Expression from '../../language/expression';
 import { Argument, MSymbol } from '../../language/symbol';
-//import * as M from './macro'
 import { MolScriptSymbolTable as MolScript } from '../../language/symbol-table';
 import Type from '../../language/type';
 
@@ -197,6 +196,7 @@ export const SymbolTable = [
             Alias(MolScript.structureQuery.atomProperty.core.y, 'atom.y'),
             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.atomKey, 'atom.key'),
             Alias(MolScript.structureQuery.atomProperty.core.bondCount, 'atom.bond-count'),
 
@@ -343,4 +343,4 @@ export function transpileMolScript(expr: Expression) {
 //     if (a.length === b.length) return (a < b) as any;
 //     return a.length - b.length;
 // });
-//export default [...sortedSymbols, ...NamedArgs.map(a => ':' + a), ...Constants];
+// export default [...sortedSymbols, ...NamedArgs.map(a => ':' + a), ...Constants];
\ No newline at end of file