86 lines
2.4 KiB
Rust
86 lines
2.4 KiB
Rust
//! DeepONet (Deep Operator Network) architecture.
|
|
//!
|
|
//! DeepONet learns operators by decomposing them into two neural networks:
|
|
//! - **Branch network**: Encodes the input function
|
|
//! - **Trunk network**: Encodes the evaluation locations
|
|
//!
|
|
//! The output is computed as an inner product of branch and trunk outputs.
|
|
//!
|
|
//! ## References
|
|
//!
|
|
//! Lu, L., et al. (2021). "Learning nonlinear operators via DeepONet based on
|
|
//! the universal approximation theorem of operators." Nature Machine Intelligence.
|
|
|
|
use rtx_backend::Backend;
|
|
use rtx_nn::generic::GenericModule;
|
|
use rtx_tensor::generic::GenericTensor;
|
|
use std::fmt::Debug;
|
|
|
|
use crate::Result;
|
|
|
|
/// DeepONet operator network.
|
|
///
|
|
/// Learns mappings between function spaces using separate branch and trunk networks.
|
|
#[derive(Debug)]
|
|
pub struct DeepONet<B: Backend<FloatElem = f32>> {
|
|
_device: B::Device,
|
|
}
|
|
|
|
impl<B: Backend<FloatElem = f32>> DeepONet<B> {
|
|
/// Create a new DeepONet.
|
|
///
|
|
/// # Arguments
|
|
/// * `branch_input_dim` - Input dimension for branch network
|
|
/// * `trunk_input_dim` - Input dimension for trunk network
|
|
/// * `hidden_dim` - Hidden layer dimension
|
|
/// * `output_dim` - Output dimension
|
|
/// * `device` - Device to create the model on
|
|
pub fn new(
|
|
_branch_input_dim: usize,
|
|
_trunk_input_dim: usize,
|
|
_hidden_dim: usize,
|
|
_output_dim: usize,
|
|
device: &B::Device,
|
|
) -> Result<Self> {
|
|
Ok(Self {
|
|
_device: device.clone(),
|
|
})
|
|
}
|
|
}
|
|
|
|
impl<B: Backend<FloatElem = f32>> GenericModule<B> for DeepONet<B> {
|
|
fn forward(&self, input: &GenericTensor<B, 2>) -> GenericTensor<B, 2> {
|
|
// Placeholder: return zeros with correct shape
|
|
let shape = input.shape();
|
|
GenericTensor::zeros(shape, &self._device)
|
|
}
|
|
|
|
fn device(&self) -> &B::Device {
|
|
&self._device
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use rtx_backend_cpu::{CpuBackend, CpuDevice};
|
|
|
|
#[test]
|
|
fn test_deeponet_creation() {
|
|
let device = CpuDevice::new();
|
|
let result = DeepONet::<CpuBackend>::new(100, 2, 128, 1, &device);
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_deeponet_forward_shape() {
|
|
let device = CpuDevice::new();
|
|
let net = DeepONet::<CpuBackend>::new(100, 2, 128, 1, &device).unwrap();
|
|
|
|
let input = GenericTensor::randn([8, 100], &device);
|
|
let output = net.forward(&input);
|
|
|
|
assert_eq!(output.shape(), [8, 100]);
|
|
}
|
|
}
|