improve result by use the noise level after first iteration

package.json

@@ -44,6 +44,7 @@
     "rollup": "^3.22.0"
   },
   "dependencies": {
-    "cuthill-mckee": "^1.0.0"
+    "cuthill-mckee": "^1.0.0",
+    "ml-spectra-processing": "^14.3.0"
   }
 }

src/index.js

@@ -1,6 +1,40 @@
+import { xMultiply, xNoiseSanPlot } from 'ml-spectra-processing';
+
 import cholesky from './choleskySolver';
 import { updateSystem, getDeltaMatrix, getCloseIndex } from './utils';
 
+function getControlPoints(x, y, options = {}) {
+  const { length } = x;
+  let { controlPoints = Int8Array.from({ length }).fill(0) } = options;
+  const { zones = [], weights = Float64Array.from({ length }).fill(1) } =
+    options;
+
+  if (x.length !== y.length) {
+    throw new RangeError('Y should match the length with X');
+  } else if (controlPoints.length !== x.length) {
+    throw new RangeError('controlPoints should match the length with X');
+  } else if (weights.length !== x.length) {
+    throw new RangeError('weights should match the length with X');
+  }
+
+  zones.forEach((range) => {
+    let indexFrom = getCloseIndex(x, range.from);
+    let indexTo = getCloseIndex(x, range.to);
+    if (indexFrom > indexTo) [indexFrom, indexTo] = [indexTo, indexFrom];
+    for (let i = indexFrom; i < indexTo; i++) {
+      controlPoints[i] = 1;
+    }
+  });
+
+  return {
+    weights:
+      'controlPoints' in options || zones.length > 0
+        ? xMultiply(weights, controlPoints)
+        : weights,
+    controlPoints,
+  };
+}
+
 /**
  * Fit the baseline drift by iteratively changing weights of sum square error between the fitted baseline and original signals,
  * for further information about the parameters you can get the [paper of airPLS](https://github.com/zmzhang/airPLS/blob/master/airPLS_manuscript.pdf)
@@ -11,45 +45,27 @@ import { updateSystem, getDeltaMatrix, getCloseIndex } from './utils';
  * @param {number} [options.tolerance = 0.001] - Factor of the sum of absolute value of original data, to compute stop criterion
  * @param {Array<number>} [options.weights = [1,1,...]] - Initial weights vector, default each point has the same weight
  * @param {number} [options.lambda = 100] - Factor of weights matrix in -> [I + lambda D'D]z = x
- * @param {Array<number>} [options.controlPoints = []] - Array of x axis values to force that baseline cross those points.
- * @param {Array<number>} [options.baseLineZones = []] - Array of x axis values (as from - to), to force that baseline cross those zones.
+ * @param {Array<number>} [options.controlPoints = []] - Array of 0|1 to force the baseline cross those points.
+ * @param {Array<number>} [options.zones = []] - Array of x axis values (as from - to), to force that baseline cross those zones.
  * @returns {{corrected: Array<number>, error: number, iteration: number, baseline: Array<number>}}
  */
 export default function airPLS(x, y, options = {}) {
-  let {
-    maxIterations = 100,
-    lambda = 100,
-    tolerance = 0.001,
-    weights = new Array(y.length).fill(1),
-    controlPoints = [],
-    baseLineZones = [],
-  } = options;
-
-  if (controlPoints.length > 0) {
-    controlPoints.forEach((e, i, arr) => (arr[i] = getCloseIndex(x, e)));
-  }
-  if (baseLineZones.length > 0) {
-    baseLineZones.forEach((range) => {
-      let indexFrom = getCloseIndex(x, range.from);
-      let indexTo = getCloseIndex(x, range.to);
-      if (indexFrom > indexTo) [indexFrom, indexTo] = [indexTo, indexFrom];
-      for (let i = indexFrom; i < indexTo; i++) {
-        controlPoints.push(i);
-      }
-    });
-  }
+  const { weights, controlPoints } = getControlPoints(x, y, options);
+  let { maxIterations = 100, lambda = 10, tolerance = 0.001 } = options;
 
   let baseline, iteration;
-  let nbPoints = y.length;
-  let l = nbPoints - 1;
   let sumNegDifferences = Number.MAX_SAFE_INTEGER;
   let stopCriterion = tolerance * y.reduce((sum, e) => Math.abs(e) + sum, 0);
 
+  const { length } = y;
   let { lowerTriangularNonZeros, permutationEncodedArray } = getDeltaMatrix(
-    nbPoints,
+    length,
     lambda,
   );
 
+  let threshold = 1;
+  const l = length - 1;
+  let prevNegSum = Number.MAX_SAFE_INTEGER;
   for (
     iteration = 0;
     iteration < maxIterations && Math.abs(sumNegDifferences) > stopCriterion;
@@ -61,29 +77,40 @@ export default function airPLS(x, y, options = {}) {
       weights,
     );
 
-    let cho = cholesky(leftHandSide, nbPoints, permutationEncodedArray);
+    let cho = cholesky(leftHandSide, length, permutationEncodedArray);
 
     baseline = cho(rightHandSide);
 
     sumNegDifferences = 0;
-
     let difference = y.map(calculateError);
+    if (iteration === 1) {
+      const { positive } = xNoiseSanPlot(difference);
+      threshold = positive;
+    } else {
+      const absChange = Math.abs(prevNegSum / sumNegDifferences);
+      if (absChange < 1.01 && absChange > 0.99) {
+        break;
+      }
+    }
 
+    prevNegSum = sumNegDifferences + 0;
     let maxNegativeDiff = -1 * Number.MAX_SAFE_INTEGER;
     for (let i = 1; i < l; i++) {
       let diff = difference[i];
-      if (diff >= 0) {
+      if (controlPoints[i] < 1 && Math.abs(diff) > threshold) {
         weights[i] = 0;
       } else {
-        weights[i] = Math.exp((iteration * diff) / sumNegDifferences);
-        if (maxNegativeDiff < diff) maxNegativeDiff = diff;
+        const factor = diff > 0 ? -1 : 1;
+        weights[i] = Math.exp(
+          (factor * (iteration * diff)) / Math.abs(sumNegDifferences),
+        );
       }
+
+      if (diff < 0 && maxNegativeDiff < diff) maxNegativeDiff = diff;
     }
 
-    let value = Math.exp((iteration * maxNegativeDiff) / sumNegDifferences);
-    weights[0] = value;
-    weights[l] = value;
-    controlPoints.forEach((i) => (weights[i] = value));
+    weights[0] = 1;
+    weights[l] = 1;
   }
 
   return {