90 lines
2.6 KiB
Rust
90 lines
2.6 KiB
Rust
//! # RTX Neural Operator Library
|
|
//!
|
|
//! Neural operators for learning solution operators to PDEs.
|
|
//!
|
|
//! This crate provides specialized neural network architectures for scientific computing:
|
|
//! - **Fourier Neural Operators (FNO)**: Learn mappings between function spaces using spectral convolutions
|
|
//! - **DeepONet**: Separate branch and trunk networks for operator learning
|
|
//!
|
|
//! ## Architecture
|
|
//!
|
|
//! Neural operators differ from traditional neural networks by learning mappings between
|
|
//! infinite-dimensional function spaces rather than finite-dimensional vectors. They are
|
|
//! particularly effective for solving partial differential equations (PDEs).
|
|
//!
|
|
//! ## Example
|
|
//!
|
|
//! ```rust,ignore
|
|
//! use rtx_neural_operator::{SpectralConv2d, FNO2d};
|
|
//! use rtx_backend_cpu::CpuBackend;
|
|
//!
|
|
//! let device = CpuBackend::default_device();
|
|
//! let fno = FNO2d::<CpuBackend>::new(1, 1, 32, 4, &device)?;
|
|
//! ```
|
|
|
|
#![warn(missing_docs)]
|
|
#![allow(clippy::module_name_repetitions)]
|
|
|
|
pub mod deeponet;
|
|
pub mod fno;
|
|
pub mod layers;
|
|
pub mod spectral;
|
|
pub mod weights;
|
|
|
|
// Re-export main types
|
|
pub use deeponet::DeepONet;
|
|
pub use fno::{FNO1d, FNO2d};
|
|
pub use layers::{GridPositionalEncoding, Lifting, LiftingMLP, Projection};
|
|
pub use spectral::{SpectralConv1d, SpectralConv2d};
|
|
pub use weights::{
|
|
DType, FNO2dConfig, FNO2dWeights, RawWeight, SafeTensorsFile, SpectralConvWeights, TensorInfo,
|
|
WeightError, load_config, load_fno2d_weights,
|
|
};
|
|
|
|
use thiserror::Error;
|
|
|
|
/// Result type for neural operator operations
|
|
pub type Result<T> = std::result::Result<T, NeuralOperatorError>;
|
|
|
|
/// Errors that can occur in neural operator operations
|
|
#[derive(Error, Debug)]
|
|
pub enum NeuralOperatorError {
|
|
/// Error from tensor operations
|
|
#[error("Tensor error: {0}")]
|
|
Tensor(#[from] rtx_tensor::TensorError),
|
|
|
|
/// Error from neural network operations
|
|
#[error("Neural network error: {0}")]
|
|
NeuralNetwork(String),
|
|
|
|
/// Invalid shape for operation
|
|
#[error("Invalid shape: {0}")]
|
|
InvalidShape(String),
|
|
|
|
/// Invalid parameter
|
|
#[error("Invalid parameter: {0}")]
|
|
InvalidParameter(String),
|
|
|
|
/// FFT operation failed
|
|
#[error("FFT error: {0}")]
|
|
FFT(String),
|
|
|
|
/// Mode truncation error
|
|
#[error("Mode truncation error: {0}")]
|
|
ModeTruncation(String),
|
|
|
|
/// Configuration error
|
|
#[error("Configuration error: {0}")]
|
|
Config(String),
|
|
|
|
/// Weight loading error
|
|
#[error("Weight loading error: {0}")]
|
|
Weight(#[from] WeightError),
|
|
}
|
|
|
|
impl From<rtx_nn::NNError> for NeuralOperatorError {
|
|
fn from(err: rtx_nn::NNError) -> Self {
|
|
Self::NeuralNetwork(err.to_string())
|
|
}
|
|
}
|