Ensure all functions require explicitly providing a prng

Author: cannoneyed

src/heap.ts

@@ -57,7 +57,9 @@
  * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  */
 
+import { RandomFn } from './umap';
 import * as utils from './utils';
+
 export type Heap = number[][][];
 
 /**
@@ -92,7 +94,7 @@ export function makeHeap(nPoints: number, size: number): Heap {
 export function rejectionSample(
   nSamples: number,
   poolSize: number,
-  random: () => number
+  random: RandomFn
 ) {
   const result = utils.zeros(nSamples);
   for (let i = 0; i < nSamples; i++) {
@@ -227,7 +229,7 @@ export function buildCandidates(
   nVertices: number,
   nNeighbors: number,
   maxCandidates: number,
-  random: () => number
+  random: RandomFn
 ) {
   const candidateNeighbors = makeHeap(nVertices, maxCandidates);
   for (let i = 0; i < nVertices; i++) {

src/nn_descent.ts

@@ -61,12 +61,12 @@ import * as heap from './heap';
 import * as matrix from './matrix';
 import * as tree from './tree';
 import * as utils from './utils';
-import { Vectors, DistanceFn } from './umap';
+import { RandomFn, Vectors, DistanceFn } from './umap';
 
 /**
  * Create a version of nearest neighbor descent.
  */
-export function makeNNDescent(distanceFn: DistanceFn, random: () => number) {
+export function makeNNDescent(distanceFn: DistanceFn, random: RandomFn) {
   return function nNDescent(
     data: Vectors,
     leafArray: Vectors,
@@ -150,25 +150,28 @@ export type InitFromRandomFn = (
   nNeighbors: number,
   data: Vectors,
   queryPoints: Vectors,
-  _heap: heap.Heap
+  _heap: heap.Heap,
+  random: RandomFn
 ) => void;
 
 export type InitFromTreeFn = (
   _tree: tree.FlatTree,
   data: Vectors,
   queryPoints: Vectors,
-  _heap: heap.Heap
+  _heap: heap.Heap,
+  random: RandomFn
 ) => void;
 
 export function makeInitializations(distanceFn: DistanceFn) {
   function initFromRandom(
     nNeighbors: number,
     data: Vectors,
     queryPoints: Vectors,
-    _heap: heap.Heap
+    _heap: heap.Heap,
+    random: RandomFn
   ) {
     for (let i = 0; i < queryPoints.length; i++) {
-      const indices = utils.rejectionSample(nNeighbors, data.length);
+      const indices = utils.rejectionSample(nNeighbors, data.length, random);
       for (let j = 0; j < indices.length; j++) {
         if (indices[j] < 0) {
           continue;
@@ -183,10 +186,11 @@ export function makeInitializations(distanceFn: DistanceFn) {
     _tree: tree.FlatTree,
     data: Vectors,
     queryPoints: Vectors,
-    _heap: heap.Heap
+    _heap: heap.Heap,
+    random: RandomFn
   ) {
     for (let i = 0; i < queryPoints.length; i++) {
-      const indices = tree.searchFlatTree(queryPoints[i], _tree);
+      const indices = tree.searchFlatTree(queryPoints[i], _tree, random);
 
       for (let j = 0; j < indices.length; j++) {
         if (indices[j] < 0) {
@@ -252,13 +256,14 @@ export function initializeSearch(
   queryPoints: Vectors,
   nNeighbors: number,
   initFromRandom: InitFromRandomFn,
-  initFromTree: InitFromTreeFn
+  initFromTree: InitFromTreeFn,
+  random: RandomFn
 ) {
   const results = heap.makeHeap(queryPoints.length, nNeighbors);
-  initFromRandom(nNeighbors, data, queryPoints, results);
+  initFromRandom(nNeighbors, data, queryPoints, results, random);
   if (forest) {
     for (let tree of forest) {
-      initFromTree(tree, data, queryPoints, results);
+      initFromTree(tree, data, queryPoints, results, random);
     }
   }
   return results;

src/tree.ts

@@ -58,7 +58,7 @@
  */
 
 import * as utils from './utils';
-import { Vector, Vectors } from './umap';
+import { RandomFn, Vector, Vectors } from './umap';
 
 /**
  * Tree functionality for approximating nearest neighbors
@@ -88,7 +88,7 @@ export function makeForest(
   data: Vectors,
   nNeighbors: number,
   nTrees: number,
-  random: () => number
+  random: RandomFn
 ) {
   const leafSize = Math.max(10, nNeighbors);
 
@@ -108,7 +108,7 @@ function makeTree(
   data: Vectors,
   leafSize = 30,
   n: number,
-  random: () => number
+  random: RandomFn
 ): RandomProjectionTreeNode {
   const indices = utils.range(data.length);
   const tree = makeEuclideanTree(data, indices, leafSize, n, random);
@@ -120,7 +120,7 @@ function makeEuclideanTree(
   indices: number[],
   leafSize = 30,
   q: number,
-  random: () => number
+  random: RandomFn
 ): RandomProjectionTreeNode {
   if (indices.length > leafSize) {
     const splitResults = euclideanRandomProjectionSplit(data, indices, random);
@@ -160,7 +160,7 @@ function makeEuclideanTree(
 function euclideanRandomProjectionSplit(
   data: Vectors,
   indices: number[],
-  random: () => number
+  random: RandomFn
 ) {
   const dim = data[0].length;
 
@@ -343,14 +343,19 @@ export function makeLeafArray(rpForest: FlatTree[]): number[][] {
 /**
  * Selects the side of the tree to search during flat tree search.
  */
-function selectSide(hyperplane: number[], offset: number, point: Vector) {
+function selectSide(
+  hyperplane: number[],
+  offset: number,
+  point: Vector,
+  random: RandomFn
+) {
   let margin = offset;
   for (let d = 0; d < point.length; d++) {
     margin += hyperplane[d] * point[d];
   }
 
   if (margin === 0) {
-    const side = utils.tauRandInt(2);
+    const side = utils.tauRandInt(2, random);
     return side;
   } else if (margin > 0) {
     return 0;
@@ -362,10 +367,19 @@ function selectSide(hyperplane: number[], offset: number, point: Vector) {
 /**
  * Searches a flattened rp-tree for a point.
  */
-export function searchFlatTree(point: Vector, tree: FlatTree) {
+export function searchFlatTree(
+  point: Vector,
+  tree: FlatTree,
+  random: RandomFn
+) {
   let node = 0;
   while (tree.children[node][0] > 0) {
-    const side = selectSide(tree.hyperplanes[node], tree.offsets[node], point);
+    const side = selectSide(
+      tree.hyperplanes[node],
+      tree.offsets[node],
+      point,
+      random
+    );
     if (side === 0) {
       node = tree.children[node][0];
     } else {

src/umap.ts

@@ -65,6 +65,7 @@ import * as utils from './utils';
 import LM from 'ml-levenberg-marquardt';
 
 export type DistanceFn = (x: Vector, y: Vector) => number;
+export type RandomFn = () => number;
 export type EpochCallback = (epoch: number) => boolean | void;
 export type Vector = number[];
 export type Vectors = Vector[];
@@ -142,7 +143,7 @@ export interface UMAPParameters {
    * The pseudo-random number generator used by the stochastic parts of the
    * algorithm.
    */
-  random?: () => number;
+  random?: RandomFn;
   /**
    * Interpolate between (fuzzy) union and intersection as the set operation
    * used to combine local fuzzy simplicial sets to obtain a global fuzzy
@@ -413,7 +414,8 @@ export class UMAP {
       toTransform,
       nNeighbors,
       this.initFromRandom,
-      this.initFromTree
+      this.initFromTree,
+      this.random
     );
 
     const result = this.search(rawData, this.searchGraph, init, toTransform);

src/utils.ts

@@ -17,17 +17,19 @@
  * ==============================================================================
  */
 
+import { RandomFn } from './umap';
+
 /**
  * Simple random integer function
  */
-export function tauRandInt(n: number, random = Math.random) {
+export function tauRandInt(n: number, random: RandomFn) {
   return Math.floor(random() * n);
 }
 
 /**
  * Simple random float function
  */
-export function tauRand(random = Math.random) {
+export function tauRand(random: RandomFn) {
   return random();
 }
 /**
@@ -132,12 +134,16 @@ export function max2d(input: number[][]): number {
  * integer is selected twice. The duplication constraint is achieved via
  * rejection sampling.
  */
-export function rejectionSample(nSamples: number, poolSize: number): number[] {
+export function rejectionSample(
+  nSamples: number,
+  poolSize: number,
+  random: RandomFn
+): number[] {
   const result = zeros(nSamples);
   for (let i = 0; i < nSamples; i++) {
     let rejectSample = true;
     while (rejectSample) {
-      const j = tauRandInt(poolSize);
+      const j = tauRandInt(poolSize, random);
       let broken = false;
       for (let k = 0; k < i; k++) {
         if (j === result[k]) {

test/umap.test.ts

@@ -21,6 +21,7 @@ import {
   UMAP,
   findABParams,
   euclidean,
+  RandomFn,
   TargetMetric,
   Vector,
 } from '../src/umap';
@@ -32,12 +33,11 @@ import {
   testLabels,
   testResults2D,
   testResults3D,
-  transformResult2d,
 } from './test_data';
 import Prando from 'prando';
 
 describe('UMAP', () => {
-  let random: () => number;
+  let random: RandomFn;
 
   // Expected "clustering" ratios, representing inter-cluster distance vs mean
   // distance to other points.

test/utils.test.ts

@@ -18,8 +18,12 @@
  */
 
 import * as utils from '../src/utils';
+import Prando from 'prando';
 
 describe('umap utils', () => {
+  const prando = new Prando(42);
+  const random = () => prando.next();
+
   test('norm function', () => {
     const results = utils.norm([1, 2, 3, 4]);
     expect(results).toEqual(Math.sqrt(30));
@@ -81,7 +85,7 @@ describe('umap utils', () => {
   });
 
   test('rejection sample', () => {
-    const results = utils.rejectionSample(5, 10);
+    const results = utils.rejectionSample(5, 10, random);
     const entries = new Set<number>();
     for (const r of results) {
       expect(entries.has(r)).toBe(false);