340 lines
10 KiB
Rust
340 lines
10 KiB
Rust
//! Fast inference engine for temperature field visualization
|
|
//!
|
|
//! Provides optimized methods for querying the trained PINN
|
|
//! at many points for real-time visualization.
|
|
|
|
use crate::network::{Coordinate4D, ThermalPinn};
|
|
use bioheat_shared::{BoundingBox3D, Point3D, SliceAxis, SliceData, TemperatureField};
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// Inference engine for fast temperature queries
|
|
#[derive(Debug, Clone)]
|
|
pub struct InferenceEngine {
|
|
/// The trained network
|
|
pinn: ThermalPinn,
|
|
/// Domain bounds
|
|
bounds: BoundingBox3D,
|
|
}
|
|
|
|
impl InferenceEngine {
|
|
/// Create a new inference engine from a trained PINN
|
|
#[must_use]
|
|
pub fn new(pinn: ThermalPinn, bounds: BoundingBox3D) -> Self {
|
|
Self { pinn, bounds }
|
|
}
|
|
|
|
/// Query temperature at a single point
|
|
#[must_use]
|
|
pub fn query_point(&self, point: &Point3D, t: f32) -> f32 {
|
|
let coord = Coordinate4D::from_point_and_time(*point, t);
|
|
self.pinn.forward(&coord) as f32
|
|
}
|
|
|
|
/// Query temperature at multiple points (same time)
|
|
#[must_use]
|
|
pub fn query_points(&self, points: &[Point3D], t: f32) -> Vec<f32> {
|
|
points.iter().map(|p| self.query_point(p, t)).collect()
|
|
}
|
|
|
|
/// Generate a 3D temperature field on a regular grid
|
|
#[must_use]
|
|
pub fn generate_field(&self, resolution: (usize, usize, usize), t: f32) -> TemperatureField {
|
|
let (nx, ny, nz) = resolution;
|
|
let size = self.bounds.size();
|
|
|
|
let dx = size.x / (nx - 1).max(1) as f32;
|
|
let dy = size.y / (ny - 1).max(1) as f32;
|
|
let dz = size.z / (nz - 1).max(1) as f32;
|
|
|
|
let mut values = Vec::with_capacity(nx * ny * nz);
|
|
|
|
for k in 0..nz {
|
|
for j in 0..ny {
|
|
for i in 0..nx {
|
|
let point = Point3D::new(
|
|
self.bounds.min.x + i as f32 * dx,
|
|
self.bounds.min.y + j as f32 * dy,
|
|
self.bounds.min.z + k as f32 * dz,
|
|
);
|
|
values.push(self.query_point(&point, t));
|
|
}
|
|
}
|
|
}
|
|
|
|
TemperatureField {
|
|
resolution,
|
|
values,
|
|
bounds: self.bounds,
|
|
}
|
|
}
|
|
|
|
/// Generate a 2D slice at the specified position
|
|
#[must_use]
|
|
pub fn generate_slice(
|
|
&self,
|
|
axis: SliceAxis,
|
|
position: f32,
|
|
resolution: (usize, usize),
|
|
t: f32,
|
|
) -> SliceData {
|
|
let (nu, nv) = resolution;
|
|
let size = self.bounds.size();
|
|
|
|
// Determine the slice bounds based on axis
|
|
let (u_range, v_range, pos_clamped) = match axis {
|
|
SliceAxis::X => {
|
|
let pos = position.clamp(self.bounds.min.x, self.bounds.max.x);
|
|
(
|
|
(self.bounds.min.y, self.bounds.max.y),
|
|
(self.bounds.min.z, self.bounds.max.z),
|
|
pos,
|
|
)
|
|
}
|
|
SliceAxis::Y => {
|
|
let pos = position.clamp(self.bounds.min.y, self.bounds.max.y);
|
|
(
|
|
(self.bounds.min.x, self.bounds.max.x),
|
|
(self.bounds.min.z, self.bounds.max.z),
|
|
pos,
|
|
)
|
|
}
|
|
SliceAxis::Z => {
|
|
let pos = position.clamp(self.bounds.min.z, self.bounds.max.z);
|
|
(
|
|
(self.bounds.min.x, self.bounds.max.x),
|
|
(self.bounds.min.y, self.bounds.max.y),
|
|
pos,
|
|
)
|
|
}
|
|
};
|
|
|
|
let du = (u_range.1 - u_range.0) / (nu - 1).max(1) as f32;
|
|
let dv = (v_range.1 - v_range.0) / (nv - 1).max(1) as f32;
|
|
|
|
let mut values = Vec::with_capacity(nu * nv);
|
|
|
|
for iv in 0..nv {
|
|
for iu in 0..nu {
|
|
let u = u_range.0 + iu as f32 * du;
|
|
let v = v_range.0 + iv as f32 * dv;
|
|
|
|
let point = match axis {
|
|
SliceAxis::X => Point3D::new(pos_clamped, u, v),
|
|
SliceAxis::Y => Point3D::new(u, pos_clamped, v),
|
|
SliceAxis::Z => Point3D::new(u, v, pos_clamped),
|
|
};
|
|
|
|
values.push(self.query_point(&point, t));
|
|
}
|
|
}
|
|
|
|
// Calculate index from position
|
|
let index = match axis {
|
|
SliceAxis::X => ((pos_clamped - self.bounds.min.x) / size.x * 100.0) as usize,
|
|
SliceAxis::Y => ((pos_clamped - self.bounds.min.y) / size.y * 100.0) as usize,
|
|
SliceAxis::Z => ((pos_clamped - self.bounds.min.z) / size.z * 100.0) as usize,
|
|
};
|
|
|
|
SliceData {
|
|
axis,
|
|
position: pos_clamped,
|
|
index,
|
|
resolution,
|
|
values,
|
|
bounds_2d: (u_range.0, u_range.1, v_range.0, v_range.1),
|
|
}
|
|
}
|
|
|
|
/// Generate all three orthogonal slices at a point
|
|
#[must_use]
|
|
pub fn generate_orthogonal_slices(
|
|
&self,
|
|
center: &Point3D,
|
|
resolution: (usize, usize),
|
|
t: f32,
|
|
) -> (SliceData, SliceData, SliceData) {
|
|
let x_slice = self.generate_slice(SliceAxis::X, center.x, resolution, t);
|
|
let y_slice = self.generate_slice(SliceAxis::Y, center.y, resolution, t);
|
|
let z_slice = self.generate_slice(SliceAxis::Z, center.z, resolution, t);
|
|
(x_slice, y_slice, z_slice)
|
|
}
|
|
|
|
/// Find the location of maximum temperature
|
|
#[must_use]
|
|
pub fn find_max_temperature(
|
|
&self,
|
|
resolution: (usize, usize, usize),
|
|
t: f32,
|
|
) -> (Point3D, f32) {
|
|
let field = self.generate_field(resolution, t);
|
|
let (nx, ny, _nz) = resolution;
|
|
|
|
let mut max_temp = f32::NEG_INFINITY;
|
|
let mut max_idx = 0;
|
|
|
|
for (idx, &temp) in field.values.iter().enumerate() {
|
|
if temp > max_temp {
|
|
max_temp = temp;
|
|
max_idx = idx;
|
|
}
|
|
}
|
|
|
|
// Convert index back to (i, j, k)
|
|
let i = max_idx % nx;
|
|
let j = (max_idx / nx) % ny;
|
|
let k = max_idx / (nx * ny);
|
|
|
|
let point = field.coords_at(i, j, k);
|
|
(point, max_temp)
|
|
}
|
|
|
|
/// Query temperature along a line (for 1D profile plots)
|
|
#[must_use]
|
|
pub fn query_line(
|
|
&self,
|
|
start: &Point3D,
|
|
end: &Point3D,
|
|
num_points: usize,
|
|
t: f32,
|
|
) -> Vec<(f32, f32)> {
|
|
let mut results = Vec::with_capacity(num_points);
|
|
|
|
for i in 0..num_points {
|
|
let frac = i as f32 / (num_points - 1).max(1) as f32;
|
|
let point = Point3D::new(
|
|
start.x + frac * (end.x - start.x),
|
|
start.y + frac * (end.y - start.y),
|
|
start.z + frac * (end.z - start.z),
|
|
);
|
|
|
|
let distance = start.distance(&point);
|
|
let temp = self.query_point(&point, t);
|
|
results.push((distance, temp));
|
|
}
|
|
|
|
results
|
|
}
|
|
|
|
/// Get the bounds
|
|
#[must_use]
|
|
pub fn bounds(&self) -> &BoundingBox3D {
|
|
&self.bounds
|
|
}
|
|
}
|
|
|
|
/// Performance metrics for inference
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct InferenceMetrics {
|
|
/// Number of points queried
|
|
pub num_points: usize,
|
|
/// Total time in milliseconds
|
|
pub total_time_ms: f32,
|
|
/// Time per point in microseconds
|
|
pub time_per_point_us: f32,
|
|
/// Throughput in points per second
|
|
pub throughput: f32,
|
|
}
|
|
|
|
impl InferenceMetrics {
|
|
/// Create from measurement
|
|
#[must_use]
|
|
pub fn from_measurement(num_points: usize, elapsed_ms: f32) -> Self {
|
|
let time_per_point_us = if num_points > 0 {
|
|
(elapsed_ms * 1000.0) / num_points as f32
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
let throughput = if elapsed_ms > 0.0 {
|
|
num_points as f32 / (elapsed_ms / 1000.0)
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
Self {
|
|
num_points,
|
|
total_time_ms: elapsed_ms,
|
|
time_per_point_us,
|
|
throughput,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::config::NetworkConfig;
|
|
|
|
fn create_test_engine() -> InferenceEngine {
|
|
let pinn = ThermalPinn::new(NetworkConfig {
|
|
fourier_features: 16,
|
|
hidden_layers: vec![32, 32],
|
|
..Default::default()
|
|
});
|
|
let bounds = BoundingBox3D::from_dimensions(0.1, 0.1, 0.1);
|
|
InferenceEngine::new(pinn, bounds)
|
|
}
|
|
|
|
#[test]
|
|
fn test_query_point() {
|
|
let engine = create_test_engine();
|
|
let temp = engine.query_point(&Point3D::origin(), 0.0);
|
|
assert!(temp.is_finite());
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_field() {
|
|
let engine = create_test_engine();
|
|
let field = engine.generate_field((8, 8, 8), 100.0);
|
|
assert_eq!(field.len(), 512);
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_slice() {
|
|
let engine = create_test_engine();
|
|
let slice = engine.generate_slice(SliceAxis::Z, 0.0, (32, 32), 50.0);
|
|
assert_eq!(slice.resolution, (32, 32));
|
|
assert_eq!(slice.values.len(), 1024);
|
|
}
|
|
|
|
#[test]
|
|
fn test_orthogonal_slices() {
|
|
let engine = create_test_engine();
|
|
let (x_slice, y_slice, z_slice) =
|
|
engine.generate_orthogonal_slices(&Point3D::origin(), (16, 16), 0.0);
|
|
|
|
assert_eq!(x_slice.axis, SliceAxis::X);
|
|
assert_eq!(y_slice.axis, SliceAxis::Y);
|
|
assert_eq!(z_slice.axis, SliceAxis::Z);
|
|
}
|
|
|
|
#[test]
|
|
fn test_query_line() {
|
|
let engine = create_test_engine();
|
|
let start = Point3D::new(-0.05, 0.0, 0.0);
|
|
let end = Point3D::new(0.05, 0.0, 0.0);
|
|
|
|
let profile = engine.query_line(&start, &end, 10, 0.0);
|
|
assert_eq!(profile.len(), 10);
|
|
|
|
// First point should be at distance 0
|
|
assert!(profile[0].0.abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_find_max() {
|
|
let engine = create_test_engine();
|
|
let (point, temp) = engine.find_max_temperature((8, 8, 8), 0.0);
|
|
|
|
assert!(engine.bounds().contains(&point));
|
|
assert!(temp.is_finite());
|
|
}
|
|
|
|
#[test]
|
|
fn test_metrics() {
|
|
let metrics = InferenceMetrics::from_measurement(1000, 100.0);
|
|
assert_eq!(metrics.num_points, 1000);
|
|
assert!((metrics.throughput - 10000.0).abs() < 100.0);
|
|
}
|
|
}
|