chore: change inputs format

mljs · May 7, 2024 · b023cba · b023cba
1 parent 6121522
commit b023cba
Show file tree

Hide file tree

Showing 5 changed files with 98 additions and 100 deletions.
diff --git a/README.md b/README.md
@@ -16,49 +16,29 @@ Polynomial Regression.
 ```js
 import { PolynomialRegression2D } from 'ml-regression-polynomial-2d';
 
-const x = [
-  [0, 10],
-  [1, 11],
-  [2, 12],
-  [3, 13],
-  [4, 14],
-  [5, 15],
-  [6, 16],
-  [7, 17],
-  [8, 18],
-  [9, 19],
-  [10, 20],
-  [11, 21],
-  [12, 22],
-  [13, 23],
-  [14, 24],
-  [15, 25],
-  [16, 26],
-  [17, 27],
-  [18, 28],
-  [19, 29],
-  [20, 30],
-];
-const y = [
+const inputs = {
+  x: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
+  y: [
+    10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28,
+    29, 30,
+  ],
+};
+const z = [
   20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38,
   39, 40,
 ];
 const order = 2; // setup the maximum degree of the polynomial
 
-const regression = new PolynomialRegression2D(x, y, { order });
+const regression = new PolynomialRegression2D(inputs, z, { order });
 
 //prediction
-console.log(regression.predict([0.5, 1.5])); // Apply the model to some x tuple.
-console.log(
-  regression.predict([
-    [0.5, 1.5],
-    [1.5, 2.5],
-  ]),
-); // Apply the model to an array of x tuple.
+console.log(regression.predict({ x: 0.5, y: 1.5 })); // Apply the model to some x tuple.
+console.log(regression.predict({ x: [0.5, 1.5], y: [1.5, 2.5] })); // Apply the model to an array of points tuple.
 console.log(regression.coefficients); // Prints the coefficients in increasing order of power (from 0 to degree).
 console.log(regression.toString(3)); // Prints a human-readable version of the function.
 console.log(regression.toLaTeX());
-console.log(regression.score(x, y));
+console.log(regression.score); // { r, r2, chi2, rmsd } statistical scores
+console.log(regression.getScore(x, y)); // calculate the score for another database
 ```
 
 ## License

diff --git a/src/BaseRegression2D.ts b/src/BaseRegression2D.ts
@@ -1,4 +1,4 @@
-import { type NumberArray } from 'cheminfo-types';
+import { DataXY, PointXY, type NumberArray } from 'cheminfo-types';
 import { isAnyArray } from 'is-any-array';
 
 import checkArrayLength from './checkArrayLength';
@@ -16,24 +16,25 @@ export default class BaseRegression2D {
     }
   }
 
-  predict(x: NumberArray): number;
-  predict(x: NumberArray[]): number[];
-  predict(x: NumberArray | NumberArray[]) {
-    if (isOnePoint(x)) {
-      return this._predict(x);
-    } else if (isAnyArray(x[0])) {
-      const y = [];
-      for (const xVal of x) {
-        y.push(this._predict(xVal));
+  predict(inputs: PointXY): number;
+  predict(inputs: DataXY): NumberArray;
+  predict(inputs: PointXY | DataXY): number | NumberArray {
+    if (isOnePoint(inputs)) {
+      return this._predict(inputs);
+    } else if (isAnyArray(inputs.x)) {
+      const { x, y } = inputs;
+      const result = new Float64Array(x.length);
+      for (let i = 0; i < x.length; i++) {
+        result[i] = this._predict({ x: x[i], y: y[i] });
       }
-      return y;
+      return result;
     } else {
       throw new TypeError('x must be a number or array');
     }
   }
 
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
-  _predict(x: NumberArray): number {
+  _predict(x: PointXY): number {
     throw new Error('_predict must be implemented');
   }
 
@@ -57,14 +58,9 @@ export default class BaseRegression2D {
    * @param y - response variable
    * @return - Object with further statistics.
    */
-  score(x: NumberArray[], y: NumberArray): RegressionScore {
-    checkArrayLength(x, y);
-
-    const n = x.length;
-    const y2 = new Float64Array(n);
-    for (let i = 0; i < n; i++) {
-      y2[i] = this._predict(x[i]);
-    }
+  getScore(input: DataXY, z: NumberArray): RegressionScore {
+    checkArrayLength(input, z);
+    const y2 = this.predict(input);
 
     let xSum = 0;
     let ySum = 0;
@@ -73,16 +69,17 @@ export default class BaseRegression2D {
     let xSquared = 0;
     let ySquared = 0;
     let xY = 0;
+    const n = z.length;
     for (let i = 0; i < n; i++) {
       xSum += y2[i];
-      ySum += y[i];
+      ySum += z[i];
       xSquared += y2[i] * y2[i];
-      ySquared += y[i] * y[i];
-      xY += y2[i] * y[i];
-      if (y[i] !== 0) {
-        chi2 += ((y[i] - y2[i]) * (y[i] - y2[i])) / y[i];
+      ySquared += z[i] * z[i];
+      xY += y2[i] * z[i];
+      if (z[i] !== 0) {
+        chi2 += ((z[i] - y2[i]) * (z[i] - y2[i])) / z[i];
       }
-      rmsd += (y[i] - y2[i]) * (y[i] - y2[i]);
+      rmsd += (z[i] - y2[i]) * (z[i] - y2[i]);
     }
 
     const r =
@@ -98,7 +95,7 @@ export default class BaseRegression2D {
   }
 }
 
-function isOnePoint(x: NumberArray | NumberArray[]): x is NumberArray {
-  return !isAnyArray(x[0]);
+function isOnePoint(input: PointXY | DataXY): input is PointXY {
+  return !isAnyArray(input.x);
 }
 export { checkArrayLength };
diff --git a/src/PolynomialRegression2D.ts b/src/PolynomialRegression2D.ts
@@ -1,4 +1,4 @@
-import { type NumberArray } from 'cheminfo-types';
+import { DataXY, PointXY, type NumberArray } from 'cheminfo-types';
 import { Matrix, SVD } from 'ml-matrix';
 import { maybeToPrecision } from 'ml-regression-base';
 
@@ -12,11 +12,18 @@ export interface PolynomialRegression2DOptions {
   order?: number;
 }
 
+interface Score {
+  r: number;
+  r2: number;
+  chi2: number;
+  rmsd: number;
+}
 // Implements the Kernel ridge regression algorithm.
 // http://www.ics.uci.edu/~welling/classnotes/papers_class/Kernel-Ridge.pdf
 export class PolynomialRegression2D extends BaseRegression2D {
   order: number;
   coefficients: Matrix;
+  score: Score;
   /**
    * Constructor for the 2D polynomial fitting
    *
@@ -25,7 +32,7 @@ export class PolynomialRegression2D extends BaseRegression2D {
    * @constructor
    */
   constructor(
-    inputs: NumberArray[],
+    inputs: DataXY,
     outputs: NumberArray,
     options: PolynomialRegression2DOptions = {},
   ) {
@@ -36,29 +43,30 @@ export class PolynomialRegression2D extends BaseRegression2D {
       this.coefficients = Matrix.columnVector(outputs.coefficients);
       // @ts-expect-error internal use only
       this.order = outputs.order;
+      // @ts-expect-error internal use only
+      this.score = outputs.score;
     } else {
       checkArrayLength(inputs, outputs);
       const { order = 2 } = options;
       this.order = order;
       this.coefficients = train(inputs, outputs, order);
+      this.score = this.getScore(inputs, outputs);
     }
   }
 
-  _predict(newInputs: NumberArray) {
-    const x1 = newInputs[0];
-    const x2 = newInputs[1];
+  _predict(newInputs: PointXY) {
+    const { x, y } = newInputs;
 
-    let y = 0;
+    let z = 0;
     let column = 0;
-
     for (let i = 0; i <= this.order; i++) {
       for (let j = 0; j <= this.order - i; j++) {
-        y += x1 ** i * x2 ** j * this.coefficients.get(column, 0);
+        z += x ** i * y ** j * this.coefficients.get(column, 0);
         column++;
       }
     }
 
-    return y;
+    return z;
   }
 
   toString(precision: number) {
@@ -116,6 +124,7 @@ export class PolynomialRegression2D extends BaseRegression2D {
     return {
       name: 'polyfit2D',
       order: this.order,
+      score: this.score,
       coefficients: this.coefficients,
     };
   }
@@ -152,14 +161,13 @@ function powColVector(x: Matrix, power: number) {
  * @param x - A matrix with n rows and 2 columns.
  * @param y - A vector of the prediction values.
  */
-function train(
-  x: NumberArray[] | Matrix,
-  y: NumberArray | Matrix,
-  order: number,
-) {
-  if (!Matrix.isMatrix(x)) x = new Matrix(x);
+function train(input: DataXY, y: NumberArray | Matrix, order: number) {
   if (!Matrix.isMatrix(y)) y = Matrix.columnVector(y);
 
+  const x = new Matrix(y.rows, 2);
+  x.setColumn(0, input.x);
+  x.setColumn(1, input.y);
+
   if (y.rows !== x.rows) {
     y = y.transpose();
   }

diff --git a/src/__tests__/PolynomialRegression2D.test.ts b/src/__tests__/PolynomialRegression2D.test.ts
@@ -1,20 +1,21 @@
-import { NumberArray } from 'cheminfo-types';
 import { describe, it, expect } from 'vitest';
 
 import { PolynomialRegression2D } from '../PolynomialRegression2D';
 
 describe('2D polinomial fit', () => {
-  const X = new Array(21);
+  const x = new Array(21);
   const y = new Array(21);
+  const z = new Array(21);
   for (let i = 0; i < 21; ++i) {
-    X[i] = [i, i + 10];
-    y[i] = i + 20;
+    x[i] = i;
+    y[i] = i + 10;
+    z[i] = i + 20;
   }
 
-  const pf = new PolynomialRegression2D(X, y, {
+  const pf = new PolynomialRegression2D({ x, y }, z, {
     order: 2,
   });
-
+  console.log({ x, y });
   it('Training coefficients', () => {
     const estimatedCoefficients = [
       1.5587e1, 3.8873e-1, 5.2582e-3, 4.8498e-1, 2.1127e-3, -7.3709e-3,
@@ -28,13 +29,15 @@ describe('2D polinomial fit', () => {
   });
 
   it('Prediction', () => {
-    const test: NumberArray[] = new Array(11);
+    const xTest = new Float64Array(11);
+    const yTest = new Float64Array(11);
     let val = 0.5;
     for (let i = 0; i < 11; ++i) {
-      test[i] = [val, val + 10];
+      xTest[i] = val;
+      yTest[i] = val + 10;
       val++;
     }
-    const y = pf.predict(test);
+    const y = pf.predict({ x: xTest, y: yTest });
 
     let j = 0;
     for (let i = 20.5; i < 30.5; i++, j++) {
@@ -50,29 +53,35 @@ describe('2D polinomial fit', () => {
 
     const len = 21;
 
-    const X = new Array(len);
+    const x = new Array(len);
     let val = 5.0;
     const y = new Array(len);
+    const z = new Array(len);
     for (let i = 0; i < len; ++i, val += 0.5) {
-      X[i] = [val, val];
-      y[i] = val * val + val * val;
+      x[i] = val;
+      y[i] = val;
+      z[i] = val * val + val * val;
     }
 
-    const polynomialRegression2D = new PolynomialRegression2D(X, y, {
+    const polynomialRegression2D = new PolynomialRegression2D({ x, y }, z, {
       order: 2,
     });
 
-    const test = 10;
+    const length = 10;
     let x1 = -4.75;
     let x2 = 4.75;
-    const X1: NumberArray[] = new Array(test);
-    for (let i = 0; i < test; ++i) {
-      X1[i] = [x1, x2];
+    const testData = {
+      x: Float64Array.from({ length }),
+      y: Float64Array.from({ length }),
+    };
+    for (let i = 0; i < length; ++i) {
+      testData.x[i] = x1;
+      testData.y[i] = x2;
       x1++;
       x2--;
     }
 
-    const predict = polynomialRegression2D.predict(X1);
+    const predict = polynomialRegression2D.predict(testData);
     for (let i = 0; i < testValues.length; ++i) {
       expect(predict[i]).toBeCloseTo(testValues[i], 1e-2);
     }

diff --git a/src/checkArrayLength.ts b/src/checkArrayLength.ts
@@ -1,22 +1,26 @@
-import { NumberArray } from 'cheminfo-types';
+import { DataXY, NumberArray } from 'cheminfo-types';
 import { isAnyArray } from 'is-any-array';
 /**
  * Check that x and y are arrays with the same length.
  * @param x - first array
  * @param y - second array
  * @throws if x or y are not the same length, or if they are not arrays
  */
-export default function checkArrayLength(x: NumberArray[], y: NumberArray) {
-  if ((!isAnyArray(x) && !isAnyArray(x[0])) || !isAnyArray(y)) {
-    throw new TypeError('x and y must be arrays');
+export default function checkArrayLength(input: DataXY, output: NumberArray) {
+  if (!isAnyArray(input.x) || !isAnyArray(input.y) || !isAnyArray(output)) {
+    throw new TypeError('x, y and outputs must be arrays');
   }
-  if (x[0].length < 2) {
+  if (input.x.length < 2) {
     throw new RangeError(
       'explanatory variable should be two element per point',
     );
   }
 
-  if (x.length !== y.length) {
+  if (input.x.length !== input.y.length) {
     throw new RangeError('x and y data must have the same length');
   }
+
+  if (input.x.length !== output.length) {
+    throw new RangeError('input and outputs must have the same length');
+  }
 }