refactor: change exposed methods and accepted types

README.md

@@ -15,15 +15,21 @@ Compares two spectra using a tree similarity.
 ```js
 import { createTree, treeSimilarity } from 'ml-tree-similarity';
 
-const a = [[1, 2, 3, 4, 5, 6, 7], [0.3, 0.7, 4, 0.3, 0.2, 5, 0.3]];
-const b = [[1, 2, 3, 4, 5, 6, 7], [0.3, 4, 0.7, 0.3, 5, 0.2, 0.3]];
+const a = {
+  x: [1, 2, 3, 4, 5, 6, 7],
+  y: [0.3, 0.7, 4, 0.3, 0.2, 5, 0.3],
+};
+const b = {
+  x: [1, 2, 3, 4, 5, 6, 7],
+  y: [0.3, 4, 0.7, 0.3, 5, 0.2, 0.3],
+};
 
 // create a tree
 const options = { from: 1, to: 7 };
-const A = createTree(a, options);
+const aTree = createTree(a, options);
+const bTree = createTree(b, options);
 
-// a pre-calculated tree is also a valid input
-const ans = treeSimilarity(A, b, options);
+const ans = treeSimilarity(aTree, bTree, options);
 ```
 
 ## [API Documentation](https://mljs.github.io/tree-similarity/)

package.json

@@ -40,6 +40,7 @@
   },
   "dependencies": {
     "binary-search": "^1.3.6",
+    "cheminfo-types": "^1.7.2",
     "num-sort": "^3.0.0"
   }
 }

src/__tests__/createTree.test.js

@@ -6,7 +6,7 @@ describe('simple trees', () => {
   it('two peaks, same height', () => {
     let x = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
     let y = [0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0];
-    let tree = createTree([x, y]);
+    let tree = createTree({ x, y });
 
     expect(tree.center).toBe(5);
     expect(tree.sum).toBe(2);
@@ -34,7 +34,7 @@ describe('simple trees', () => {
     y[20] = 20;
     y[80] = 20;
 
-    let tree = createTree([x, y]);
+    let tree = createTree({ x, y });
 
     expect(tree.center).toBe(50);
     expect(tree.sum).toBe(40);

src/__tests__/tree.test.js

@@ -1,28 +1,21 @@
 import { describe, it, expect } from 'vitest';
 
-import { treeSimilarity, getFunction } from '..';
+import { createTree, treeSimilarity } from '../index';
 
-let a = [
-  [1, 2, 3, 4, 5, 6, 7],
-  [0.3, 0.7, 4, 0.3, 0.2, 5, 0.3],
-];
-let b = [
-  [1, 2, 3, 4, 5, 6, 7],
-  [0.3, 4, 0.7, 0.3, 5, 0.2, 0.3],
-];
+const a = {
+  x: [1, 2, 3, 4, 5, 6, 7],
+  y: [0.3, 0.7, 4, 0.3, 0.2, 5, 0.3],
+};
+const b = {
+  x: [1, 2, 3, 4, 5, 6, 7],
+  y: [0.3, 4, 0.7, 0.3, 5, 0.2, 0.3],
+};
 
 describe('Tree similarity', () => {
   it('should work with two arrays', () => {
-    expect(treeSimilarity(a, b)).toBeCloseTo(0.653354, 4);
-  });
-
-  it('should currify the options', () => {
-    let options = {
-      alpha: 0.4,
-      beta: 0.5,
-      gamma: 0.001,
-    };
-    let func = getFunction(options);
-    expect(func(a, b)).toBe(treeSimilarity(a, b, options));
+    expect(treeSimilarity(createTree(a), createTree(b))).toBeCloseTo(
+      0.653354,
+      4,
+    );
   });
 });

src/__tests__/getSimilarity.test.js → src/__tests__/treeSimilarity.test.js

@@ -1,15 +1,15 @@
 import { describe, it, expect } from 'vitest';
 
 import { createTree } from '../createTree';
-import { getSimilarity } from '../getSimilarity';
+import { treeSimilarity } from '../treeSimilarity';
 
 describe('simple trees', () => {
   it('same tree', () => {
     let x = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
     let y = [0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0];
-    let tree1 = createTree([x, y]);
-    let tree2 = createTree([x, y]);
+    let tree1 = createTree({ x, y });
+    let tree2 = createTree({ x, y });
 
-    expect(getSimilarity(tree1, tree2, { beta: 1 })).toBe(1);
+    expect(treeSimilarity(tree1, tree2, { beta: 1 })).toBe(1);
   });
 });

src/createTree.js

@@ -1,53 +1,50 @@
 import binarySearch from 'binary-search';
 import { numberSortAscending } from 'num-sort';
 
+/**
+ * @typedef {import("../types").Tree} Tree
+ * @typedef {import("../types").CreateTreeOptions} CreateTreeOptions
+ * @typedef {import("../types").Spectrum} Spectrum
+ */
+
 /**
  * Function that creates the tree
- * @param {Array<Array<number>>} spectrum
- * @param {object} [options]
- * @return {Tree|null}
- * left and right have the same structure than the parent,
- * or are null if they are leaves
+ * @param {Spectrum} spectrum
+ * @param {CreateTreeOptions} [options]
+ * @return { Tree | null }
  */
 export function createTree(spectrum, options = {}) {
-  const X = spectrum[0];
+  const { x, y } = spectrum;
   const {
     minWindow = 0.16,
     threshold = 0.01,
-    from = X[0],
-    to = X[X.length - 1],
+    from = x[0],
+    to = x[x.length - 1],
   } = options;
 
-  return mainCreateTree(
-    spectrum[0],
-    spectrum[1],
-    from,
-    to,
-    minWindow,
-    threshold,
-  );
+  return mainCreateTree(x, y, from, to, minWindow, threshold);
 }
 
-function mainCreateTree(X, Y, from, to, minWindow, threshold) {
+function mainCreateTree(x, y, from, to, minWindow, threshold) {
   if (to - from < minWindow) {
     return null;
   }
 
   // search first point
-  let start = binarySearch(X, from, numberSortAscending);
+  let start = binarySearch(x, from, numberSortAscending);
   if (start < 0) {
     start = ~start;
   }
 
   // stop at last point
   let sum = 0;
   let center = 0;
-  for (let i = start; i < X.length; i++) {
-    if (X[i] >= to) {
+  for (let i = start; i < x.length; i++) {
+    if (x[i] >= to) {
       break;
     }
-    sum += Y[i];
-    center += X[i] * Y[i];
+    sum += y[i];
+    center += x[i] * y[i];
   }
 
   if (sum < threshold) {
@@ -59,24 +56,15 @@ function mainCreateTree(X, Y, from, to, minWindow, threshold) {
     return null;
   }
   if (center - from < minWindow / 4) {
-    return mainCreateTree(X, Y, center, to, minWindow, threshold);
+    return mainCreateTree(x, y, center, to, minWindow, threshold);
   } else if (to - center < minWindow / 4) {
-    return mainCreateTree(X, Y, from, center, minWindow, threshold);
+    return mainCreateTree(x, y, from, center, minWindow, threshold);
   } else {
-    return new Tree(
+    return {
       sum,
       center,
-      mainCreateTree(X, Y, from, center, minWindow, threshold),
-      mainCreateTree(X, Y, center, to, minWindow, threshold),
-    );
-  }
-}
-
-class Tree {
-  constructor(sum, center, left, right) {
-    this.sum = sum;
-    this.center = center;
-    this.left = left;
-    this.right = right;
+      left: mainCreateTree(x, y, from, center, minWindow, threshold),
+      right: mainCreateTree(x, y, center, to, minWindow, threshold),
+    };
   }
 }

src/index.js

@@ -1,11 +1,2 @@
-import { getSimilarity } from './getSimilarity';
-
+export { treeSimilarity } from './treeSimilarity';
 export { createTree } from './createTree';
-
-export function treeSimilarity(A, B, options = {}) {
-  return getSimilarity(A, B, options);
-}
-
-export function getFunction(options = {}) {
-  return (A, B) => getSimilarity(A, B, options);
-}

src/treeSimilarity.js

@@ -0,0 +1,42 @@
+/**
+ * @typedef {import("../types").Tree} Tree
+ * @typedef {import("../types").TreeSimilarityOptions} TreeSimilarityOptions
+ */
+/**
+ * Similarity between two nodes
+ * @param {Tree | null} a - tree A node
+ * @param {Tree | null} b - tree B node
+ * @param {TreeSimilarityOptions} [options]
+ * @return {number} similarity measure between tree nodes
+ */
+export function treeSimilarity(a, b, options = {}) {
+  const { alpha = 0.1, beta = 0.33, gamma = 0.001 } = options;
+
+  if (a === null || b === null) {
+    return 0;
+  }
+
+  if (!isTree(a) || !isTree(b)) {
+    throw new Error('tree similarity expects tree as inputs');
+  }
+
+  if (a.sum === 0 && b.sum === 0) {
+    return 1;
+  }
+
+  const C =
+    (alpha * Math.min(a.sum, b.sum)) / Math.max(a.sum, b.sum) +
+    (1 - alpha) * Math.exp(-gamma * Math.abs(a.center - b.center));
+
+  return (
+    beta * C +
+    ((1 - beta) *
+      (treeSimilarity(a.left, b.left, options) +
+        treeSimilarity(a.right, b.right, options))) /
+      2
+  );
+}
+
+function isTree(a) {
+  return ['sum', 'center', 'left', 'right'].every((key) => key in a);
+}

types.d.ts

@@ -0,0 +1,45 @@
+import { NumberArray } from "cheminfo-types";
+
+export interface Tree {
+  sum: number;
+  center: number;
+  /**
+   * left and right have the same structure than the parent,
+   * or are null if they are leaves
+   */
+  left: Tree | null;
+  right: Tree | null;
+}
+
+export interface CreateTreeOptions {
+  /**
+   * low limit of the tree
+   * @default x[0]
+   */
+  from?: number
+  /**
+   * high limit of the tree
+   * @default x.at(-1)
+   */
+  to?: number
+  /**
+   * minimal sum value to accept a node
+   * @default 0.01 
+   */ 
+  threshold?: number;
+  /**
+   * minimal window width to create a node
+   * @default 0.16
+   */
+  minWindow?: number
+}
+
+export interface TreeSimilarityOptions {
+  alpha?: number;
+  beta?: number;
+  gamma?: number;
+}
+export interface Spectrum {
+  x: NumberArray;
+  y: NumberArray;
+}