diff --git a/src/mol-math/geometry/lookup3d/grid.ts b/src/mol-math/geometry/lookup3d/grid.ts
index db2409a1ba38052418d98605c79918063bf61db2..e91b0aff68af67792644b4ce9f4348b9186d0fb5 100644
--- a/src/mol-math/geometry/lookup3d/grid.ts
+++ b/src/mol-math/geometry/lookup3d/grid.ts
@@ -12,15 +12,20 @@ import { PositionData } from '../common';
import { Vec3 } from '../../linear-algebra';
import { OrderedSet } from 'mol-data/int';
-function GridLookup3D(data: PositionData): Lookup3D {
- return new GridLookup3DImpl(data);
+interface GridLookup3D<T = number> extends Lookup3D<T> {
+ readonly buckets: { readonly offset: ArrayLike<number>, readonly count: ArrayLike<number>, readonly array: ArrayLike<number> }
+}
+
+function GridLookup3D(data: PositionData, cellSize?: Vec3): GridLookup3D {
+ return new GridLookup3DImpl(data, cellSize);
}
export { GridLookup3D }
-class GridLookup3DImpl implements Lookup3D<number> {
+class GridLookup3DImpl implements GridLookup3D<number> {
private ctx: QueryContext;
boundary: Lookup3D['boundary'];
+ buckets: GridLookup3D['buckets'];
find(x: number, y: number, z: number, radius: number): Result<number> {
this.ctx.x = x;
@@ -31,6 +36,7 @@ class GridLookup3DImpl implements Lookup3D<number> {
query(this.ctx);
return this.ctx.result;
}
+
check(x: number, y: number, z: number, radius: number): boolean {
this.ctx.x = x;
this.ctx.y = y;
@@ -40,10 +46,11 @@ class GridLookup3DImpl implements Lookup3D<number> {
return query(this.ctx);
}
- constructor(data: PositionData) {
- const structure = build(data);
+ constructor(data: PositionData, cellSize?: Vec3) {
+ const structure = build(data, cellSize);
this.ctx = createContext(structure);
this.boundary = { box: structure.boundingBox, sphere: structure.boundingSphere };
+ this.buckets = { offset: structure.bucketOffset, count: structure.bucketCounts, array: structure.bucketArray };
}
}
@@ -82,18 +89,18 @@ interface BuildState {
expandedBox: Box3D,
boundingBox: Box3D,
boundingSphere: Sphere3D,
- count: number
+ elementCount: number
}
function _build(state: BuildState): Grid3D {
- const { expandedBox, size: [sX, sY, sZ], data: { x: px, y: py, z: pz, radius, indices }, count, delta } = state;
+ const { expandedBox, size: [sX, sY, sZ], data: { x: px, y: py, z: pz, radius, indices }, elementCount, delta } = state;
const n = sX * sY * sZ;
const { min: [minX, minY, minZ] } = expandedBox;
let maxRadius = 0;
let bucketCount = 0;
const grid = new Uint32Array(n);
- const bucketIndex = new Int32Array(count);
+ const bucketIndex = new Int32Array(elementCount);
for (let t = 0, _t = OrderedSet.size(indices); t < _t; t++) {
const i = OrderedSet.getAt(indices, t);
const x = Math.floor((px[i] - minX) / delta[0]);
@@ -124,14 +131,14 @@ function _build(state: BuildState): Grid3D {
}
}
- const bucketOffset = new Uint32Array(count);
- for (let i = 1; i < count; ++i) {
+ const bucketOffset = new Uint32Array(bucketCount);
+ for (let i = 1; i < bucketCount; ++i) {
bucketOffset[i] += bucketOffset[i - 1] + bucketCounts[i - 1];
}
const bucketFill = new Int32Array(bucketCount);
- const bucketArray = new Int32Array(count);
- for (let i = 0; i < count; i++) {
+ const bucketArray = new Int32Array(elementCount);
+ for (let i = 0; i < elementCount; i++) {
const bucketIdx = grid[bucketIndex[i]]
if (bucketIdx > 0) {
const k = bucketIdx - 1;
@@ -156,7 +163,7 @@ function _build(state: BuildState): Grid3D {
}
}
-function build(data: PositionData) {
+function build(data: PositionData, cellSize?: Vec3) {
const boundingBox = Box3D.computeBounding(data);
// need to expand the grid bounds to avoid rounding errors
const expandedBox = Box3D.expand(Box3D.empty(), boundingBox, Vec3.create(0.5, 0.5, 0.5));
@@ -168,7 +175,10 @@ function build(data: PositionData) {
const elementCount = OrderedSet.size(indices);
- if (elementCount > 0) {
+ if (cellSize) {
+ size = [Math.ceil(S[0] / cellSize[0]), Math.ceil(S[1] / cellSize[1]), Math.ceil(S[2] / cellSize[2])];
+ delta = cellSize;
+ } else if (elementCount > 0) {
// size of the box
// required "grid volume" so that each cell contains on average 32 elements.
const V = Math.ceil(elementCount / 32);
@@ -194,7 +204,7 @@ function build(data: PositionData) {
expandedBox,
boundingBox,
boundingSphere,
- count: elementCount,
+ elementCount: elementCount,
delta
}
diff --git a/src/mol-model/structure/structure/structure.ts b/src/mol-model/structure/structure/structure.ts
index 29a2be47685bcf25bc74255f0335db6c3f306f08..e22d4b3fba06f57d2ecd2d6dd101ec171e2bfbad 100644
--- a/src/mol-model/structure/structure/structure.ts
+++ b/src/mol-model/structure/structure/structure.ts
@@ -18,11 +18,12 @@ import { InterUnitBonds, computeInterUnitBonds } from './unit/links';
import { PairRestraints, CrossLinkRestraint, extractCrossLinkRestraints } from './unit/pair-restraints';
import StructureSymmetry from './symmetry';
import StructureProperties from './properties';
-import { ResidueIndex } from '../model/indexing';
+import { ResidueIndex, ChainIndex } from '../model/indexing';
import { Carbohydrates } from './carbohydrates/data';
import { computeCarbohydrates } from './carbohydrates/compute';
import { Vec3 } from 'mol-math/linear-algebra';
import { idFactory } from 'mol-util/id-factory';
+import { GridLookup3D } from 'mol-math/geometry';
class Structure {
/** Maps unit.id to unit */
@@ -175,7 +176,12 @@ namespace Structure {
}
const elements = SortedArray.ofBounds(start as ElementIndex, chains.offsets[c + 1] as ElementIndex);
- builder.addUnit(Unit.Kind.Atomic, model, SymmetryOperator.Default, elements);
+
+ if (isWaterChain(model, c as ChainIndex, elements)) {
+ partitionAtomicUnit(model, elements, builder);
+ } else {
+ builder.addUnit(Unit.Kind.Atomic, model, SymmetryOperator.Default, elements);
+ }
}
const cs = model.coarseHierarchy;
@@ -191,6 +197,26 @@ namespace Structure {
return builder.getStructure();
}
+ function isWaterChain(model: Model, chainIndex: ChainIndex, indices: SortedArray) {
+ const e = model.atomicHierarchy.index.getEntityFromChain(chainIndex);
+ return model.entities.data.type.value(e) === 'water';
+ }
+
+ function partitionAtomicUnit(model: Model, indices: SortedArray, builder: StructureBuilder) {
+ const { x, y, z } = model.atomicConformation;
+ const lookup = GridLookup3D({ x, y, z, indices }, Vec3.create(64, 64, 64));
+ const { offset, count, array } = lookup.buckets;
+
+ for (let i = 0, _i = offset.length; i < _i; i++) {
+ const start = offset[i];
+ const set = new Int32Array(count[i]);
+ for (let j = 0, _j = count[i]; j < _j; j++) {
+ set[j] = indices[array[start + j]];
+ }
+ builder.addUnit(Unit.Kind.Atomic, model, SymmetryOperator.Default, SortedArray.ofSortedArray(set));
+ }
+ }
+
function addCoarseUnits(builder: StructureBuilder, model: Model, elements: CoarseElements, kind: Unit.Kind) {
const { chainElementSegments } = elements;
for (let cI = 0; cI < chainElementSegments.count; cI++) {