Compute bonds on unit level

src/apps/structure-info/index.ts

@@ -10,8 +10,9 @@ require('util.promisify').shim();
 
 // import { Table } from 'mol-data/db'
 import CIF from 'mol-io/reader/cif'
-import { Model } from 'mol-model/structure'
+import { Model, Structure, ElementSet, Unit } from 'mol-model/structure'
 import { Run, Progress } from 'mol-task'
+import { OrderedSet } from 'mol-data/int';
 
 async function parseCif(data: string|Uint8Array) {
     const comp = CIF.parse(data);
@@ -37,24 +38,34 @@ export function atomLabel(model: Model, aI: number) {
 }
 
 
-function printBonds(model: Model) {
-    // TODO: do bonds again
-    // const { count, offset, neighbor } = Model.bonds(model)
-    // for (let i = 0; i < count; ++i) {
-    //     const start = offset[i];
-    //     const end = offset[i + 1];
-    //     for (let bI = start; bI < end; bI++) {
-    //         console.log(`${atomLabel(model, i)} -- ${atomLabel(model, neighbor[bI])}`)
-    //     }
-    // }
+function printBonds(structure: Structure) {
+
+    const { units, elements } = structure;
+    const unitIds = ElementSet.unitIndices(elements);
+
+    for (let i = 0, _i = OrderedSet.size(unitIds); i < _i; i++) {
+        const unit = units[OrderedSet.getAt(unitIds, i)];
+        const group = ElementSet.groupFromUnitIndex(elements, OrderedSet.getAt(unitIds, i));
+
+        const { count, offset, neighbor } = Unit.getGroupBonds(unit, group);
+        const { model }  = unit;
+
+        for (let j = 0; j < count; ++j) {
+            const start = offset[j];
+            const end = offset[j + 1];
+            for (let bI = start; bI < end; bI++) {
+                console.log(`${atomLabel(model, j)} -- ${atomLabel(model, neighbor[bI])}`)
+            }
+        }
+    }
 }
 
 async function run(pdb: string) {
     const mmcif = await getPdb(pdb)
     const models = Model.create({ kind: 'mmCIF', data: mmcif });
-    // const structure = Structure.ofModel(models[0])
+    const structure = Structure.ofModel(models[0])
     // console.log(structure)
-    printBonds(models[0])
+    printBonds(structure)
 }
 
 const parser = new argparse.ArgumentParser({

src/mol-math/geometry/lookup3d/grid.ts

@@ -256,7 +256,7 @@ function query(ctx: QueryContext): boolean {
                     if (distSq <= rSq) {
                         if (maxRadius > 0 && Math.sqrt(distSq) - radius![idx] > inputRadius) continue;
                         if (isCheck) return true;
-                        Result.add(result, bucketArray[i], distSq);
+                        Result.add(result, idx, distSq);
                     }
                 }
             }

src/mol-model/structure/model/formats/mmcif/bonds.ts

@@ -6,12 +6,12 @@
  */
 
 import Model from '../../model'
-import Bonds from '../../properties/bonds'
 import { BondType } from '../../types'
 import { findEntityIdByAsymId, findAtomIndexByLabelName } from './util'
 import { Column } from 'mol-data/db'
+import { GroupBonds } from '../../../structure/element/properties/bonds/group-data';
 
-export class StructConn implements Bonds.StructConn {
+export class StructConn implements GroupBonds.StructConn {
     private _residuePairIndex: Map<string, StructConn.Entry[]> | undefined = void 0;
     private _atomIndex: Map<number, StructConn.Entry[]> | undefined = void 0;
 
@@ -71,7 +71,7 @@ export class StructConn implements Bonds.StructConn {
 }
 
 export namespace StructConn {
-    export interface Entry extends Bonds.StructConnEntry {
+    export interface Entry extends GroupBonds.StructConnEntry {
         distance: number,
         order: number,
         flags: number,
@@ -181,7 +181,7 @@ export namespace StructConn {
     }
 }
 
-export class ComponentBondInfo implements Bonds.ComponentBondInfo {
+export class ComponentBondInfo implements GroupBonds.ComponentBondInfo {
     entries: Map<string, ComponentBondInfo.Entry> = new Map();
 
     newEntry(id: string) {
@@ -192,7 +192,7 @@ export class ComponentBondInfo implements Bonds.ComponentBondInfo {
 }
 
 export namespace ComponentBondInfo {
-    export class Entry implements Bonds.ComponentBondInfoEntry {
+    export class Entry implements GroupBonds.ComponentBondInfoEntry {
         map: Map<string, Map<string, { order: number, flags: number }>> = new Map();
 
         add(a: string, b: string, order: number, flags: number, swap = true) {

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

@@ -7,13 +7,15 @@
 import { OrderedSet } from 'mol-data/int'
 import { Lookup3D } from 'mol-math/geometry'
 import Unit from '../unit'
+import { GroupBonds } from './properties/bonds/group-data';
 
 interface ElementGroup {
     elements: OrderedSet,
     // Unique identifier of the group, usable as partial key for various "caches".
     key: number,
 
-    __lookup3d__?: Lookup3D
+    __lookup3d__?: Lookup3D,
+    __bonds__?: GroupBonds
 }
 
 namespace ElementGroup {

src/mol-model/structure/model/properties/bonds.ts → src/mol-model/structure/structure/element/properties/bonds/group-data.ts

@@ -5,9 +5,9 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { BondType } from '../types'
+import { BondType } from '../../../../model/types'
 
-interface Bonds {
+interface GroupBonds {
     /**
      * Where bonds for atom A start and end.
      * Start offset at idx, end at idx + 1
@@ -21,8 +21,8 @@ interface Bonds {
     count: number
 }
 
-namespace Bonds {
-    export function createEmpty(): Bonds {
+namespace GroupBonds {
+    export function createEmpty(): GroupBonds {
         return { offset: [], neighbor: [], order: [], flags: [], count: 0 }
     }
     export function isCovalent(flags: number) {
@@ -46,6 +46,4 @@ namespace Bonds {
     }
 }
 
-
-
-export default Bonds
+export { GroupBonds }

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

@@ -8,6 +8,7 @@ import { SymmetryOperator } from 'mol-math/geometry/symmetry-operator'
 import ElementGroup from './element/group'
 import { Model } from '../model'
 import { GridLookup3D } from 'mol-math/geometry'
+import { computeUnitBonds } from './element/properties/bonds/group-compute';
 
 // A building block of a structure that corresponds to an atomic or a coarse grained representation
 // 'conveniently grouped together'.
@@ -95,6 +96,16 @@ namespace Unit {
 
         throw 'not implemented';
     }
+
+    export function getGroupBonds(unit: Unit, group: ElementGroup) {
+        if (group.__bonds__) return group.__bonds__;
+        if (Unit.isAtomic(unit)) {
+            group.__bonds__ = computeUnitBonds(unit, group);
+            return group.__bonds__;
+        }
+        
+        throw 'not implemented';
+    }
 }
 
 export default Unit;