Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,405 @@
|
||||
//! Neural operator inference engine
|
||||
//!
|
||||
//! Provides the main demo functionality for FNO inference.
|
||||
|
||||
use std::path::Path;
|
||||
use std::time::Instant;
|
||||
|
||||
use rtx_backend_cpu::{CpuBackend, CpuDevice};
|
||||
use rtx_neural_operator::{FNO2d, load_fno2d_weights};
|
||||
use rtx_neural_operator_shared::{
|
||||
config::{PDEConfig, PDEType},
|
||||
error::{NeuralOperatorError, Result},
|
||||
ipc::{ModelInfo, PerformanceMetrics, SolutionData},
|
||||
};
|
||||
use rtx_nn::GenericModule4D;
|
||||
use rtx_tensor::GenericTensor;
|
||||
|
||||
/// Result from a solve operation
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SolveResult {
|
||||
/// Solution field (flattened [H, W])
|
||||
pub solution: Vec<f32>,
|
||||
/// Inference time in milliseconds
|
||||
pub inference_time_ms: f64,
|
||||
}
|
||||
|
||||
/// Neural Operator Demo engine
|
||||
///
|
||||
/// Manages FNO model loading and inference for interactive PDE solving.
|
||||
pub struct NeuralOperatorDemo {
|
||||
model: Option<FNO2d<CpuBackend>>,
|
||||
config: Option<PDEConfig>,
|
||||
metrics: PerformanceMetrics,
|
||||
last_solution: Option<Vec<f32>>,
|
||||
}
|
||||
|
||||
impl NeuralOperatorDemo {
|
||||
/// Creates a new demo instance
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
model: None,
|
||||
config: None,
|
||||
metrics: PerformanceMetrics::new(),
|
||||
last_solution: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Initializes the demo with a PDE configuration
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `config` - PDE configuration specifying type and resolution
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns error if weight file not found or model loading fails.
|
||||
pub fn initialize(&mut self, config: PDEConfig) -> Result<()> {
|
||||
tracing::info!(
|
||||
"Initializing neural operator for {:?} at {}x{} resolution",
|
||||
config.pde_type,
|
||||
config.resolution,
|
||||
config.resolution
|
||||
);
|
||||
|
||||
// For now, we'll create a model with default weights
|
||||
// In production, this would load from weight files
|
||||
self.config = Some(config);
|
||||
self.metrics = PerformanceMetrics::new();
|
||||
self.last_solution = None;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Initializes with weights from a file path
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `config` - PDE configuration
|
||||
/// * `weights_path` - Path to `SafeTensors` weight file
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns error if weight loading fails.
|
||||
pub fn initialize_with_weights(
|
||||
&mut self,
|
||||
config: PDEConfig,
|
||||
weights_path: impl AsRef<Path>,
|
||||
) -> Result<()> {
|
||||
tracing::info!("Loading FNO weights from {:?}", weights_path.as_ref());
|
||||
|
||||
let weights = load_fno2d_weights(weights_path.as_ref())
|
||||
.map_err(|e| NeuralOperatorError::weight_load(e.to_string()))?;
|
||||
|
||||
let device = CpuDevice::default();
|
||||
let model = FNO2d::from_weights(&weights, &device)
|
||||
.map_err(|e| NeuralOperatorError::weight_load(e.to_string()))?;
|
||||
|
||||
self.model = Some(model);
|
||||
self.config = Some(config);
|
||||
self.metrics = PerformanceMetrics::new();
|
||||
self.last_solution = None;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Solves the PDE with the given input field
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `input` - Input field (flattened [H, W] array)
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns error if model not initialized or inference fails.
|
||||
pub fn solve(&mut self, input: &[f32]) -> Result<SolveResult> {
|
||||
let config = self
|
||||
.config
|
||||
.as_ref()
|
||||
.ok_or(NeuralOperatorError::NotInitialized)?;
|
||||
|
||||
let expected_size = (config.resolution * config.resolution) as usize;
|
||||
if input.len() != expected_size {
|
||||
return Err(NeuralOperatorError::invalid_dimensions(
|
||||
format!("[{expected_size}]"),
|
||||
format!("[{}]", input.len()),
|
||||
));
|
||||
}
|
||||
|
||||
let start = Instant::now();
|
||||
|
||||
// If we have a loaded model, use it
|
||||
let solution = if let Some(ref model) = self.model {
|
||||
// Create input tensor [1, 1, H, W]
|
||||
let h = config.resolution as usize;
|
||||
let w = config.resolution as usize;
|
||||
let device = CpuDevice::default();
|
||||
|
||||
let tensor: GenericTensor<CpuBackend, 4> =
|
||||
GenericTensor::from_slice(input, [1, 1, h, w], &device);
|
||||
|
||||
// Run inference
|
||||
let output = model.forward_4d(&tensor);
|
||||
|
||||
// Extract result
|
||||
output.to_vec()
|
||||
} else {
|
||||
// No model loaded - return a simple demo output
|
||||
// This is a placeholder that shows the demo UI works
|
||||
generate_demo_solution(input, config)
|
||||
};
|
||||
|
||||
let inference_time_ms = start.elapsed().as_secs_f64() * 1000.0;
|
||||
|
||||
// Record metrics
|
||||
self.metrics.record_inference(inference_time_ms);
|
||||
|
||||
// Cache solution
|
||||
self.last_solution = Some(solution.clone());
|
||||
|
||||
Ok(SolveResult {
|
||||
solution,
|
||||
inference_time_ms,
|
||||
})
|
||||
}
|
||||
|
||||
/// Creates solution data from a solve result
|
||||
#[must_use]
|
||||
pub fn create_solution_data(&self, result: &SolveResult) -> Option<SolutionData> {
|
||||
let config = self.config.as_ref()?;
|
||||
Some(SolutionData::new(
|
||||
result.solution.clone(),
|
||||
config.resolution,
|
||||
config.resolution,
|
||||
result.inference_time_ms,
|
||||
))
|
||||
}
|
||||
|
||||
/// Returns the last computed solution
|
||||
#[must_use]
|
||||
pub fn last_solution(&self) -> Option<&Vec<f32>> {
|
||||
self.last_solution.as_ref()
|
||||
}
|
||||
|
||||
/// Returns current performance metrics
|
||||
#[must_use]
|
||||
pub fn metrics(&self) -> &PerformanceMetrics {
|
||||
&self.metrics
|
||||
}
|
||||
|
||||
/// Returns model information
|
||||
#[must_use]
|
||||
pub fn model_info(&self) -> Option<ModelInfo> {
|
||||
let config = self.config.as_ref()?;
|
||||
Some(ModelInfo::new(
|
||||
format!("FNO2d-{}", config.pde_type.name()),
|
||||
config.pde_type.name(),
|
||||
config.resolution,
|
||||
config.n_modes,
|
||||
config.model_width,
|
||||
config.n_layers,
|
||||
))
|
||||
}
|
||||
|
||||
/// Returns the current configuration
|
||||
#[must_use]
|
||||
pub fn config(&self) -> Option<&PDEConfig> {
|
||||
self.config.as_ref()
|
||||
}
|
||||
|
||||
/// Returns whether the model is initialized
|
||||
#[must_use]
|
||||
pub fn is_initialized(&self) -> bool {
|
||||
self.config.is_some()
|
||||
}
|
||||
|
||||
/// Returns whether a trained model is loaded
|
||||
#[must_use]
|
||||
pub fn has_model(&self) -> bool {
|
||||
self.model.is_some()
|
||||
}
|
||||
|
||||
/// Resets the demo to initial state
|
||||
pub fn reset(&mut self) {
|
||||
self.model = None;
|
||||
self.config = None;
|
||||
self.metrics = PerformanceMetrics::new();
|
||||
self.last_solution = None;
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for NeuralOperatorDemo {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Generates a demo solution when no model is loaded
|
||||
///
|
||||
/// This creates a visually interesting output that demonstrates the UI
|
||||
/// without requiring a trained model.
|
||||
fn generate_demo_solution(input: &[f32], config: &PDEConfig) -> Vec<f32> {
|
||||
let n = config.resolution as usize;
|
||||
let mut solution = vec![0.0; n * n];
|
||||
|
||||
// Create a simple diffusion-like response to the input
|
||||
// This mimics what a PDE solver would produce
|
||||
for i in 0..n {
|
||||
for j in 0..n {
|
||||
let idx = i * n + j;
|
||||
let x = j as f32 / n as f32;
|
||||
let y = i as f32 / n as f32;
|
||||
|
||||
// Combine input with a smooth basis function
|
||||
let input_val = input[idx];
|
||||
let smooth = (std::f32::consts::PI * x).sin() * (std::f32::consts::PI * y).sin();
|
||||
|
||||
// Different responses for different PDE types
|
||||
let response = match config.pde_type {
|
||||
PDEType::DarcyFlow => {
|
||||
// Pressure-like response
|
||||
input_val * smooth * 0.5 + (1.0 - x) * 0.3
|
||||
}
|
||||
PDEType::HeatEquation => {
|
||||
// Temperature diffusion
|
||||
input_val * smooth.powi(2) * 0.8
|
||||
}
|
||||
PDEType::Poisson => {
|
||||
// Potential field
|
||||
input_val * smooth * 0.6 + smooth * 0.2
|
||||
}
|
||||
PDEType::NavierStokes => {
|
||||
// Velocity-like field
|
||||
let vortex = ((x - 0.5).powi(2) + (y - 0.5).powi(2)).sqrt();
|
||||
input_val * (1.0 - vortex).max(0.0) * 0.7
|
||||
}
|
||||
};
|
||||
|
||||
solution[idx] = response;
|
||||
}
|
||||
}
|
||||
|
||||
solution
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_demo_creation() {
|
||||
let demo = NeuralOperatorDemo::new();
|
||||
assert!(!demo.is_initialized());
|
||||
assert!(!demo.has_model());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_initialize() {
|
||||
let mut demo = NeuralOperatorDemo::new();
|
||||
let config = PDEConfig::darcy(64);
|
||||
demo.initialize(config).unwrap();
|
||||
|
||||
assert!(demo.is_initialized());
|
||||
assert!(!demo.has_model()); // No weights loaded
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_solve_without_model() {
|
||||
let mut demo = NeuralOperatorDemo::new();
|
||||
let config = PDEConfig::darcy(32);
|
||||
demo.initialize(config).unwrap();
|
||||
|
||||
let input = vec![1.0; 32 * 32];
|
||||
let result = demo.solve(&input).unwrap();
|
||||
|
||||
assert_eq!(result.solution.len(), 32 * 32);
|
||||
assert!(result.inference_time_ms > 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_solve_dimension_mismatch() {
|
||||
let mut demo = NeuralOperatorDemo::new();
|
||||
let config = PDEConfig::darcy(64);
|
||||
demo.initialize(config).unwrap();
|
||||
|
||||
let input = vec![1.0; 32 * 32]; // Wrong size
|
||||
let result = demo.solve(&input);
|
||||
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_solve_not_initialized() {
|
||||
let mut demo = NeuralOperatorDemo::new();
|
||||
let input = vec![1.0; 64 * 64];
|
||||
let result = demo.solve(&input);
|
||||
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_metrics_recording() {
|
||||
let mut demo = NeuralOperatorDemo::new();
|
||||
let config = PDEConfig::darcy(16);
|
||||
demo.initialize(config).unwrap();
|
||||
|
||||
let input = vec![1.0; 16 * 16];
|
||||
demo.solve(&input).unwrap();
|
||||
demo.solve(&input).unwrap();
|
||||
|
||||
let metrics = demo.metrics();
|
||||
assert_eq!(metrics.inference_count, 2);
|
||||
assert!(metrics.avg_inference_time_ms > 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_model_info() {
|
||||
let mut demo = NeuralOperatorDemo::new();
|
||||
let config = PDEConfig::darcy(64)
|
||||
.with_modes(12, 12)
|
||||
.with_width(32)
|
||||
.with_layers(4);
|
||||
demo.initialize(config).unwrap();
|
||||
|
||||
let info = demo.model_info().unwrap();
|
||||
assert_eq!(info.resolution, 64);
|
||||
assert_eq!(info.n_modes, (12, 12));
|
||||
assert_eq!(info.model_width, 32);
|
||||
assert_eq!(info.n_layers, 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reset() {
|
||||
let mut demo = NeuralOperatorDemo::new();
|
||||
let config = PDEConfig::darcy(64);
|
||||
demo.initialize(config).unwrap();
|
||||
|
||||
let input = vec![1.0; 64 * 64];
|
||||
demo.solve(&input).unwrap();
|
||||
|
||||
demo.reset();
|
||||
|
||||
assert!(!demo.is_initialized());
|
||||
assert!(demo.last_solution().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_different_pde_types() {
|
||||
let mut demo = NeuralOperatorDemo::new();
|
||||
let input = vec![1.0; 16 * 16];
|
||||
|
||||
for pde_type in [
|
||||
PDEType::DarcyFlow,
|
||||
PDEType::HeatEquation,
|
||||
PDEType::Poisson,
|
||||
PDEType::NavierStokes,
|
||||
] {
|
||||
let config = PDEConfig::new(pde_type, 16);
|
||||
demo.initialize(config).unwrap();
|
||||
|
||||
let result = demo.solve(&input).unwrap();
|
||||
assert_eq!(result.solution.len(), 16 * 16);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user