Skip to content
Snippets Groups Projects
Commit cbf86ece authored by Alexander Rose's avatar Alexander Rose
Browse files

molecular surface improvements

parent e87abc47
No related branches found
No related tags found
No related merge requests found
......@@ -10,7 +10,6 @@
import { fillUniform } from 'mol-util/array';
import { Vec3, Tensor } from 'mol-math/linear-algebra';
import { ParamDefinition as PD } from 'mol-util/param-definition';
import { Lookup3D, Result } from './lookup3d/common';
import { RuntimeContext } from 'mol-task';
import { OrderedSet } from 'mol-data/int';
import { PositionData } from './common';
......@@ -44,6 +43,23 @@ function getAngleTables (probePositions: number): AnglesTables {
return { cosTable, sinTable}
}
//
export const MolecularSurfaceCalculationParams = {
resolution: PD.Numeric(0.5, { min: 0.01, max: 10, step: 0.01 }),
probeRadius: PD.Numeric(1.4, { min: 0, max: 10, step: 0.1 }),
probePositions: PD.Numeric(30, { min: 12, max: 90, step: 1 }),
}
export const DefaultMolecularSurfaceCalculationProps = PD.getDefaultValues(MolecularSurfaceCalculationParams)
export type MolecularSurfaceCalculationProps = typeof DefaultMolecularSurfaceCalculationProps
export async function calcMolecularSurface(ctx: RuntimeContext, position: Required<PositionData>, maxRadius: number, props: MolecularSurfaceCalculationProps) {
// Field generation method adapted from AstexViewer (Mike Hartshorn) by Fred Ludlow.
// Other parts based heavily on NGL (Alexander Rose) EDT Surface class
let lastClip = -1
/**
* Is the point at x,y,z obscured by any of the atoms specifeid by indices in neighbours.
* Ignore indices a and b (these are the relevant atoms in projectPoints/Torii)
......@@ -52,20 +68,20 @@ function getAngleTables (probePositions: number): AnglesTables {
*
* `a` and `b` must be resolved indices
*/
function obscured (state: MolSurfCalcState, x: number, y: number, z: number, a: number, b: number) {
if (state.lastClip !== -1) {
const ai = state.lastClip
if (ai !== a && ai !== b && singleAtomObscures(state, ai, x, y, z)) {
function obscured(x: number, y: number, z: number, a: number, b: number) {
if (lastClip !== -1) {
const ai = lastClip
if (ai !== a && ai !== b && singleAtomObscures(ai, x, y, z)) {
return ai
} else {
state.lastClip = -1
lastClip = -1
}
}
for (let j = 0, jl = state.neighbours.count; j < jl; ++j) {
const ai = OrderedSet.getAt(state.position.indices, state.neighbours.indices[j])
if (ai !== a && ai !== b && singleAtomObscures(state, ai, x, y, z)) {
state.lastClip = ai
for (let j = 0, jl = neighbours.count; j < jl; ++j) {
const ai = OrderedSet.getAt(position.indices, neighbours.indices[j])
if (ai !== a && ai !== b && singleAtomObscures(ai, x, y, z)) {
lastClip = ai
return ai
}
}
......@@ -76,11 +92,11 @@ function obscured (state: MolSurfCalcState, x: number, y: number, z: number, a:
/**
* `ai` must be a resolved index
*/
function singleAtomObscures (state: MolSurfCalcState, ai: number, x: number, y: number, z: number) {
const r = state.position.radius[ai]
const dx = state.position.x[ai] - x
const dy = state.position.y[ai] - y
const dz = state.position.z[ai] - z
function singleAtomObscures(ai: number, x: number, y: number, z: number) {
const r = radius[ai]
const dx = px[ai] - x
const dy = py[ai] - y
const dz = pz[ai] - z
const dSq = dx * dx + dy * dy + dz * dz
return dSq < (r * r)
}
......@@ -95,31 +111,22 @@ function singleAtomObscures (state: MolSurfCalcState, ai: number, x: number, y:
* Calculate delta distance and set grid value to minimum of
* itself and delta
*/
async function projectPoints (ctx: RuntimeContext, state: MolSurfCalcState) {
const { position, lookup3d, min, space, data, idData, scaleFactor, updateChunk } = state
const { gridx, gridy, gridz } = state
const { indices, x, y, z, radius } = position
const n = OrderedSet.size(indices)
const [ dimX, dimY, dimZ ] = space.dimensions
const iu = dimZ, iv = dimY, iuv = iu * iv
for (let i = 0; i < n; ++i) {
function projectPointsRange (begI: number, endI: number) {
for (let i = begI; i < endI; ++i) {
const j = OrderedSet.getAt(indices, i)
const vx = x[j], vy = y[j], vz = z[j]
const vx = px[j], vy = py[j], vz = pz[j]
const rad = radius[j]
const rSq = rad * rad
state.neighbours = lookup3d.find(vx, vy, vz, rad)
lookup3d.find(vx, vy, vz, rad)
// Number of grid points, round this up...
const ng = Math.ceil(rad * scaleFactor)
// Center of the atom, mapped to grid points (take floor)
const iax = Math.floor(scaleFactor * (vx - min[0]))
const iay = Math.floor(scaleFactor * (vy - min[1]))
const iaz = Math.floor(scaleFactor * (vz - min[2]))
const iax = Math.floor(scaleFactor * (vx - minX))
const iay = Math.floor(scaleFactor * (vy - minY))
const iaz = Math.floor(scaleFactor * (vz - minZ))
// Extents of grid to consider for this atom
const begX = Math.max(0, iax - ng)
......@@ -158,7 +165,7 @@ async function projectPoints (ctx: RuntimeContext, state: MolSurfCalcState) {
const spy = dy * ap + vy
const spz = dz * ap + vz
if (obscured(state, spx, spy, spz, j, -1) === -1) {
if (obscured(spx, spy, spz, j, -1) === -1) {
const dd = rad - d
if (dd < data[idx]) {
data[idx] = dd
......@@ -169,8 +176,14 @@ async function projectPoints (ctx: RuntimeContext, state: MolSurfCalcState) {
}
}
}
}
}
if (i % updateChunk === 0 && ctx.shouldUpdate) {
async function projectPoints() {
for (let i = 0; i < n; i += updateChunk) {
projectPointsRange(i, Math.min(i + updateChunk, n))
if (ctx.shouldUpdate) {
await ctx.update({ message: 'projecting points', current: i, max: n })
}
}
......@@ -184,21 +197,12 @@ const n2 = Vec3()
/**
* `a` and `b` must be resolved indices
*/
function projectTorus (state: MolSurfCalcState, a: number, b: number) {
const { position, min, space, data, idData } = state
const { cosTable, sinTable, probePositions, probeRadius, scaleFactor } = state
const { gridx, gridy, gridz } = state
const [ dimX, dimY, dimZ ] = space.dimensions
const iu = dimZ, iv = dimY, iuv = iu * iv
const ng = Math.max(5, 2 + Math.floor(probeRadius * scaleFactor))
const rA = position.radius[a]
const rB = position.radius[b]
const dx = atob[0] = position.x[b] - position.x[a]
const dy = atob[1] = position.y[b] - position.y[a]
const dz = atob[2] = position.z[b] - position.z[a]
function projectTorus(a: number, b: number) {
const rA = radius[a]
const rB = radius[b]
const dx = atob[0] = px[b] - px[a]
const dy = atob[1] = py[b] - py[a]
const dz = atob[2] = pz[b] - pz[a]
const dSq = dx * dx + dy * dy + dz * dz
// This check now redundant as already done in AVHash.withinRadii
......@@ -230,11 +234,11 @@ function projectTorus (state: MolSurfCalcState, a: number, b: number) {
Vec3.scale(n2, n2, rInt)
Vec3.scale(atob, atob, dmp)
mid[0] = atob[0] + position.x[a]
mid[1] = atob[1] + position.y[a]
mid[2] = atob[2] + position.z[a]
mid[0] = atob[0] + px[a]
mid[1] = atob[1] + py[a]
mid[2] = atob[2] + pz[a]
state.lastClip = -1
lastClip = -1
for (let i = 0; i < probePositions; ++i) {
const cost = cosTable[i]
......@@ -244,18 +248,18 @@ function projectTorus (state: MolSurfCalcState, a: number, b: number) {
const py = mid[1] + cost * n1[1] + sint * n2[1]
const pz = mid[2] + cost * n1[2] + sint * n2[2]
if (obscured(state, px, py, pz, a, b) === -1) {
const iax = Math.floor(scaleFactor * (px - min[0]))
const iay = Math.floor(scaleFactor * (py - min[1]))
const iaz = Math.floor(scaleFactor * (pz - min[2]))
if (obscured(px, py, pz, a, b) === -1) {
const iax = Math.floor(scaleFactor * (px - minX))
const iay = Math.floor(scaleFactor * (py - minY))
const iaz = Math.floor(scaleFactor * (pz - minZ))
const begX = Math.max(0, iax - ng)
const begY = Math.max(0, iay - ng)
const begZ = Math.max(0, iaz - ng)
const begX = Math.max(0, iax - ngTorus)
const begY = Math.max(0, iay - ngTorus)
const begZ = Math.max(0, iaz - ngTorus)
const endX = Math.min(dimX, iax + ng + 2)
const endY = Math.min(dimY, iay + ng + 2)
const endZ = Math.min(dimZ, iaz + ng + 2)
const endX = Math.min(dimX, iax + ngTorus + 2)
const endY = Math.min(dimY, iay + ngTorus + 2)
const endZ = Math.min(dimZ, iaz + ngTorus + 2)
for (let xi = begX; xi < endX; ++xi) {
const dx = px - gridx[xi]
......@@ -287,15 +291,13 @@ function projectTorus (state: MolSurfCalcState, a: number, b: number) {
}
}
async function projectTorii (ctx: RuntimeContext, state: MolSurfCalcState) {
const { n, lookup3d, position, updateChunk } = state
const { x, y, z, indices, radius } = position
async function projectTorii () {
for (let i = 0; i < n; ++i) {
const k = OrderedSet.getAt(indices, i)
state.neighbours = lookup3d.find(x[k], y[k], z[k], radius[k])
for (let j = 0, jl = state.neighbours.count; j < jl; ++j) {
const l = OrderedSet.getAt(indices, state.neighbours.indices[j])
if (k < l) projectTorus(state, k, l)
lookup3d.find(px[k], py[k], pz[k], radius[k])
for (let j = 0, jl = neighbours.count; j < jl; ++j) {
const l = OrderedSet.getAt(indices, neighbours.indices[j])
if (k < l) projectTorus(k, l)
}
if (i % updateChunk === 0 && ctx.shouldUpdate) {
......@@ -304,60 +306,30 @@ async function projectTorii (ctx: RuntimeContext, state: MolSurfCalcState) {
}
}
//
interface MolSurfCalcState {
/** Cached last value for obscured test */
lastClip: number
/** Neighbours as transient result array from lookup3d */
neighbours: Result<number>
lookup3d: Lookup3D
position: Required<PositionData>
min: Vec3
updateChunk: number
maxRadius: number
n: number
resolution: number
scaleFactor: number
probeRadius: number
/** Angle lookup tables */
cosTable: Float32Array
sinTable: Float32Array
probePositions: number
/** grid lookup tables */
gridx: Float32Array
gridy: Float32Array
gridz: Float32Array
expandedBox: Box3D
space: Tensor.Space
data: Tensor.Data
idData: Tensor.Data
}
async function createState(ctx: RuntimeContext, position: Required<PositionData>, maxRadius: number, props: MolecularSurfaceCalculationProps): Promise<MolSurfCalcState> {
// console.time('MolecularSurface')
// console.time('MolecularSurface createState')
const { resolution, probeRadius, probePositions } = props
const scaleFactor = 1 / resolution
const ngTorus = Math.max(5, 2 + Math.floor(probeRadius * scaleFactor))
const cellSize = Vec3.create(maxRadius, maxRadius, maxRadius)
Vec3.scale(cellSize, cellSize, 2)
const lookup3d = GridLookup3D(position, cellSize)
const neighbours = lookup3d.result
const box = lookup3d.boundary.box
const { indices } = position
const { indices, x: px, y: py, z: pz, radius } = position
const n = OrderedSet.size(indices)
const pad = maxRadius * 2 + resolution
const expandedBox = Box3D.expand(Box3D(), box, Vec3.create(pad, pad, pad));
const min = expandedBox.min
const [ minX, minY, minZ ] = expandedBox.min
const scaledBox = Box3D.scale(Box3D(), expandedBox, scaleFactor)
const dim = Box3D.size(Vec3(), scaledBox)
Vec3.ceil(dim, dim)
const [ dimX, dimY, dimZ ] = dim
const iu = dimZ, iv = dimY, iuv = iu * iv
const { cosTable, sinTable } = getAngleTables(probePositions)
const space = Tensor.Space(dim, [0, 1, 2], Float32Array)
......@@ -367,81 +339,28 @@ async function createState(ctx: RuntimeContext, position: Required<PositionData>
fillUniform(data, -1001.0)
fillUniform(idData, -1)
const gridx = fillGridDim(dim[0], min[0], resolution)
const gridy = fillGridDim(dim[1], min[1], resolution)
const gridz = fillGridDim(dim[2], min[2], resolution)
const updateChunk = Math.ceil(1000000 / (Math.pow(maxRadius, 3) * resolution))
return {
lastClip: -1,
neighbours: lookup3d.find(0, 0, 0, 0),
lookup3d,
position,
min,
updateChunk,
maxRadius,
n,
resolution,
scaleFactor,
probeRadius,
cosTable,
sinTable,
probePositions,
gridx,
gridy,
gridz,
expandedBox,
space,
data,
idData,
}
}
//
export const MolecularSurfaceCalculationParams = {
resolution: PD.Numeric(0.5, { min: 0.01, max: 10, step: 0.01 }),
probeRadius: PD.Numeric(1.4, { min: 0, max: 10, step: 0.1 }),
probePositions: PD.Numeric(30, { min: 12, max: 90, step: 1 }),
}
export const DefaultMolecularSurfaceCalculationProps = PD.getDefaultValues(MolecularSurfaceCalculationParams)
export type MolecularSurfaceCalculationProps = typeof DefaultMolecularSurfaceCalculationProps
export async function calcMolecularSurface(ctx: RuntimeContext, position: Required<PositionData>, maxRadius: number, props: MolecularSurfaceCalculationProps) {
// Field generation method adapted from AstexViewer (Mike Hartshorn) by Fred Ludlow.
// Other parts based heavily on NGL (Alexander Rose) EDT Surface class
console.time('MolecularSurface')
console.time('MolecularSurface createState')
const state = await createState(ctx, position, maxRadius, props)
console.timeEnd('MolecularSurface createState')
const gridx = fillGridDim(dimX, minX, resolution)
const gridy = fillGridDim(dimY, minY, resolution)
const gridz = fillGridDim(dimZ, minZ, resolution)
console.time('MolecularSurface projectPoints')
await projectPoints(ctx, state)
console.timeEnd('MolecularSurface projectPoints')
const updateChunk = Math.ceil(100000 / ((Math.pow(Math.pow(maxRadius, 3), 3) * scaleFactor)))
// console.timeEnd('MolecularSurface createState')
console.time('MolecularSurface projectTorii')
await projectTorii(ctx, state)
console.timeEnd('MolecularSurface projectTorii')
// console.time('MolecularSurface projectPoints')
await projectPoints()
// console.timeEnd('MolecularSurface projectPoints')
console.timeEnd('MolecularSurface')
// console.time('MolecularSurface projectTorii')
await projectTorii()
// console.timeEnd('MolecularSurface projectTorii')
// console.timeEnd('MolecularSurface')
const field = Tensor.create(state.space, state.data)
const idField = Tensor.create(state.space, state.idData)
const field = Tensor.create(space, data)
const idField = Tensor.create(space, idData)
const { resolution, expandedBox } = state
const transform = Mat4.identity()
Mat4.fromScaling(transform, Vec3.create(resolution, resolution, resolution))
Mat4.setTranslation(transform, expandedBox.min)
console.log({ field, idField, transform, state })
// console.log({ field, idField, transform, updateChunk })
return { field, idField, transform }
}
\ No newline at end of file
......@@ -27,7 +27,6 @@ export type MolecularSurfaceMeshParams = typeof MolecularSurfaceMeshParams
//
async function createMolecularSurfaceMesh(ctx: VisualContext, unit: Unit, structure: Structure, theme: Theme, props: MolecularSurfaceCalculationProps, mesh?: Mesh): Promise<Mesh> {
console.log(props)
const { transform, field, idField } = await computeUnitMolecularSurface(unit, props).runInContext(ctx.runtime)
const params = {
......
......@@ -11,13 +11,11 @@ import { PositionData, DensityData } from 'mol-math/geometry';
import { MolecularSurfaceCalculationProps, calcMolecularSurface } from 'mol-math/geometry/molecular-surface';
import { OrderedSet } from 'mol-data/int';
export function computeUnitMolecularSurface(unit: Unit, props: MolecularSurfaceCalculationProps) {
function getPositionDataAndMaxRadius(unit: Unit, props: MolecularSurfaceCalculationProps) {
const { position, radius } = getUnitConformationAndRadius(unit)
return Task.create('Molecular Surface', async ctx => {
const { indices } = position
const n = OrderedSet.size(indices)
const radii = new Float32Array(OrderedSet.max(indices))
const radii = new Float32Array(OrderedSet.end(indices))
let maxRadius = 0
for (let i = 0; i < n; ++i) {
......@@ -25,13 +23,15 @@ export function computeUnitMolecularSurface(unit: Unit, props: MolecularSurfaceC
const r = radius(j)
if (maxRadius < r) maxRadius = r
radii[j] = r + props.probeRadius
if (i % 100000 === 0 && ctx.shouldUpdate) {
await ctx.update({ message: 'calculating max radius', current: i, max: n })
}
return { position: { ...position, radius: radii }, maxRadius }
}
return await MolecularSurface(ctx, { ...position, radius: radii }, maxRadius, props);
export function computeUnitMolecularSurface(unit: Unit, props: MolecularSurfaceCalculationProps) {
const { position, maxRadius } = getPositionDataAndMaxRadius(unit, props)
return Task.create('Molecular Surface', async ctx => {
return await MolecularSurface(ctx, position, maxRadius, props);
});
}
......
......@@ -84,43 +84,59 @@ async function init() {
console.timeEnd('computeModelDSSP');
(models[0].properties as any).secondaryStructure = secondaryStructure
const structure = await getStructure(models[0])
const show = {
cartoon: false,
ballAndStick: true,
molecularSurface: true,
gaussianSurface: false,
}
const cartoonRepr = getCartoonRepr()
const ballAndStick = getBallAndStickRepr()
const ballAndStickRepr = getBallAndStickRepr()
const molecularSurfaceRepr = getMolecularSurfaceRepr()
const gaussianSurfaceRepr = getGaussianSurfaceRepr()
// cartoonRepr.setTheme({
// color: reprCtx.colorThemeRegistry.create('secondary-structure', { structure }),
// size: reprCtx.sizeThemeRegistry.create('uniform', { structure })
// })
// await cartoonRepr.createOrUpdate({ ...CartoonRepresentationProvider.defaultValues, quality: 'auto' }, structure).run()
if (show.cartoon) {
cartoonRepr.setTheme({
color: reprCtx.colorThemeRegistry.create('secondary-structure', { structure }),
size: reprCtx.sizeThemeRegistry.create('uniform', { structure })
})
await cartoonRepr.createOrUpdate({ ...CartoonRepresentationProvider.defaultValues, quality: 'auto' }, structure).run()
}
// ballAndStick.setTheme({
// color: reprCtx.colorThemeRegistry.create('secondary-structure', { structure }),
// size: reprCtx.sizeThemeRegistry.create('uniform', { structure })
// })
// await ballAndStick.createOrUpdate({ ...BallAndStickRepresentationProvider.defaultValues, quality: 'auto' }, structure).run()
if (show.ballAndStick) {
ballAndStickRepr.setTheme({
color: reprCtx.colorThemeRegistry.create('secondary-structure', { structure }),
size: reprCtx.sizeThemeRegistry.create('uniform', { structure })
})
await ballAndStickRepr.createOrUpdate({ ...BallAndStickRepresentationProvider.defaultValues, quality: 'auto' }, structure).run()
}
if (show.molecularSurface) {
molecularSurfaceRepr.setTheme({
color: reprCtx.colorThemeRegistry.create('secondary-structure', { structure }),
size: reprCtx.sizeThemeRegistry.create('physical', { structure })
})
console.time('molecular surface')
await molecularSurfaceRepr.createOrUpdate({ ...MolecularSurfaceRepresentationProvider.defaultValues, quality: 'custom', alpha: 1.0, flatShaded: true, doubleSided: true, resolution: 0.3 }, structure).run()
await molecularSurfaceRepr.createOrUpdate({ ...MolecularSurfaceRepresentationProvider.defaultValues, quality: 'custom', alpha: 0.5, flatShaded: true, doubleSided: true, resolution: 0.3 }, structure).run()
console.timeEnd('molecular surface')
}
if (show.gaussianSurface) {
gaussianSurfaceRepr.setTheme({
color: reprCtx.colorThemeRegistry.create('secondary-structure', { structure }),
size: reprCtx.sizeThemeRegistry.create('physical', { structure })
})
console.time('gaussian surface')
await gaussianSurfaceRepr.createOrUpdate({ ...GaussianSurfaceRepresentationProvider.defaultValues, quality: 'custom', alpha: 1.0, flatShaded: true, doubleSided: true, resolution: 0.3 }, structure).run()
console.timeEnd('gaussian surface')
}
// gaussianSurfaceRepr.setTheme({
// color: reprCtx.colorThemeRegistry.create('secondary-structure', { structure }),
// size: reprCtx.sizeThemeRegistry.create('physical', { structure })
// })
// console.time('gaussian surface')
// await gaussianSurfaceRepr.createOrUpdate({ ...GaussianSurfaceRepresentationProvider.defaultValues, quality: 'custom', alpha: 1.0, flatShaded: true, doubleSided: true, resolution: 0.3 }, structure).run()
// console.timeEnd('gaussian surface')
// canvas3d.add(cartoonRepr)
// canvas3d.add(ballAndStick)
canvas3d.add(molecularSurfaceRepr)
// canvas3d.add(gaussianSurfaceRepr)
if (show.cartoon) canvas3d.add(cartoonRepr)
if (show.ballAndStick) canvas3d.add(ballAndStickRepr)
if (show.molecularSurface) canvas3d.add(molecularSurfaceRepr)
if (show.gaussianSurface) canvas3d.add(gaussianSurfaceRepr)
canvas3d.resetCamera()
}
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment