Initial commit

2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
@@ -0,0 +1,641 @@
+//! SeismicAI: Earthquake Simulation & Early Warning with Neural Operators
+//!
+//! This demo showcases neural operator-based seismic wave propagation for
+//! earthquake simulation and early warning systems. It demonstrates:
+//! - Fourier Neural Operators for elastic wave propagation
+//! - Ground motion prediction equations (GMPEs)
+//! - Real-time earthquake early warning
+//! - Site effects and attenuation modeling
+
+pub mod attenuation;
+pub mod sample_data;
+pub mod wave_operator;
+
+use attenuation::{AttenuationModel, SiteEffects, GMPE};
+use seismic_shared::{
+    AlertLevel, EarthquakeSource, EarthquakeWarning, GeoLocation, GroundMotion, SimulationConfig,
+    SimulationResult, StationConfig, VelocityModel, WarningConfig, WaveFieldSnapshot,
+};
+use wave_operator::{WaveFNO, WaveField};
+
+// ============================================================================
+// SeismicAI System
+// ============================================================================
+
+/// Main SeismicAI system for earthquake simulation and early warning.
+#[derive(Debug)]
+pub struct SeismicAI {
+    /// Wave operator for full waveform simulation.
+    wave_operator: WaveFNO,
+    /// Source encoder for earthquake parametrization.
+    source_encoder: SourceEncoder,
+    /// Arrival predictor for travel time estimation.
+    arrival_predictor: ArrivalPredictor,
+    /// Attenuation model for ground motion prediction.
+    attenuation_model: AttenuationModel,
+    /// Simulation configuration.
+    config: SimulationConfig,
+    /// Warning configuration.
+    warning_config: WarningConfig,
+    /// RNG state for simulations.
+    rng_state: u64,
+}
+
+/// Encodes earthquake source parameters for neural network input.
+#[derive(Debug)]
+pub struct SourceEncoder {
+    /// Hidden dimension.
+    hidden_dim: usize,
+    /// Encoding weights.
+    weights: Vec<f32>,
+}
+
+impl SourceEncoder {
+    pub fn new(hidden_dim: usize) -> Self {
+        Self {
+            hidden_dim,
+            weights: vec![0.01; hidden_dim * 10], // 10 source parameters
+        }
+    }
+
+    /// Encode earthquake source to feature vector.
+    pub fn encode(&self, source: &EarthquakeSource) -> Vec<f32> {
+        let mut features = Vec::with_capacity(self.hidden_dim);
+
+        // Extract key source parameters
+        let lat = source.hypocenter.location.latitude as f32;
+        let lon = source.hypocenter.location.longitude as f32;
+        let depth = source.hypocenter.depth_km as f32;
+        let mag = source.magnitude as f32;
+        let strike = source.mechanism.strike as f32 / 360.0;
+        let dip = source.mechanism.dip as f32 / 90.0;
+        let rake = (source.mechanism.rake as f32 + 180.0) / 360.0;
+
+        // Simple linear encoding
+        for i in 0..self.hidden_dim {
+            let w_idx = i * 7;
+            let feature = self.weights[w_idx] * lat
+                + self.weights[w_idx + 1] * lon
+                + self.weights[w_idx + 2] * depth
+                + self.weights[w_idx + 3] * mag
+                + self.weights[w_idx + 4] * strike
+                + self.weights[w_idx + 5] * dip
+                + self.weights[w_idx + 6] * rake;
+
+            features.push(feature.tanh());
+        }
+
+        features
+    }
+}
+
+/// Predicts P and S wave arrival times.
+#[derive(Debug)]
+pub struct ArrivalPredictor {
+    /// Average P-wave velocity (km/s).
+    avg_vp: f64,
+    /// Average S-wave velocity (km/s).
+    avg_vs: f64,
+}
+
+impl ArrivalPredictor {
+    pub fn new() -> Self {
+        Self {
+            avg_vp: 6.0,
+            avg_vs: 3.5,
+        }
+    }
+
+    /// Update velocities from a velocity model.
+    pub fn from_velocity_model(model: &VelocityModel) -> Self {
+        let n = model.layers.len() as f64;
+        let avg_vp = model.layers.iter().map(|l| l.vp).sum::<f64>() / n;
+        let avg_vs = model.layers.iter().map(|l| l.vs).sum::<f64>() / n;
+        Self { avg_vp, avg_vs }
+    }
+
+    /// Predict P-wave arrival time.
+    pub fn predict_p_arrival(&self, distance_km: f64, depth_km: f64) -> f64 {
+        let path = (distance_km.powi(2) + depth_km.powi(2)).sqrt();
+        path / self.avg_vp
+    }
+
+    /// Predict S-wave arrival time.
+    pub fn predict_s_arrival(&self, distance_km: f64, depth_km: f64) -> f64 {
+        let path = (distance_km.powi(2) + depth_km.powi(2)).sqrt();
+        path / self.avg_vs
+    }
+
+    /// Predict S-P time (useful for distance estimation).
+    pub fn predict_sp_time(&self, distance_km: f64, depth_km: f64) -> f64 {
+        self.predict_s_arrival(distance_km, depth_km)
+            - self.predict_p_arrival(distance_km, depth_km)
+    }
+
+    /// Estimate distance from S-P time.
+    pub fn estimate_distance(&self, sp_time: f64) -> f64 {
+        // Simplified: assumes horizontal distance >> depth
+        sp_time / (1.0 / self.avg_vs - 1.0 / self.avg_vp)
+    }
+}
+
+impl Default for ArrivalPredictor {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl SeismicAI {
+    /// Create a new SeismicAI system.
+    pub fn new(config: SimulationConfig) -> Self {
+        Self {
+            wave_operator: WaveFNO::new(config.fno_modes, config.hidden_dim, 4),
+            source_encoder: SourceEncoder::new(config.hidden_dim),
+            arrival_predictor: ArrivalPredictor::new(),
+            attenuation_model: AttenuationModel::nga_west2(),
+            config,
+            warning_config: WarningConfig::default(),
+            rng_state: 42,
+        }
+    }
+
+    /// Create with specific warning configuration.
+    pub fn with_warning_config(mut self, warning_config: WarningConfig) -> Self {
+        self.warning_config = warning_config;
+        self
+    }
+
+    /// Set the attenuation model.
+    pub fn with_attenuation_model(mut self, model: AttenuationModel) -> Self {
+        self.attenuation_model = model;
+        self
+    }
+
+    /// Simulate seismic wave propagation.
+    pub fn simulate(
+        &mut self,
+        source: &EarthquakeSource,
+        stations: &[StationConfig],
+        velocity_model: &VelocityModel,
+    ) -> SimulationResult {
+        let start = std::time::Instant::now();
+
+        // Update arrival predictor with velocity model
+        self.arrival_predictor = ArrivalPredictor::from_velocity_model(velocity_model);
+
+        // Calculate grid dimensions
+        let (dx, dy, dz) = (self.config.dx, self.config.dx, self.config.dx);
+        let (lx, ly, lz) = self.config.domain_size;
+        let nx = (lx / dx) as usize;
+        let ny = (ly / dy) as usize;
+        let nz = (lz / dz) as usize;
+
+        // Initialize wave field
+        let mut wave_field = WaveField::new(nx, ny, nz);
+
+        // Inject source
+        self.wave_operator
+            .inject_source(&mut wave_field, source, &self.config, velocity_model);
+
+        // Propagate waves
+        let num_steps = (self.config.duration / self.config.dt) as usize;
+        let fields = self
+            .wave_operator
+            .propagate(&wave_field, self.config.dt, num_steps);
+
+        // Compute seismograms at stations
+        let seismograms =
+            self.wave_operator
+                .compute_seismograms(&fields, stations, &self.config, velocity_model);
+
+        // Predict ground motion at each station
+        let ground_motions = self.predict_ground_motion(source, stations, velocity_model);
+
+        // Generate warning
+        let warning = self.generate_warning(source, stations, velocity_model);
+
+        // Create snapshots (sample every N steps)
+        let snapshot_interval = num_steps / 10;
+        let snapshots: Vec<WaveFieldSnapshot> = fields
+            .iter()
+            .enumerate()
+            .filter(|(i, _)| i % snapshot_interval == 0)
+            .take(10)
+            .map(|(_, field)| self.create_snapshot(field))
+            .collect();
+
+        let computation_time_ms = start.elapsed().as_secs_f64() * 1000.0;
+
+        SimulationResult {
+            source: source.clone(),
+            ground_motions,
+            seismograms: Some(seismograms),
+            snapshots: Some(snapshots),
+            warning: Some(warning),
+            computation_time_ms,
+        }
+    }
+
+    /// Predict ground motion at stations using GMPE.
+    pub fn predict_ground_motion(
+        &self,
+        source: &EarthquakeSource,
+        stations: &[StationConfig],
+        velocity_model: &VelocityModel,
+    ) -> Vec<(StationConfig, GroundMotion)> {
+        stations
+            .iter()
+            .map(|station| {
+                // Get base prediction from GMPE
+                let mut gm = self
+                    .attenuation_model
+                    .predict(source, station, velocity_model);
+
+                // Apply site effects
+                if self.config.include_site_effects {
+                    let site_effects = SiteEffects::from_station(station);
+                    gm = site_effects.apply(&gm);
+                }
+
+                (station.clone(), gm)
+            })
+            .collect()
+    }
+
+    /// Generate earthquake early warning.
+    pub fn generate_warning(
+        &mut self,
+        source: &EarthquakeSource,
+        stations: &[StationConfig],
+        velocity_model: &VelocityModel,
+    ) -> EarthquakeWarning {
+        // Find stations that would detect the event
+        let detecting_stations: Vec<_> = stations
+            .iter()
+            .filter(|s| {
+                let dist = source.hypocenter.location.distance_km(&s.location);
+                dist < self.warning_config.alert_radius_km
+            })
+            .collect();
+
+        if detecting_stations.len() < self.warning_config.min_stations
+            || source.magnitude < self.warning_config.min_magnitude
+        {
+            return EarthquakeWarning::default();
+        }
+
+        // Estimate magnitude (with uncertainty)
+        let estimated_magnitude = source.magnitude + (self.random() - 0.5) * 0.3;
+
+        // Estimate location (with uncertainty)
+        let estimated_location = GeoLocation::new(
+            source.hypocenter.location.latitude + (self.random() - 0.5) * 0.1,
+            source.hypocenter.location.longitude + (self.random() - 0.5) * 0.1,
+        );
+
+        let estimated_depth = source.hypocenter.depth_km + (self.random() - 0.5) * 5.0;
+
+        // Calculate warning time for a typical target location
+        let target_distance = 50.0; // km
+        let p_arrival = self
+            .arrival_predictor
+            .predict_p_arrival(target_distance, source.hypocenter.depth_km);
+        let s_arrival = self
+            .arrival_predictor
+            .predict_s_arrival(target_distance, source.hypocenter.depth_km);
+
+        // Time to shaking at target
+        let processing_delay = 3.0; // seconds for detection and processing
+        let time_to_shaking = (s_arrival - p_arrival - processing_delay).max(0.0);
+
+        // Predict ground motion at target
+        let target_station = StationConfig::new(
+            "TARGET",
+            source.hypocenter.location.latitude + 0.45,
+            source.hypocenter.location.longitude,
+        );
+        let expected_gm = self
+            .attenuation_model
+            .predict(source, &target_station, velocity_model);
+
+        // Determine alert level
+        let alert_level = if expected_gm.mmi >= 7.0 {
+            AlertLevel::Severe
+        } else if expected_gm.mmi >= 5.0 {
+            AlertLevel::Warning
+        } else if expected_gm.mmi >= 4.0 {
+            AlertLevel::Watch
+        } else if expected_gm.mmi >= 3.0 {
+            AlertLevel::Advisory
+        } else {
+            AlertLevel::None
+        };
+
+        // Calculate confidence based on number of detecting stations
+        let confidence = (detecting_stations.len() as f64 / 10.0).min(0.95);
+
+        EarthquakeWarning {
+            event_id: source.event_id.clone(),
+            alert_level,
+            estimated_magnitude,
+            estimated_location,
+            estimated_depth_km: estimated_depth,
+            detecting_stations: detecting_stations.len(),
+            time_to_shaking,
+            expected_ground_motion: expected_gm,
+            warning_time: p_arrival + processing_delay,
+            origin_time: 0.0,
+            confidence,
+        }
+    }
+
+    /// Create a wave field snapshot for visualization.
+    fn create_snapshot(&self, field: &WaveField) -> WaveFieldSnapshot {
+        let (nx, ny, _) = field.dimensions;
+
+        // Extract surface (z=0) displacement
+        let mut disp_x = vec![vec![0.0_f32; ny]; nx];
+        let mut disp_y = vec![vec![0.0_f32; ny]; nx];
+        let mut disp_z = vec![vec![0.0_f32; ny]; nx];
+
+        for i in 0..nx {
+            for j in 0..ny {
+                let (dx, dy, dz) = field.displacement_at(i, j, 0);
+                disp_x[i][j] = dx;
+                disp_y[i][j] = dy;
+                disp_z[i][j] = dz;
+            }
+        }
+
+        // Grid coordinates
+        let grid_x: Vec<f32> = (0..nx).map(|i| i as f32 * self.config.dx as f32).collect();
+        let grid_y: Vec<f32> = (0..ny).map(|j| j as f32 * self.config.dx as f32).collect();
+
+        WaveFieldSnapshot {
+            time: field.time,
+            displacement_x: disp_x,
+            displacement_y: disp_y,
+            displacement_z: disp_z,
+            grid_x,
+            grid_y,
+        }
+    }
+
+    /// 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
+    }
+}
+
+// ============================================================================
+// Demo Entry Point
+// ============================================================================
+
+/// Run the SeismicAI demo.
+pub fn run_demo() -> SimulationResult {
+    // Create demo scenario
+    let (source, stations, velocity_model, config) = sample_data::local_test_scenario();
+
+    // Create SeismicAI system
+    let mut seismic_ai = SeismicAI::new(config)
+        .with_warning_config(sample_data::urban_warning_config())
+        .with_attenuation_model(AttenuationModel::california());
+
+    // Run simulation
+    seismic_ai.simulate(&source, &stations, &velocity_model)
+}
+
+/// Run a major earthquake scenario.
+pub fn run_major_scenario() -> SimulationResult {
+    let (source, stations, velocity_model, config) = sample_data::hayward_scenario();
+
+    let mut seismic_ai =
+        SeismicAI::new(config).with_warning_config(sample_data::urban_warning_config());
+
+    seismic_ai.simulate(&source, &stations, &velocity_model)
+}
+
+// ============================================================================
+// Tests
+// ============================================================================
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_seismic_ai_creation() {
+        let config = SimulationConfig::default();
+        let seismic_ai = SeismicAI::new(config);
+        assert_eq!(seismic_ai.arrival_predictor.avg_vp, 6.0);
+    }
+
+    #[test]
+    fn test_source_encoder() {
+        let encoder = SourceEncoder::new(32);
+        let source = sample_data::local_earthquake();
+        let features = encoder.encode(&source);
+
+        assert_eq!(features.len(), 32);
+        // All features should be in tanh range [-1, 1]
+        for f in &features {
+            assert!(*f >= -1.0 && *f <= 1.0);
+        }
+    }
+
+    #[test]
+    fn test_arrival_predictor() {
+        let predictor = ArrivalPredictor::new();
+
+        let p_time = predictor.predict_p_arrival(60.0, 10.0);
+        let s_time = predictor.predict_s_arrival(60.0, 10.0);
+
+        // S should arrive after P
+        assert!(s_time > p_time);
+
+        // S-P time increases with distance
+        let sp_near = predictor.predict_sp_time(30.0, 10.0);
+        let sp_far = predictor.predict_sp_time(100.0, 10.0);
+        assert!(sp_far > sp_near);
+    }
+
+    #[test]
+    fn test_arrival_predictor_from_model() {
+        let model = sample_data::california_velocity_model();
+        let predictor = ArrivalPredictor::from_velocity_model(&model);
+
+        // Should use model velocities
+        assert!(predictor.avg_vp > 4.0 && predictor.avg_vp < 8.0);
+        assert!(predictor.avg_vs > 2.0 && predictor.avg_vs < 5.0);
+    }
+
+    #[test]
+    fn test_distance_estimation() {
+        let predictor = ArrivalPredictor::new();
+
+        // Round-trip test
+        let true_distance = 80.0;
+        let sp_time = predictor.predict_sp_time(true_distance, 10.0);
+        let estimated = predictor.estimate_distance(sp_time);
+
+        // Should be approximately correct (ignoring depth effects)
+        assert!((estimated - true_distance).abs() < 20.0);
+    }
+
+    #[test]
+    fn test_ground_motion_prediction() {
+        let config = sample_data::fast_config();
+        let seismic_ai = SeismicAI::new(config);
+
+        let source = sample_data::local_earthquake();
+        let stations = sample_data::sparse_network();
+        let velocity_model = sample_data::generic_velocity_model();
+
+        let gm = seismic_ai.predict_ground_motion(&source, &stations, &velocity_model);
+
+        assert_eq!(gm.len(), stations.len());
+
+        for (station, motion) in &gm {
+            assert!(motion.pga > 0.0);
+            assert!(motion.p_arrival_time > 0.0);
+            assert!(motion.s_arrival_time > motion.p_arrival_time);
+            assert!(motion.mmi >= 1.0);
+            // Station should be returned
+            assert!(!station.code.is_empty());
+        }
+    }
+
+    #[test]
+    fn test_warning_generation() {
+        let config = sample_data::fast_config();
+        let mut seismic_ai = SeismicAI::new(config).with_warning_config(WarningConfig {
+            min_magnitude: 3.0,
+            min_stations: 2,
+            ..Default::default()
+        });
+
+        let source = sample_data::local_earthquake();
+        let stations = sample_data::sparse_network();
+        let velocity_model = sample_data::generic_velocity_model();
+
+        let warning = seismic_ai.generate_warning(&source, &stations, &velocity_model);
+
+        // Should generate a warning for this magnitude
+        assert!(!warning.event_id.is_empty());
+        assert!(warning.detecting_stations >= 2);
+        assert!(warning.confidence > 0.0);
+    }
+
+    #[test]
+    fn test_no_warning_for_small_event() {
+        let config = sample_data::fast_config();
+        let mut seismic_ai = SeismicAI::new(config).with_warning_config(WarningConfig {
+            min_magnitude: 5.0, // Higher threshold
+            ..Default::default()
+        });
+
+        let mut source = sample_data::local_earthquake();
+        source.magnitude = 3.0; // Below threshold
+
+        let stations = sample_data::sparse_network();
+        let velocity_model = sample_data::generic_velocity_model();
+
+        let warning = seismic_ai.generate_warning(&source, &stations, &velocity_model);
+
+        // Should not generate alert
+        assert_eq!(warning.alert_level, AlertLevel::None);
+    }
+
+    #[test]
+    fn test_simulation() {
+        let (source, stations, velocity_model, config) = sample_data::local_test_scenario();
+        let mut seismic_ai = SeismicAI::new(config);
+
+        let result = seismic_ai.simulate(&source, &stations, &velocity_model);
+
+        // Check result completeness
+        assert_eq!(result.ground_motions.len(), stations.len());
+        assert!(result.seismograms.is_some());
+        assert!(result.snapshots.is_some());
+        assert!(result.warning.is_some());
+        assert!(result.computation_time_ms > 0.0);
+    }
+
+    #[test]
+    fn test_seismogram_output() {
+        let (source, stations, velocity_model, mut config) = sample_data::local_test_scenario();
+        config.duration = 10.0; // Shorter for testing
+        let mut seismic_ai = SeismicAI::new(config);
+
+        let result = seismic_ai.simulate(&source, &stations, &velocity_model);
+
+        let seismograms = result.seismograms.unwrap();
+        assert_eq!(seismograms.len(), stations.len());
+
+        for seis in &seismograms {
+            assert!(!seis.station_code.is_empty());
+            assert!(!seis.time.is_empty());
+            assert_eq!(seis.east.len(), seis.time.len());
+        }
+    }
+
+    #[test]
+    fn test_wave_field_snapshot() {
+        let (source, stations, velocity_model, mut config) = sample_data::local_test_scenario();
+        config.duration = 5.0;
+        let mut seismic_ai = SeismicAI::new(config);
+
+        let result = seismic_ai.simulate(&source, &stations, &velocity_model);
+
+        let snapshots = result.snapshots.unwrap();
+        assert!(!snapshots.is_empty());
+
+        for snap in &snapshots {
+            assert!(!snap.displacement_x.is_empty());
+            assert!(!snap.grid_x.is_empty());
+            assert!(snap.time >= 0.0);
+        }
+    }
+
+    #[test]
+    fn test_run_demo() {
+        let result = run_demo();
+        assert!(result.computation_time_ms > 0.0);
+        assert!(!result.ground_motions.is_empty());
+    }
+
+    #[test]
+    fn test_attenuation_model_config() {
+        let config = sample_data::fast_config();
+        let seismic_ai = SeismicAI::new(config).with_attenuation_model(AttenuationModel::japan());
+
+        assert_eq!(seismic_ai.attenuation_model.name(), "Japan");
+    }
+
+    #[test]
+    fn test_warning_config() {
+        let config = sample_data::fast_config();
+        let warning_config = WarningConfig {
+            min_magnitude: 4.0,
+            min_stations: 5,
+            ..Default::default()
+        };
+        let seismic_ai = SeismicAI::new(config).with_warning_config(warning_config);
+
+        assert_eq!(seismic_ai.warning_config.min_magnitude, 4.0);
+        assert_eq!(seismic_ai.warning_config.min_stations, 5);
+    }
+
+    #[test]
+    fn test_major_scenario() {
+        // Just test it runs without panicking
+        let (source, stations, velocity_model, _) = sample_data::hayward_scenario();
+        assert!(source.magnitude >= 6.5);
+        assert!(stations.len() >= 50);
+        assert!(!velocity_model.layers.is_empty());
+    }
+}