mol-script: added more link props, support different element reference for atomic props

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

@@ -7,7 +7,6 @@
 import { Structure, StructureElement, Unit } from '../structure';
 import { now } from '../../../mol-util/now';
 import { ElementIndex } from '../model';
-import { Link } from '../structure/unit/links';
 import { LinkType } from '../model/types';
 import { StructureSelection } from './selection';
 
@@ -32,7 +31,7 @@ export class QueryContext implements QueryContextView {
     currentStructure: Structure = void 0 as any;
 
     /** Current link between atoms */
-    readonly atomicLink = QueryContextLinkInfo.empty<Unit.Atomic>();
+    readonly atomicLink = new QueryContextLinkInfo<Unit.Atomic>();
 
     /** Supply this from the outside. Used by the internal.generator.current symbol */
     currentSelection: StructureSelection | undefined = void 0;
@@ -54,7 +53,7 @@ export class QueryContext implements QueryContextView {
 
     pushCurrentLink() {
         if (this.atomicLink) this.currentAtomicLinkStack.push(this.atomicLink);
-        (this.atomicLink as QueryContextLinkInfo<Unit.Atomic>) = QueryContextLinkInfo.empty();
+        (this.atomicLink as QueryContextLinkInfo<Unit.Atomic>) = new QueryContextLinkInfo();
         return this.atomicLink;
     }
 
@@ -112,14 +111,15 @@ export interface QueryContextOptions {
 export interface QueryPredicate { (ctx: QueryContext): boolean }
 export interface QueryFn<T = any> { (ctx: QueryContext): T }
 
-export interface QueryContextLinkInfo<U extends Unit = Unit> {
-    link: Link.Location<U>,
-    type: LinkType,
-    order: number
-}
+export class QueryContextLinkInfo<U extends Unit = Unit> {
+    a: StructureElement.Location<U> = StructureElement.Location.create();
+    aIndex: StructureElement.UnitIndex = 0 as StructureElement.UnitIndex;
+    b: StructureElement.Location<U> = StructureElement.Location.create();
+    bIndex: StructureElement.UnitIndex = 0 as StructureElement.UnitIndex;
+    type: LinkType = LinkType.Flag.None;
+    order: number = 0;
 
-export namespace QueryContextLinkInfo {
-    export function empty<U extends Unit = Unit>(): QueryContextLinkInfo<U> {
-        return { link: Link.Location() as Link.Location<U>, type: LinkType.Flag.None, order: 0 };
+    get length() {
+        return StructureElement.Location.distance(this.a, this. b);
     }
 }

src/mol-model/structure/query/queries/filters.ts

@@ -250,7 +250,7 @@ function checkConnected(ctx: IsConnectedToCtx, structure: Structure) {
         const linkedUnits = interLinks.getLinkedUnits(unit);
         const luCount = linkedUnits.length;
 
-        atomicLink.link.aUnit = inputUnit;
+        atomicLink.a.unit = inputUnit;
 
         const srcElements = unit.elements;
         const inputElements = inputUnit.elements;
@@ -258,14 +258,16 @@ function checkConnected(ctx: IsConnectedToCtx, structure: Structure) {
         for (let i = 0 as StructureElement.UnitIndex, _i = srcElements.length; i < _i; i++) {
             const inputIndex = SortedArray.indexOf(inputElements, srcElements[i]) as StructureElement.UnitIndex;
 
-            atomicLink.link.aIndex = inputIndex;
-            atomicLink.link.bUnit = inputUnit;
+            atomicLink.aIndex = inputIndex;
+            atomicLink.a.element = srcElements[i];
+            atomicLink.b.unit = inputUnit;
 
             tElement.unit = unit;
             for (let l = offset[inputIndex], _l = offset[inputIndex + 1]; l < _l; l++) {
                 tElement.element = inputElements[b[l]];
                 if (!target.hasElement(tElement)) continue;
-                atomicLink.link.bIndex = b[l] as StructureElement.UnitIndex;
+                atomicLink.bIndex = b[l] as StructureElement.UnitIndex;
+                atomicLink.b.element = inputUnit.elements[b[l]];
                 atomicLink.type = flags[l];
                 atomicLink.order = order[l];
                 if (linkTest(queryCtx)) return true;
@@ -274,14 +276,15 @@ function checkConnected(ctx: IsConnectedToCtx, structure: Structure) {
             for (let li = 0; li < luCount; li++) {
                 const lu = linkedUnits[li];
                 tElement.unit = lu.unitB;
-                atomicLink.link.bUnit = lu.unitB;
+                atomicLink.b.unit = lu.unitB;
                 const bElements = lu.unitB.elements;
                 const bonds = lu.getBonds(inputIndex);
                 for (let bi = 0, _bi = bonds.length; bi < _bi; bi++) {
                     const bond = bonds[bi];
                     tElement.element = bElements[bond.indexB];
                     if (!target.hasElement(tElement)) continue;
-                    atomicLink.link.bIndex = bond.indexB;
+                    atomicLink.bIndex = bond.indexB;
+                    atomicLink.b.element = tElement.element;
                     atomicLink.type = bond.flag;
                     atomicLink.order = bond.order;
                     if (linkTest(queryCtx)) return true;

src/mol-model/structure/query/queries/modifiers.ts

@@ -356,18 +356,20 @@ function expandConnected(ctx: QueryContext, structure: Structure, linkTest: Quer
         const { offset: intraLinkOffset, b: intraLinkB, edgeProps: { flags, order } } = inputUnit.links;
 
         // Process intra unit links
-        atomicLink.link.aUnit = inputUnit;
-        atomicLink.link.bUnit = inputUnit;
+        atomicLink.a.unit = inputUnit;
+        atomicLink.b.unit = inputUnit;
         for (let i = 0, _i = unit.elements.length; i < _i; i++) {
             // add the current element
             builder.addToUnit(unit.id, unit.elements[i]);
 
             const srcIndex = SortedArray.indexOf(inputUnit.elements, unit.elements[i]);
-            atomicLink.link.aIndex = srcIndex as StructureElement.UnitIndex;
+            atomicLink.aIndex = srcIndex as StructureElement.UnitIndex;
+            atomicLink.a.element = unit.elements[i];
 
             // check intra unit links
             for (let lI = intraLinkOffset[srcIndex], _lI = intraLinkOffset[srcIndex + 1]; lI < _lI; lI++) {
-                atomicLink.link.bIndex = intraLinkB[lI] as StructureElement.UnitIndex;
+                atomicLink.bIndex = intraLinkB[lI] as StructureElement.UnitIndex;
+                atomicLink.b.element = inputUnit.elements[intraLinkB[lI]];
                 atomicLink.type = flags[lI];
                 atomicLink.order = order[lI];
                 if (linkTest(ctx)) {
@@ -380,14 +382,17 @@ function expandConnected(ctx: QueryContext, structure: Structure, linkTest: Quer
         for (const linkedUnit of interLinks.getLinkedUnits(inputUnit)) {
             if (processedUnits.has(linkedUnit.unitA.id)) continue;
 
-            atomicLink.link.bUnit = linkedUnit.unitB;
+            atomicLink.b.unit = linkedUnit.unitB;
             for (const aI of linkedUnit.linkedElementIndices) {
                 // check if the element is in the expanded structure
                 if (!SortedArray.has(unit.elements, inputUnit.elements[aI])) continue;
 
-                atomicLink.link.aIndex = aI;
+                atomicLink.aIndex = aI;
+                atomicLink.a.element = inputUnit.elements[aI];
                 for (const bond of linkedUnit.getBonds(aI)) {
-                    atomicLink.link.bIndex = bond.indexB;
+                    atomicLink.bIndex = bond.indexB;
+                    // TODO: optimize the lookup?
+                    atomicLink.b.element = linkedUnit.unitB.elements[bond.indexB];
                     atomicLink.type = bond.flag;
                     atomicLink.order = bond.order;
                     if (linkTest(ctx)) {

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

@@ -7,6 +7,7 @@
 
 import { ElementIndex } from '../../model';
 import Unit from '../unit';
+import { Vec3 } from '../../../../mol-math/linear-algebra';
 
 export { Location }
 
@@ -18,7 +19,7 @@ interface Location<U = Unit> {
 }
 
 namespace Location {
-    export function create(unit?: Unit, element?: ElementIndex): Location {
+    export function create<U extends Unit>(unit?: U, element?: ElementIndex): Location<U> {
         return { kind: 'element-location', unit: unit!, element: element || (0 as ElementIndex) };
     }
 
@@ -37,4 +38,11 @@ namespace Location {
     export function is(x: any): x is Location {
         return !!x && x.kind === 'element-location';
     }
+
+    const pA = Vec3.zero(), pB = Vec3.zero();
+    export function distance(a: Location, b: Location) {
+        a.unit.conformation.position(a.element, pA);
+        b.unit.conformation.position(b.element, pB);
+        return Vec3.distance(pA, pB);
+    }
 }

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

@@ -314,7 +314,10 @@ const linkProperty = {
     '@header': 'Link Properties',
 
     flags: linkProp(Types.LinkFlags),
-    order: linkProp(Type.Num)
+    order: linkProp(Type.Num),
+    length: linkProp(Type.Num),
+    atomA: linkProp(Types.ElementReference),
+    atomB: linkProp(Types.ElementReference)
 }
 
 function atomProp(type: Type, description?: string) {

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

@@ -194,8 +194,8 @@ const symbols = [
 
     // ============= SLOTS ================
     // TODO: slots might not be needed after all: reducer simply pushes/pops current element
-    C(MolScript.structureQuery.slot.element, (ctx, _) => ctx.element),
-    // C(MolScript.structureQuery.slot.elementSetReduce, (ctx, _) => ctx.element),
+    // C(MolScript.structureQuery.slot.element, (ctx, _) => ctx_.element),
+    // C(MolScript.structureQuery.slot.elementSetReduce, (ctx, _) => ctx_.element),
 
     // ============= FILTERS ================
     D(MolScript.structureQuery.filter.pick, (ctx, xs) => Queries.filters.pick(xs[0] as any, xs['test'])(ctx)),
@@ -256,9 +256,9 @@ const symbols = [
 
     // ~~~ CORE ~~~
     D(MolScript.structureQuery.atomProperty.core.elementSymbol, atomProp(StructureProperties.atom.type_symbol)),
-    D(MolScript.structureQuery.atomProperty.core.vdw, (ctx, _) => VdwRadius(StructureProperties.atom.type_symbol(ctx.element))),
-    D(MolScript.structureQuery.atomProperty.core.mass, (ctx, _) => AtomWeight(StructureProperties.atom.type_symbol(ctx.element))),
-    D(MolScript.structureQuery.atomProperty.core.atomicNumber, (ctx, _) => AtomNumber(StructureProperties.atom.type_symbol(ctx.element))),
+    D(MolScript.structureQuery.atomProperty.core.vdw, (ctx, xs) => VdwRadius(StructureProperties.atom.type_symbol((xs && xs[0] && xs[0](ctx) as any) || ctx.element))),
+    D(MolScript.structureQuery.atomProperty.core.mass, (ctx, xs) => AtomWeight(StructureProperties.atom.type_symbol((xs && xs[0] && xs[0](ctx) as any) || ctx.element))),
+    D(MolScript.structureQuery.atomProperty.core.atomicNumber, (ctx, xs) => AtomNumber(StructureProperties.atom.type_symbol((xs && xs[0] && xs[0](ctx) as any) || ctx.element))),
     D(MolScript.structureQuery.atomProperty.core.x, atomProp(StructureProperties.atom.x)),
     D(MolScript.structureQuery.atomProperty.core.y, atomProp(StructureProperties.atom.y)),
     D(MolScript.structureQuery.atomProperty.core.z, atomProp(StructureProperties.atom.z)),
@@ -266,7 +266,10 @@ const symbols = [
     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)),
+    D(MolScript.structureQuery.atomProperty.core.atomKey, (ctx, xs) => {
+        const e = (xs && xs[0] && xs[0](ctx) as any) || ctx.element;
+        return cantorPairing(e.unit.id, e.element)
+    }),
 
     // TODO:
     // D(MolScript.structureQuery.atomProperty.core.bondCount, (ctx, _) => ),
@@ -283,13 +286,13 @@ const symbols = [
     // authResidueId: prop((env, v) => ResidueIdentifier.authOfResidueIndex(env.context.model, getAddress(env, v).residue)),
 
     // keys
-    D(MolScript.structureQuery.atomProperty.macromolecular.residueKey, (ctx, _) => StructureElement.residueIndex(ctx.element)),
-    D(MolScript.structureQuery.atomProperty.macromolecular.chainKey, (ctx, _) => StructureElement.chainIndex(ctx.element)),
-    D(MolScript.structureQuery.atomProperty.macromolecular.entityKey, (ctx, _) => StructureElement.entityIndex(ctx.element)),
+    D(MolScript.structureQuery.atomProperty.macromolecular.residueKey, (ctx, xs) => StructureElement.residueIndex((xs && xs[0] && xs[0](ctx) as any) || ctx.element)),
+    D(MolScript.structureQuery.atomProperty.macromolecular.chainKey, (ctx, xs) => StructureElement.chainIndex((xs && xs[0] && xs[0](ctx) as any) || ctx.element)),
+    D(MolScript.structureQuery.atomProperty.macromolecular.entityKey, (ctx, xs) => StructureElement.entityIndex((xs && xs[0] && xs[0](ctx) as any) || ctx.element)),
 
     // mmCIF
     D(MolScript.structureQuery.atomProperty.macromolecular.id, atomProp(StructureProperties.atom.id)),
-    D(MolScript.structureQuery.atomProperty.macromolecular.isHet, (ctx, _) => StructureProperties.residue.group_PDB(ctx.element) !== 'ATOM'),
+    D(MolScript.structureQuery.atomProperty.macromolecular.isHet, (ctx, xs) => StructureProperties.residue.group_PDB((xs && xs[0] && xs[0](ctx) as any) || ctx.element) !== 'ATOM'),
 
     D(MolScript.structureQuery.atomProperty.macromolecular.label_atom_id, atomProp(StructureProperties.atom.label_atom_id)),
     D(MolScript.structureQuery.atomProperty.macromolecular.label_alt_id, atomProp(StructureProperties.atom.label_alt_id)),
@@ -322,6 +325,9 @@ const symbols = [
     // ============= BOND PROPERTIES ================
     D(MolScript.structureQuery.linkProperty.order, (ctx, xs) => ctx.atomicLink.order),
     D(MolScript.structureQuery.linkProperty.flags, (ctx, xs) => ctx.atomicLink.type),
+    D(MolScript.structureQuery.linkProperty.atomA, (ctx, xs) => ctx.atomicLink.a),
+    D(MolScript.structureQuery.linkProperty.atomB, (ctx, xs) => ctx.atomicLink.b),
+    D(MolScript.structureQuery.linkProperty.length, (ctx, xs) => ctx.atomicLink.length),
 
 
     ////////////////////////////////////
@@ -331,8 +337,8 @@ const symbols = [
     D(MolScript.internal.generator.current, function internal_generator_current(ctx, xs) { return ctx.tryGetCurrentSelection() }),
 ];
 
-function atomProp(p: (e: StructureElement.Location) => any): (ctx: QueryContext, _: any) => any {
-    return (ctx, _) => p(ctx.element);
+function atomProp(p: (e: StructureElement.Location) => any): (ctx: QueryContext, xs: any) => any {
+    return (ctx, xs) => p((xs && xs[0] && xs[0](ctx) as any) || ctx.element);
 }
 
 function linkFlag(current: LinkType, f: string): LinkType {

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

@@ -248,6 +248,9 @@ export const SymbolTable = [
         [
             'Link Properties',
             Alias(MolScript.structureQuery.linkProperty.order, 'link.order'),
+            Alias(MolScript.structureQuery.linkProperty.length, 'link.length'),
+            Alias(MolScript.structureQuery.linkProperty.atomA, 'link.atom-a'),
+            Alias(MolScript.structureQuery.linkProperty.atomB, 'link.atom-b'),
             Macro(MSymbol('link.is', Arguments.List(StructureQueryTypes.LinkFlag), Type.Bool,
                 `Test if the current link has at least one (or all if partial = false) of the specified flags: ${Type.oneOfValues(StructureQueryTypes.LinkFlag).join(', ')}`),
             args => B.core.flags.hasAny([B.struct.linkProperty.flags(), B.struct.type.linkFlags(getPositionalArgs(args))])),