565 lines
19 KiB
Rust
565 lines
19 KiB
Rust
//! Icosahedral mesh encoder for weather prediction.
|
|
//!
|
|
//! This module implements the mesh structure and encoding/decoding
|
|
//! for converting between regular grid data and icosahedral mesh
|
|
//! representation used in GraphCast-style weather models.
|
|
|
|
use weathercast_shared::{AtmosphericState, MeshConfig};
|
|
|
|
// ============================================================================
|
|
// Icosahedral Mesh Structure
|
|
// ============================================================================
|
|
|
|
/// Icosahedral mesh for global weather modeling.
|
|
#[derive(Debug, Clone)]
|
|
pub struct IcosahedralMesh {
|
|
/// Mesh vertices as (x, y, z) on unit sphere.
|
|
pub vertices: Vec<(f32, f32, f32)>,
|
|
/// Mesh edges as vertex index pairs.
|
|
pub edges: Vec<(usize, usize)>,
|
|
/// Mesh faces as vertex index triplets.
|
|
pub faces: Vec<(usize, usize, usize)>,
|
|
/// Refinement level.
|
|
pub refinement_level: usize,
|
|
/// Average edge length in km.
|
|
pub avg_edge_length: f32,
|
|
}
|
|
|
|
impl IcosahedralMesh {
|
|
/// Create a new icosahedral mesh with given configuration.
|
|
pub fn new(config: &MeshConfig) -> Self {
|
|
let (vertices, edges, faces) = Self::generate_icosahedral_mesh(config.refinement_levels);
|
|
|
|
Self {
|
|
vertices,
|
|
edges,
|
|
faces,
|
|
refinement_level: config.refinement_levels,
|
|
avg_edge_length: config.avg_edge_length,
|
|
}
|
|
}
|
|
|
|
/// Generate icosahedral mesh at given refinement level.
|
|
fn generate_icosahedral_mesh(
|
|
refinement: usize,
|
|
) -> (
|
|
Vec<(f32, f32, f32)>,
|
|
Vec<(usize, usize)>,
|
|
Vec<(usize, usize, usize)>,
|
|
) {
|
|
// Start with base icosahedron (12 vertices, 20 faces, 30 edges)
|
|
let phi = f32::midpoint(1.0, 5.0_f32.sqrt()); // Golden ratio
|
|
|
|
let mut vertices = vec![
|
|
(-1.0, phi, 0.0),
|
|
(1.0, phi, 0.0),
|
|
(-1.0, -phi, 0.0),
|
|
(1.0, -phi, 0.0),
|
|
(0.0, -1.0, phi),
|
|
(0.0, 1.0, phi),
|
|
(0.0, -1.0, -phi),
|
|
(0.0, 1.0, -phi),
|
|
(phi, 0.0, -1.0),
|
|
(phi, 0.0, 1.0),
|
|
(-phi, 0.0, -1.0),
|
|
(-phi, 0.0, 1.0),
|
|
];
|
|
|
|
// Normalize to unit sphere
|
|
for v in &mut vertices {
|
|
let len = (v.0 * v.0 + v.1 * v.1 + v.2 * v.2).sqrt();
|
|
v.0 /= len;
|
|
v.1 /= len;
|
|
v.2 /= len;
|
|
}
|
|
|
|
// Base faces (20 triangular faces of icosahedron)
|
|
let mut faces = vec![
|
|
(0, 11, 5),
|
|
(0, 5, 1),
|
|
(0, 1, 7),
|
|
(0, 7, 10),
|
|
(0, 10, 11),
|
|
(1, 5, 9),
|
|
(5, 11, 4),
|
|
(11, 10, 2),
|
|
(10, 7, 6),
|
|
(7, 1, 8),
|
|
(3, 9, 4),
|
|
(3, 4, 2),
|
|
(3, 2, 6),
|
|
(3, 6, 8),
|
|
(3, 8, 9),
|
|
(4, 9, 5),
|
|
(2, 4, 11),
|
|
(6, 2, 10),
|
|
(8, 6, 7),
|
|
(9, 8, 1),
|
|
];
|
|
|
|
// Refine mesh
|
|
for _ in 0..refinement {
|
|
let (new_vertices, new_faces) = Self::subdivide(&vertices, &faces);
|
|
vertices = new_vertices;
|
|
faces = new_faces;
|
|
}
|
|
|
|
// Generate edges from faces
|
|
let mut edges = Vec::new();
|
|
let mut edge_set = std::collections::HashSet::new();
|
|
|
|
for &(a, b, c) in &faces {
|
|
for (v1, v2) in [(a, b), (b, c), (c, a)] {
|
|
let edge = if v1 < v2 { (v1, v2) } else { (v2, v1) };
|
|
if edge_set.insert(edge) {
|
|
edges.push(edge);
|
|
}
|
|
}
|
|
}
|
|
|
|
(vertices, edges, faces)
|
|
}
|
|
|
|
/// Subdivide mesh faces (each triangle becomes 4 triangles).
|
|
fn subdivide(
|
|
vertices: &[(f32, f32, f32)],
|
|
faces: &[(usize, usize, usize)],
|
|
) -> (Vec<(f32, f32, f32)>, Vec<(usize, usize, usize)>) {
|
|
let mut new_vertices = vertices.to_vec();
|
|
let mut new_faces = Vec::new();
|
|
let mut midpoint_cache: std::collections::HashMap<(usize, usize), usize> =
|
|
std::collections::HashMap::new();
|
|
|
|
let get_midpoint = |v1: usize,
|
|
v2: usize,
|
|
vertices: &[(f32, f32, f32)],
|
|
new_vertices: &mut Vec<(f32, f32, f32)>,
|
|
cache: &mut std::collections::HashMap<(usize, usize), usize>|
|
|
-> usize {
|
|
let key = if v1 < v2 { (v1, v2) } else { (v2, v1) };
|
|
|
|
if let Some(&idx) = cache.get(&key) {
|
|
return idx;
|
|
}
|
|
|
|
let p1 = vertices[v1];
|
|
let p2 = vertices[v2];
|
|
|
|
// Calculate midpoint and project to sphere
|
|
let mut mid = (
|
|
f32::midpoint(p1.0, p2.0),
|
|
f32::midpoint(p1.1, p2.1),
|
|
f32::midpoint(p1.2, p2.2),
|
|
);
|
|
let len = (mid.0 * mid.0 + mid.1 * mid.1 + mid.2 * mid.2).sqrt();
|
|
mid.0 /= len;
|
|
mid.1 /= len;
|
|
mid.2 /= len;
|
|
|
|
let idx = new_vertices.len();
|
|
new_vertices.push(mid);
|
|
cache.insert(key, idx);
|
|
idx
|
|
};
|
|
|
|
for &(a, b, c) in faces {
|
|
let ab = get_midpoint(a, b, vertices, &mut new_vertices, &mut midpoint_cache);
|
|
let bc = get_midpoint(b, c, vertices, &mut new_vertices, &mut midpoint_cache);
|
|
let ca = get_midpoint(c, a, vertices, &mut new_vertices, &mut midpoint_cache);
|
|
|
|
new_faces.push((a, ab, ca));
|
|
new_faces.push((b, bc, ab));
|
|
new_faces.push((c, ca, bc));
|
|
new_faces.push((ab, bc, ca));
|
|
}
|
|
|
|
(new_vertices, new_faces)
|
|
}
|
|
|
|
/// Get vertex latitude/longitude.
|
|
pub fn vertex_latlon(&self, idx: usize) -> Option<(f32, f32)> {
|
|
self.vertices.get(idx).map(|&(x, y, z)| {
|
|
let lat = y.asin().to_degrees();
|
|
let lon = x.atan2(z).to_degrees();
|
|
(lat, lon)
|
|
})
|
|
}
|
|
|
|
/// Find nearest vertex to given lat/lon.
|
|
pub fn nearest_vertex(&self, lat: f32, lon: f32) -> usize {
|
|
let lat_rad = lat.to_radians();
|
|
let lon_rad = lon.to_radians();
|
|
|
|
// Convert to Cartesian
|
|
let x = lat_rad.cos() * lon_rad.sin();
|
|
let y = lat_rad.sin();
|
|
let z = lat_rad.cos() * lon_rad.cos();
|
|
|
|
// Find nearest vertex
|
|
self.vertices
|
|
.iter()
|
|
.enumerate()
|
|
.min_by(|(_, v1), (_, v2)| {
|
|
let d1 = (v1.0 - x).powi(2) + (v1.1 - y).powi(2) + (v1.2 - z).powi(2);
|
|
let d2 = (v2.0 - x).powi(2) + (v2.1 - y).powi(2) + (v2.2 - z).powi(2);
|
|
d1.partial_cmp(&d2).unwrap_or(std::cmp::Ordering::Equal)
|
|
})
|
|
.map_or(0, |(i, _)| i)
|
|
}
|
|
|
|
/// Get neighbors of a vertex.
|
|
pub fn vertex_neighbors(&self, idx: usize) -> Vec<usize> {
|
|
self.edges
|
|
.iter()
|
|
.filter_map(|&(a, b)| {
|
|
if a == idx {
|
|
Some(b)
|
|
} else if b == idx {
|
|
Some(a)
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
.collect()
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Mesh Encoder
|
|
// ============================================================================
|
|
|
|
/// Encoder for converting between grid and mesh representations.
|
|
#[derive(Debug)]
|
|
pub struct MeshEncoder {
|
|
/// Number of mesh nodes.
|
|
#[allow(dead_code)]
|
|
num_nodes: usize,
|
|
/// Hidden dimension.
|
|
hidden_dim: usize,
|
|
/// Encoder weights.
|
|
encoder_weights: Vec<f32>,
|
|
/// Decoder weights.
|
|
decoder_weights: Vec<f32>,
|
|
/// RNG state.
|
|
rng_state: u64,
|
|
}
|
|
|
|
impl MeshEncoder {
|
|
/// Create a new mesh encoder.
|
|
pub fn new(num_nodes: usize, hidden_dim: usize) -> Self {
|
|
let mut encoder = Self {
|
|
num_nodes,
|
|
hidden_dim,
|
|
encoder_weights: Vec::new(),
|
|
decoder_weights: Vec::new(),
|
|
rng_state: 42,
|
|
};
|
|
encoder.initialize_weights();
|
|
encoder
|
|
}
|
|
|
|
/// Initialize encoder/decoder weights.
|
|
fn initialize_weights(&mut self) {
|
|
// Encoder: input_dim -> hidden_dim
|
|
let input_dim = 8 * 5 + 5; // 8 pressure levels * 5 vars + 5 surface vars
|
|
self.encoder_weights = (0..input_dim * self.hidden_dim)
|
|
.map(|_| self.random() as f32 * 0.1 - 0.05)
|
|
.collect();
|
|
|
|
// Decoder: hidden_dim -> output_dim
|
|
self.decoder_weights = (0..self.hidden_dim * input_dim)
|
|
.map(|_| self.random() as f32 * 0.1 - 0.05)
|
|
.collect();
|
|
}
|
|
|
|
/// Encode atmospheric state to mesh representation.
|
|
pub fn encode_state(&mut self, state: &AtmosphericState, mesh: &IcosahedralMesh) -> Vec<f32> {
|
|
let num_nodes = mesh.vertices.len();
|
|
let mut mesh_state = vec![0.0f32; num_nodes * self.hidden_dim];
|
|
|
|
// Flatten state to feature vector
|
|
let features = self.state_to_features(state);
|
|
|
|
// Apply encoder at each mesh node
|
|
for node in 0..num_nodes {
|
|
let (lat, lon) = mesh.vertex_latlon(node).unwrap_or((0.0, 0.0));
|
|
|
|
// Modulate features by position
|
|
let lat_factor = lat.to_radians().cos();
|
|
let lon_factor = (lon.to_radians() * 2.0).sin() * 0.1;
|
|
|
|
for h in 0..self.hidden_dim {
|
|
let mut val = 0.0f32;
|
|
|
|
// Linear projection with positional modulation
|
|
for (f_idx, &f) in features.iter().enumerate() {
|
|
let w_idx = f_idx * self.hidden_dim + h;
|
|
let weight = self.encoder_weights.get(w_idx).copied().unwrap_or(0.01);
|
|
val += f * weight;
|
|
}
|
|
|
|
// Add positional encoding
|
|
val += (lat_factor * (h as f32 * 0.1)).sin() * 0.5;
|
|
val += lon_factor * (h as f32 * 0.1).cos() * 0.3;
|
|
|
|
// Apply non-linearity
|
|
mesh_state[node * self.hidden_dim + h] = self.gelu(val);
|
|
}
|
|
}
|
|
|
|
mesh_state
|
|
}
|
|
|
|
/// Decode mesh representation back to atmospheric state.
|
|
pub fn decode_state(
|
|
&mut self,
|
|
mesh_state: &[f32],
|
|
mesh: &IcosahedralMesh,
|
|
timestamp: f32,
|
|
) -> AtmosphericState {
|
|
let num_nodes = mesh.vertices.len();
|
|
|
|
// Average over all nodes (simplified global mean)
|
|
let mut avg_hidden = vec![0.0f32; self.hidden_dim];
|
|
for node in 0..num_nodes {
|
|
for h in 0..self.hidden_dim {
|
|
let idx = node * self.hidden_dim + h;
|
|
avg_hidden[h] += mesh_state.get(idx).copied().unwrap_or(0.0) / num_nodes as f32;
|
|
}
|
|
}
|
|
|
|
// Decode to output features
|
|
let output_dim = 8 * 5 + 5;
|
|
let mut output = vec![0.0f32; output_dim];
|
|
|
|
for o in 0..output_dim {
|
|
for h in 0..self.hidden_dim {
|
|
let w_idx = h * output_dim + o;
|
|
let weight = self.decoder_weights.get(w_idx).copied().unwrap_or(0.01);
|
|
output[o] += avg_hidden[h] * weight;
|
|
}
|
|
}
|
|
|
|
// Reconstruct atmospheric state with physical constraints
|
|
self.features_to_state(&output, timestamp)
|
|
}
|
|
|
|
/// Convert atmospheric state to feature vector.
|
|
fn state_to_features(&self, state: &AtmosphericState) -> Vec<f32> {
|
|
let mut features = Vec::new();
|
|
|
|
// Normalize and add pressure level variables
|
|
for i in 0..8 {
|
|
features.push(state.temperature.get(i).copied().unwrap_or(250.0) / 300.0);
|
|
features.push(state.geopotential.get(i).copied().unwrap_or(5000.0) / 10000.0);
|
|
features.push(state.humidity.get(i).copied().unwrap_or(0.005) * 100.0);
|
|
features.push(state.wind_u.get(i).copied().unwrap_or(0.0) / 50.0);
|
|
features.push(state.wind_v.get(i).copied().unwrap_or(0.0) / 50.0);
|
|
}
|
|
|
|
// Surface variables
|
|
features.push(state.surface_pressure / 101325.0);
|
|
features.push(state.temperature_2m / 300.0);
|
|
features.push(state.wind_u_10m / 20.0);
|
|
features.push(state.wind_v_10m / 20.0);
|
|
features.push(state.precipitation / 10.0);
|
|
|
|
features
|
|
}
|
|
|
|
/// Convert feature vector back to atmospheric state.
|
|
fn features_to_state(&mut self, features: &[f32], timestamp: f32) -> AtmosphericState {
|
|
let pressure_levels = vec![1000.0, 925.0, 850.0, 700.0, 500.0, 300.0, 200.0, 50.0];
|
|
|
|
let mut temperature = Vec::new();
|
|
let mut geopotential = Vec::new();
|
|
let mut humidity = Vec::new();
|
|
let mut wind_u = Vec::new();
|
|
let mut wind_v = Vec::new();
|
|
|
|
for i in 0..8 {
|
|
let base = i * 5;
|
|
// Clamp temperature to physically reasonable range (150K to 350K)
|
|
let temp = (features.get(base).copied().unwrap_or(0.83) * 300.0).clamp(150.0, 350.0);
|
|
temperature.push(temp);
|
|
geopotential.push(features.get(base + 1).copied().unwrap_or(0.5) * 10000.0);
|
|
humidity.push((features.get(base + 2).copied().unwrap_or(0.5) / 100.0).max(0.0));
|
|
wind_u.push(features.get(base + 3).copied().unwrap_or(0.0) * 50.0);
|
|
wind_v.push(features.get(base + 4).copied().unwrap_or(0.0) * 50.0);
|
|
}
|
|
|
|
// Add some stochastic variability (small enough to stay in range)
|
|
for t in &mut temperature {
|
|
*t += (self.random() as f32 - 0.5) * 2.0;
|
|
*t = t.clamp(150.0, 350.0); // Ensure still in range after noise
|
|
}
|
|
|
|
AtmosphericState {
|
|
temperature,
|
|
geopotential,
|
|
humidity,
|
|
wind_u,
|
|
wind_v,
|
|
pressure_levels,
|
|
surface_pressure: features.get(40).copied().unwrap_or(1.0) * 101325.0,
|
|
temperature_2m: (features.get(41).copied().unwrap_or(0.96) * 300.0).clamp(150.0, 350.0),
|
|
wind_u_10m: features.get(42).copied().unwrap_or(0.0) * 20.0,
|
|
wind_v_10m: features.get(43).copied().unwrap_or(0.0) * 20.0,
|
|
precipitation: (features.get(44).copied().unwrap_or(0.0) * 10.0).max(0.0),
|
|
timestamp,
|
|
}
|
|
}
|
|
|
|
/// GELU activation function.
|
|
fn gelu(&self, x: f32) -> f32 {
|
|
0.5 * x * (1.0 + ((2.0 / std::f32::consts::PI).sqrt() * (x + 0.044715 * x.powi(3))).tanh())
|
|
}
|
|
|
|
/// Random number generator.
|
|
fn random(&mut self) -> f64 {
|
|
self.rng_state = self
|
|
.rng_state
|
|
.wrapping_mul(6364136223846793005)
|
|
.wrapping_add(1442695040888963407);
|
|
(self.rng_state >> 11) as f64 / (1u64 << 53) as f64
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Tests
|
|
// ============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_mesh_creation() {
|
|
let config = MeshConfig::from_refinement(3);
|
|
let mesh = IcosahedralMesh::new(&config);
|
|
|
|
assert!(!mesh.vertices.is_empty());
|
|
assert!(!mesh.edges.is_empty());
|
|
assert!(!mesh.faces.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_mesh_refinement_levels() {
|
|
// Test different refinement levels
|
|
for level in 0..4 {
|
|
let config = MeshConfig::from_refinement(level);
|
|
let mesh = IcosahedralMesh::new(&config);
|
|
|
|
// Icosahedral mesh: V = 10 * 4^n + 2
|
|
let expected_vertices = 10 * 4_usize.pow(level as u32) + 2;
|
|
assert_eq!(mesh.vertices.len(), expected_vertices);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_vertices_on_unit_sphere() {
|
|
let config = MeshConfig::from_refinement(2);
|
|
let mesh = IcosahedralMesh::new(&config);
|
|
|
|
for &(x, y, z) in &mesh.vertices {
|
|
let len = (x * x + y * y + z * z).sqrt();
|
|
assert!((len - 1.0).abs() < 0.001, "Vertex not on unit sphere");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_vertex_latlon() {
|
|
let config = MeshConfig::from_refinement(2);
|
|
let mesh = IcosahedralMesh::new(&config);
|
|
|
|
for idx in 0..mesh.vertices.len() {
|
|
let latlon = mesh.vertex_latlon(idx);
|
|
assert!(latlon.is_some());
|
|
|
|
let (lat, lon) = latlon.unwrap();
|
|
assert!(lat >= -90.0 && lat <= 90.0);
|
|
assert!(lon >= -180.0 && lon <= 180.0);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_nearest_vertex() {
|
|
let config = MeshConfig::from_refinement(3);
|
|
let mesh = IcosahedralMesh::new(&config);
|
|
|
|
// Find vertex near New York City
|
|
let nyc_vertex = mesh.nearest_vertex(40.7, -74.0);
|
|
let (lat, _lon) = mesh.vertex_latlon(nyc_vertex).unwrap();
|
|
|
|
// Should be reasonably close
|
|
assert!((lat - 40.7).abs() < 10.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_vertex_neighbors() {
|
|
let config = MeshConfig::from_refinement(2);
|
|
let mesh = IcosahedralMesh::new(&config);
|
|
|
|
let neighbors = mesh.vertex_neighbors(0);
|
|
|
|
// Each vertex should have multiple neighbors (typically 5-6 for icosahedral mesh)
|
|
assert!(!neighbors.is_empty());
|
|
assert!(neighbors.len() >= 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_mesh_encoder_creation() {
|
|
let encoder = MeshEncoder::new(100, 64);
|
|
assert_eq!(encoder.num_nodes, 100);
|
|
assert_eq!(encoder.hidden_dim, 64);
|
|
}
|
|
|
|
#[test]
|
|
fn test_encode_decode_roundtrip() {
|
|
let config = MeshConfig::from_refinement(2);
|
|
let mesh = IcosahedralMesh::new(&config);
|
|
let mut encoder = MeshEncoder::new(mesh.vertices.len(), 64);
|
|
|
|
let original_state = weathercast_shared::sample_atmospheric_state();
|
|
let encoded = encoder.encode_state(&original_state, &mesh);
|
|
let decoded = encoder.decode_state(&encoded, &mesh, 0.0);
|
|
|
|
// Check that decoded state has same structure
|
|
assert_eq!(decoded.temperature.len(), original_state.temperature.len());
|
|
assert_eq!(
|
|
decoded.pressure_levels.len(),
|
|
original_state.pressure_levels.len()
|
|
);
|
|
|
|
// Values should be in reasonable ranges
|
|
for t in &decoded.temperature {
|
|
assert!(*t > 100.0 && *t < 400.0, "Temperature out of range: {}", t);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_encoder_mesh_state_size() {
|
|
let config = MeshConfig::from_refinement(2);
|
|
let mesh = IcosahedralMesh::new(&config);
|
|
let hidden_dim = 64;
|
|
let mut encoder = MeshEncoder::new(mesh.vertices.len(), hidden_dim);
|
|
|
|
let state = weathercast_shared::sample_atmospheric_state();
|
|
let encoded = encoder.encode_state(&state, &mesh);
|
|
|
|
assert_eq!(encoded.len(), mesh.vertices.len() * hidden_dim);
|
|
}
|
|
|
|
#[test]
|
|
fn test_gelu_activation() {
|
|
let encoder = MeshEncoder::new(10, 32);
|
|
|
|
// GELU(0) should be approximately 0
|
|
assert!(encoder.gelu(0.0).abs() < 0.01);
|
|
|
|
// GELU should be monotonically increasing for positive x
|
|
assert!(encoder.gelu(1.0) > encoder.gelu(0.5));
|
|
assert!(encoder.gelu(0.5) > encoder.gelu(0.0));
|
|
}
|
|
}
|