Initial commit

crates/specialized/rtx-science/src/physics/conservation.rs

@@ -0,0 +1,681 @@
+//! Conservation Laws for Physics-Informed Neural Networks
+//!
+//! This module implements various conservation laws that can be enforced
+//! in PINNs to ensure physical consistency of the learned solutions.
+
+use crate::Tensor;
+use crate::error::{ConservationLaw, Result, ScienceError};
+use crate::variable_extensions::VariableExt;
+use async_trait::async_trait;
+use rtx_autograd::Variable;
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+
+/// Trait for conservation law constraints in PINNs
+#[async_trait]
+pub trait ConservationLoss: Send + Sync {
+    /// Compute the conservation law violation
+    async fn compute_loss(
+        &self,
+        coordinates: &Tensor, // Input coordinates [x, t, ...]
+        solution: &Variable,  // Network output u(x,t)
+        du_dx: &Variable,     // ∂u/∂x
+        du_dt: &Variable,     // ∂u/∂t
+    ) -> Result<f32>;
+
+    /// Get the name of the conservation law
+    fn name(&self) -> &str;
+
+    /// Get the conservation law type
+    fn law_type(&self) -> ConservationLaw;
+
+    /// Get tolerance for violation detection
+    fn tolerance(&self) -> f64;
+
+    /// Validate conservation at given points
+    async fn validate_conservation(
+        &self,
+        coordinates: &Tensor,
+        solution: &Variable,
+        du_dx: &Variable,
+        du_dt: &Variable,
+    ) -> Result<ValidationResult>;
+}
+
+/// Conservation law validation result
+#[derive(Debug, Clone)]
+pub struct ValidationResult {
+    /// Maximum violation magnitude
+    pub max_violation: f64,
+    /// Mean violation magnitude
+    pub mean_violation: f64,
+    /// Points where tolerance is exceeded
+    pub violation_points: Vec<usize>,
+    /// Whether conservation is satisfied within tolerance
+    pub is_satisfied: bool,
+    /// Detailed diagnostics
+    pub diagnostics: HashMap<String, f64>,
+}
+
+/// Mass conservation law: ∂ρ/∂t + ∇·(ρv) = 0
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct MassConservation {
+    /// Density field identifier
+    pub density_field: String,
+    /// Velocity field identifiers [u, v, w]
+    pub velocity_fields: Vec<String>,
+    /// Tolerance for conservation violation
+    pub tolerance: f64,
+    /// Weight in loss function
+    pub weight: f64,
+}
+
+/// Momentum conservation law: ∂(ρv)/∂t + ∇·(ρv⊗v) = -∇p + μ∇²v + f
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct MomentumConservation {
+    /// Density field
+    pub density_field: String,
+    /// Velocity components
+    pub velocity_fields: Vec<String>,
+    /// Pressure field
+    pub pressure_field: String,
+    /// Viscosity coefficient
+    pub viscosity: f64,
+    /// External force fields
+    pub force_fields: Vec<String>,
+    /// Tolerance for conservation violation
+    pub tolerance: f64,
+    /// Weight in loss function  
+    pub weight: f64,
+}
+
+/// Energy conservation law: ∂E/∂t + ∇·(Ev + pv - k∇T) = Q
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct EnergyConservation {
+    /// Total energy field
+    pub energy_field: String,
+    /// Velocity fields
+    pub velocity_fields: Vec<String>,
+    /// Pressure field
+    pub pressure_field: String,
+    /// Temperature field
+    pub temperature_field: String,
+    /// Thermal conductivity
+    pub thermal_conductivity: f64,
+    /// Heat source term
+    pub heat_source: f64,
+    /// Tolerance for conservation violation
+    pub tolerance: f64,
+    /// Weight in loss function
+    pub weight: f64,
+}
+
+/// Angular momentum conservation: L = r × p = constant
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AngularMomentumConservation {
+    /// Position fields [x, y, z]
+    pub position_fields: Vec<String>,
+    /// Momentum fields [px, py, pz]  
+    pub momentum_fields: Vec<String>,
+    /// Reference point for angular momentum calculation
+    pub reference_point: Vec<f64>,
+    /// Tolerance for conservation violation
+    pub tolerance: f64,
+    /// Weight in loss function
+    pub weight: f64,
+}
+
+/// Charge conservation: ∂ρ/∂t + ∇·J = 0
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChargeConservation {
+    /// Charge density field
+    pub charge_density_field: String,
+    /// Current density fields [Jx, Jy, Jz]
+    pub current_density_fields: Vec<String>,
+    /// Tolerance for conservation violation
+    pub tolerance: f64,
+    /// Weight in loss function
+    pub weight: f64,
+}
+
+/// Probability conservation (quantum mechanics): ∫|ψ|²dx = 1
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ProbabilityConservation {
+    /// Wavefunction field (complex)
+    pub wavefunction_field: String,
+    /// Integration domain bounds
+    pub domain_bounds: Vec<(f64, f64)>,
+    /// Tolerance for conservation violation
+    pub tolerance: f64,
+    /// Weight in loss function
+    pub weight: f64,
+}
+
+/// Conservation validator for comprehensive checking
+pub struct ConservationValidator {
+    /// List of conservation laws to check
+    conservation_laws: Vec<Box<dyn ConservationLoss + Send + Sync>>,
+    /// Global tolerance settings
+    global_tolerance: f64,
+    /// Validation frequency (every N training steps)
+    validation_frequency: usize,
+    /// Current validation step
+    current_step: usize,
+}
+
+impl std::fmt::Debug for ConservationValidator {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ConservationValidator")
+            .field("num_conservation_laws", &self.conservation_laws.len())
+            .field("global_tolerance", &self.global_tolerance)
+            .field("validation_frequency", &self.validation_frequency)
+            .field("current_step", &self.current_step)
+            .finish()
+    }
+}
+
+impl MassConservation {
+    /// Create a new mass conservation law
+    #[must_use]
+    pub fn new(tolerance: f64) -> Self {
+        Self {
+            density_field: "rho".to_string(),
+            velocity_fields: vec!["u".to_string(), "v".to_string()],
+            tolerance,
+            weight: 1.0,
+        }
+    }
+
+    /// Set velocity field names
+    #[must_use]
+    pub fn with_velocity_fields(mut self, fields: Vec<String>) -> Self {
+        self.velocity_fields = fields;
+        self
+    }
+
+    /// Set loss weight
+    #[must_use]
+    pub fn with_weight(mut self, weight: f64) -> Self {
+        self.weight = weight;
+        self
+    }
+}
+
+#[async_trait]
+impl ConservationLoss for MassConservation {
+    async fn compute_loss(
+        &self,
+        _coordinates: &Tensor,
+        _solution: &Variable,
+        du_dx: &Variable,
+        du_dt: &Variable,
+    ) -> Result<f32> {
+        // For simplified mass conservation: ∂ρ/∂t + ∂(ρu)/∂x = 0
+        // Assuming solution represents density ρ and we have velocity field
+
+        // ∂ρ/∂t term
+        let drho_dt = du_dt.clone();
+
+        // ∂(ρu)/∂x term - simplified to ρ∂u/∂x + u∂ρ/∂x
+        // For this example, assume constant velocity or treat solution as ρu
+        let drhou_dx = du_dx.clone();
+
+        // Mass conservation residual: ∂ρ/∂t + ∂(ρu)/∂x = 0
+        let residual = drho_dt.add(&drhou_dx)?;
+
+        // Return mean squared residual as f32
+        let squared_residual = residual.multiply(&residual)?;
+        let loss_var = squared_residual.mean_square()?;
+        let loss_value = loss_var.value().to_scalar::<f32>()?;
+
+        Ok(loss_value)
+    }
+
+    fn name(&self) -> &'static str {
+        "Mass Conservation"
+    }
+
+    fn law_type(&self) -> ConservationLaw {
+        ConservationLaw::Mass
+    }
+
+    fn tolerance(&self) -> f64 {
+        self.tolerance
+    }
+
+    async fn validate_conservation(
+        &self,
+        coordinates: &Tensor,
+        solution: &Variable,
+        du_dx: &Variable,
+        du_dt: &Variable,
+    ) -> Result<ValidationResult> {
+        // Compute conservation residual
+        let residual = self
+            .compute_loss(coordinates, solution, du_dx, du_dt)
+            .await?;
+        let residual_values = [residual]; // f32 scalar converted to vec
+
+        let violations: Vec<f64> = residual_values
+            .iter()
+            .map(|&x| f64::from(x.abs()))
+            .collect();
+
+        let max_violation = violations.iter().copied().fold(0.0, f64::max);
+        let mean_violation = violations.iter().sum::<f64>() / violations.len() as f64;
+
+        let violation_points: Vec<usize> = violations
+            .iter()
+            .enumerate()
+            .filter_map(|(i, &v)| if v > self.tolerance { Some(i) } else { None })
+            .collect();
+
+        let is_satisfied = max_violation <= self.tolerance;
+
+        let mut diagnostics = HashMap::new();
+        diagnostics.insert("max_violation".to_string(), max_violation);
+        diagnostics.insert("mean_violation".to_string(), mean_violation);
+        diagnostics.insert(
+            "violation_rate".to_string(),
+            violation_points.len() as f64 / violations.len() as f64,
+        );
+
+        Ok(ValidationResult {
+            max_violation,
+            mean_violation,
+            violation_points,
+            is_satisfied,
+            diagnostics,
+        })
+    }
+}
+
+impl MomentumConservation {
+    /// Create a new momentum conservation law
+    #[must_use]
+    pub fn new(viscosity: f64, tolerance: f64) -> Self {
+        Self {
+            density_field: "rho".to_string(),
+            velocity_fields: vec!["u".to_string(), "v".to_string()],
+            pressure_field: "p".to_string(),
+            viscosity,
+            force_fields: Vec::new(),
+            tolerance,
+            weight: 1.0,
+        }
+    }
+
+    /// Add external force field
+    #[must_use]
+    pub fn with_force_field(mut self, force_field: String) -> Self {
+        self.force_fields.push(force_field);
+        self
+    }
+}
+
+#[async_trait]
+impl ConservationLoss for MomentumConservation {
+    async fn compute_loss(
+        &self,
+        _coordinates: &Tensor,
+        _solution: &Variable,
+        du_dx: &Variable,
+        du_dt: &Variable,
+    ) -> Result<f32> {
+        // Simplified momentum conservation: ∂(ρu)/∂t + ∂p/∂x = μ∂²u/∂x²
+        // Assuming solution represents momentum ρu
+
+        let momentum_time_derivative = du_dt.clone();
+
+        // For now, assume pressure gradient is represented by some component
+        // In a full implementation, this would be computed from pressure field
+        let pressure_gradient = du_dx.multiply_scalar(0.1)?; // Placeholder
+
+        // Viscous term (simplified)
+        let viscous_term = du_dx.multiply_scalar(self.viscosity as f32)?;
+
+        // Momentum conservation residual
+        let residual = momentum_time_derivative
+            .add(&pressure_gradient)?
+            .subtract(&viscous_term)?;
+
+        let squared_residual = residual.multiply(&residual)?;
+        let loss_var = squared_residual.mean_square()?;
+        let loss_value = loss_var.value().to_scalar::<f32>()?;
+
+        Ok(loss_value)
+    }
+
+    fn name(&self) -> &'static str {
+        "Momentum Conservation"
+    }
+
+    fn law_type(&self) -> ConservationLaw {
+        ConservationLaw::Momentum
+    }
+
+    fn tolerance(&self) -> f64 {
+        self.tolerance
+    }
+
+    async fn validate_conservation(
+        &self,
+        coordinates: &Tensor,
+        solution: &Variable,
+        du_dx: &Variable,
+        du_dt: &Variable,
+    ) -> Result<ValidationResult> {
+        let residual = self
+            .compute_loss(coordinates, solution, du_dx, du_dt)
+            .await?;
+        let residual_values = [residual]; // f32 scalar converted to vec
+
+        let violations: Vec<f64> = residual_values
+            .iter()
+            .map(|&x| f64::from(x.abs()))
+            .collect();
+
+        let max_violation = violations.iter().copied().fold(0.0, f64::max);
+        let mean_violation = violations.iter().sum::<f64>() / violations.len() as f64;
+
+        let violation_points: Vec<usize> = violations
+            .iter()
+            .enumerate()
+            .filter_map(|(i, &v)| if v > self.tolerance { Some(i) } else { None })
+            .collect();
+
+        let is_satisfied = max_violation <= self.tolerance;
+
+        let mut diagnostics = HashMap::new();
+        diagnostics.insert("max_violation".to_string(), max_violation);
+        diagnostics.insert("mean_violation".to_string(), mean_violation);
+        diagnostics.insert("viscosity".to_string(), self.viscosity);
+
+        Ok(ValidationResult {
+            max_violation,
+            mean_violation,
+            violation_points,
+            is_satisfied,
+            diagnostics,
+        })
+    }
+}
+
+impl EnergyConservation {
+    /// Create a new energy conservation law
+    #[must_use]
+    pub fn new(thermal_conductivity: f64, tolerance: f64) -> Self {
+        Self {
+            energy_field: "E".to_string(),
+            velocity_fields: vec!["u".to_string(), "v".to_string()],
+            pressure_field: "p".to_string(),
+            temperature_field: "T".to_string(),
+            thermal_conductivity,
+            heat_source: 0.0,
+            tolerance,
+            weight: 1.0,
+        }
+    }
+
+    /// Set heat source term
+    #[must_use]
+    pub fn with_heat_source(mut self, heat_source: f64) -> Self {
+        self.heat_source = heat_source;
+        self
+    }
+}
+
+#[async_trait]
+impl ConservationLoss for EnergyConservation {
+    async fn compute_loss(
+        &self,
+        _coordinates: &Tensor,
+        _solution: &Variable,
+        du_dx: &Variable,
+        du_dt: &Variable,
+    ) -> Result<f32> {
+        // Simplified energy conservation: ∂E/∂t = -∂q/∂x + Q
+        // where q is heat flux and Q is heat source
+
+        let energy_time_derivative = du_dt.clone();
+
+        // Heat flux term (simplified as thermal conductivity * temperature gradient)
+        let heat_flux_gradient = du_dx.multiply_scalar(self.thermal_conductivity as f32)?;
+
+        // Heat source term
+        let du_dt_tensor = du_dt.value();
+        let source_shape = &[du_dt_tensor.shape().dims()[0]];
+        let source_tensor =
+            Tensor::full(source_shape, self.heat_source as f32, du_dt_tensor.device())?;
+        let source_term = Variable::constant(source_tensor);
+
+        // Energy conservation residual
+        let residual = energy_time_derivative
+            .add(&heat_flux_gradient)?
+            .subtract(&source_term)?;
+
+        let squared_residual = residual.multiply(&residual)?;
+        let loss_var = squared_residual.mean_square()?;
+        let loss_value = loss_var.value().to_scalar::<f32>()?;
+
+        Ok(loss_value)
+    }
+
+    fn name(&self) -> &'static str {
+        "Energy Conservation"
+    }
+
+    fn law_type(&self) -> ConservationLaw {
+        ConservationLaw::Energy
+    }
+
+    fn tolerance(&self) -> f64 {
+        self.tolerance
+    }
+
+    async fn validate_conservation(
+        &self,
+        coordinates: &Tensor,
+        solution: &Variable,
+        du_dx: &Variable,
+        du_dt: &Variable,
+    ) -> Result<ValidationResult> {
+        let residual = self
+            .compute_loss(coordinates, solution, du_dx, du_dt)
+            .await?;
+        let residual_values = [residual]; // f32 scalar converted to vec
+
+        let violations: Vec<f64> = residual_values
+            .iter()
+            .map(|&x| f64::from(x.abs()))
+            .collect();
+
+        let max_violation = violations.iter().copied().fold(0.0, f64::max);
+        let mean_violation = violations.iter().sum::<f64>() / violations.len() as f64;
+
+        let violation_points: Vec<usize> = violations
+            .iter()
+            .enumerate()
+            .filter_map(|(i, &v)| if v > self.tolerance { Some(i) } else { None })
+            .collect();
+
+        let is_satisfied = max_violation <= self.tolerance;
+
+        let mut diagnostics = HashMap::new();
+        diagnostics.insert("max_violation".to_string(), max_violation);
+        diagnostics.insert("mean_violation".to_string(), mean_violation);
+        diagnostics.insert(
+            "thermal_conductivity".to_string(),
+            self.thermal_conductivity,
+        );
+        diagnostics.insert("heat_source".to_string(), self.heat_source);
+
+        Ok(ValidationResult {
+            max_violation,
+            mean_violation,
+            violation_points,
+            is_satisfied,
+            diagnostics,
+        })
+    }
+}
+
+impl ConservationValidator {
+    /// Create a new conservation validator
+    #[must_use]
+    pub fn new(global_tolerance: f64, validation_frequency: usize) -> Self {
+        Self {
+            conservation_laws: Vec::new(),
+            global_tolerance,
+            validation_frequency,
+            current_step: 0,
+        }
+    }
+
+    /// Add a conservation law to validate
+    #[must_use]
+    pub fn add_conservation_law(mut self, law: Box<dyn ConservationLoss + Send + Sync>) -> Self {
+        self.conservation_laws.push(law);
+        self
+    }
+
+    /// Validate all conservation laws
+    pub async fn validate_all(
+        &mut self,
+        coordinates: &Tensor,
+        solution: &Variable,
+        du_dx: &Variable,
+        du_dt: &Variable,
+    ) -> Result<Vec<ValidationResult>> {
+        self.current_step += 1;
+
+        // Only validate at specified frequency
+        if !self.current_step.is_multiple_of(self.validation_frequency) {
+            return Ok(Vec::new());
+        }
+
+        let mut results = Vec::new();
+
+        for law in &self.conservation_laws {
+            let result = law
+                .validate_conservation(coordinates, solution, du_dx, du_dt)
+                .await?;
+
+            // Log critical violations
+            if !result.is_satisfied && result.max_violation > self.global_tolerance * 10.0 {
+                tracing::warn!(
+                    "Critical conservation violation in {}: max = {:.2e}, tolerance = {:.2e}",
+                    law.name(),
+                    result.max_violation,
+                    law.tolerance()
+                );
+            }
+
+            results.push(result);
+        }
+
+        Ok(results)
+    }
+
+    /// Get summary of all conservation laws
+    #[must_use]
+    pub fn summary(&self) -> HashMap<String, String> {
+        let mut summary = HashMap::new();
+
+        summary.insert(
+            "num_laws".to_string(),
+            self.conservation_laws.len().to_string(),
+        );
+        summary.insert(
+            "global_tolerance".to_string(),
+            self.global_tolerance.to_string(),
+        );
+        summary.insert(
+            "validation_frequency".to_string(),
+            self.validation_frequency.to_string(),
+        );
+        summary.insert("current_step".to_string(), self.current_step.to_string());
+
+        for (i, law) in self.conservation_laws.iter().enumerate() {
+            summary.insert(format!("law_{i}_name"), law.name().to_string());
+            summary.insert(format!("law_{i}_tolerance"), law.tolerance().to_string());
+        }
+
+        summary
+    }
+}
+
+/// Factory function for creating conservation laws
+pub fn create_conservation_law(
+    law_type: ConservationLaw,
+    tolerance: f64,
+    parameters: &HashMap<String, f64>,
+) -> Result<Box<dyn ConservationLoss + Send + Sync>> {
+    match law_type {
+        ConservationLaw::Mass => Ok(Box::new(MassConservation::new(tolerance))),
+        ConservationLaw::Momentum => {
+            let viscosity = parameters.get("viscosity").copied().unwrap_or(1e-3);
+            Ok(Box::new(MomentumConservation::new(viscosity, tolerance)))
+        }
+        ConservationLaw::Energy => {
+            let thermal_conductivity = parameters
+                .get("thermal_conductivity")
+                .copied()
+                .unwrap_or(0.1);
+            let mut energy_law = EnergyConservation::new(thermal_conductivity, tolerance);
+            if let Some(&heat_source) = parameters.get("heat_source") {
+                energy_law = energy_law.with_heat_source(heat_source);
+            }
+            Ok(Box::new(energy_law))
+        }
+        _ => Err(ScienceError::physics(
+            format!("Unsupported conservation law: {law_type:?}"),
+            crate::error::PhysicsDomain::FluidDynamics,
+        )),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::Device;
+
+    #[tokio::test]
+    async fn test_mass_conservation() -> Result<()> {
+        let device = Device::cpu();
+        let conservation = MassConservation::new(1e-6);
+
+        // Create dummy data
+        let coords = Tensor::zeros(&[10, 2], &device)?;
+        let solution = Variable::new(Tensor::ones(&[10], &device)?, true);
+        let du_dx = Variable::new(Tensor::zeros(&[10], &device)?, false);
+        let du_dt = Variable::new(Tensor::zeros(&[10], &device)?, false);
+
+        let loss = conservation
+            .compute_loss(&coords, &solution, &du_dx, &du_dt)
+            .await?;
+        assert!(loss >= 0.0);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_conservation_validator() {
+        let validator = ConservationValidator::new(1e-6, 10);
+        assert_eq!(validator.conservation_laws.len(), 0);
+
+        let summary = validator.summary();
+        assert_eq!(summary["num_laws"], "0");
+    }
+
+    #[test]
+    fn test_conservation_law_factory() -> Result<()> {
+        let mut params = HashMap::new();
+        params.insert("viscosity".to_string(), 1e-3);
+
+        let law = create_conservation_law(ConservationLaw::Momentum, 1e-6, &params)?;
+        assert_eq!(law.name(), "Momentum Conservation");
+
+        Ok(())
+    }
+}