qpet_sz/

lib.rs

use std::num::NonZeroUsize;

use ndarray::{Array, ArrayView, Dimension, IxDyn};

#[derive(Clone, Debug)]
pub enum CompressionMode<'a> {
    /// Symbolic Quantity of Interest
    SymbolicQuantityOfInterest {
        /// quantity of interest expression
        qoi: &'a str,
        /// side length of the region of the quantity of interest,
        /// 1 for pointwise
        qoi_region_size: NonZeroUsize,
        /// quantity of interest error bound
        qoi_error_bound: QoIErrorBound,
        /// data error bound
        data_error_bound: DataErrorBound,
        /// positive quantity of interest c parameter (2.0 is a good default)
        qoi_c: f64,
    },
}

#[derive(Clone, Debug, Copy)]
pub enum DataErrorBound {
    Absolute(f64),
    Relative(f64),
    PSNR(f64),
    L2Norm(f64),
    AbsoluteAndRelative {
        absolute_bound: f64,
        relative_bound: f64,
    },
    AbsoluteOrRelative {
        absolute_bound: f64,
        relative_bound: f64,
    },
}

#[derive(Clone, Debug, Copy)]
pub enum QoIErrorBound {
    /// Absolute error bound
    Absolute(f64),
    /// Range-relative error bound
    Relative(f64),
}

impl DataErrorBound {
    const fn code(&self) -> u8 {
        #[allow(clippy::cast_possible_truncation)]
        match self {
            Self::Absolute(_) => const { qpet_sz_sys::QoZ_EB_EB_ABS as u8 },
            Self::Relative(_) => const { qpet_sz_sys::QoZ_EB_EB_REL as u8 },
            Self::PSNR(_) => const { qpet_sz_sys::QoZ_EB_EB_PSNR as u8 },
            Self::L2Norm(_) => const { qpet_sz_sys::QoZ_EB_EB_L2NORM as u8 },
            Self::AbsoluteAndRelative { .. } => const { qpet_sz_sys::QoZ_EB_EB_ABS_AND_REL as u8 },
            Self::AbsoluteOrRelative { .. } => const { qpet_sz_sys::QoZ_EB_EB_ABS_OR_REL as u8 },
        }
    }

    const fn abs_bound(&self) -> f64 {
        match self {
            Self::Absolute(bound) => *bound,
            Self::AbsoluteOrRelative { absolute_bound, .. }
            | Self::AbsoluteAndRelative { absolute_bound, .. } => *absolute_bound,
            _ => 0.0,
        }
    }

    const fn rel_bound(&self) -> f64 {
        match self {
            Self::Relative(bound) => *bound,
            Self::AbsoluteOrRelative { relative_bound, .. }
            | Self::AbsoluteAndRelative { relative_bound, .. } => *relative_bound,
            _ => 0.0,
        }
    }

    const fn l2norm_bound(&self) -> f64 {
        match self {
            Self::L2Norm(bound) => *bound,
            _ => 0.0,
        }
    }

    const fn psnr_bound(&self) -> f64 {
        match self {
            Self::PSNR(bound) => *bound,
            _ => 0.0,
        }
    }
}

impl QoIErrorBound {
    const fn code(&self) -> u8 {
        #[allow(clippy::cast_possible_truncation)]
        match self {
            Self::Absolute(_) => const { qpet_sz_sys::QoZ_EB_EB_ABS as u8 },
            Self::Relative(_) => const { qpet_sz_sys::QoZ_EB_EB_REL as u8 },
        }
    }

    const fn bound(&self) -> f64 {
        match self {
            Self::Absolute(bound) | Self::Relative(bound) => *bound,
        }
    }
}

#[derive(Clone, Debug)]
pub struct Config<'a> {
    mode: CompressionMode<'a>,
    openmp: bool,
}

pub trait Element: private::Element {}
impl Element for f32 {}
impl Element for f64 {}

mod private {
    pub trait Element: Copy {
        unsafe fn compress(
            config: qpet_sz_sys::QoZ_Config,
            data: *const Self,
            len: *mut usize,
        ) -> *mut u8;

        unsafe fn decompress(
            compressed_data: *const u8,
            compressed_len: usize,
            decompressed_data: *mut *mut Self,
        ) -> qpet_sz_sys::QoZ_Config;

        unsafe fn dealloc(data: *mut Self);
    }

    impl Element for f32 {
        unsafe fn compress(
            config: qpet_sz_sys::QoZ_Config,
            data: *const Self,
            len: *mut usize,
        ) -> *mut u8 {
            qpet_sz_sys::compress_float(config, data, len).cast()
        }

        unsafe fn decompress(
            compressed_data: *const u8,
            compressed_len: usize,
            decompressed_data: *mut *mut Self,
        ) -> qpet_sz_sys::QoZ_Config {
            qpet_sz_sys::decompress_float(compressed_data.cast(), compressed_len, decompressed_data)
        }

        unsafe fn dealloc(data: *mut Self) {
            qpet_sz_sys::dealloc_float(data);
        }
    }

    impl Element for f64 {
        unsafe fn compress(
            config: qpet_sz_sys::QoZ_Config,
            data: *const Self,
            len: *mut usize,
        ) -> *mut u8 {
            qpet_sz_sys::compress_double(config, data, len).cast()
        }

        unsafe fn decompress(
            compressed_data: *const u8,
            compressed_len: usize,
            decompressed_data: *mut *mut Self,
        ) -> qpet_sz_sys::QoZ_Config {
            qpet_sz_sys::decompress_double(
                compressed_data.cast(),
                compressed_len,
                decompressed_data,
            )
        }

        unsafe fn dealloc(data: *mut Self) {
            qpet_sz_sys::dealloc_double(data);
        }
    }
}

impl<'a> Config<'a> {
    #[must_use]
    pub const fn new(mode: CompressionMode<'a>) -> Self {
        Self {
            mode,
            openmp: false,
        }
    }

    #[cfg(feature = "openmp")]
    #[must_use]
    pub const fn openmp(mut self, openmp: bool) -> Self {
        self.openmp = openmp;
        self
    }
}

#[allow(clippy::needless_pass_by_value)]
#[allow(clippy::missing_panics_doc)] // FIXME
pub fn compress<T: Element, D: Dimension>(data: ArrayView<T, D>, config: &Config) -> Vec<u8> {
    let Config {
        mode:
            CompressionMode::SymbolicQuantityOfInterest {
                qoi,
                qoi_region_size,
                qoi_error_bound,
                data_error_bound,
                qoi_c,
            },
        openmp,
    } = config;

    let mut qoi_string = Vec::from(qoi.as_bytes());
    qoi_string.push(b'\0');

    #[allow(clippy::unwrap_used)] // FIXME
    let raw_config = qpet_sz_sys::QoZ_Config {
        N: data.ndim().try_into().unwrap(),
        dims: data.shape().as_ptr().cast_mut(),
        num: data.len(),
        errorBoundMode: data_error_bound.code(),
        absErrorBound: data_error_bound.abs_bound(),
        relErrorBound: data_error_bound.rel_bound(),
        l2normErrorBound: data_error_bound.l2norm_bound(),
        psnrErrorBound: data_error_bound.psnr_bound(),
        openmp: *openmp,
        qoi: 14, // symbolic QoI
        qoiEB: qoi_error_bound.bound(),
        qoiEBMode: qoi_error_bound.code(),
        qoi_string: qoi_string.as_mut_ptr().cast(),
        qoiRegionSize: qoi_region_size.get().try_into().unwrap(),
        qoiRegionMode: match *qoi_region_size {
            NonZeroUsize::MIN => 0, // pointwise
            _ => 1,                 // regional average
        },
        error_std_rate: *qoi_c,
        analytical: true,
    };

    // FIXME: remove debug print
    eprintln!(
        "compress config N={} num={} shape={:?}",
        raw_config.N,
        raw_config.num,
        data.shape()
    );

    let mut len: usize = 0;

    let compressed_ptr = unsafe { T::compress(raw_config, data.as_ptr(), &raw mut len) };

    let compressed_data = unsafe { std::slice::from_raw_parts(compressed_ptr, len) };
    let compressed_data = Vec::from(compressed_data);

    unsafe { qpet_sz_sys::dealloc_char(compressed_ptr.cast()) };

    compressed_data
}

#[must_use]
#[allow(clippy::missing_panics_doc)] // FIXME
pub fn decompress<T: Element>(compressed_data: &[u8]) -> Array<T, IxDyn> {
    let mut decompressed_ptr = std::ptr::null_mut();

    let config = unsafe {
        T::decompress(
            compressed_data.as_ptr(),
            compressed_data.len(),
            &raw mut decompressed_ptr,
        )
    };

    let decompressed_data = unsafe { std::slice::from_raw_parts(decompressed_ptr, config.num) };
    let decompressed_data = Vec::from(decompressed_data);

    unsafe { T::dealloc(decompressed_ptr) };

    #[allow(clippy::unwrap_used)] // FIXME
    let shape: Vec<usize> = (0..config.N)
        .map(|i| unsafe { std::ptr::read(config.dims.add(i.try_into().unwrap())) })
        .collect();

    // FIXME: remove debug print
    eprintln!(
        "decompress config N={} num={} shape={:?}",
        config.N, config.num, shape
    );

    unsafe {
        qpet_sz_sys::dealloc_size_t(config.dims);
    }

    unsafe {
        qpet_sz_sys::dealloc_char(config.qoi_string);
    }

    #[allow(clippy::unwrap_used)] // FIXME
    let decompressed_data = Array::from_shape_vec(shape, decompressed_data).unwrap();

    decompressed_data
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;

    #[test]
    fn log10_decode() {
        let encoded = compress(
            Array::<f32, _>::logspace(2.0, 0.0, 100.0, 721 * 1440)
                .into_shape_clone((721, 1440))
                .unwrap()
                .view(),
            &Config::new(CompressionMode::SymbolicQuantityOfInterest {
                qoi: "log(x, 10)",
                qoi_region_size: NonZeroUsize::MIN,
                qoi_error_bound: QoIErrorBound::Absolute(0.25),
                data_error_bound: DataErrorBound::Absolute(f64::MAX),
                qoi_c: 2.0,
            }),
        );
        let decoded = decompress::<f32>(&encoded);

        assert_eq!(decoded.len(), 721 * 1440);
        assert_eq!(decoded.shape(), &[721, 1440]);
    }
}