use ndarray::{Array2, Array3};
use crate::psf::{Kernel2D, Kernel3D};
use crate::{Error, Result};
use super::operator::{LinearOperator2D, LinearOperator3D};
use super::validate::{finite_real_2d, finite_real_3d};
pub(crate) struct Convolution2D {
kernel: Array2<f32>,
adjoint_kernel: Array2<f32>,
}
impl Convolution2D {
pub(crate) fn new(kernel: &Kernel2D) -> Result<Self> {
let kernel = kernel.as_array().as_standard_layout().to_owned();
finite_real_2d(&kernel)?;
let adjoint_kernel = flip_2d(&kernel);
Ok(Self {
kernel,
adjoint_kernel,
})
}
}
impl LinearOperator2D for Convolution2D {
fn apply(&self, input: &Array2<f32>) -> Result<Array2<f32>> {
convolve_2d_same(input, &self.kernel)
}
fn adjoint(&self, input: &Array2<f32>) -> Result<Array2<f32>> {
convolve_2d_same(input, &self.adjoint_kernel)
}
}
pub(crate) struct Convolution3D {
kernel: Array3<f32>,
adjoint_kernel: Array3<f32>,
}
impl Convolution3D {
pub(crate) fn new(kernel: &Kernel3D) -> Result<Self> {
let kernel = kernel.as_array().as_standard_layout().to_owned();
finite_real_3d(&kernel)?;
let adjoint_kernel = flip_3d(&kernel);
Ok(Self {
kernel,
adjoint_kernel,
})
}
}
impl LinearOperator3D for Convolution3D {
fn apply(&self, input: &Array3<f32>) -> Result<Array3<f32>> {
convolve_3d_same(input, &self.kernel)
}
fn adjoint(&self, input: &Array3<f32>) -> Result<Array3<f32>> {
convolve_3d_same(input, &self.adjoint_kernel)
}
}
fn convolve_2d_same(input: &Array2<f32>, kernel: &Array2<f32>) -> Result<Array2<f32>> {
finite_real_2d(input)?;
finite_real_2d(kernel)?;
let (height, width) = input.dim();
let (kernel_h, kernel_w) = kernel.dim();
let center_y = to_i64(kernel_h / 2)?;
let center_x = to_i64(kernel_w / 2)?;
let height_i64 = to_i64(height)?;
let width_i64 = to_i64(width)?;
let mut output = Array2::zeros((height, width));
for y in 0..height {
let y_i64 = y as i64;
for x in 0..width {
let x_i64 = x as i64;
let mut acc = 0.0_f32;
for ky in 0..kernel_h {
let ky_i64 = ky as i64;
let iy = y_i64 + ky_i64 - center_y;
if iy < 0 || iy >= height_i64 {
continue;
}
let iy_usize = iy as usize;
for kx in 0..kernel_w {
let kx_i64 = kx as i64;
let ix = x_i64 + kx_i64 - center_x;
if ix < 0 || ix >= width_i64 {
continue;
}
let ix_usize = ix as usize;
acc += input[[iy_usize, ix_usize]] * kernel[[ky, kx]];
}
}
if !acc.is_finite() {
return Err(Error::NonFiniteInput);
}
output[[y, x]] = acc;
}
}
Ok(output)
}
fn convolve_3d_same(input: &Array3<f32>, kernel: &Array3<f32>) -> Result<Array3<f32>> {
finite_real_3d(input)?;
finite_real_3d(kernel)?;
let (depth, height, width) = input.dim();
let (kernel_d, kernel_h, kernel_w) = kernel.dim();
let center_d = to_i64(kernel_d / 2)?;
let center_y = to_i64(kernel_h / 2)?;
let center_x = to_i64(kernel_w / 2)?;
let depth_i64 = to_i64(depth)?;
let height_i64 = to_i64(height)?;
let width_i64 = to_i64(width)?;
let mut output = Array3::zeros((depth, height, width));
for d in 0..depth {
let d_i64 = d as i64;
for y in 0..height {
let y_i64 = y as i64;
for x in 0..width {
let x_i64 = x as i64;
let mut acc = 0.0_f32;
for kd in 0..kernel_d {
let kd_i64 = kd as i64;
let id = d_i64 + kd_i64 - center_d;
if id < 0 || id >= depth_i64 {
continue;
}
let id_usize = id as usize;
for ky in 0..kernel_h {
let ky_i64 = ky as i64;
let iy = y_i64 + ky_i64 - center_y;
if iy < 0 || iy >= height_i64 {
continue;
}
let iy_usize = iy as usize;
for kx in 0..kernel_w {
let kx_i64 = kx as i64;
let ix = x_i64 + kx_i64 - center_x;
if ix < 0 || ix >= width_i64 {
continue;
}
let ix_usize = ix as usize;
acc += input[[id_usize, iy_usize, ix_usize]] * kernel[[kd, ky, kx]];
}
}
}
if !acc.is_finite() {
return Err(Error::NonFiniteInput);
}
output[[d, y, x]] = acc;
}
}
}
Ok(output)
}
fn flip_2d(kernel: &Array2<f32>) -> Array2<f32> {
let (height, width) = kernel.dim();
let mut flipped = Array2::zeros((height, width));
for y in 0..height {
for x in 0..width {
flipped[[height - 1 - y, width - 1 - x]] = kernel[[y, x]];
}
}
flipped
}
fn flip_3d(kernel: &Array3<f32>) -> Array3<f32> {
let (depth, height, width) = kernel.dim();
let mut flipped = Array3::zeros((depth, height, width));
for d in 0..depth {
for y in 0..height {
for x in 0..width {
flipped[[depth - 1 - d, height - 1 - y, width - 1 - x]] = kernel[[d, y, x]];
}
}
}
flipped
}
fn to_i64(value: usize) -> Result<i64> {
i64::try_from(value).map_err(|_| Error::DimensionMismatch)
}
#[cfg(test)]
mod tests {
use ndarray::{array, Array2, Array3};
use super::{Convolution2D, Convolution3D};
use crate::core::operator::{
inner_product_2d, inner_product_3d, LinearOperator2D, LinearOperator3D,
};
use crate::Kernel2D;
use crate::Kernel3D;
#[test]
fn convolution_2d_passes_adjoint_consistency() {
let kernel = Kernel2D::new(array![
[0.0_f32, 1.0_f32, 0.0_f32],
[1.0_f32, 2.0_f32, 1.0_f32],
[0.0_f32, 1.0_f32, 0.0_f32]
])
.unwrap();
let op = Convolution2D::new(&kernel).unwrap();
let x: Array2<f32> = array![[1.0_f32, 2.0_f32, 3.0_f32], [4.0_f32, 5.0_f32, 6.0_f32]];
let y: Array2<f32> = array![[0.5_f32, 1.5_f32, 0.5_f32], [2.0_f32, 1.0_f32, 0.0_f32]];
let ax = op.apply(&x).unwrap();
let aty = op.adjoint(&y).unwrap();
let lhs = inner_product_2d(&ax, &y).unwrap();
let rhs = inner_product_2d(&x, &aty).unwrap();
assert!((lhs - rhs).abs() < 1e-4);
}
#[test]
fn convolution_3d_passes_adjoint_consistency() {
let kernel = Kernel3D::new(array![[
[0.0_f32, 0.0_f32, 0.0_f32],
[0.0_f32, 1.0_f32, 0.0_f32],
[0.0_f32, 0.0_f32, 0.0_f32]
]])
.unwrap();
let op = Convolution3D::new(&kernel).unwrap();
let x: Array3<f32> = array![
[[1.0_f32, 2.0_f32], [3.0_f32, 4.0_f32]],
[[5.0_f32, 6.0_f32], [7.0_f32, 8.0_f32]]
];
let y: Array3<f32> = array![
[[0.5_f32, 0.0_f32], [1.0_f32, 1.5_f32]],
[[2.0_f32, 0.5_f32], [1.0_f32, 0.0_f32]]
];
let ax = op.apply(&x).unwrap();
let aty = op.adjoint(&y).unwrap();
let lhs = inner_product_3d(&ax, &y).unwrap();
let rhs = inner_product_3d(&x, &aty).unwrap();
assert!((lhs - rhs).abs() < 1e-4);
}
}