1use axonml::autograd::Variable;
18use axonml::nn::{Dropout, GRU, LayerNorm, Linear, Module, Parameter, ReLU, Softmax};
19use axonml::tensor::Tensor;
20
21#[derive(Debug, Clone)]
27pub struct HvacConfig {
28 pub num_features: usize,
30 pub seq_len: usize,
32 pub hidden_size: usize,
34 pub num_layers: usize,
36 pub num_classes: usize,
38 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
55pub 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
99pub struct HvacPredictor {
110 config: HvacConfig,
111
112 input_proj: Linear,
114 input_norm: LayerNorm,
115 input_relu: ReLU,
116
117 gru: GRU,
119
120 head_imminent: PredictionHead, head_warning: PredictionHead, head_early: PredictionHead, softmax: Softmax,
127}
128
129#[derive(Debug)]
131pub struct HvacOutput {
132 pub imminent_logits: Variable,
134 pub warning_logits: Variable,
136 pub early_logits: Variable,
138}
139
140impl HvacPredictor {
141 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 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 let values = data.to_vec();
174
175 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 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); 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 let encoded = self.gru.forward(&proj);
211
212 let pooled = self.mean_pool(&encoded);
214
215 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 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 pub fn config(&self) -> &HvacConfig {
244 &self.config
245 }
246
247 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 let output = self.forward_multi(x);
260 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
275fn 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
306pub 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
330pub 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
362fn main() {
367 println!("╔════════════════════════════════════════════════════════════╗");
368 println!("║ HVAC Multi-Horizon Predictor - AxonML Native ║");
369 println!("╚════════════════════════════════════════════════════════════╝");
370 println!();
371
372 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 let batch_size = 2;
390 let mut input_data = vec![0.5f32; batch_size * config.seq_len * config.num_features];
391
392 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 for i in 0..6 {
398 input_data[base + i] = 0.5;
399 }
400 input_data[base + 6] = 0.83; input_data[base + 7] = 0.375; 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 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}