feat: add reference shift for predictors

package.json

@@ -49,6 +49,7 @@
   "dependencies": {
     "common-spectrum": "2.7.0",
     "ensure-string": "^1.2.0",
+    "md5": "^2.3.0",
     "ml-gsd": "^12.1.6",
     "openchemlib-utils": "^5.21.1",
     "vamas": "^0.3.0"

src/utilities/prediction/__tests__/predictPolymer.test.js

@@ -1,10 +1,10 @@
-import { Molecule } from 'openchemlib'
+import { Molecule } from 'openchemlib';
 
-import { predictPolymer } from '../predictPolymer'
+import { predictPolymer } from '../predictPolymer';
 
 test('predictPolymer', async () => {
-  const idCode = "gJQDHODnJRmT@@" // [R2]OCC[R1]
-  const molecule = Molecule.fromIDCode(idCode)
-  const polymer = await predictPolymer(molecule)
-  console.log(polymer)
-})
+  const idCode = 'gJQDHODnJRmT@@'; // [R2]OCC[R1]
+  const molecule = Molecule.fromIDCode(idCode);
+  const polymer = await predictPolymer(molecule);
+  console.log(polymer);
+});

src/utilities/prediction/__tests__/predictUsingHoseCodes.test.js

@@ -5,12 +5,25 @@ import { predictUsingHoseCodes } from '../predictUsingHoseCodes.js';
 test('predictUsingHoseCodes', async () => {
   const molecule = OCL.Molecule.fromSmiles('CCO');
 
-  const result = await predictUsingHoseCodes(molecule);
-  expect(Object.keys(result)).toHaveLength(4);
-  expect(Object.keys(result)).toStrictEqual([
+  const resultNoRef = await predictUsingHoseCodes(molecule, {
+    energyReference: 'none',
+  });
+  expect(Object.keys(resultNoRef)).toHaveLength(4);
+  expect(Object.keys(resultNoRef)).toStrictEqual([
     'grouped',
     'spectrum',
     'annotations',
     'peaks',
   ]);
+
+  expect(resultNoRef.peaks[0]).toStrictEqual({
+    transition: 'C1s',
+    x: 290.6834,
+    y: 1,
+  });
+
+  const resultSolidRef = await predictUsingHoseCodes(molecule, {
+    energyReference: 'solid',
+  });
+  expect(resultSolidRef.peaks[0].x).toBeCloseTo(285.30823, 4);
 });

src/utilities/prediction/predictPolymer.js

@@ -1,22 +1,23 @@
-import { createPolymer, } from 'openchemlib-utils';
+import { createPolymer } from 'openchemlib-utils';
 
 import { predictUsingHoseCodes } from './predictUsingHoseCodes.js';
 
-
 export async function predictPolymer(monomer, options = {}) {
-  const { Molecule } = monomer.getOCL()
+  const { Molecule } = monomer.getOCL();
   const polymer = createPolymer(monomer, {
     count: 5, // we create 5 monomers and we use the middle one
     markMonomer: true,
     ...getAlphaGamma(Molecule),
   });
 
-  const prediction = await predictUsingHoseCodes(polymer, { atomMapNo: 3, ...options });
+  const prediction = await predictUsingHoseCodes(polymer, {
+    atomMapNo: 3,
+    ...options,
+  });
 
-  return prediction
+  return prediction;
 }
 
-
 function getAlphaGamma(Molecule) {
   const r1AtomicNo = Molecule.getAtomicNoFromLabel(
     'R1',
@@ -32,5 +33,3 @@ function getAlphaGamma(Molecule) {
   gamma.setAtomicNo(0, r2AtomicNo);
   return { alpha, gamma };
 }
-
-

src/utilities/prediction/predictUsingAI.js

@@ -1,8 +1,8 @@
+import md5 from 'md5';
 import { generateSpectrum } from 'spectrum-generator';
 
 const cache = {};
 
-
 /**
  *
  * @param {import('openchemlib').Molecule} [molecule]
@@ -11,16 +11,15 @@ const cache = {};
  * @param {import('spectrum-generator').GenerateSpectrumOptions} [options.spectrum]
  */
 export async function predictUsingAI(molecule, options = {}) {
-  const {
-    spectrum: spectrumOptions,
-  } = options;
+  const { spectrum: spectrumOptions, energyReference = 'solid' } = options;
 
   const moleculeWithH = molecule.getCompactCopy();
   moleculeWithH.addImplicitHydrogens();
 
   const molfile = moleculeWithH.toMolfile();
+  const cacheKey = md5(molfile + energyReference);
 
-  if (!cache[molfile]) {
+  if (!cache[cacheKey]) {
     const response = await fetch(
       'https://xps-service.cheminfo.org/v1/predict_binding_energies',
       {
@@ -32,13 +31,14 @@ export async function predictUsingAI(molecule, options = {}) {
           molFile: molfile,
           method: 'GFN2xTB',
           fmax: 0.01,
+          energyReference,
         }),
       },
     );
-    cache[molfile] = await response.json();
+    cache[cacheKey] = await response.json();
   }
 
-  const predictions = cache[molfile].predictions;
+  const predictions = cache[cacheKey].predictions;
 
   const peaks = [];
   const xyz = [predictions.length, 'XTB optimized geometry'];

src/utilities/prediction/predictUsingHoseCodes.js

@@ -5,18 +5,18 @@ let data;
 
 const xShiftFcts = {
   solid: {
-    c1s: (x) => x,
-    o1s: (x) => x,
-  }
-}
+    C1s: (x) => -2501.002792622873 + 18.12807741 * x - 0.02938832 * x ** 2,
+    O1s: (x) => -14914.578039681715 + 5.67166123e1 * x - 5.20506217e-2 * x ** 2,
+  },
+};
 
 /**
  *
  * @param {import('openchemlib').Molecule} [molecule]
  * @param {object} options
  * @param {string} [options.statsKey='gw']
  * @param {number|undefined} [options.atomMapNo] // allows to filter for specific atoms
- * @param {string} [options.energyReference='solid']
+ * @param {'solid'|'none'} [options.energyReference='solid']
  * @param {import('spectrum-generator').GenerateSpectrumOptions} [options.spectrum]
  */
 export async function predictUsingHoseCodes(molecule, options = {}) {
@@ -28,8 +28,6 @@ export async function predictUsingHoseCodes(molecule, options = {}) {
   } = options;
   await ensureSpheres();
 
-  const xShiftFct = undefined;
-
   const spheres = data.spheres[statsKey];
   const diaIDs = getDiastereotopicAtomIDsAndH(molecule);
   const annotations = [];
@@ -45,9 +43,12 @@ export async function predictUsingHoseCodes(molecule, options = {}) {
       atomMapNo: molecule.getAtomMapNo(index),
       _highlight: [diaIDs[index].oclID],
     }))
-    .filter((value) => value.atomLabel && value.atomLabel !== 'H' && value.atomLabel !== '?')
+    .filter(
+      (value) =>
+        value.atomLabel && value.atomLabel !== 'H' && value.atomLabel !== '?',
+    )
     .filter((value) => value.hoses)
-    .filter((value) => (!atomMapNo) || (value.atomMapNo === atomMapNo));
+    .filter((value) => !atomMapNo || value.atomMapNo === atomMapNo);
 
   for (const hoseCode of values) {
     for (let i = hoseCode.hoses.length - 1; i >= 0; i--) {
@@ -66,10 +67,15 @@ export async function predictUsingHoseCodes(molecule, options = {}) {
 
   const peaks = values
     .filter((value) => value.prediction)
-    .map((value) => ({ x: value.prediction.boxplot.median, y: 1 }));
+    .map((value) => ({
+      x: value.prediction.boxplot.median,
+      y: 1,
+      transition: `${value.atomLabel}1s`,
+    }));
 
-  if (xShiftFct) {
-    for (const peak of peaks) {
+  for (const peak of peaks) {
+    const xShiftFct = xShiftFcts[energyReference]?.[peak.transition];
+    if (xShiftFct) {
       peak.x = xShiftFct(peak.x);
     }
   }
@@ -81,7 +87,6 @@ export async function predictUsingHoseCodes(molecule, options = {}) {
   // during this operation we should take care to combine _highlights
   const uniqueValuesObject = {};
   for (const value of values) {
-    console.log(value)
     const key = `${value.prediction.idCode}`;
     if (uniqueValuesObject[key]) {
       uniqueValuesObject[key].atomNumbers.push(value.atomNumber);
@@ -90,7 +95,11 @@ export async function predictUsingHoseCodes(molecule, options = {}) {
         uniqueValuesObject[key]._highlight.push(highlight);
       }
     } else {
-      const { atomNumber, ...uniqueValueObject } = { ...value, atomNumbers: [value.atomNumber] }
+      // eslint-disable-next-line no-unused-vars
+      const { atomNumber, ...uniqueValueObject } = {
+        ...value,
+        atomNumbers: [value.atomNumber],
+      };
       uniqueValuesObject[key] = uniqueValueObject;
     }
   }