LoRaWAN-Signal-Mapping/server/train-ai/AITrainer.mjs

"use strict";

import path from 'path';
import fs from 'fs';

import brain from 'brain.js';

class AITrainer {
	constructor({ ansi, settings, log, root_dir, GatewayRepo, DatasetFetcher }) {
		this.a = ansi;
		this.settings = settings;
		this.root_dir = root_dir;
		this.l = log;
		this.dataset_fetcher = DatasetFetcher;
		this.repo_gateway = GatewayRepo;
	}
	
	async train_unified() {
		let filepath = path.join(
			this.root_dir,
			"..",
			this.settings.ai.output_directory,
			"ai.json"
		);
		
		if(!fs.existsSync(path.dirname(filepath)))
			await fs.promises.mkdir(path.dirname(filepath), { recursive: true });
		
		let training_result = await this.train_gateway(null, filepath);
		
		let { latitude, longitude } = this.repo_gateway.get_average_location();
		
		await fs.promises.writeFile(
			path.join(
				path.dirname(this.root_dir),
				this.settings.ai.output_directory,
				"index.json"
			),
			JSON.stringify({
				properties: {
					rssi_min: -150,
					rssi_max:  0,
					training_mode: "unified"
				},
				index: [{
					id: "Unified AI (not a real gateway)",
					filename: "ai.json",
					latitude, longitude,
					net_settings: training_result.net_settings
				}]
			}, null, "\t")
		);
	}
	
	async train_all() {
		let index = [];
		for(let gateway of this.repo_gateway.iterate()) {
			let filepath = path.join(this.root_dir, "..", this.settings.ai.output_directory, `${gateway.id}.json`);
			
			if(!fs.existsSync(path.dirname(filepath)))
				await fs.promises.mkdir(path.dirname(filepath), { recursive: true });
			
			let result = await this.train_gateway(gateway.id, filepath);
			if(!result || result.success === false) {
				this.l.warn(`Warning: Failed to train AI for ${gateway.id}.`);
				continue;
			}
			
			
			this.l.log(`Saved to ${filepath}`);
			
			index.push({
				id: gateway.id,
				filename: path.basename(filepath),
				latitude: gateway.latitude,
				longitude: gateway.longitude,
				net_settings: result.net_settings
			});
		}
		
		await fs.promises.writeFile(
			path.join(
				path.dirname(this.root_dir),
				this.settings.ai.output_directory,
				"index.json"
			),
			JSON.stringify({
				properties: {
					rssi_min: -150,
					rssi_max:  0,
					training_mode: "split"
				},
				index
			})
		);
	}
	
	/**
	 * Trains an AI to predict the coverage of a specific gateway.
	 * @param	{string}	gateway_id				The id of the gateway to train an AI for.
	 * @param	{string}	destination_filename	The absolute path to the file to serialise the trained to. Required because we can't serialise and return a TensorFlow model, it has to be sent somewhere because the API is backwards and upside-down :-/
	 * @return	{Promise}	A promise that resolves when training and serialisation is complete.
	 */
	async train_gateway(gateway_id, destination_filename) {
		this.l.log(`${this.a.fgreen}${this.a.hicol}Training AI for gateway ${gateway_id}${this.a.reset}`);
		
		// Create the neural network
		let net_settings = {
			hiddenLayers: this.settings.ai.network_arch,
			activation: "sigmoid"
		};
		let net = new brain.NeuralNetwork(net_settings);
		
		// Fetch the dataset
		let dataset = this.dataset_fetcher.fetch_all(gateway_id);
		
		await net.trainAsync(dataset, {
			iterations: this.settings.ai.epochs,
			errorThresh: this.settings.ai.error_threshold,
			
			learningRate: this.settings.ai.learning_rate,
			momentum: this.settings.ai.momentum,
			
			log: (log_line) => this.l.log(`[brain.js] ${log_line}`),
			logPeriod: 50,
			
			timeout: Infinity
		});
		
		await fs.promises.writeFile(destination_filename, JSON.stringify(net.toJSON()), null, "\t");
		// console.log(result);
		
		return {
			success: true,
			net_settings
		};
	}
}

export default AITrainer;