rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

Initial commit

Browse Source
This commit is contained in:
redclawsystems
2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
demos/rtx-weathercast-demo/src/sample_data.rs
+565
View File
@@ -0,0 +1,565 @@
//! Sample data and configurations for WeatherCast demo.
//!
//! Provides pre-configured scenarios for different types of weather
//! predictions including global forecasts, regional high-resolution
//! forecasts, and ensemble predictions.
use weathercast_shared::{
AtmosphericState, GridPoint, MeshConfig, PredictionConfig, TrainingConfig,
};
// ============================================================================
// Global Forecast Configurations
// ============================================================================
/// Create a global medium-range forecast configuration.
#[must_use]
pub fn global_forecast() -> (MeshConfig, PredictionConfig) {
let mesh_config = MeshConfig::from_refinement(5); // ~40km resolution
let pred_config = PredictionConfig {
lead_time: 240.0, // 10 days
ensemble_size: 1,
resolution: 0.25,
time_step: 6.0,
num_steps: 40,
multi_scale: true,
variables: vec![
"temperature".to_string(),
"geopotential".to_string(),
"humidity".to_string(),
"wind_u".to_string(),
"wind_v".to_string(),
],
pressure_levels: vec![1000.0, 925.0, 850.0, 700.0, 500.0, 300.0, 200.0, 50.0],
};
(mesh_config, pred_config)
}
/// Create a short-range global forecast (0-72h).
#[must_use]
pub fn global_short_range() -> (MeshConfig, PredictionConfig) {
let mesh_config = MeshConfig::from_refinement(5);
let pred_config = PredictionConfig {
lead_time: 72.0,
time_step: 3.0,
num_steps: 24,
..PredictionConfig::default()
};
(mesh_config, pred_config)
}
// ============================================================================
// Regional Forecast Configurations
// ============================================================================
/// Create a regional high-resolution forecast for North America.
#[must_use]
pub fn regional_forecast() -> (MeshConfig, PredictionConfig) {
let mesh_config = MeshConfig {
refinement_levels: 6,
num_nodes: 40962,
num_edges: 122880,
num_faces: 81920,
avg_edge_length: 20.0, // ~20km
multi_scale: true,
regional_focus: Some((40.0, -100.0, 2000.0)), // Central US, 2000km radius
};
let pred_config = PredictionConfig {
lead_time: 48.0, // 2 days
time_step: 1.0, // Hourly
num_steps: 48,
resolution: 0.1, // ~10km effective
ensemble_size: 1,
multi_scale: true,
..PredictionConfig::default()
};
(mesh_config, pred_config)
}
/// Create a European regional forecast.
#[must_use]
pub fn european_regional() -> (MeshConfig, PredictionConfig) {
let mesh_config = MeshConfig {
refinement_levels: 6,
num_nodes: 40962,
num_edges: 122880,
num_faces: 81920,
avg_edge_length: 15.0,
multi_scale: true,
regional_focus: Some((50.0, 10.0, 1500.0)), // Central Europe
};
let pred_config = PredictionConfig {
lead_time: 72.0,
time_step: 1.0,
num_steps: 72,
resolution: 0.1,
..PredictionConfig::default()
};
(mesh_config, pred_config)
}
/// Create an Asian regional forecast.
#[must_use]
pub fn asian_regional() -> (MeshConfig, PredictionConfig) {
let mesh_config = MeshConfig::regional(35.0, 135.0, 2000.0); // Japan/East Asia
let pred_config = PredictionConfig {
lead_time: 72.0,
time_step: 3.0,
num_steps: 24,
..PredictionConfig::default()
};
(mesh_config, pred_config)
}
// ============================================================================
// Ensemble Forecast Configurations
// ============================================================================
/// Create ensemble forecast configuration (50 members).
#[must_use]
pub fn ensemble_forecast() -> (MeshConfig, PredictionConfig) {
let mesh_config = MeshConfig::from_refinement(4); // Coarser for ensemble
let pred_config = PredictionConfig {
lead_time: 360.0, // 15 days
ensemble_size: 50,
time_step: 6.0,
num_steps: 60,
resolution: 0.5,
multi_scale: false,
..PredictionConfig::default()
};
(mesh_config, pred_config)
}
/// Create large ensemble (100 members) for probabilistic forecasting.
#[must_use]
pub fn large_ensemble() -> (MeshConfig, PredictionConfig) {
let mesh_config = MeshConfig::from_refinement(3); // Even coarser
let pred_config = PredictionConfig {
lead_time: 240.0,
ensemble_size: 100,
time_step: 12.0,
num_steps: 20,
resolution: 1.0,
..PredictionConfig::default()
};
(mesh_config, pred_config)
}
// ============================================================================
// High-Resolution Configurations
// ============================================================================
/// Create ultra high-resolution forecast (convection-permitting).
#[must_use]
pub fn high_resolution_forecast() -> (MeshConfig, PredictionConfig) {
let mesh_config = MeshConfig {
refinement_levels: 7,
num_nodes: 163842,
num_edges: 491520,
num_faces: 327680,
avg_edge_length: 5.0, // ~5km (convection-permitting)
multi_scale: true,
regional_focus: Some((40.0, -100.0, 500.0)), // Small regional domain
};
let pred_config = PredictionConfig {
lead_time: 24.0, // 1 day
time_step: 0.5, // 30-minute steps
num_steps: 48,
resolution: 0.05,
ensemble_size: 1,
multi_scale: true,
variables: vec![
"temperature".to_string(),
"geopotential".to_string(),
"humidity".to_string(),
"wind_u".to_string(),
"wind_v".to_string(),
"precipitation".to_string(),
"cloud_cover".to_string(),
],
pressure_levels: vec![
1000.0, 925.0, 850.0, 700.0, 600.0, 500.0, 400.0, 300.0, 250.0, 200.0, 150.0, 100.0,
50.0,
],
};
(mesh_config, pred_config)
}
// ============================================================================
// Sample Atmospheric States
// ============================================================================
/// Create initial state for mid-latitude cyclone.
#[must_use]
pub fn midlatitude_cyclone() -> AtmosphericState {
AtmosphericState {
temperature: vec![280.0, 278.0, 275.0, 265.0, 248.0, 225.0, 215.0, 200.0],
geopotential: vec![0.0, 800.0, 1500.0, 3100.0, 5600.0, 9200.0, 11800.0, 20500.0],
humidity: vec![0.015, 0.012, 0.008, 0.003, 0.001, 0.0002, 0.00005, 0.000001],
wind_u: vec![-5.0, -8.0, -15.0, -25.0, -40.0, -50.0, -45.0, -25.0],
wind_v: vec![10.0, 15.0, 20.0, 25.0, 20.0, 10.0, 5.0, 0.0],
pressure_levels: vec![1000.0, 925.0, 850.0, 700.0, 500.0, 300.0, 200.0, 50.0],
surface_pressure: 99000.0, // Low pressure center
temperature_2m: 283.0,
wind_u_10m: -5.0,
wind_v_10m: 8.0,
precipitation: 5.0,
timestamp: 0.0,
}
}
/// Create initial state for tropical cyclone.
#[must_use]
pub fn tropical_cyclone() -> AtmosphericState {
AtmosphericState {
temperature: vec![299.0, 297.0, 292.0, 280.0, 260.0, 228.0, 217.0, 195.0],
geopotential: vec![0.0, 750.0, 1500.0, 3000.0, 5800.0, 9500.0, 12200.0, 20800.0],
humidity: vec![
0.020, 0.018, 0.015, 0.008, 0.002, 0.0001, 0.00001, 0.0000001,
],
wind_u: vec![-30.0, -40.0, -50.0, -55.0, -45.0, -20.0, -5.0, 5.0],
wind_v: vec![30.0, 40.0, 50.0, 55.0, 45.0, 20.0, 5.0, -5.0],
pressure_levels: vec![1000.0, 925.0, 850.0, 700.0, 500.0, 300.0, 200.0, 50.0],
surface_pressure: 94000.0, // Very low for hurricane
temperature_2m: 300.0,
wind_u_10m: -35.0,
wind_v_10m: 35.0,
precipitation: 50.0,
timestamp: 0.0,
}
}
/// Create initial state for high pressure system (anticyclone).
#[must_use]
pub fn anticyclone() -> AtmosphericState {
AtmosphericState {
temperature: vec![295.0, 292.0, 287.0, 275.0, 255.0, 228.0, 218.0, 205.0],
geopotential: vec![0.0, 820.0, 1550.0, 3150.0, 5750.0, 9400.0, 12100.0, 20600.0],
humidity: vec![
0.008, 0.006, 0.004, 0.002, 0.0005, 0.0001, 0.00002, 0.000001,
],
wind_u: vec![2.0, 3.0, 5.0, 8.0, 12.0, 15.0, 10.0, 5.0],
wind_v: vec![-2.0, -3.0, -5.0, -8.0, -10.0, -8.0, -5.0, -2.0],
pressure_levels: vec![1000.0, 925.0, 850.0, 700.0, 500.0, 300.0, 200.0, 50.0],
surface_pressure: 103000.0, // High pressure
temperature_2m: 298.0,
wind_u_10m: 2.0,
wind_v_10m: -1.0,
precipitation: 0.0,
timestamp: 0.0,
}
}
/// Create initial state for winter storm.
#[must_use]
pub fn winter_storm() -> AtmosphericState {
AtmosphericState {
temperature: vec![268.0, 265.0, 260.0, 250.0, 235.0, 215.0, 208.0, 195.0],
geopotential: vec![0.0, 700.0, 1350.0, 2800.0, 5200.0, 8800.0, 11500.0, 20000.0],
humidity: vec![
0.003, 0.003, 0.002, 0.001, 0.0005, 0.0001, 0.00002, 0.000001,
],
wind_u: vec![-15.0, -20.0, -30.0, -45.0, -55.0, -60.0, -55.0, -30.0],
wind_v: vec![5.0, 8.0, 12.0, 15.0, 12.0, 5.0, 0.0, -5.0],
pressure_levels: vec![1000.0, 925.0, 850.0, 700.0, 500.0, 300.0, 200.0, 50.0],
surface_pressure: 98500.0,
temperature_2m: 268.0,
wind_u_10m: -12.0,
wind_v_10m: 5.0,
precipitation: 15.0, // Snow equivalent
timestamp: 0.0,
}
}
// ============================================================================
// Sample Locations
// ============================================================================
/// Create sample verification locations.
#[must_use]
pub fn verification_locations() -> Vec<GridPoint> {
vec![
GridPoint {
lat: 40.7128,
lon: -74.0060,
elevation: 10.0,
land_mask: 1.0,
}, // New York
GridPoint {
lat: 34.0522,
lon: -118.2437,
elevation: 71.0,
land_mask: 1.0,
}, // Los Angeles
GridPoint {
lat: 51.5074,
lon: -0.1278,
elevation: 11.0,
land_mask: 1.0,
}, // London
GridPoint {
lat: 35.6762,
lon: 139.6503,
elevation: 40.0,
land_mask: 1.0,
}, // Tokyo
GridPoint {
lat: -33.8688,
lon: 151.2093,
elevation: 58.0,
land_mask: 1.0,
}, // Sydney
GridPoint {
lat: 55.7558,
lon: 37.6173,
elevation: 156.0,
land_mask: 1.0,
}, // Moscow
GridPoint {
lat: -22.9068,
lon: -43.1729,
elevation: 11.0,
land_mask: 1.0,
}, // Rio de Janeiro
GridPoint {
lat: 1.3521,
lon: 103.8198,
elevation: 15.0,
land_mask: 1.0,
}, // Singapore
]
}
/// Create sample ocean locations for marine forecasting.
#[must_use]
pub fn marine_locations() -> Vec<GridPoint> {
vec![
GridPoint {
lat: 30.0,
lon: -60.0,
elevation: 0.0,
land_mask: 0.0,
}, // North Atlantic
GridPoint {
lat: 0.0,
lon: -160.0,
elevation: 0.0,
land_mask: 0.0,
}, // Central Pacific
GridPoint {
lat: -45.0,
lon: 20.0,
elevation: 0.0,
land_mask: 0.0,
}, // Southern Ocean
GridPoint {
lat: 20.0,
lon: 120.0,
elevation: 0.0,
land_mask: 0.0,
}, // South China Sea
]
}
// ============================================================================
// Training Configurations
// ============================================================================
/// Create quick training configuration for testing.
#[must_use]
pub fn quick_training() -> TrainingConfig {
TrainingConfig {
epochs: 10,
batch_size: 8,
learning_rate: 1e-3,
weight_decay: 1e-5,
num_ar_steps: 2,
gradient_accumulation: 1,
curriculum: false,
variable_weights: vec![
("temperature".to_string(), 1.0),
("geopotential".to_string(), 1.0),
],
noise_level: 0.01,
}
}
/// Create full training configuration.
#[must_use]
pub fn full_training() -> TrainingConfig {
TrainingConfig {
epochs: 200,
batch_size: 32,
learning_rate: 1e-4,
weight_decay: 1e-5,
num_ar_steps: 12,
gradient_accumulation: 4,
curriculum: true,
variable_weights: vec![
("temperature".to_string(), 1.0),
("geopotential".to_string(), 0.5),
("humidity".to_string(), 1.0),
("wind_u".to_string(), 1.0),
("wind_v".to_string(), 1.0),
],
noise_level: 0.005,
}
}
/// Create fine-tuning configuration.
#[must_use]
pub fn finetune_training() -> TrainingConfig {
TrainingConfig {
epochs: 50,
batch_size: 16,
learning_rate: 1e-5,
weight_decay: 1e-6,
num_ar_steps: 6,
gradient_accumulation: 2,
curriculum: false,
..TrainingConfig::default()
}
}
// ============================================================================
// Tests
// ============================================================================
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_global_forecast_config() {
let (mesh, pred) = global_forecast();
assert_eq!(mesh.refinement_levels, 5);
assert_eq!(pred.lead_time, 240.0);
}
#[test]
fn test_regional_forecast_config() {
let (mesh, pred) = regional_forecast();
assert!(mesh.regional_focus.is_some());
assert!(mesh.multi_scale);
assert_eq!(pred.time_step, 1.0);
}
#[test]
fn test_ensemble_config() {
let (_, pred) = ensemble_forecast();
assert_eq!(pred.ensemble_size, 50);
}
#[test]
fn test_large_ensemble() {
let (_, pred) = large_ensemble();
assert_eq!(pred.ensemble_size, 100);
}
#[test]
fn test_high_resolution_config() {
let (mesh, pred) = high_resolution_forecast();
assert!(mesh.avg_edge_length < 10.0);
assert!(pred.variables.len() > 5);
}
#[test]
fn test_midlatitude_cyclone() {
let state = midlatitude_cyclone();
assert!(state.surface_pressure < 101325.0); // Low pressure
assert!(state.precipitation > 0.0);
}
#[test]
fn test_tropical_cyclone() {
let state = tropical_cyclone();
assert!(state.surface_pressure < 95000.0); // Very low
assert!(state.wind_speed(0) > 30.0); // Strong winds
}
#[test]
fn test_anticyclone() {
let state = anticyclone();
assert!(state.surface_pressure > 101325.0); // High pressure
assert!(state.precipitation == 0.0);
}
#[test]
fn test_winter_storm() {
let state = winter_storm();
assert!(state.temperature_2m < 273.0); // Below freezing
}
#[test]
fn test_verification_locations() {
let locations = verification_locations();
assert!(!locations.is_empty());
for loc in &locations {
assert!(loc.lat >= -90.0 && loc.lat <= 90.0);
assert!(loc.lon >= -180.0 && loc.lon <= 180.0);
}
}
#[test]
fn test_marine_locations() {
let locations = marine_locations();
for loc in &locations {
assert_eq!(loc.land_mask, 0.0); // All should be ocean
}
}
#[test]
fn test_training_configs() {
let quick = quick_training();
let full = full_training();
assert!(quick.epochs < full.epochs);
assert!(quick.learning_rate > full.learning_rate);
}
#[test]
fn test_european_regional() {
let (mesh, _) = european_regional();
let (lat, lon, _) = mesh.regional_focus.unwrap();
assert!((lat - 50.0).abs() < 1.0);
assert!((lon - 10.0).abs() < 1.0);
}
#[test]
fn test_asian_regional() {
let (mesh, _) = asian_regional();
assert!(mesh.regional_focus.is_some());
}
#[test]
fn test_global_short_range() {
let (_, pred) = global_short_range();
assert_eq!(pred.lead_time, 72.0);
assert_eq!(pred.time_step, 3.0);
}
#[test]
fn test_finetune_training() {
let config = finetune_training();
assert!(config.learning_rate < 1e-4);
}
}