MoleculeToRepresentation

Source

index.js
import macro from 'vtk.js/Sources/macro';
import vtkPolyData from 'vtk.js/Sources/Common/DataModel/PolyData';
import vtkDataArray from 'vtk.js/Sources/Common/Core/DataArray';
import * as vtkMath from 'vtk.js/Sources/Common/Core/Math';

import atomElem from 'vtk.js/Utilities/XMLConverter/chemistry/elements.json';

const { vtkErrorMacro, vtkDebugMacro } = macro;

// ----------------------------------------------------------------------------
// Globals
// ----------------------------------------------------------------------------

const ATOMS = {};
atomElem.atoms.forEach((a) => {
ATOMS[a.atomicNumber] = a;
});

// ----------------------------------------------------------------------------
// vtkMoleculeToRepresentation methods
// ----------------------------------------------------------------------------

function vtkMoleculeToRepresentation(publicAPI, model) {
const bondPositionData = [];
const bondScaleData = [];
const bondOrientationData = [];
const bondColorData = [];

// Set our className
model.classHierarchy.push('vtkMoleculeToRepresentation');

function addBond(
position,
orientation,
length,
color = [1.0, 1.0, 1.0],
radius = model.bondRadius
) {
bondScaleData.push(length);
bondScaleData.push(radius);

bondOrientationData.push(orientation[0]);
bondOrientationData.push(orientation[1]);
bondOrientationData.push(orientation[2]);

bondPositionData.push(position[0]);
bondPositionData.push(position[1]);
bondPositionData.push(position[2]);

for (let i = 0; i < color.length; ++i) {
bondColorData.push(color[i] * 255);
}
}

publicAPI.requestData = (inData, outData) => {
// input
const moleculedata = inData[0];

if (!moleculedata) {
vtkErrorMacro('Invalid or missing input');
return 1;
}

// output
const SphereData = vtkPolyData.newInstance();
const StickData = vtkPolyData.newInstance();

// Fetch from input molecule data
let numPts = 0;
let numBonds = 0;
let pointsArray = null;
let atomicNumber = null;
let bondIndex = null;
let bondOrder = null;

// Empty arrays
bondPositionData.length = 0;
bondScaleData.length = 0;
bondOrientationData.length = 0;
bondColorData.length = 0;

if (moleculedata.getAtoms()) {
if (moleculedata.getAtoms().coords !== undefined) {
if (moleculedata.getAtoms().coords['3d'] !== undefined) {
pointsArray = moleculedata.getAtoms().coords['3d'];
numPts = pointsArray.length / 3;
}
}
if (moleculedata.getAtoms().elements !== undefined) {
if (moleculedata.getAtoms().elements.number !== undefined) {
atomicNumber = moleculedata.getAtoms().elements.number;
}
}
}
if (moleculedata.getBonds()) {
if (moleculedata.getBonds().connections !== undefined) {
if (moleculedata.getBonds().connections.index !== undefined) {
bondIndex = moleculedata.getBonds().connections.index;
numBonds = bondIndex.length / 2;
}
}
if (moleculedata.getBonds().order !== undefined) {
bondOrder = moleculedata.getBonds().order;
}
}

const pointsData = [];
const scaleData = [];
const colorData = [];

const radiusArray = [];
const covalentArray = [];
const colorArray = [];

vtkDebugMacro('Checking for bonds with tolerance ', model.tolerance);

// go through each points and fill from elements.json
/* eslint-disable no-continue */
let ptsIdx = 0;
for (let i = 0; i < numPts; i++) {
// fetch from elements.json
if (atomicNumber) {
radiusArray.push(ATOMS[atomicNumber[i]][model.radiusType]);
covalentArray.push(ATOMS[atomicNumber[i]].radiusCovalent);
colorArray.push(ATOMS[atomicNumber[i]].elementColor[0]);
colorArray.push(ATOMS[atomicNumber[i]].elementColor[1]);
colorArray.push(ATOMS[atomicNumber[i]].elementColor[2]);
}

// skip atoms specified by hideElements
// model.hideHydrogen = false; // show hydrogen
if (model.hideElements.indexOf(ATOMS[atomicNumber[i]].id) !== -1) {
continue;
}

// points
ptsIdx = i * 3;
pointsData.push(pointsArray[ptsIdx]);
pointsData.push(pointsArray[ptsIdx + 1]);
pointsData.push(pointsArray[ptsIdx + 2]);

// radius
if (radiusArray.length > 0) {
scaleData.push(radiusArray[i] * model.atomicRadiusScaleFactor);
}

// colors
if (colorArray.length > 0) {
ptsIdx = i * 3;
colorData.push(colorArray[ptsIdx] * 255);
colorData.push(colorArray[ptsIdx + 1] * 255);
colorData.push(colorArray[ptsIdx + 2] * 255);
}
}

// if we don't have Bonds provided
// we fill up a bondIndex and a bondOrder
if (!bondIndex) {
bondIndex = [];
bondOrder = [];
// default bond display
/* eslint-disable no-continue */
for (let i = 0; i < numPts; i++) {
for (let j = i + 1; j < numPts; j++) {
const cutoff = covalentArray[i] + covalentArray[j] + model.tolerance;
const jPtsIdx = j * 3;
const iPtsIdx = i * 3;
const diff = [
pointsArray[jPtsIdx],
pointsArray[jPtsIdx + 1],
pointsArray[jPtsIdx + 2],
];
diff[0] -= pointsArray[iPtsIdx];
diff[1] -= pointsArray[iPtsIdx + 1];
diff[2] -= pointsArray[iPtsIdx + 2];

if (
Math.abs(diff[0]) > cutoff ||
Math.abs(diff[1]) > cutoff ||
Math.abs(diff[2]) > cutoff
) {
continue;
}

// Check radius and add bond if needed
const cutoffSq = cutoff * cutoff;
const diffsq =
diff[0] * diff[0] + diff[1] * diff[1] + diff[2] * diff[2];
if (diffsq < cutoffSq && diffsq > 0.1) {
// appendBond between i and j
bondIndex.push(i);
bondIndex.push(j);
bondOrder.push(1);
}
}
}
numBonds = bondIndex.length / 2;
}

// now we have the bonds, draw them
for (let index = 0; index < numBonds; index++) {
// appendBond between i and j
const i = bondIndex[index * 2];
const j = bondIndex[index * 2 + 1];

// Do not append if i or j belong to element to not display
if (
model.hideElements.indexOf(ATOMS[atomicNumber[i]].id) !== -1 ||
model.hideElements.indexOf(ATOMS[atomicNumber[j]].id) !== -1
) {
continue;
}

const jPtsIdx = j * 3;
const iPtsIdx = i * 3;
const diff = [
pointsArray[jPtsIdx],
pointsArray[jPtsIdx + 1],
pointsArray[jPtsIdx + 2],
];
diff[0] -= pointsArray[iPtsIdx];
diff[1] -= pointsArray[iPtsIdx + 1];
diff[2] -= pointsArray[iPtsIdx + 2];
const diffsq = diff[0] * diff[0] + diff[1] * diff[1] + diff[2] * diff[2];

let bondDelta = (2 + model.deltaBondFactor) * model.bondRadius; // distance between 2 bonds

// scale bonds if total distance from bonds is bigger than 2r*factor with r = min(r_i, r_j)
const r = Math.min(
radiusArray[i] * model.atomicRadiusScaleFactor,
radiusArray[j] * model.atomicRadiusScaleFactor
);
const t = (bondOrder[index] - 1) * bondDelta + 2 * model.bondRadius;
if (t > 2 * r * 0.6) {
model.bondRadius *= (2 * r * 0.6) / t;
// recompute bondDelta
bondDelta = (2 + model.deltaBondFactor) * model.bondRadius; // distance between 2 bonds
}

// Display multiple bond
// loop such as 0 11 22 if odd order / 00 11 22 33 if even order
// To make: 0 22 44 66 88 ... 11 33 55 77 ....
// because the offset has to be:
// (with bd= bondDelta. Note the minus is added just before creating bondPos)
// - odd order: 0 2bd/2 -2bd/2 4bd/2 -4bd/2 ...
// - even order: 1bd/2 -1bd/2 3bd/2 -3bd/2 ...
// Then, to transform loop to offset we have:
// - odd order: x * 2 <=> x * 2 + 1 - 1
// - even order: x * 2 + 1
// (with x the loop <=> floor(k/2))
const oddOrEven = bondOrder[index] % 2; // zero if even order / one if odd order
for (let k = oddOrEven; k < bondOrder[index] + oddOrEven; k++) {
// dist from center to bond depending of number of bond
let offset = ((Math.floor(k / 2) * 2 + 1 - oddOrEven) * bondDelta) / 2;

const vectUnitJI = [
diff[0] / Math.sqrt(diffsq),
diff[1] / Math.sqrt(diffsq),
diff[2] / Math.sqrt(diffsq),
];
const vectUnitJIperp = [0, 0, 0];
// Search perp to vectUnitJI: find axis != 0 to create vectUnitJIperp such as dot(vectUnitJIperp,vectUnitJI) = 0
for (let coord = 0; coord < 3; coord++) {
if (Math.abs(vectUnitJI[coord]) < 0.000001) {
continue;
}
vectUnitJIperp[coord] =
-(
vectUnitJI[(coord + 2) % 3] * vectUnitJI[(coord + 2) % 3] +
vectUnitJI[(coord + 1) % 3] * vectUnitJI[(coord + 1) % 3]
) / vectUnitJI[coord];
vectUnitJIperp[(coord + 1) % 3] = vectUnitJI[(coord + 1) % 3];
vectUnitJIperp[(coord + 2) % 3] = vectUnitJI[(coord + 2) % 3];
vtkMath.normalize(vectUnitJIperp);
break;
}

offset *= (-1) ** (k % 2);
/*
If atoms have a color associated, and if the atoms involved in the bond
are different species, then each bond will be represented by
two sticks, so that they can be colored with the same color as the atoms
involved in the bond.
*/
let bondPos;

if (
atomicNumber &&
atomicNumber[i] !== atomicNumber[j] &&
colorArray.length > 0
) {
const bondLength = Math.sqrt(diffsq) / 2.0;

bondPos = [
pointsArray[jPtsIdx] -
(bondLength * vectUnitJI[0]) / 2.0 +
offset * vectUnitJIperp[0],
pointsArray[jPtsIdx + 1] -
(bondLength * vectUnitJI[1]) / 2.0 +
offset * vectUnitJIperp[1],
pointsArray[jPtsIdx + 2] -
(bondLength * vectUnitJI[2]) / 2.0 +
offset * vectUnitJIperp[2],
];

addBond(
bondPos,
vectUnitJI,
bondLength,
colorArray.slice(jPtsIdx, jPtsIdx + 3)
);

bondPos = [
pointsArray[iPtsIdx] +
(bondLength * vectUnitJI[0]) / 2.0 +
offset * vectUnitJIperp[0],
pointsArray[iPtsIdx + 1] +
(bondLength * vectUnitJI[1]) / 2.0 +
offset * vectUnitJIperp[1],
pointsArray[iPtsIdx + 2] +
(bondLength * vectUnitJI[2]) / 2.0 +
offset * vectUnitJIperp[2],
];
addBond(
bondPos,
vectUnitJI,
bondLength,
colorArray.slice(iPtsIdx, iPtsIdx + 3)
);
} else {
const bondLength = Math.sqrt(diffsq);

bondPos = [
pointsArray[jPtsIdx] - diff[0] / 2.0 + offset * vectUnitJIperp[0],
pointsArray[jPtsIdx + 1] -
diff[1] / 2.0 +
offset * vectUnitJIperp[1],
pointsArray[jPtsIdx + 2] -
diff[2] / 2.0 +
offset * vectUnitJIperp[2],
];
if (colorArray.length > 0) {
addBond(
bondPos,
vectUnitJI,
bondLength,
colorArray.slice(iPtsIdx, iPtsIdx + 3)
);
} else {
addBond(bondPos, vectUnitJI, bondLength);
}
}
}
}

SphereData.getPoints().setData(pointsData, 3);

if (radiusArray) {
const scales = vtkDataArray.newInstance({
numberOfComponents: 1,
values: scaleData,
name: publicAPI.getSphereScaleArrayName(),
});
SphereData.getPointData().addArray(scales);
}

if (colorArray.length > 0) {
const colors = vtkDataArray.newInstance({
numberOfComponents: 3,
values: Uint8Array.from(colorData),
name: 'colors',
});
SphereData.getPointData().setScalars(colors);
}

StickData.getPoints().setData(bondPositionData, 3);

const stickScales = vtkDataArray.newInstance({
numberOfComponents: 2,
values: bondScaleData,
name: 'stickScales',
});
StickData.getPointData().addArray(stickScales);

const orientation = vtkDataArray.newInstance({
numberOfComponents: 3,
values: bondOrientationData,
name: 'orientation',
});
StickData.getPointData().addArray(orientation);

if (colorArray.length > 0) {
const bondColors = vtkDataArray.newInstance({
numberOfComponents: 3,
values: Uint8Array.from(bondColorData),
name: 'colors',
});
StickData.getPointData().setScalars(bondColors);
}

// Update output
outData[0] = SphereData;
outData[1] = StickData;

return 1;
};
}

// ----------------------------------------------------------------------------
// Object factory
// ----------------------------------------------------------------------------

const DEFAULT_VALUES = {
sphereScaleArrayName: 'radius',
tolerance: 0.45,
atomicRadiusScaleFactor: 0.3,
bondRadius: 0.075,
deltaBondFactor: 0.6,
radiusType: 'radiusVDW',
hideElements: '',
};

// ----------------------------------------------------------------------------

export function extend(publicAPI, model, initialValues = {}) {
Object.assign(model, DEFAULT_VALUES, initialValues);

// Build VTK API
macro.obj(publicAPI, model);
macro.setGet(publicAPI, model, [
'atomicRadiusScaleFactor',
'bondRadius',
'deltaBondFactor',
'hideElements',
'radiusType',
'sphereScaleArrayName',
'tolerance',
]);
macro.algo(publicAPI, model, 1, 2);
vtkMoleculeToRepresentation(publicAPI, model);
}

// ----------------------------------------------------------------------------

export const newInstance = macro.newInstance(
extend,
'vtkMoleculeToRepresentation'
);

// ----------------------------------------------------------------------------

export default { newInstance, extend };