Files
rustytorch/demos/rtx-aeroflow-demo/src/geometry.rs
T
2026-03-04 00:08:42 +00:00

401 lines
13 KiB
Rust

//! Geometry encoding for neural operators.
use aeroflow_shared::{AirfoilGeometry, Point2D, WingGeometry};
/// Geometry encoder for converting airfoil/wing shapes to feature vectors.
#[derive(Debug)]
pub struct GeometryEncoder {
/// Hidden dimension for encoded features.
hidden_dim: usize,
/// Number of Chebyshev nodes for interpolation.
#[allow(dead_code)]
num_nodes: usize,
}
impl GeometryEncoder {
/// Create a new geometry encoder.
pub fn new(hidden_dim: usize) -> Self {
Self {
hidden_dim,
num_nodes: 64,
}
}
/// Encode an airfoil geometry to feature vector.
pub fn encode_airfoil(&self, airfoil: &AirfoilGeometry) -> Vec<f32> {
let mut features = Vec::with_capacity(self.hidden_dim);
// Global shape parameters
features.push(airfoil.max_thickness);
features.push(airfoil.max_camber);
features.push(airfoil.max_camber_position);
features.push(airfoil.chord);
// Encode surface shape using Chebyshev coefficients
let upper_coeffs = self.chebyshev_coefficients(&airfoil.upper_surface, 10);
let lower_coeffs = self.chebyshev_coefficients(&airfoil.lower_surface, 10);
features.extend(upper_coeffs);
features.extend(lower_coeffs);
// Curvature features
let upper_curvature = self.compute_curvature(&airfoil.upper_surface);
let lower_curvature = self.compute_curvature(&airfoil.lower_surface);
features.push(upper_curvature.iter().sum::<f32>() / upper_curvature.len() as f32);
features.push(lower_curvature.iter().sum::<f32>() / lower_curvature.len() as f32);
// Leading edge radius
features.push(self.estimate_le_radius(airfoil));
// Trailing edge angle
features.push(self.compute_te_angle(airfoil));
// Pad or truncate to hidden_dim
features.resize(self.hidden_dim, 0.0);
features
}
/// Encode a wing geometry to feature vector.
pub fn encode_wing(&self, wing: &WingGeometry) -> Vec<f32> {
let mut features = Vec::with_capacity(self.hidden_dim);
// Encode root airfoil
let root_features = self.encode_airfoil(&wing.root_airfoil);
features.extend(root_features.iter().take(20));
// Planform parameters
features.push(wing.span);
features.push(wing.root_chord);
features.push(wing.tip_chord);
features.push(wing.tip_chord / wing.root_chord); // Taper ratio
features.push(wing.sweep_angle.to_radians());
features.push(wing.dihedral_angle.to_radians());
features.push(wing.twist_angle.to_radians());
// Aspect ratio
let avg_chord = f32::midpoint(wing.root_chord, wing.tip_chord);
let wing_area = wing.span * avg_chord;
let aspect_ratio = wing.span.powi(2) / wing_area;
features.push(aspect_ratio);
// Pad to hidden_dim
features.resize(self.hidden_dim, 0.0);
features
}
/// Compute Chebyshev coefficients for surface representation.
fn chebyshev_coefficients(&self, points: &[Point2D], num_coeffs: usize) -> Vec<f32> {
if points.is_empty() {
return vec![0.0; num_coeffs];
}
let n = points.len();
let mut coeffs = vec![0.0; num_coeffs];
for k in 0..num_coeffs {
let mut sum = 0.0;
for (j, point) in points.iter().enumerate() {
let x = 2.0 * j as f32 / (n - 1) as f32 - 1.0; // Map to [-1, 1]
let tk = self.chebyshev_t(k, x);
sum += point.y * tk;
}
coeffs[k] = sum / n as f32;
}
coeffs
}
/// Chebyshev polynomial of the first kind.
fn chebyshev_t(&self, n: usize, x: f32) -> f32 {
match n {
0 => 1.0,
1 => x,
_ => 2.0 * x * self.chebyshev_t(n - 1, x) - self.chebyshev_t(n - 2, x),
}
}
/// Compute curvature along a surface.
fn compute_curvature(&self, points: &[Point2D]) -> Vec<f32> {
if points.len() < 3 {
return vec![0.0; points.len()];
}
let mut curvature = vec![0.0; points.len()];
for i in 1..points.len() - 1 {
let p0 = &points[i - 1];
let p1 = &points[i];
let p2 = &points[i + 1];
// First derivatives
let dx1 = p1.x - p0.x;
let dy1 = p1.y - p0.y;
let dx2 = p2.x - p1.x;
let dy2 = p2.y - p1.y;
// Second derivatives
let ddx = dx2 - dx1;
let ddy = dy2 - dy1;
// Curvature formula: κ = |x'y'' - y'x''| / (x'^2 + y'^2)^(3/2)
let dx_avg = f32::midpoint(dx1, dx2);
let dy_avg = f32::midpoint(dy1, dy2);
let denom = (dx_avg.powi(2) + dy_avg.powi(2)).powf(1.5).max(1e-10);
curvature[i] = (dx_avg * ddy - dy_avg * ddx).abs() / denom;
}
curvature
}
/// Estimate leading edge radius.
fn estimate_le_radius(&self, airfoil: &AirfoilGeometry) -> f32 {
// Approximate LE radius from thickness distribution
// For NACA 4-digit: r_LE ≈ 1.1019 * (t/c)^2
1.1019 * airfoil.max_thickness.powi(2)
}
/// Compute trailing edge angle.
fn compute_te_angle(&self, airfoil: &AirfoilGeometry) -> f32 {
if airfoil.upper_surface.len() < 2 || airfoil.lower_surface.len() < 2 {
return 0.0;
}
let upper_last = &airfoil.upper_surface[airfoil.upper_surface.len() - 1];
let upper_prev = &airfoil.upper_surface[airfoil.upper_surface.len() - 2];
let lower_last = &airfoil.lower_surface[airfoil.lower_surface.len() - 1];
let lower_prev = &airfoil.lower_surface[airfoil.lower_surface.len() - 2];
let upper_slope = (upper_last.y - upper_prev.y) / (upper_last.x - upper_prev.x).max(1e-6);
let lower_slope = (lower_last.y - lower_prev.y) / (lower_last.x - lower_prev.x).max(1e-6);
// TE angle is the angle between upper and lower surface slopes
(upper_slope.atan() - lower_slope.atan()).abs()
}
/// Interpolate surface at given x positions.
pub fn interpolate_surface(&self, points: &[Point2D], x_positions: &[f32]) -> Vec<f32> {
x_positions
.iter()
.map(|&x| {
// Find bracketing points
let mut y = 0.0;
for i in 1..points.len() {
if points[i].x >= x {
let t = (x - points[i - 1].x) / (points[i].x - points[i - 1].x).max(1e-10);
y = points[i - 1].y + t * (points[i].y - points[i - 1].y);
break;
}
}
y
})
.collect()
}
/// Compute area enclosed by airfoil.
pub fn compute_area(&self, airfoil: &AirfoilGeometry) -> f32 {
let mut area = 0.0;
// Shoelace formula for upper surface
for i in 1..airfoil.upper_surface.len() {
let p0 = &airfoil.upper_surface[i - 1];
let p1 = &airfoil.upper_surface[i];
area += (p1.x - p0.x) * (p0.y + p1.y) / 2.0;
}
// Subtract lower surface (traversed in reverse)
for i in 1..airfoil.lower_surface.len() {
let p0 = &airfoil.lower_surface[i - 1];
let p1 = &airfoil.lower_surface[i];
area -= (p1.x - p0.x) * (p0.y + p1.y) / 2.0;
}
area.abs()
}
}
/// CST (Class-Shape Transformation) parameterization.
#[derive(Debug, Clone)]
pub struct CSTAirfoil {
/// Class function exponents.
pub n1: f32,
pub n2: f32,
/// Upper surface coefficients.
pub upper_coeffs: Vec<f32>,
/// Lower surface coefficients.
pub lower_coeffs: Vec<f32>,
/// Trailing edge thickness.
pub te_thickness: f32,
}
impl Default for CSTAirfoil {
fn default() -> Self {
Self {
n1: 0.5,
n2: 1.0,
upper_coeffs: vec![0.2, 0.3, 0.2, 0.1, 0.1],
lower_coeffs: vec![-0.2, -0.1, -0.1, -0.05, -0.05],
te_thickness: 0.0,
}
}
}
impl CSTAirfoil {
/// Generate airfoil geometry from CST parameters.
pub fn to_airfoil(&self, num_points: usize) -> AirfoilGeometry {
let mut upper = Vec::with_capacity(num_points);
let mut lower = Vec::with_capacity(num_points);
for i in 0..num_points {
let psi = 1.0 - (std::f32::consts::PI * i as f32 / (num_points - 1) as f32).cos();
let psi = psi / 2.0; // Map to [0, 1]
// Class function
let class_func = psi.powf(self.n1) * (1.0 - psi).powf(self.n2);
// Shape functions (Bernstein polynomials)
let upper_shape = self.bernstein_sum(&self.upper_coeffs, psi);
let lower_shape = self.bernstein_sum(&self.lower_coeffs, psi);
let y_upper = class_func * upper_shape + psi * self.te_thickness / 2.0;
let y_lower = class_func * lower_shape - psi * self.te_thickness / 2.0;
upper.push(Point2D::new(psi, y_upper));
lower.push(Point2D::new(psi, y_lower));
}
let max_thickness = upper
.iter()
.zip(lower.iter())
.map(|(u, l)| u.y - l.y)
.fold(0.0f32, f32::max);
AirfoilGeometry {
name: "CST Airfoil".to_string(),
upper_surface: upper,
lower_surface: lower,
chord: 1.0,
max_thickness,
max_camber: 0.0,
max_camber_position: 0.0,
}
}
/// Evaluate Bernstein polynomial sum.
fn bernstein_sum(&self, coeffs: &[f32], t: f32) -> f32 {
let n = coeffs.len() - 1;
let mut sum = 0.0;
for (k, &coeff) in coeffs.iter().enumerate() {
sum += coeff * self.bernstein(n, k, t);
}
sum
}
/// Bernstein basis polynomial.
fn bernstein(&self, n: usize, k: usize, t: f32) -> f32 {
let binom = self.binomial(n, k) as f32;
binom * t.powi(k as i32) * (1.0 - t).powi((n - k) as i32)
}
/// Binomial coefficient.
fn binomial(&self, n: usize, k: usize) -> usize {
if k > n {
return 0;
}
let mut result = 1;
for i in 0..k {
result = result * (n - i) / (i + 1);
}
result
}
}
#[cfg(test)]
mod tests {
use super::*;
use aeroflow_shared::naca_4digit;
#[test]
fn test_encoder_creation() {
let encoder = GeometryEncoder::new(64);
assert_eq!(encoder.hidden_dim, 64);
}
#[test]
fn test_encode_airfoil() {
let encoder = GeometryEncoder::new(32);
let airfoil = naca_4digit("0012", 50).unwrap();
let features = encoder.encode_airfoil(&airfoil);
assert_eq!(features.len(), 32);
assert!((features[0] - 0.12).abs() < 0.001); // max_thickness
}
#[test]
fn test_encode_wing() {
let encoder = GeometryEncoder::new(64);
let wing = WingGeometry::default();
let features = encoder.encode_wing(&wing);
assert_eq!(features.len(), 64);
}
#[test]
fn test_chebyshev() {
let encoder = GeometryEncoder::new(32);
// T_0(x) = 1
assert!((encoder.chebyshev_t(0, 0.5) - 1.0).abs() < 0.001);
// T_1(x) = x
assert!((encoder.chebyshev_t(1, 0.5) - 0.5).abs() < 0.001);
// T_2(x) = 2x^2 - 1
assert!((encoder.chebyshev_t(2, 0.5) - (-0.5)).abs() < 0.001);
}
#[test]
fn test_curvature() {
let encoder = GeometryEncoder::new(32);
let points = vec![
Point2D::new(0.0, 0.0),
Point2D::new(0.5, 0.1),
Point2D::new(1.0, 0.0),
];
let curvature = encoder.compute_curvature(&points);
assert_eq!(curvature.len(), 3);
assert!(curvature[1] > 0.0); // Should have some curvature
}
#[test]
fn test_area() {
let encoder = GeometryEncoder::new(32);
let airfoil = naca_4digit("0012", 50).unwrap();
let area = encoder.compute_area(&airfoil);
assert!(area > 0.0);
assert!(area < 0.2); // Reasonable area for NACA 0012
}
#[test]
fn test_cst_airfoil() {
let cst = CSTAirfoil::default();
let airfoil = cst.to_airfoil(50);
assert_eq!(airfoil.upper_surface.len(), 50);
assert_eq!(airfoil.lower_surface.len(), 50);
assert!(airfoil.max_thickness > 0.0);
}
#[test]
fn test_bernstein() {
let cst = CSTAirfoil::default();
// B_{0,0}(0.5) = 1
assert!((cst.bernstein(0, 0, 0.5) - 1.0).abs() < 0.001);
// B_{2,1}(0.5) = 2 * 0.5 * 0.5 = 0.5
assert!((cst.bernstein(2, 1, 0.5) - 0.5).abs() < 0.001);
}
}