Skip to main content

hvac_model/
hvac_model.rs

1//! HVAC Multi-Horizon Predictor - Native AxonML Implementation
2//!
3//! # File
4//! `crates/axonml/examples/hvac_model.rs`
5//!
6//! # Author
7//! Andrew Jewell Sr - AutomataNexus
8//!
9//! # Updated
10//! March 8, 2026
11//!
12//! # Disclaimer
13//! Use at own risk. This software is provided "as is", without warranty of any
14//! kind, express or implied. The author and AutomataNexus shall not be held
15//! liable for any damages arising from the use of this software.
16
17use axonml::autograd::Variable;
18use axonml::nn::{Dropout, GRU, LayerNorm, Linear, Module, Parameter, ReLU, Softmax};
19use axonml::tensor::Tensor;
20
21// =============================================================================
22// Model Configuration
23// =============================================================================
24
25/// HVAC Model Configuration
26#[derive(Debug, Clone)]
27pub struct HvacConfig {
28    /// Number of input features (sensor readings)
29    pub num_features: usize,
30    /// Sequence length (time steps)
31    pub seq_len: usize,
32    /// Hidden dimension size
33    pub hidden_size: usize,
34    /// Number of GRU layers
35    pub num_layers: usize,
36    /// Number of failure classes
37    pub num_classes: usize,
38    /// Dropout rate
39    pub dropout: f32,
40}
41
42impl Default for HvacConfig {
43    fn default() -> Self {
44        Self {
45            num_features: 28,
46            seq_len: 120,
47            hidden_size: 128,
48            num_layers: 2,
49            num_classes: 20,
50            dropout: 0.1,
51        }
52    }
53}
54
55// =============================================================================
56// Prediction Head
57// =============================================================================
58
59/// Classification head for a single prediction horizon
60pub struct PredictionHead {
61    fc1: Linear,
62    fc2: Linear,
63    fc3: Linear,
64    relu: ReLU,
65    dropout: Dropout,
66}
67
68impl PredictionHead {
69    pub fn new(hidden_size: usize, num_classes: usize, dropout: f32) -> Self {
70        Self {
71            fc1: Linear::new(hidden_size, hidden_size),
72            fc2: Linear::new(hidden_size, 64),
73            fc3: Linear::new(64, num_classes),
74            relu: ReLU,
75            dropout: Dropout::new(dropout),
76        }
77    }
78}
79
80impl Module for PredictionHead {
81    fn forward(&self, x: &Variable) -> Variable {
82        let x = self.fc1.forward(x);
83        let x = self.relu.forward(&x);
84        let x = self.dropout.forward(&x);
85        let x = self.fc2.forward(&x);
86        let x = self.relu.forward(&x);
87        let x = self.dropout.forward(&x);
88        self.fc3.forward(&x)
89    }
90
91    fn parameters(&self) -> Vec<Parameter> {
92        let mut params = self.fc1.parameters();
93        params.extend(self.fc2.parameters());
94        params.extend(self.fc3.parameters());
95        params
96    }
97}
98
99// =============================================================================
100// HVAC Multi-Horizon Predictor
101// =============================================================================
102
103/// HVAC Multi-Horizon Failure Predictor
104///
105/// Predicts potential HVAC system failures at 3 time horizons:
106/// - Imminent: 5 minutes
107/// - Warning: 15 minutes
108/// - Early: 30 minutes
109pub struct HvacPredictor {
110    config: HvacConfig,
111
112    // Input projection
113    input_proj: Linear,
114    input_norm: LayerNorm,
115    input_relu: ReLU,
116
117    // Temporal encoder
118    gru: GRU,
119
120    // Prediction heads
121    head_imminent: PredictionHead, // 5 min
122    head_warning: PredictionHead,  // 15 min
123    head_early: PredictionHead,    // 30 min
124
125    // Output
126    softmax: Softmax,
127}
128
129/// Output from the HVAC predictor
130#[derive(Debug)]
131pub struct HvacOutput {
132    /// Imminent (5 min) class predictions
133    pub imminent_logits: Variable,
134    /// Warning (15 min) class predictions
135    pub warning_logits: Variable,
136    /// Early (30 min) class predictions
137    pub early_logits: Variable,
138}
139
140impl HvacPredictor {
141    /// Creates a new HVAC predictor with the given configuration
142    pub fn new(config: HvacConfig) -> Self {
143        Self {
144            input_proj: Linear::new(config.num_features, config.hidden_size),
145            input_norm: LayerNorm::new(vec![config.hidden_size]),
146            input_relu: ReLU,
147            gru: GRU::new(config.hidden_size, config.hidden_size, config.num_layers),
148            head_imminent: PredictionHead::new(
149                config.hidden_size,
150                config.num_classes,
151                config.dropout,
152            ),
153            head_warning: PredictionHead::new(
154                config.hidden_size,
155                config.num_classes,
156                config.dropout,
157            ),
158            head_early: PredictionHead::new(config.hidden_size, config.num_classes, config.dropout),
159            softmax: Softmax::new(-1),
160            config,
161        }
162    }
163
164    /// Mean pooling over sequence dimension
165    fn mean_pool(&self, x: &Variable) -> Variable {
166        let data = x.data();
167        let shape = data.shape();
168        let batch_size = shape[0];
169        let seq_len = shape[1];
170        let hidden = shape[2];
171
172        // Reshape to [batch * seq, hidden] then back
173        let values = data.to_vec();
174
175        // Calculate mean over sequence dimension
176        let mut pooled = vec![0.0f32; batch_size * hidden];
177        for b in 0..batch_size {
178            for h in 0..hidden {
179                let mut sum = 0.0;
180                for s in 0..seq_len {
181                    let idx = b * seq_len * hidden + s * hidden + h;
182                    sum += values[idx];
183                }
184                pooled[b * hidden + h] = sum / seq_len as f32;
185            }
186        }
187
188        let pooled_tensor = Tensor::from_vec(pooled, &[batch_size, hidden])
189            .expect("Failed to create pooled tensor");
190        Variable::new(pooled_tensor, x.requires_grad())
191    }
192
193    /// Forward pass returning logits for all 3 horizons
194    pub fn forward_multi(&self, x: &Variable) -> HvacOutput {
195        let x_data = x.data();
196        let shape = x_data.shape();
197        let batch_size = shape[0];
198        let seq_len = shape[1];
199        drop(x_data); // Release borrow
200
201        // Input projection: [batch, seq, features] -> [batch, seq, hidden]
202        // Reshape for linear: [batch * seq, features]
203        let x_flat = x.reshape(&[batch_size * seq_len, self.config.num_features]);
204        let proj = self.input_proj.forward(&x_flat);
205        let proj = self.input_norm.forward(&proj);
206        let proj = self.input_relu.forward(&proj);
207        let proj = proj.reshape(&[batch_size, seq_len, self.config.hidden_size]);
208
209        // GRU encoding: [batch, seq, hidden] -> [batch, seq, hidden]
210        let encoded = self.gru.forward(&proj);
211
212        // Mean pooling: [batch, seq, hidden] -> [batch, hidden]
213        let pooled = self.mean_pool(&encoded);
214
215        // Prediction heads
216        let imminent_logits = self.head_imminent.forward(&pooled);
217        let warning_logits = self.head_warning.forward(&pooled);
218        let early_logits = self.head_early.forward(&pooled);
219
220        HvacOutput {
221            imminent_logits,
222            warning_logits,
223            early_logits,
224        }
225    }
226
227    /// Get predicted classes (argmax of logits)
228    pub fn predict(&self, x: &Variable) -> (Vec<usize>, Vec<usize>, Vec<usize>) {
229        let output = self.forward_multi(x);
230
231        let imminent_probs = self.softmax.forward(&output.imminent_logits);
232        let warning_probs = self.softmax.forward(&output.warning_logits);
233        let early_probs = self.softmax.forward(&output.early_logits);
234
235        (
236            argmax_batch(&imminent_probs),
237            argmax_batch(&warning_probs),
238            argmax_batch(&early_probs),
239        )
240    }
241
242    /// Returns the model configuration
243    pub fn config(&self) -> &HvacConfig {
244        &self.config
245    }
246
247    /// Returns the number of trainable parameters
248    pub fn num_parameters(&self) -> usize {
249        self.parameters()
250            .iter()
251            .map(|p| p.variable().data().numel())
252            .sum()
253    }
254}
255
256impl Module for HvacPredictor {
257    fn forward(&self, x: &Variable) -> Variable {
258        // Return concatenated logits for all horizons
259        let output = self.forward_multi(x);
260        // For single output, return imminent predictions
261        output.imminent_logits
262    }
263
264    fn parameters(&self) -> Vec<Parameter> {
265        let mut params = self.input_proj.parameters();
266        params.extend(self.input_norm.parameters());
267        params.extend(self.gru.parameters());
268        params.extend(self.head_imminent.parameters());
269        params.extend(self.head_warning.parameters());
270        params.extend(self.head_early.parameters());
271        params
272    }
273}
274
275// =============================================================================
276// Helper Functions
277// =============================================================================
278
279/// Get argmax for each sample in batch
280fn argmax_batch(x: &Variable) -> Vec<usize> {
281    let data = x.data();
282    let shape = data.shape();
283    let batch_size = shape[0];
284    let num_classes = shape[1];
285    let values = data.to_vec();
286
287    let mut results = Vec::with_capacity(batch_size);
288    for b in 0..batch_size {
289        let start = b * num_classes;
290        let end = start + num_classes;
291        let slice = &values[start..end];
292
293        let mut max_idx = 0;
294        let mut max_val = slice[0];
295        for (i, &v) in slice.iter().enumerate() {
296            if v > max_val {
297                max_val = v;
298                max_idx = i;
299            }
300        }
301        results.push(max_idx);
302    }
303    results
304}
305
306/// Failure type names
307pub const FAILURE_TYPES: [&str; 20] = [
308    "normal",
309    "pump_failure_hw_5",
310    "pump_failure_hw_6",
311    "pump_failure_cw_3",
312    "pump_failure_cw_4",
313    "pump_failure_2pipe_a",
314    "pump_failure_2pipe_b",
315    "pressure_low_hw",
316    "pressure_high_hw",
317    "pressure_low_cw",
318    "pressure_high_cw",
319    "temp_anomaly_hw_supply",
320    "temp_anomaly_cw_supply",
321    "temp_anomaly_space",
322    "valve_stuck_1_3",
323    "valve_stuck_2_3",
324    "vfd_fault",
325    "sensor_drift",
326    "chiller_fault",
327    "interlock_violation",
328];
329
330/// Feature names for the 28 sensor inputs
331pub const FEATURE_NAMES: [&str; 28] = [
332    "hw_pump_5_current",
333    "hw_pump_6_current",
334    "cw_pump_3_current",
335    "cw_pump_4_current",
336    "2pipe_pump_a_current",
337    "2pipe_pump_b_current",
338    "hw_supply_4pipe_temp",
339    "cw_supply_4pipe_temp",
340    "hw_supply_2pipe_temp",
341    "cw_return_2pipe_temp",
342    "outdoor_air_temp",
343    "mech_room_temp",
344    "space_sensor_1_temp",
345    "space_sensor_2_temp",
346    "hw_pressure_4pipe",
347    "cw_pressure_4pipe",
348    "hw_pump_5_vfd_speed",
349    "hw_pump_6_vfd_speed",
350    "cw_pump_3_vfd_speed",
351    "cw_pump_4_vfd_speed",
352    "2pipe_pump_a_vfd_speed",
353    "2pipe_pump_b_vfd_speed",
354    "steam_valve_1_3_pos",
355    "steam_valve_2_3_pos",
356    "summer_winter_mode",
357    "hw_lead_pump_id",
358    "cw_lead_pump_id",
359    "2pipe_lead_pump_id",
360];
361
362// =============================================================================
363// Main
364// =============================================================================
365
366fn main() {
367    println!("╔════════════════════════════════════════════════════════════╗");
368    println!("║     HVAC Multi-Horizon Predictor - AxonML Native           ║");
369    println!("╚════════════════════════════════════════════════════════════╝");
370    println!();
371
372    // Create model with default config
373    let config = HvacConfig::default();
374    println!("Model Configuration:");
375    println!("  Input features: {}", config.num_features);
376    println!("  Sequence length: {}", config.seq_len);
377    println!("  Hidden size: {}", config.hidden_size);
378    println!("  GRU layers: {}", config.num_layers);
379    println!("  Output classes: {}", config.num_classes);
380    println!("  Dropout: {}", config.dropout);
381    println!();
382
383    let model = HvacPredictor::new(config.clone());
384    println!("Model created!");
385    println!("  Total parameters: {}", model.num_parameters());
386    println!();
387
388    // Create sample input
389    let batch_size = 2;
390    let mut input_data = vec![0.5f32; batch_size * config.seq_len * config.num_features];
391
392    // Simulate normal HVAC readings
393    for b in 0..batch_size {
394        for t in 0..config.seq_len {
395            let base = (b * config.seq_len + t) * config.num_features;
396            // Pump currents ~25A (normalized)
397            for i in 0..6 {
398                input_data[base + i] = 0.5;
399            }
400            // Temperatures (normalized)
401            input_data[base + 6] = 0.83; // HW supply ~180F
402            input_data[base + 7] = 0.375; // CW supply ~55F
403            // VFD speeds ~60%
404            for i in 16..22 {
405                input_data[base + i] = 0.6;
406            }
407        }
408    }
409
410    let input = Tensor::from_vec(
411        input_data,
412        &[batch_size, config.seq_len, config.num_features],
413    )
414    .expect("Failed to create input tensor");
415
416    let input_var = Variable::new(input, false);
417    println!("Input shape: {:?}", input_var.data().shape());
418
419    // Run inference
420    println!();
421    println!("Running inference...");
422    let (imminent, warning, early) = model.predict(&input_var);
423
424    println!();
425    println!("Predictions:");
426    println!("────────────────────────────────────────────────────────────");
427    for b in 0..batch_size {
428        println!("Sample {}:", b);
429        println!(
430            "  5 min (Imminent): {} - {}",
431            imminent[b], FAILURE_TYPES[imminent[b]]
432        );
433        println!(
434            "  15 min (Warning): {} - {}",
435            warning[b], FAILURE_TYPES[warning[b]]
436        );
437        println!(
438            "  30 min (Early):   {} - {}",
439            early[b], FAILURE_TYPES[early[b]]
440        );
441    }
442    println!("────────────────────────────────────────────────────────────");
443    println!();
444    println!("Model ready for training with your HVAC sensor data!");
445}