/**
* Main class for the creation of machine learning models.
* Acts as a thin client/proxy to the Web Worker where the actual models reside.
* @class
* @name nn
*/
export default class nn {
static worker = null;
static pending = new Map();
static msgId = 0;
static {
if (typeof window !== 'undefined' && typeof Worker !== 'undefined') {
const workerCode = `
importScripts("https://cdnjs.cloudflare.com/ajax/libs/tensorflow/3.10.0/tf.min.js");
/**
* Fat Worker for Neural Network operations.
* Holds all model instances and performs all heavy computation.
*/
const modelRegistry = new Map();
self.onmessage = async function (e) {
const { id, action, payload } = e.data;
try {
let result;
switch (action) {
case "CREATE":
result = await createModel(payload);
break;
case "TRAIN":
result = await trainModel(payload);
break;
case "PREDICT":
result = await predictModel(payload);
break;
case "SAVE":
result = await saveModel(payload);
break;
case "LOAD":
result = await loadModel(payload);
break;
case "DISPOSE":
result = disposeModel(payload);
break;
default:
throw new Error(\`Unknown action: \${action}\`);
}
// Send success response
self.postMessage({ id, status: "SUCCESS", payload: result });
} catch (error) {
self.postMessage({ id, status: "ERROR", payload: error.message });
}
};
// --- Model Creation ---
async function createModel({ modelId, type, config }) {
let model;
switch (type.toLowerCase()) {
case "cnn":
model = createCNN(config);
break;
case "lstm":
model = createLSTM(config);
break;
case "transformer":
model = createTransformer(config);
break;
case "segmentation":
model = createSegmentationModel(config);
break;
case "dense":
default:
model = createDense(config);
break;
}
modelRegistry.set(modelId, model);
return { modelId, type, status: "CREATED" };
}
function createDense(config) {
const model = tf.sequential();
model.add(tf.layers.dense({
inputShape: [config.inputs],
units: config.inputs,
activation: config.activation || 'relu'
}));
model.add(tf.layers.dense({
units: config.neurons,
activation: config.activation || 'relu'
}));
model.add(tf.layers.dense({
units: config.outputs,
activation: config.outputActivation || 'linear' // Default to linear for regression
}));
return model;
}
function createCNN(config) {
const model = tf.sequential();
model.add(tf.layers.conv2d({
inputShape: config.inputShape,
kernelSize: config.kernelSize || 3,
filters: config.filters || 16,
activation: 'relu'
}));
model.add(tf.layers.maxPooling2d({ poolSize: [2, 2] }));
model.add(tf.layers.flatten());
model.add(tf.layers.dense({ units: config.classes || 1, activation: 'softmax' }));
return model;
}
function createSegmentationModel(config) {
const model = tf.sequential();
// Encoder
model.add(tf.layers.conv2d({
inputShape: config.inputShape, // [height, width, channels]
filters: 16,
kernelSize: 3,
padding: 'same',
activation: 'relu'
}));
model.add(tf.layers.conv2d({
filters: 32,
kernelSize: 3,
padding: 'same',
activation: 'relu'
}));
// Decoder / Output
model.add(tf.layers.conv2d({
filters: config.classes || 2,
kernelSize: 1,
padding: 'same',
activation: 'softmax' // Pixel-wise classification
}));
return model;
}
function createLSTM(config) {
const model = tf.sequential();
model.add(tf.layers.lstm({
units: config.units || 50,
inputShape: [config.timeSteps, config.features],
returnSequences: false
}));
model.add(tf.layers.dense({ units: config.outputs || 1 }));
return model;
}
function createTransformer(config) {
const input = tf.input({ shape: [config.timeSteps, config.features] });
const d_model = config.d_model || 64;
let x = tf.layers.conv1d({ filters: d_model, kernelSize: 1, padding: 'same', activation: 'linear' }).apply(input);
const query = tf.layers.dense({ units: d_model }).apply(x);
const key = tf.layers.dense({ units: d_model }).apply(x);
const value = tf.layers.dense({ units: d_model }).apply(x);
x = tf.layers.dense({ units: d_model * 2, activation: 'relu' }).apply(x);
x = tf.layers.dense({ units: d_model }).apply(x);
x = tf.layers.globalAveragePooling1d().apply(x);
const output = tf.layers.dense({ units: config.outputs || 1 }).apply(x);
const model = tf.model({ inputs: input, outputs: output });
return model;
}
// --- Training ---
async function trainModel({ modelId, data, config }) {
const model = modelRegistry.get(modelId);
if (!model) throw new Error(\`Model \${modelId} not found\`);
const xs = tf.tensor(data.inputs, data.inputShape);
const ys = tf.tensor(data.outputs, data.outputShape);
const { epochs, batchSize, validationSplit } = config;
model.compile({
optimizer: config.optimizer || "adam",
loss: config.loss || "meanSquaredError",
metrics: config.metrics || ["mse"],
});
await model.fit(xs, ys, {
epochs: epochs || 50,
batchSize: batchSize || 32,
validationSplit: validationSplit || 0.1,
callbacks: {
onEpochEnd: (epoch, logs) => {
self.postMessage({
action: "EPOCH_END",
payload: { modelId, epoch, logs },
});
},
},
});
xs.dispose();
ys.dispose();
return { modelId, status: "TRAINED" };
}
// --- Prediction ---
async function predictModel({ modelId, inputs, inputShape }) {
const model = modelRegistry.get(modelId);
if (!model) throw new Error(\`Model \${modelId} not found\`);
const xs = tf.tensor(inputs, inputShape);
const prediction = model.predict(xs);
const result = await prediction.array();
xs.dispose();
prediction.dispose();
return result;
}
// --- Utilities ---
async function saveModel({ modelId, name }) {
const model = modelRegistry.get(modelId);
if (!model) throw new Error(\`Model \${modelId} not found\`);
await model.save(\`downloads://\${name}\`);
return { status: "SAVED", name };
}
async function loadModel({ modelId, url }) {
const model = await tf.loadLayersModel(url);
modelRegistry.set(modelId, model);
return { modelId, status: "LOADED" };
}
function disposeModel({ modelId }) {
const model = modelRegistry.get(modelId);
if (model) {
model.dispose();
modelRegistry.delete(modelId);
}
return { status: "DISPOSED" };
}
`;
const blob = new Blob([workerCode], { type: 'application/javascript' });
this.worker = new Worker(URL.createObjectURL(blob));
}
}
static {
if (this.worker) {
this.worker.onmessage = (e) => {
const { id, status, payload, action } = e.data;
// Handle progress updates (e.g., EPOCH_END) without resolving the promise yet
if (action === "EPOCH_END") {
console.log(`[Training] Epoch ${payload.epoch}: loss = ${payload.logs.loss}`);
return;
}
if (this.pending.has(id)) {
const { resolve, reject } = this.pending.get(id);
this.pending.delete(id);
if (status === "SUCCESS") resolve(payload);
else reject(new Error(payload));
}
};
}
}
static _send(action, payload) {
const id = this.msgId++;
return new Promise((resolve, reject) => {
this.pending.set(id, { resolve, reject });
this.worker.postMessage({ id, action, payload });
});
}
/**
* Create a machine learning model for hydrological prediction and analysis
* Supports 5 model architectures: Dense, LSTM, CNN, Transformer, Segmentation
* Models run in Web Worker for non-blocking performance
*
* @function createModel
* @memberof nn
* @async
* @param {Object} options - Function options
* @param {Object} options.params - Model parameters
* @param {string} options.params.type - Model type: 'dense', 'lstm', 'cnn', 'transformer', 'segmentation'
* @param {Object} options.args - Model configuration (varies by type)
* @param {Object} [options.data] - Not used
* @returns {Promise<ModelProxy>} Model proxy object for training and prediction
*
* @example
* // 1. Dense Neural Network - Simple regression/classification
* const denseModel = await hydro.analyze.nn.createModel({
* params: { type: 'dense' },
* args: {
* inputShape: [5], // 5 input features
* units: [32, 16, 1], // 3 layers: 32→16→1 neurons
* activation: 'relu', // Activation function
* outputActivation: 'linear' // Linear output for regression
* }
* });
* // Use case: Predict streamflow from rainfall, temperature, etc.
*
* @example
* // 2. LSTM - Time series forecasting
* const lstmModel = await hydro.analyze.nn.createModel({
* params: { type: 'lstm' },
* args: {
* timeSteps: 30, // 30 days lookback
* features: 3, // 3 variables (precip, temp, flow)
* units: [64, 32], // 2 LSTM layers
* outputUnits: 1, // Predict 1 day ahead
* dropout: 0.2 // 20% dropout for regularization
* }
* });
* // Use case: Forecast streamflow 1 day ahead using 30-day history
*
* @example
* // 3. CNN - Spatial pattern recognition
* const cnnModel = await hydro.analyze.nn.createModel({
* params: { type: 'cnn' },
* args: {
* inputShape: [64, 64, 3], // 64x64 RGB image
* filters: [32, 64, 128], // 3 conv layers with increasing filters
* kernelSize: 3, // 3x3 convolution kernels
* poolSize: 2, // 2x2 max pooling
* denseUnits: [128, 10], // Dense layers after convolution
* activation: 'relu'
* }
* });
* // Use case: Classify land cover from satellite imagery
*
* @example
* // 4. Transformer - Advanced time series with attention
* const transformerModel = await hydro.analyze.nn.createModel({
* params: { type: 'transformer' },
* args: {
* timeSteps: 60, // 60-step sequence
* features: 5, // 5 input features
* numHeads: 4, // 4 attention heads
* dModel: 128, // Model dimension
* numLayers: 2, // 2 transformer blocks
* dff: 256, // Feedforward dimension
* outputUnits: 1
* }
* });
* // Use case: Complex multivariate forecasting with attention mechanisms
*
* @example
* // 5. Segmentation - Image segmentation (U-Net architecture)
* const segModel = await hydro.analyze.nn.createModel({
* params: { type: 'segmentation' },
* args: {
* inputShape: [256, 256, 3], // Input image size
* filters: [64, 128, 256], // Encoder filters
* numClasses: 3, // 3 classes (water, land, vegetation)
* kernelSize: 3,
* poolSize: 2
* }
* });
* // Use case: Segment water bodies from aerial imagery
*
* @example
* // Complete workflow: Create → Train → Predict
* const model = await hydro.analyze.nn.createModel({
* params: { type: 'lstm' },
* args: { timeSteps: 7, features: 2, units: [32], outputUnits: 1 }
* });
*
* // Train the model
* await model.train({
* params: { epochs: 50, batchSize: 32 },
* data: [trainingInputs, trainingOutputs]
* });
*
* // Make predictions
* const predictions = await model.predict({ data: testInputs });
*/
static async createModel({ params = {}, args = {}, data } = {}) {
const modelId = `model_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
const type = params.type || "dense";
const config = args; // Direct assignment, NO backward compatibility
await this._send("CREATE", { modelId, type, config });
// Return a proxy object that remembers its ID and config
return new ModelProxy(modelId, config);
}
/**
* Load previously saved machine learning model
* Restores architecture and weights from storage
* Supports URL or IndexedDB loading
*
* @function loadModel
* @memberof nn
* @async
* @param {Object} options - Loading options
* @param {Object} options.params - Parameters
* @param {string} options.params.url - Path to model.json or storage key
* @param {Object} [options.args] - Additional arguments
* @param {Object} [options.data] - Not used
* @returns {Promise<ModelProxy>} Loaded model proxy
*
* @example
* // Load from URL
* const model = await hydro.analyze.nn.loadModel({
* params: { url: 'https://example.com/models/streamflow/model.json' }
* });
* const predictions = await model.predict({ data: testData });
*
* @example
* // Load from IndexedDB
* const model = await hydro.analyze.nn.loadModel({
* params: { url: 'indexeddb://my-trained-model' }
* });
*
* @example
* // Complete save/load cycle
* // Training session:
* const model1 = await hydro.analyze.nn.createModel({/* config});
* await model1.train({ data: [X, y], params: { epochs: 100 } });
* await model1.save({ params: { name: 'flood-model', location: 'indexeddb' } });
*
* // Later session:
* const model2 = await hydro.analyze.nn.loadModel({
* params: { url: 'indexeddb://flood-model' }
* });
* const results = await model2.predict({ data: newData });
*/
static async loadModel({ params = {}, args, data } = {}) {
const modelId = `model_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
const url = params.url;
if (!url) {
throw new Error("URL is required to load a model.");
}
await this._send("LOAD", { modelId, url });
return new ModelProxy(modelId, {});
}
}
/**
* Proxy class representing a model living in the worker.
*/
class ModelProxy {
constructor(modelId, config = {}) {
this.id = modelId;
this.config = config;
}
/**
* Preprocesses 1D data into LSTM-ready 3D format.
* @method preprocess
* @private
* @param {Object} params - Contains: inputWindow (timeSteps), outputWindow (predictions).
* @param {Array} data - 1D array of values.
* @returns {Array} [X, y] where X is 3D array and y is 2D array.
*/
preprocess({ params = {}, data } = {}) {
const inputWindow = params.inputWindow || this.config.timeSteps || 1;
const outputWindow = params.outputWindow || 1;
const series = data;
if (!Array.isArray(series)) {
throw new Error("Data must be a 1D array.");
}
const X = [];
const y = [];
// Sliding window
for (let i = 0; i <= series.length - inputWindow - outputWindow; i++) {
const inputSeq = series.slice(i, i + inputWindow).map(val => [val]); // Shape: [timeSteps, features=1]
const outputSeq = series.slice(i + inputWindow, i + inputWindow + outputWindow); // Shape: [outputWindow]
X.push(inputSeq);
y.push(outputSeq);
}
return [X, y];
}
/**
* Train the neural network model with hydrological data
* Supports multiple data formats and training configurations
* Training runs in Web Worker for non-blocking UI
*
* @function train
* @memberof ModelProxy
* @async
* @param {Object} options - Training options
* @param {Object} [options.params] - Training parameters
* @param {number} [options.params.epochs=50] - Number of training epochs
* @param {number} [options.params.batchSize=32] - Batch size for training
* @param {number} [options.params.validationSplit=0.2] - Validation data fraction (0-1)
* @param {boolean} [options.params.shuffle=true] - Shuffle training data
* @param {Function} [options.params.onEpochEnd] - Callback after each epoch
* @param {Object|Array} options.data - Training data: [inputs, outputs] or {inputs, outputs}
* @returns {Promise<Object>} Training history with loss and metrics
*
* @example
* // Basic training: Streamflow prediction from rainfall
* const model = await hydro.analyze.nn.createModel({
* params: { type: 'dense' },
* args: { inputShape: [3], units: [16, 8, 1] }
* });
*
* const inputs = [[5, 20, 15], [3, 18, 12], [8, 25, 18]]; // [precip, temp, humidity]
* const outputs = [[120], [95], [145]]; // streamflow
*
* const history = await model.train({
* params: { epochs: 100, batchSize: 16 },
* data: [inputs, outputs]
* });
* console.log(`Final loss: ${history.loss[history.loss.length - 1]}`);
*
* @example
* // LSTM training: Time series forecasting
* const lstmModel = await hydro.analyze.nn.createModel({
* params: { type: 'lstm' },
* args: { timeSteps: 7, features: 2, units: [32], outputUnits: 1 }
* });
*
* // Training data: 7 days of [rainfall, temperature] → next day flow
* const sequences = [
* [[[10,20],[12,22],[8,21],[15,23],[5,19],[7,18],[9,20]]], // Input sequence
* [[125]] // Target: next day flow
* ];
*
* await lstmModel.train({
* params: { epochs: 50, validationSplit: 0.2 },
* data: sequences
* });
*
* @example
* // Training with callbacks and monitoring
* await model.train({
* params: {
* epochs: 100,
* batchSize: 32,
* onEpochEnd: (epoch, logs) => {
* console.log(`Epoch ${epoch}: loss=${logs.loss.toFixed(4)}, val_loss=${logs.val_loss?.toFixed(4)}`);
* // Early stopping logic could go here
* }
* },
* data: [X_train, y_train]
* });
*
* @example
* // Object format (alternative to array format)
* await model.train({
* params: { epochs: 50 },
* data: {
* inputs: flowObservations,
* outputs: nextDayFlow
* }
* });
*/
async train({ params = {}, args, data } = {}) {
let inputs, outputs;
// Handle [X, y] format
if (Array.isArray(data) && data.length === 2) {
inputs = data[0];
outputs = data[1];
} else if (typeof data === 'object' && data.inputs && data.outputs) {
inputs = data.inputs;
outputs = data.outputs;
} else {
// Fallback or cache key handling (omitted helper for brevity in this strict refactor unless needed, assuming direct data for now based on request)
// If user passes string, we could still support it, but let's stick to the requested simplification first.
// Re-adding cache support briefly to be safe.
if (typeof data === 'string') {
const cache = globalThis.hydro?.cache;
if (cache) {
const cached = await cache.get(data);
if (cached) {
inputs = cached.data.inputs;
outputs = cached.data.outputs;
}
}
}
}
if (!inputs || !outputs) {
throw new Error("Invalid data format. Provide [inputs, outputs] or {inputs, outputs}.");
}
return nn._send("TRAIN", {
modelId: this.id,
data: { inputs, outputs },
config: params
});
}
/**
* Generate predictions from trained model
* Supports auto-reshaping for LSTM inputs
* Returns predictions for hydrological forecasting or classification
*
* @function predict
* @memberof ModelProxy
* @async
* @param {Object} options - Prediction options
* @param {Object} [options.params] - Prediction parameters
* @param {Array} [options.params.inputShape] - Override input shape
* @param {Array} options.data - Input data for prediction
* @returns {Promise<Array>} Model predictions
*
* @example
* // Dense model: Predict streamflow from current conditions
* const model = await hydro.analyze.nn.createModel({
* params: { type: 'dense' },
* args: { inputShape: [3], units: [16, 1] }
* });
* await model.train({ params: { epochs: 50 }, data: [X_train, y_train] });
*
* // Predict for new conditions
* const newConditions = [[6, 22, 14]]; // [precip, temp, humidity]
* const flowPrediction = await model.predict({ data: newConditions });
* console.log(`Predicted flow: ${flowPrediction[0][0].toFixed(2)} m³/s`);
*
* @example
* // LSTM: Forecast next day flow (auto-reshaping)
* const lstmModel = await hydro.analyze.nn.createModel({
* params: { type: 'lstm' },
* args: { timeSteps: 7, features: 1, units: [32], outputUnits: 1 }
* });
* await lstmModel.train({ params: { epochs: 50 }, data: [sequences, targets] });
*
* // 1D input automatically reshaped to [1, 7, 1] for LSTM
* const last7Days = [100, 105, 98, 110, 115, 108, 112];
* const nextDayPrediction = await lstmModel.predict({ data: last7Days });
* console.log(`Tomorrow's flow: ${nextDayPrediction[0][0].toFixed(2)} m³/s`);
*
* @example
* // Batch predictions: Multiple scenarios
* const scenarios = [
* [5, 20, 15], // Scenario 1: light rain
* [15, 22, 18], // Scenario 2: heavy rain
* [2, 25, 12] // Scenario 3: minimal rain
* ];
* const predictions = await model.predict({ data: scenarios });
* scenarios.forEach((scenario, i) => {
* console.log(`Scenario ${i+1}: ${predictions[i][0].toFixed(2)} m³/s`);
* });
*
* @example
* // Classification: Flood risk category
* const classifier = await hydro.analyze.nn.createModel({
* params: { type: 'dense' },
* args: { inputShape: [4], units: [32, 16, 3], outputActivation: 'softmax' }
* });
* // Train with flood categories: 0=Normal, 1=Warning, 2=Danger
*
* const currentConditions = [[120, 8, 25, 85]]; // [flow, precip, temp, saturation]
* const probabilities = await classifier.predict({ data: currentConditions });
* const category = probabilities[0].indexOf(Math.max(...probabilities[0]));
* const categories = ['Normal', 'Warning', 'Danger'];
* console.log(`Flood Risk: ${categories[category]} (${(probabilities[0][category]*100).toFixed(1)}%)`);
*/
async predict({ params = {}, args, data } = {}) {
let inputs = data;
let inputShape = params.inputShape;
// Auto-reshape if 1D and we know the timeSteps
if (Array.isArray(data) && !Array.isArray(data[0]) && this.config.timeSteps) {
if (data.length === this.config.timeSteps) {
// Reshape to [1, timeSteps, 1]
inputs = [data.map(v => [v])];
// We don't strictly need to pass inputShape if the tensor inference works,
// but let's pass it if we can derive it.
// inputShape = [1, this.config.timeSteps, 1];
}
}
return nn._send("PREDICT", {
modelId: this.id,
inputs: inputs,
inputShape: inputShape
});
}
/**
* Saves the model to the local file system (browser downloads).
* @method save
* @memberof ModelProxy
* @param {Object} params - Contains: name (filename for the model).
* @param {Object} args - Not used by this function.
* @param {Object} data - Not used by this function.
* @returns {Promise<Object>} Save status.
*
* @example
* await model.save({
* params: { name: 'my-trained-model' }
* });
*/
async save({ params = {}, args, data } = {}) {
return nn._send("SAVE", {
modelId: this.id,
name: params.name || 'model'
});
}
/**
* Disposes the model from memory in the worker.
* @method dispose
* @memberof ModelProxy
* @param {Object} params - Not used by this function.
* @param {Object} args - Not used by this function.
* @param {Object} data - Not used by this function.
* @returns {void}
*
* @example
* model.dispose();
*/
dispose({ params, args, data } = {}) {
nn._send("DISPOSE", { modelId: this.id });
}
}