gspx 0.1.2

Sparse graph signal processing and spectral graph wavelets in Rust
Documentation
use std::fmt;
use std::fs::File;
use std::path::{Path, PathBuf};

use ndarray::Array2;
use serde_json::Value;

use crate::{ChebyKernel, Laplacian, VfKernel, error::GspError};

use super::{load_laplacian, load_ply_laplacian, load_ply_xyz, load_signal};

/// Friendly handle to the named example datasets used throughout the examples.
#[derive(Debug, Clone)]
pub struct ExampleData {
    root: PathBuf,
}

impl ExampleData {
    /// Uses a caller-provided local example-data directory.
    ///
    /// # Errors
    /// Returns an error when the directory does not exist or is not a directory.
    pub fn from_dir<P: Into<PathBuf>>(root: P) -> Result<Self, GspError> {
        let root = root.into();
        if !root.is_dir() {
            return Err(GspError::ResourceNotFound(format!(
                "example-data directory not found: {}",
                root.display()
            )));
        }
        Ok(Self { root })
    }

    /// Returns the underlying data root.
    #[must_use]
    pub fn root(&self) -> &Path {
        &self.root
    }

    /// Loads a graph Laplacian for a dataset and graph kind.
    ///
    /// # Errors
    /// Returns an error when the selected combination is unavailable or malformed.
    pub fn graph(&self, dataset: Dataset, kind: GraphKind) -> Result<Laplacian, GspError> {
        let Some(path) = dataset.graph_path(&self.root, kind) else {
            return Err(GspError::ResourceNotFound(format!(
                "{dataset} does not provide a {kind} graph"
            )));
        };
        if !path.exists() {
            return Err(GspError::ResourceNotFound(format!(
                "graph file not found: {}",
                path.display()
            )));
        }
        match kind {
            GraphKind::Mesh => load_ply_laplacian(&path),
            GraphKind::Delay | GraphKind::Impedance | GraphKind::Length => load_laplacian(&path),
        }
    }

    /// Loads coordinates for a dataset when available.
    ///
    /// # Errors
    /// Returns an error when the dataset does not provide coordinates or the resource is malformed.
    pub fn coords(&self, dataset: Dataset) -> Result<Array2<f64>, GspError> {
        let Some(source) = dataset.coords_source(&self.root) else {
            return Err(GspError::ResourceNotFound(format!(
                "{dataset} does not provide coordinates"
            )));
        };
        let path = source.path();
        if !path.exists() {
            return Err(GspError::ResourceNotFound(format!(
                "coordinate file not found: {}",
                path.display()
            )));
        }
        match source {
            CoordSource::Mat(path) => load_signal(&path),
            CoordSource::Ply(path) => load_ply_xyz(&path),
        }
    }

    /// Loads a named kernel preset as a validated VF kernel.
    ///
    /// # Errors
    /// Returns an error when the preset is unavailable or invalid.
    pub fn vf_kernel(&self, preset: KernelPreset) -> Result<VfKernel, GspError> {
        let value = load_kernel_json(&preset.path(&self.root))?;
        VfKernel::from_json_value(&value)
    }

    /// Loads a named kernel preset as a validated Chebyshev kernel.
    ///
    /// # Errors
    /// Returns an error when the preset is unavailable or invalid.
    pub fn cheby_kernel(&self, preset: KernelPreset) -> Result<ChebyKernel, GspError> {
        let value = load_kernel_json(&preset.path(&self.root))?;
        ChebyKernel::from_json_value(&value)
    }
}

/// Named datasets exposed by the local example data.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Dataset {
    East,
    EastWest,
    Hawaii,
    Neiso,
    Texas,
    Usa,
    Wecc,
    Bunny,
    Horse,
    Lbrain,
    Rbrain,
}

impl Dataset {
    fn graph_file(self, kind: GraphKind) -> Option<(&'static str, &'static str)> {
        match (self, kind) {
            (Self::East, GraphKind::Delay) => Some(("DELAY", "EAST.mat")),
            (Self::EastWest, GraphKind::Delay) => Some(("DELAY", "EASTWEST.mat")),
            (Self::EastWest, GraphKind::Impedance) => Some(("IMPEDANCE", "EASTWEST.mat")),
            (Self::EastWest, GraphKind::Length) => Some(("LENGTH", "EASTWEST.mat")),
            (Self::Hawaii, GraphKind::Delay) => Some(("DELAY", "HAWAII.mat")),
            (Self::Hawaii, GraphKind::Impedance) => Some(("IMPEDANCE", "HAWAII.mat")),
            (Self::Hawaii, GraphKind::Length) => Some(("LENGTH", "HAWAII.mat")),
            (Self::Neiso, GraphKind::Delay) => Some(("DELAY", "NEISO.mat")),
            (Self::Neiso, GraphKind::Length) => Some(("LENGTH", "NEISO.mat")),
            (Self::Texas, GraphKind::Delay) => Some(("DELAY", "TEXAS.mat")),
            (Self::Texas, GraphKind::Impedance) => Some(("IMPEDANCE", "TEXAS.mat")),
            (Self::Texas, GraphKind::Length) => Some(("LENGTH", "TEXAS.mat")),
            (Self::Usa, GraphKind::Delay) => Some(("DELAY", "USA.mat")),
            (Self::Usa, GraphKind::Impedance) => Some(("IMPEDANCE", "USA.mat")),
            (Self::Usa, GraphKind::Length) => Some(("LENGTH", "USA.mat")),
            (Self::Wecc, GraphKind::Delay) => Some(("DELAY", "WECC.mat")),
            (Self::Wecc, GraphKind::Impedance) => Some(("IMPEDANCE", "WECC.mat")),
            (Self::Wecc, GraphKind::Length) => Some(("LENGTH", "WECC.mat")),
            (Self::Bunny, GraphKind::Mesh) => Some(("MESH", "BUNNY.ply")),
            (Self::Horse, GraphKind::Mesh) => Some(("MESH", "HORSE.ply")),
            (Self::Lbrain, GraphKind::Mesh) => Some(("MESH", "LBRAIN.ply")),
            (Self::Rbrain, GraphKind::Mesh) => Some(("MESH", "RBRAIN.ply")),
            _ => None,
        }
    }

    fn graph_path(self, root: &Path, kind: GraphKind) -> Option<PathBuf> {
        let (dir, file) = self.graph_file(kind)?;
        Some(root.join(dir).join(file))
    }

    fn coords_source(self, root: &Path) -> Option<CoordSource> {
        match self {
            Self::EastWest => Some(CoordSource::Mat(root.join("SIGNALS").join("EASTWEST.mat"))),
            Self::Hawaii => Some(CoordSource::Mat(root.join("SIGNALS").join("HAWAII.mat"))),
            Self::Texas => Some(CoordSource::Mat(root.join("SIGNALS").join("TEXAS.mat"))),
            Self::Usa => Some(CoordSource::Mat(root.join("SIGNALS").join("USA.mat"))),
            Self::Bunny => Some(CoordSource::Ply(root.join("MESH").join("BUNNY.ply"))),
            Self::Horse => Some(CoordSource::Ply(root.join("MESH").join("HORSE.ply"))),
            Self::Lbrain => Some(CoordSource::Ply(root.join("MESH").join("LBRAIN.ply"))),
            Self::Rbrain => Some(CoordSource::Ply(root.join("MESH").join("RBRAIN.ply"))),
            Self::East | Self::Neiso | Self::Wecc => None,
        }
    }
}

impl fmt::Display for Dataset {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let label = match self {
            Self::East => "east",
            Self::EastWest => "eastwest",
            Self::Hawaii => "hawaii",
            Self::Neiso => "neiso",
            Self::Texas => "texas",
            Self::Usa => "usa",
            Self::Wecc => "wecc",
            Self::Bunny => "bunny",
            Self::Horse => "horse",
            Self::Lbrain => "lbrain",
            Self::Rbrain => "rbrain",
        };
        f.write_str(label)
    }
}

/// Graph families available for a dataset.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum GraphKind {
    Delay,
    Impedance,
    Length,
    Mesh,
}

impl fmt::Display for GraphKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let label = match self {
            Self::Delay => "delay",
            Self::Impedance => "impedance",
            Self::Length => "length",
            Self::Mesh => "mesh",
        };
        f.write_str(label)
    }
}

/// Kernel presets available through example data.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum KernelPreset {
    GaussianWavelet,
    MexicanHat,
    ModifiedMorlet,
    Shannon,
}

impl KernelPreset {
    fn file_name(self) -> &'static str {
        match self {
            Self::GaussianWavelet => "GAUSSIAN_WAV.json",
            Self::MexicanHat => "MEXICAN_HAT.json",
            Self::ModifiedMorlet => "MODIFIED_MORLET.json",
            Self::Shannon => "SHANNON.json",
        }
    }

    fn path(self, root: &Path) -> PathBuf {
        root.join("KERNELS").join(self.file_name())
    }
}

impl fmt::Display for KernelPreset {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let label = match self {
            Self::GaussianWavelet => "gaussian-wavelet",
            Self::MexicanHat => "mexican-hat",
            Self::ModifiedMorlet => "modified-morlet",
            Self::Shannon => "shannon",
        };
        f.write_str(label)
    }
}

enum CoordSource {
    Mat(PathBuf),
    Ply(PathBuf),
}

impl CoordSource {
    fn path(&self) -> &Path {
        match self {
            Self::Mat(path) | Self::Ply(path) => path,
        }
    }
}

fn load_kernel_json(path: &Path) -> Result<Value, GspError> {
    if !path.exists() {
        return Err(GspError::ResourceNotFound(format!(
            "kernel file not found: {}",
            path.display()
        )));
    }
    let file = File::open(path)?;
    Ok(serde_json::from_reader(file)?)
}