wip, MolScript query mapping

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

@@ -7,14 +7,14 @@
 
 import { Column } from 'mol-data/db'
 import { Segmentation } from 'mol-data/int';
-import { ElementIndex, ChainIndex } from '../../indexing';
+import { ElementIndex, ChainIndex, EntityIndex } from '../../indexing';
 import SortedRanges from 'mol-data/int/sorted-ranges';
 
 export interface CoarsedElementKeys {
     // assign a key to each element
-    chainKey: ArrayLike<number>,
+    chainKey: ArrayLike<ChainIndex>,
     // assign a key to each element, index to the Model.entities.data table
-    entityKey: ArrayLike<number>,
+    entityKey: ArrayLike<EntityIndex>,
 
     /** find index of the residue/feature element where seq_id is included */
     findSequenceKey(entityId: string, asym_id: string, seq_id: number): ElementIndex

src/mol-model/structure/model/properties/utils/coarse-keys.ts

@@ -7,7 +7,7 @@
 
 import { Entities } from '../common';
 import { CoarseElementData, CoarsedElementKeys } from '../coarse';
-import { ChainIndex, ElementIndex } from '../../indexing';
+import { ChainIndex, ElementIndex, EntityIndex } from '../../indexing';
 
 function getElementKey(map: Map<string, number>, key: string, counter: { index: number }) {
     if (map.has(key)) return map.get(key)!;
@@ -57,8 +57,8 @@ export function getCoarseKeys(data: CoarseElementData, entities: Entities): Coar
     const seqMaps = new Map<number, Map<number, number>>();
     const chainMaps = new Map<number, Map<string, number>>(), chainCounter = { index: 0 };
 
-    const chainKey = new Int32Array(count);
-    const entityKey = new Int32Array(count);
+    const chainKey = new Int32Array(count) as any as ChainIndex[];
+    const entityKey = new Int32Array(count) as any as EntityIndex[];
 
     for (let i = 0; i < count; i++) {
         entityKey[i] = entities.getEntityIndex(entity_id.value(i));
@@ -68,7 +68,7 @@ export function getCoarseKeys(data: CoarseElementData, entities: Entities): Coar
     for (let cI = 0; cI < chainElementSegments.count; cI++) {
         const start = chainElementSegments.offsets[cI], end = chainElementSegments.offsets[cI + 1];
         const map = getElementSubstructureKeyMap(chainMaps, entityKey[start]);
-        const key = getElementKey(map, asym_id.value(start), chainCounter);
+        const key = getElementKey(map, asym_id.value(start), chainCounter) as ChainIndex;
         for (let i = start; i < end; i++) chainKey[i] = key;
 
         // create seq_id map for the ranges defined by seq_id_begin and seq_id_end

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

@@ -7,6 +7,7 @@
 import { OrderedSet, SortedArray } from 'mol-data/int'
 import Unit from './unit'
 import { ElementIndex } from '../model';
+import { ResidueIndex, ChainIndex } from '../model/indexing';
 
 interface StructureElement<U = Unit> {
     unit: U,
@@ -58,6 +59,35 @@ namespace StructureElement {
     export function isLoci(x: any): x is Loci {
         return !!x && x.kind === 'element-loci';
     }
+
+    export function residueIndex(e: StructureElement) {
+        if (Unit.isAtomic(e.unit)) {
+            return e.unit.residueIndex[e.element];
+        } else {
+            // TODO: throw error instead?
+            return -1 as ResidueIndex;
+        }
+    }
+
+    export function chainIndex(e: StructureElement) {
+        if (Unit.isAtomic(e.unit)) {
+            return e.unit.chainIndex[e.element];
+        } else {
+            // TODO: throw error instead?
+            return -1 as ChainIndex;
+        }
+    }
+
+    export function entityIndex(l: StructureElement) {
+        switch (l.unit.kind) {
+            case Unit.Kind.Atomic:
+                return l.unit.model.atomicHierarchy.getEntityKey(l.unit.chainIndex[l.element])
+            case Unit.Kind.Spheres:
+                return l.unit.model.coarseHierarchy.spheres.entityKey[l.element]
+            case Unit.Kind.Gaussians:
+                return l.unit.model.coarseHierarchy.gaussians.entityKey[l.element]
+        }
+    }
 }
 
 export default StructureElement
\ No newline at end of file

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

@@ -89,16 +89,7 @@ const coarse = {
     gaussian_covariance_matrix: StructureElement.property(l => !Unit.isGaussians(l.unit) ? notCoarse('gaussians') : l.unit.coarseConformation.covariance_matrix[l.element])
 }
 
-function eK(l: StructureElement) {
-    switch (l.unit.kind) {
-        case Unit.Kind.Atomic:
-            return l.unit.model.atomicHierarchy.getEntityKey(l.unit.chainIndex[l.element])
-        case Unit.Kind.Spheres:
-            return l.unit.model.coarseHierarchy.spheres.entityKey[l.element]
-        case Unit.Kind.Gaussians:
-            return l.unit.model.coarseHierarchy.gaussians.entityKey[l.element]
-    }
-}
+const eK = StructureElement.entityIndex
 
 const entity = {
     key: eK,

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

@@ -12,6 +12,9 @@ export class QueryRuntimeTable {
     private map = new Map<string, QuerySymbolRuntime>();
 
     addSymbol(runtime: QuerySymbolRuntime) {
+        if (this.map.has(runtime.symbol.id)) {
+            throw new Error(`Symbol '${runtime.symbol.id}' already added.`);
+        }
         this.map.set(runtime.symbol.id, runtime);
     }
 
@@ -74,6 +77,18 @@ export interface QuerySymbolRuntime {
 export type QueryRuntimeArguments<S extends MSymbol> =
     { length?: number } & { [P in keyof S['args']['@type']]: QueryFn<S['args']['@type'][P]> }
 
+export namespace QueryRuntimeArguments {
+    export function forEachEval<S extends MSymbol, Ctx>(xs: QueryRuntimeArguments<S>, queryCtx: QueryContext, f: (arg: any, i: number, ctx: Ctx) => void, ctx: Ctx): Ctx {
+        if (typeof xs.length === 'number') {
+            for (let i = 0, _i = xs.length; i < _i; i++) f((xs as any)[i](queryCtx), i, ctx);
+        } else {
+            let i = 0;
+            for (const k of Object.keys(xs)) f((xs as any)[k](queryCtx), i++, ctx);
+        }
+        return ctx;
+    }
+}
+
 export namespace QuerySymbolRuntime {
     export function Const<S extends MSymbol<any>>(symbol: S, fn: ConstQuerySymbolFn<S>): QuerySymbolRuntime {
         return new SymbolRuntimeImpl(symbol, fn, true);

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

@@ -5,22 +5,45 @@
  */
 
 import { MolScriptSymbolTable as MolScript } from '../../language/symbol-table';
-import { DefaultQueryRuntimeTable, QuerySymbolRuntime } from './compiler';
-
+import { DefaultQueryRuntimeTable, QuerySymbolRuntime, QueryRuntimeArguments } from './compiler';
+import { Queries, StructureProperties, StructureElement, QueryContext } from 'mol-model/structure';
+import { ElementSymbol } from 'mol-model/structure/model/types';
+import { isSuperset } from 'mol-util/set';
+import toUpperCase from 'mol-util/upper-case';
+import { VdwRadius, AtomWeight, AtomNumber } from 'mol-model/structure/model/properties/atomic';
+import { cantorPairing } from 'mol-data/util';
 import C = QuerySymbolRuntime.Const
 import D = QuerySymbolRuntime.Dynamic
-import { Queries, StructureProperties } from 'mol-model/structure';
-import { ElementSymbol } from 'mol-model/structure/model/types';
 
 const symbols = [
-    C(MolScript.core.math.add, (ctx, xs) => {
-        let ret = 0;
+    // ============= TYPES =============
+
+    C(MolScript.core.type.bool, (ctx, v) => !!v[0](ctx)),
+    C(MolScript.core.type.num, (ctx, v) => +v[0](ctx)),
+    C(MolScript.core.type.str, (ctx, v) => '' + v[0](ctx)),
+    C(MolScript.core.type.list, (ctx, xs) => QueryRuntimeArguments.forEachEval(xs, ctx, (v, i, list) => list[i] = v, [] as any[])),
+    C(MolScript.core.type.set, (ctx, xs) => QueryRuntimeArguments.forEachEval(xs, ctx, (v, i, set) => set.add(v), new Set<any>())),
+    C(MolScript.core.type.regex, (ctx, v) => new RegExp(v[0](ctx), (v[1] && v[1](ctx)) || '')),
+    C(MolScript.core.type.bitflags, (ctx, v) => +v[0](ctx)),
+    C(MolScript.core.type.compositeKey, (ctx, xs) => QueryRuntimeArguments.forEachEval(xs, ctx, (v, i, list) => list[i] = '' + v, [] as string[]).join('-')),
+
+    // ============= LOGIC ================
+    C(MolScript.core.logic.not, (ctx, v) => !v[0](ctx)),
+    C(MolScript.core.logic.and, (ctx, xs) => {
         if (typeof xs.length === 'number') {
-            for (let i = 0, _i = xs.length; i < _i; i++) ret += xs[i](ctx);
+            for (let i = 0, _i = xs.length; i < _i; i++) if (!xs[i](ctx)) return false;
         } else {
-            for (const k of Object.keys(xs)) ret += xs[k](ctx);
+            for (const k of Object.keys(xs)) if (!xs[k](ctx)) return false;
         }
-        return ret;
+        return true;
+    }),
+    C(MolScript.core.logic.or, (ctx, xs) => {
+        if (typeof xs.length === 'number') {
+            for (let i = 0, _i = xs.length; i < _i; i++) if (xs[i](ctx)) return true;
+        } else {
+            for (const k of Object.keys(xs)) if (xs[k](ctx)) return true;
+        }
+        return false;
     }),
 
     // ============= RELATIONAL ================
@@ -35,8 +58,133 @@ const symbols = [
         return x >= v[1](ctx) && x <= v[2](ctx);
     }),
 
+    // ============= ARITHMETIC ================
+    C(MolScript.core.math.add, (ctx, xs) => {
+        let ret = 0;
+        if (typeof xs.length === 'number') {
+            for (let i = 0, _i = xs.length; i < _i; i++) ret += xs[i](ctx);
+        } else {
+            for (const k of Object.keys(xs)) ret += xs[k](ctx);
+        }
+        return ret;
+    }),
+    C(MolScript.core.math.sub, (ctx, xs) => {
+        let ret = 0;
+        if (typeof xs.length === 'number') {
+            if (xs.length === 1) return -xs[0](ctx);
+            ret = xs[0](ctx) || 0;
+            for (let i = 1, _i = xs.length; i < _i; i++) ret -= xs[i](ctx);
+        } else {
+            const keys = Object.keys(xs);
+            if (keys.length === 1)
+            ret = xs[keys[0]](ctx) || 0;
+            for (let i = 1, _i = keys.length; i < _i; i++) ret -= xs[keys[i]](ctx);
+        }
+        return ret;
+    }),
+    C(MolScript.core.math.mult, (ctx, xs) => {
+        let ret = 1;
+        if (typeof xs.length === 'number') {
+            for (let i = 0, _i = xs.length; i < _i; i++) ret *= xs[i](ctx);
+        } else {
+            for (const k of Object.keys(xs)) ret *= xs[k](ctx);
+        }
+        return ret;
+    }),
+    C(MolScript.core.math.div, (ctx, v) => v[0](ctx) / v[1](ctx)),
+    C(MolScript.core.math.pow, (ctx, v) => Math.pow(v[0](ctx), v[1](ctx))),
+    C(MolScript.core.math.mod, (ctx, v) => v[0](ctx) % v[1](ctx)),
+
+    C(MolScript.core.math.min, (ctx, xs) => {
+        let ret = Number.POSITIVE_INFINITY;
+        if (typeof xs.length === 'number') {
+            for (let i = 0, _i = xs.length; i < _i; i++) ret = Math.min(xs[i](ctx), ret);
+        } else {
+            for (const k of Object.keys(xs)) ret = Math.min(xs[k](ctx), ret)
+        }
+        return ret;
+    }),
+    C(MolScript.core.math.max, (ctx, xs) => {
+        let ret = Number.NEGATIVE_INFINITY;
+        if (typeof xs.length === 'number') {
+            for (let i = 0, _i = xs.length; i < _i; i++) ret = Math.max(xs[i](ctx), ret);
+        } else {
+            for (const k of Object.keys(xs)) ret = Math.max(xs[k](ctx), ret)
+        }
+        return ret;
+    }),
+
+    C(MolScript.core.math.floor, (ctx, v) => Math.floor(v[0](ctx))),
+    C(MolScript.core.math.ceil, (ctx, v) => Math.ceil(v[0](ctx))),
+    C(MolScript.core.math.roundInt, (ctx, v) => Math.round(v[0](ctx))),
+    C(MolScript.core.math.abs, (ctx, v) => Math.abs(v[0](ctx))),
+    C(MolScript.core.math.sqrt, (ctx, v) => Math.sqrt(v[0](ctx))),
+    C(MolScript.core.math.sin, (ctx, v) => Math.sin(v[0](ctx))),
+    C(MolScript.core.math.cos, (ctx, v) => Math.cos(v[0](ctx))),
+    C(MolScript.core.math.tan, (ctx, v) => Math.tan(v[0](ctx))),
+    C(MolScript.core.math.asin, (ctx, v) => Math.asin(v[0](ctx))),
+    C(MolScript.core.math.acos, (ctx, v) => Math.acos(v[0](ctx))),
+    C(MolScript.core.math.atan, (ctx, v) => Math.atan(v[0](ctx))),
+    C(MolScript.core.math.sinh, (ctx, v) => Math.sinh(v[0](ctx))),
+    C(MolScript.core.math.cosh, (ctx, v) => Math.cosh(v[0](ctx))),
+    C(MolScript.core.math.tanh, (ctx, v) => Math.tanh(v[0](ctx))),
+    C(MolScript.core.math.exp, (ctx, v) => Math.exp(v[0](ctx))),
+    C(MolScript.core.math.log, (ctx, v) => Math.log(v[0](ctx))),
+    C(MolScript.core.math.log10, (ctx, v) => Math.log10(v[0](ctx))),
+    C(MolScript.core.math.atan2, (ctx, v) => Math.atan2(v[0](ctx), v[1](ctx))),
+
+    // ============= STRING ================
+    C(MolScript.core.str.match, (ctx, v) => v[0](ctx).test(v[1](ctx))),
+    C(MolScript.core.str.concat, (ctx, xs) => {
+        let ret: string[] = [];
+        if (typeof xs.length === 'number') {
+            for (let i = 0, _i = xs.length; i < _i; i++) ret.push(xs[i](ctx).toString());
+        } else {
+            for (const k of Object.keys(xs)) ret.push(xs[k](ctx).toString());
+        }
+        return ret.join('');
+    }),
+
+    // ============= LIST ================
+    C(MolScript.core.list.getAt, (ctx, v) => v[0](ctx)[v[1](ctx)]),
+
+    // ============= SET ================
+    C(MolScript.core.set.has, (ctx, v) => v[0](ctx).has(v[1](ctx))),
+    C(MolScript.core.set.isSubset, (ctx, v) => isSuperset(v[1](ctx) as Set<any>, v[0](ctx) as Set<any>)),
+
+    // ============= FLAGS ================
+    C(MolScript.core.flags.hasAny, (ctx, v) => {
+        const test = v[1](ctx);
+        const tested = v[0](ctx);
+        if (!test) return !!tested;
+        return (tested & test) !== 0;
+    }),
+    C(MolScript.core.flags.hasAll, (ctx, v) => {
+        const test = v[1](ctx);
+        const tested = v[0](ctx);
+        if (!test) return !tested;
+        return (tested & test) === test;
+    }),
+
+    ////////////////////////////////////
+    // Structure
+
     // ============= TYPES ================
     C(MolScript.structureQuery.type.elementSymbol, (ctx, v) => ElementSymbol(v[0](ctx))),
+    C(MolScript.structureQuery.type.atomName, (ctx, v) => toUpperCase(v[0](ctx))),
+
+    // TODO:
+    // C(MolScript.structureQuery.type.bondFlags, (ctx, v) => StructureRuntime.BondProperties.createFlags(env, v)),
+    // C(MolScript.structureQuery.type.secondaryStructureFlags, (ctx, v) => StructureRuntime.AtomProperties.createSecondaryStructureFlags(env, v)),
+    // C(MolScript.structureQuery.type.entityType, (ctx, v) => StructureRuntime.Common.entityType(v[0](ctx))),
+    // C(MolScript.structureQuery.type.ringFingerprint, (ctx, v) => StructureRuntime.Common.ringFingerprint(env, v as any)),
+    // C(MolScript.structureQuery.type.authResidueId, (ctx, v) => ResidueIdentifier.auth(v[0](ctx), v[1](ctx), v[2] && v[2](ctx))),
+    // C(MolScript.structureQuery.type.labelResidueId, (ctx, v) => ResidueIdentifier.label(v[0](ctx), v[1](ctx), v[2](ctx), v[3] && v[3](ctx))),
+
+    // ============= 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),
 
     // ============= GENERATORS ================
     D(MolScript.structureQuery.generator.atomGroups, (ctx, xs) => Queries.generators.atoms({
@@ -48,10 +196,76 @@ const symbols = [
     })(ctx)),
 
     // ============= ATOM PROPERTIES ================
-    D(MolScript.structureQuery.atomProperty.macromolecular.label_comp_id, (ctx, _) => StructureProperties.residue.label_comp_id(ctx.element)),
-    D(MolScript.structureQuery.atomProperty.core.elementSymbol, (ctx, _) => StructureProperties.atom.type_symbol(ctx.element))
+
+    // ~~~ 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.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)),
+    D(MolScript.structureQuery.atomProperty.core.atomKey, (ctx, _) => cantorPairing(ctx.element.unit.id, ctx.element.element)),
+
+    // TODO:
+    // D(MolScript.structureQuery.atomProperty.core.bondCount, (ctx, _) => ),
+
+    // ~~~ TOPOLOGY ~~~
+
+    // TODO
+
+    // ~~~ MACROMOLECULAR ~~~
+
+    // TODO:
+    // // identifiers
+    // labelResidueId: prop((env, v) => ResidueIdentifier.labelOfResidueIndex(env.context.model, getAddress(env, v).residue)),
+    // 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)),
+
+    // 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.label_atom_id, atomProp(StructureProperties.atom.label_atom_id)),
+    D(MolScript.structureQuery.atomProperty.macromolecular.label_alt_id, atomProp(StructureProperties.atom.label_alt_id)),
+    D(MolScript.structureQuery.atomProperty.macromolecular.label_asym_id, atomProp(StructureProperties.chain.label_asym_id)),
+    D(MolScript.structureQuery.atomProperty.macromolecular.label_comp_id, atomProp(StructureProperties.residue.label_comp_id)),
+    D(MolScript.structureQuery.atomProperty.macromolecular.label_seq_id, atomProp(StructureProperties.residue.label_seq_id)),
+    D(MolScript.structureQuery.atomProperty.macromolecular.label_entity_id, atomProp(StructureProperties.entity.id)),
+
+    D(MolScript.structureQuery.atomProperty.macromolecular.auth_atom_id, atomProp(StructureProperties.atom.auth_atom_id)),
+    D(MolScript.structureQuery.atomProperty.macromolecular.auth_asym_id, atomProp(StructureProperties.chain.auth_asym_id)),
+    D(MolScript.structureQuery.atomProperty.macromolecular.auth_comp_id, atomProp(StructureProperties.residue.auth_comp_id)),
+    D(MolScript.structureQuery.atomProperty.macromolecular.auth_seq_id, atomProp(StructureProperties.residue.auth_seq_id)),
+
+    D(MolScript.structureQuery.atomProperty.macromolecular.pdbx_PDB_ins_code, atomProp(StructureProperties.residue.pdbx_PDB_ins_code)),
+    D(MolScript.structureQuery.atomProperty.macromolecular.pdbx_formal_charge, atomProp(StructureProperties.atom.pdbx_formal_charge)),
+    D(MolScript.structureQuery.atomProperty.macromolecular.occupancy, atomProp(StructureProperties.atom.occupancy)),
+    D(MolScript.structureQuery.atomProperty.macromolecular.B_iso_or_equiv, atomProp(StructureProperties.atom.B_iso_or_equiv)),
+
+    D(MolScript.structureQuery.atomProperty.macromolecular.entityType, atomProp(StructureProperties.entity.type)),
+
+    D(MolScript.structureQuery.atomProperty.macromolecular.isModified, (ctx, _) => ctx.element.unit.model.properties.modifiedResidues.parentId.has(StructureProperties.residue.label_comp_id(ctx.element))),
+    D(MolScript.structureQuery.atomProperty.macromolecular.modifiedParentName, (ctx, _) => {
+        const id = StructureProperties.residue.label_comp_id(ctx.element);
+        return ctx.element.unit.model.properties.modifiedResidues.parentId.get(id) || id
+    })
+
+    // TODO
+    // MolScript.structureQuery.atomProperty.macromolecular.secondaryStructureKey
+    // MolScript.structureQuery.atomProperty.macromolecular.secondaryStructureFlags
+
+    // ============= BOND PROPERTIES ================
 ];
 
+function atomProp(p: (e: StructureElement) => any): (ctx: QueryContext, _: any) => any {
+    return (ctx, _) => p(ctx.element);
+}
+
 (function () {
     for (const s of symbols) {
         DefaultQueryRuntimeTable.addSymbol(s);

src/mol-util/upper-case.ts

+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+export default function toUpperCase(value: any): string {
+    if (!value) return '';
+    return typeof value === 'string' ? value.toUpperCase() : `${value}`.toUpperCase();
+}
\ No newline at end of file