From 4708fdbd6ef27ae8be13f40d9dcdf1756a28c39a Mon Sep 17 00:00:00 2001
From: Starbeamrainbowlabs <sbrl@starbeamrainbowlabs.com>
Date: Thu, 25 Jul 2019 16:44:26 +0100
Subject: [PATCH] Refactor the AI into a web worker, but it's untested.

---
 client_src/js/ClientConfig.mjs     |   4 +-
 client_src/js/LayerAI.mjs          | 171 +++++++++++++++++++++--------
 client_src/js/Worker/AI.worker.mjs |  26 +++++
 client_src/js/Worker/AIWrapper.mjs |  84 ++++++++++++++
 package-lock.json                  |   9 ++
 package.json                       |   3 +-
 rollup.config.js                   |   5 +
 7 files changed, 251 insertions(+), 51 deletions(-)
 create mode 100644 client_src/js/Worker/AI.worker.mjs
 create mode 100644 client_src/js/Worker/AIWrapper.mjs

diff --git a/client_src/js/ClientConfig.mjs b/client_src/js/ClientConfig.mjs
index 5fa0984..3ea688c 100644
--- a/client_src/js/ClientConfig.mjs
+++ b/client_src/js/ClientConfig.mjs
@@ -10,8 +10,8 @@ export default {
 	
 	// The border around gateways that we should consult the AI on.
 	border: {
-		lat: 0.05,
-		lng: 0.1
+		lat: 0.1,
+		lng: 0.2
 	},
 	
 	// The resolution of the coverage map
diff --git a/client_src/js/LayerAI.mjs b/client_src/js/LayerAI.mjs
index 850f864..f6ec9ff 100644
--- a/client_src/js/LayerAI.mjs
+++ b/client_src/js/LayerAI.mjs
@@ -1,19 +1,20 @@
 "use strict";
 
-import path from 'path';
-
 import L from 'leaflet';
-import {
-	loadLayersModel as tf_loadLayersModel,
-	tensor as tf_tensor
-} from '@tensorflow/tfjs';
 import chroma from 'chroma-js';
 
 import GetFromUrl from './Helpers/GetFromUrl.mjs';
 import Config from './ClientConfig.mjs';
-import { normalise } from '../../common/Math.mjs';
+
+import AIWorker from './Worker/AI.worker.mjs';
+
 
 class LayerAI {
+	/**
+	 * Computes a bounding box that exactly encompasses all the gateways in 
+	 * the index.
+	 * @return	{{north:number,south:number,east:number,west:number}}	The computed bounding box
+	 */
 	get gateway_bounds() {
 		let result = {
 			east: Infinity,
@@ -31,30 +32,107 @@ class LayerAI {
 		
 		return result;
 	}
+	
+	/**
+	 * Initialises a new Leaflet AI layer instance.
+	 * @param {[type]} map [description]
+	 */
 	constructor(map) {
 		this.map = map;
-		this.gateways = new Map();
+		this.worker = new AIWorker();
+		
+		this.map_bounds = null;
 	}
 	
+	/**
+	 * Sets up the Web Worker that does the TensorFlow prediction.
+	 * Using a web worker avoids hanging the main thread.
+	 * @return	{Promise}	A Promise that resolves when setup is complete.
+	 */
+	worker_setup() {
+		// Arrow functions inherit the parent scope, including the "this" 
+		// special variable.
+		return new Promise((resolve, reject) => {
+			// Attach the listener first
+			this.worker.addEventListener("message", (event) => {
+				if(event.data.event !== "setup-complete") {
+					reject(`Error: AIWorker responded with event ${event.data.event}, but 'setup-complete' was expected.`, event.data);
+					return;
+				}
+				
+				resolve();
+			}, { once: true });
+			
+			// Ask the web worker to set itself up
+			this.worker.postMessage({
+				event: "setup",
+				bounds: this.gateway_bounds,
+				index: this.index
+			});
+		})
+	}
+	
+	/**
+	 * Uses the Web Worker to predict a row of signal strength values.
+	 * @param	{number}	latitude	The latitude for which predictions should be made.
+	 * @return	{Promise<number[]>}		A Promise returning the array of predictions calculated by the web worker.
+	 */
+	worker_predict_row(latitude) {
+		return new Promise((resolve, reject) => {
+			// Attach the event listener....
+			this.worker.addEventListener("message", (event) => {
+				if(event.data.event !== "result") {
+					reject(`Error: AIWorker responded with event ${event.data.event}, but 'result' was expected.`, event.data);
+					return;
+				}
+				
+				resolve(event.data);
+				
+			}, { once: true });
+			
+			// ....and send the request
+			this.worker.postMessage({
+				event: "predict-row",
+				latitude
+			});
+		});
+	}
+	
+	/**
+	 * Sets up the Leaflet AI visualisation layer.
+	 * @return	{Promise}	A promise that resolves when setup is complete.
+	 */
 	async setup() {
+		// Download the index file that tells us where the gateways and their 
+		// trained models are located
 		this.index = JSON.parse(
 			await GetFromUrl(Config.ai_index_file)
 		);
 		console.log(this.index);
 		
-		for(let gateway of this.index.index) {
-			this.gateways.set(
-				gateway.id,
-				await tf_loadLayersModel(`${window.location.href}/${path.dirname(Config.ai_index_file)}/${gateway.id}/model.json`)
-			);
-		}
+		// Figure out the bounds of the map we're going to generate
+		this.map_bounds = this.gateway_bounds;
+		map_bounds.north += Config.border.lat;
+		map_bounds.south -= Config.border.lat;
 		
-		this.layer = this.generate_layer();
+		map_bounds.east += Config.border.lng;
+		map_bounds.west -= Config.border.lng;
+		
+		// Setup the web worker
+		await this.worker_setup();
+		
+		// Generate the Leaflet layer
+		this.layer = await this.generate_layer();
 		this.layer.addTo(this.map);
 		console.log("[Layer/AI] Complete");
 	}
 	
-	generate_layer() {
+	/**
+	 * Generates and returns the Leaflet layer containing the AI-predicted 
+	 * values.
+	 * @return {Promise} A Promise that resolves to the generated Leaflet layer.
+	 */
+	async generate_layer() {
 		console.log("[Layer/AI] Rendering map");
 		let map = this.render_map();
 		console.log("[Layer/AI] Passing to Leaflet");
@@ -70,15 +148,12 @@ class LayerAI {
 		});
 	}
 	
-	render_map() {
-		// FUTURE: Do this in a web worker?
-		let map_bounds = this.gateway_bounds;
-		map_bounds.north += Config.border.lat;
-		map_bounds.south -= Config.border.lat;
-		
-		map_bounds.east += Config.border.lng;
-		map_bounds.west -= Config.border.lng;
-		
+	/**
+	 * Uses a Web Worker and pre-trained AIs to generate a GeoJSON map of 
+	 * signal strength for a bounding box around the known gateways.
+	 * @return	{Promise}	A Promise that resolves to a GeoJSON array representing the map.
+	 */
+	async render_map() {
 		let coverage = [],
 			colour_scale = chroma.scale([
 				Config.colour_scale.min,
@@ -88,28 +163,24 @@ class LayerAI {
 				this.index.properties.rssi_max
 			);
 		
-		for(let lat = map_bounds.south; lat < map_bounds.north; lat += Config.step.lat) {
-			for(let lng = map_bounds.west; lng < map_bounds.east; lng += Config.step.lng) {
-				let max_predicted_rssi = -Infinity;
-				
-				for(let [, ai] of this.gateways) {
-					let next_prediction = ai.predict(
-						tf_tensor([ lat, lng ], [1, 2])
-					);
-					max_predicted_rssi = Math.max(
-						max_predicted_rssi,
-						next_prediction.arraySync()[0][0]
-					);
-				}
-				
-				max_predicted_rssi = normalise(max_predicted_rssi,
-					{ min: 0, max: 1 },
-					{
-						min: this.index.properties.rssi_min,
-						max: this.index.properties.rssi_max
-					}
-				);
+		let stats = {
+			rssi_min: Infinity,
+			rssi_max: -Infinity
+		};
+		
+		for(let lat = this.map_bounds.south; lat < this.map_bounds.north; lat += Config.step.lat) {
+			let next_row = await worker_predict_row(lat);
+			let lng = this.map_bounds.west;
+			
+			
+			for(let value of next_row) {
+				// Keep up with the statistics
+				if(value > stats.rssi_max)
+					stats.rssi_max = value;
+				if(value < stats.rssi_min)
+					stats.rssi_min = value;
 				
+				// Generate the GeoJSON feature for this cell of the map
 				coverage.push({
 					type: "Feature",
 					geometry: {
@@ -125,12 +196,16 @@ class LayerAI {
 						]
 					},
 					properties: {
-						colour: colour_scale(max_predicted_rssi).toString()
+						colour: colour_scale(value).toString()
 					}
-				})
+				});
+				
+				lng += Config.step.lng;
 			}
 		}
 		
+		console.log(stats);
+		
 		return coverage;
 	}
 }
diff --git a/client_src/js/Worker/AI.worker.mjs b/client_src/js/Worker/AI.worker.mjs
new file mode 100644
index 0000000..ef60dd2
--- /dev/null
+++ b/client_src/js/Worker/AI.worker.mjs
@@ -0,0 +1,26 @@
+
+export default function(self) {
+	let ai_wrapper = new AIWrapper();
+	
+	self.addEventListener("message", async (event) => {
+		console.log(event.data);
+		switch(event.data.event) {
+			case "setup":
+				await ai_wrapper.setup(event.data.setup_info);
+				self.postMessage({
+					"event": "setup-complete"
+				});
+				break;
+			
+			case "predict-row":
+				let message = await ai_wrapper.predict_row(event.data.latitude);
+				message.event = "result";
+				self.postMessage(message);
+				break;
+			
+			case "end":
+				self.close();
+				break;
+		}
+	});
+}
diff --git a/client_src/js/Worker/AIWrapper.mjs b/client_src/js/Worker/AIWrapper.mjs
new file mode 100644
index 0000000..4b75f66
--- /dev/null
+++ b/client_src/js/Worker/AIWrapper.mjs
@@ -0,0 +1,84 @@
+"use strict";
+
+import path from 'path';
+
+import {
+	loadLayersModel as tf_loadLayersModel,
+	tensor as tf_tensor
+} from '@tensorflow/tfjs';
+
+import { normalise } from '../../common/Math.mjs';
+import Config from '../ClientConfig.mjs';
+
+class AIWrapper {
+	constructor() {
+		this.setup_complete = false;
+		
+		this.map_bounds = null;
+		this.index = null;
+		
+		this.gateways = new Map();
+	}
+	
+	async setup({ bounds, index }) {
+		this.map_bounds = bounds;
+		this.index = index;
+		
+		for(let gateway of this.index.index) {
+			this.gateways.set(
+				gateway.id,
+				await tf_loadLayersModel(`${window.location.href}/${path.dirname(Config.ai_index_file)}/${gateway.id}/model.json`)
+			);
+		}
+		
+		this.setup_complete = true;
+	}
+	
+	predict_row(lat) {
+		if(!setup_complete)
+			throw new Error("Error: Can't do predictions until the setup is complete.");
+			
+		let results = [],
+			stats = {
+				rssi_min: Infinity,
+				rssi_max: -Infinity
+			};
+		
+		for(let lng = this.map_bounds.west; lng < this.map_bounds.east; lng += Config.step.lng) {
+			let max_predicted_rssi = -Infinity;
+			
+			for(let [, ai] of this.gateways) {
+				let next_prediction = this.predict_value(lat, lng)
+				max_predicted_rssi = Math.max(
+					max_predicted_rssi,
+					next_prediction
+				);
+			}
+			
+			max_predicted_rssi = normalise(max_predicted_rssi,
+				{ min: 0, max: 1 },
+				{
+					min: this.index.properties.rssi_min,
+					max: this.index.properties.rssi_max
+				}
+			);
+			
+			if(max_predicted_rssi > stats.rssi_max)
+				stats.rssi_max = max_predicted_rssi;
+			if(max_predicted_rssi < stats.rssi_min)
+				stats.rssi_min = max_predicted_rssi;
+			
+			result.push(max_predicted_rssi);
+		}
+		
+		return { result, stats };
+	}
+	
+	predict_value(latitude, longitude) {
+		return ai.predict(
+			tf_tensor([ latitude, longitude ], [1, 2])
+		).arraySync()[0][0];
+	}
+}
+
+export default AIWrapper;
diff --git a/package-lock.json b/package-lock.json
index 8961e99..c1e7708 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -3892,6 +3892,15 @@
         "terser": "^4.1.0"
       }
     },
+    "rollup-plugin-webworkify": {
+      "version": "0.0.4",
+      "resolved": "https://registry.npmjs.org/rollup-plugin-webworkify/-/rollup-plugin-webworkify-0.0.4.tgz",
+      "integrity": "sha512-GvDPdz7qeakeB26cinVpxPgOBWHx/p3xua/itRXZ/30gGOK5QZeb/iu/9V4M7Zuhwmo2UPR1rViVX8OCqmhi1g==",
+      "dev": true,
+      "requires": {
+        "rollup-pluginutils": "^2"
+      }
+    },
     "rollup-pluginutils": {
       "version": "2.8.1",
       "resolved": "https://registry.npmjs.org/rollup-pluginutils/-/rollup-pluginutils-2.8.1.tgz",
diff --git a/package.json b/package.json
index 9e6220d..04d0abb 100644
--- a/package.json
+++ b/package.json
@@ -43,6 +43,7 @@
 		"rollup-plugin-commonjs": "^10.0.1",
 		"rollup-plugin-node-resolve": "^5.2.0",
 		"rollup-plugin-postcss": "^2.0.3",
-		"rollup-plugin-replace": "^2.2.0"
+		"rollup-plugin-replace": "^2.2.0",
+		"rollup-plugin-webworkify": "0.0.4"
 	}
 }
diff --git a/rollup.config.js b/rollup.config.js
index a849993..9144143 100644
--- a/rollup.config.js
+++ b/rollup.config.js
@@ -8,6 +8,7 @@ import postcss from 'rollup-plugin-postcss';
 import { terser } from "rollup-plugin-terser";
 import replace from 'rollup-plugin-replace';
 import builtins from '@joseph184/rollup-plugin-node-builtins';
+import webworkify from 'rollup-plugin-webworkify';
 // import json from 'rollup-plugin-json';
 
 import postcss_import from 'postcss-import';
@@ -55,6 +56,10 @@ let plugins = [
 		
 	}),
 	
+	webworkify({
+		pattern: '**/*.worker.mjs'
+	}),
+	
 	postcss({
 		plugins: [
 			postcss_import({}),
