xdmf 0.2.0

//! A library for writing XDMF files, which are commonly used in scientific simulations for visualizing datasets on meshes, for example with [Paraview](https://www.paraview.org/).
//!
//! The [XDMF](https://www.xdmf.org/) (e**X**tensible **D**ata **M**odel and **F**ormat) stores the metadata in XML files and the actual data in different formats, most commonly in HDF5 files.
use std::{
    path::{Path, PathBuf},
    str::FromStr,
};

use serde::{Deserialize, Serialize};
use xdmf_elements::{
    attribute,
    data_item::{DataContent, Format},
};

mod ascii_writer;
mod binary_writer;
mod error;
#[cfg(feature = "hdf5")]
mod hdf5_writer;
mod paraview;
mod reader;
mod time_series_writer;
mod values;
pub mod xdmf_elements;

// Re-export types used in the public API
pub use error::{Error, Result};
pub use reader::{DataInfo, TimeSeriesReader, ValueType};
pub use time_series_writer::{SubmeshCells, TimeSeriesDataWriter, TimeSeriesWriter, TimeStep};
pub use values::{ConnectivityIndex, Coordinate, Values};
pub use xdmf_elements::CellType;

/// Type of storage used for the heavy data (e.g. ASCII or HDF5)
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
pub enum DataStorage {
    /// store the data in ASCII format, each set of data is stored in a separate file.
    Ascii,
    /// store the data in ASCII format, but inline in the XDMF file. This is only recommended for small datasets.
    AsciiInline,
    /// store the data in HDF5 format, all data in a single HDF5 file.
    Hdf5SingleFile {
        /// zlib/deflate compression level applied to every dataset (0 = none, 9 = max).
        /// `None` uses the library default (currently 3, chosen for write speed over
        /// compression ratio).
        deflate_level: Option<u8>,
    },
    /// store the data in HDF5 format, one file per time step.
    Hdf5MultipleFiles {
        /// zlib/deflate compression level applied to every dataset (0 = none, 9 = max).
        /// `None` uses the library default (currently 3, chosen for write speed over
        /// compression ratio).
        deflate_level: Option<u8>,
    },
    /// store the data in uncompressed raw binary format, each set of data is stored in a separate file.
    Binary,
}

impl FromStr for DataStorage {
    type Err = String;

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "ascii" => Ok(Self::Ascii),
            "asciiinline" | "ascii_inline" | "ascii-inline" => Ok(Self::AsciiInline),
            "hdf5singlefile" | "hdf5_single_file" | "hdf5-single-file" => {
                Ok(Self::Hdf5SingleFile {
                    deflate_level: None,
                })
            }
            "hdf5multiplefiles" | "hdf5_multiple_files" | "hdf5-multiple-files" => {
                Ok(Self::Hdf5MultipleFiles {
                    deflate_level: None,
                })
            }
            "binary" => Ok(Self::Binary),
            _ => Err(format!(
                "Invalid DataStorage variant: '{s}'. Valid options are: 'Ascii', 'AsciiInline', 'Hdf5SingleFile', 'Hdf5MultipleFiles', 'Binary'"
            )),
        }
    }
}

/// this trait defines the interface used to write the heavy data
///
/// `Send + Sync` so the Python bindings can hold a writer in a `#[pyclass]` and release the GIL
/// for the duration of a write
pub(crate) trait DataWriter: Send + Sync {
    fn format(&self) -> Format;

    fn data_storage(&self) -> DataStorage;

    /// Write one point-coordinate array: the mesh's own when `submesh` is `None`, or that of the
    /// submesh at that position otherwise.
    ///
    /// Called once per mesh or once per submesh, before the matching
    /// [`write_connectivity`](Self::write_connectivity), and never both ways for one mesh: a
    /// submesh carries only the points its own cells use, so that a viewer holds each block's
    /// geometry once rather than the whole mesh's once per block.
    fn write_points(&mut self, submesh: Option<usize>, points: &Values<'_>) -> Result<DataContent>;

    /// Write one component of the mesh's own coordinates: every X value, then every Y, then every
    /// Z.
    ///
    /// Called instead of [`write_points`](Self::write_points) for a mesh whose submeshes select
    /// their points out of the mesh's rather than carrying a copy of them, and then exactly three
    /// times, in order. Split by component because that is what lets all three of a submesh's
    /// selections share the one index list naming its points -- an interleaved array would need
    /// one index per coordinate instead. Only reachable from a backend that answers
    /// [`supports_selections`](Self::supports_selections) with `true`.
    fn write_point_component(
        &mut self,
        _component: usize,
        _coordinates: &Values<'_>,
    ) -> Result<DataContent> {
        Err(Error::Internal(
            "this storage cannot be selected out of, so its submeshes carry their own points",
        ))
    }

    /// Write one connectivity array: the mesh's own when `submesh` is `None`, or that of the
    /// submesh at that position otherwise.
    ///
    /// A mesh written with submeshes has no connectivity of its own -- each submesh carries the
    /// cells it contains, indexed into its own points -- so a backend sees either one `None` call
    /// or one `Some` call per submesh, never both.
    fn write_connectivity(
        &mut self,
        submesh: Option<usize>,
        cells: &Values<'_>,
    ) -> Result<DataContent>;

    /// Write one submesh's global cell indices, as the position of each of its cells in the
    /// mesh the submesh was cut out of.
    ///
    /// Called once per scattered submesh, at mesh-write time; a contiguous one is a start and a
    /// length and needs no array. Nothing in the light data references what this writes -- it is
    /// there to read the file back with, not for `ParaView`.
    fn write_submesh_cells(&mut self, submesh: usize, cells: &Values<'_>) -> Result<DataContent>;

    /// Write one submesh's global point indices, as the position of each of its points in the
    /// mesh the submesh was cut out of.
    ///
    /// The counterpart of [`write_submesh_cells`](Self::write_submesh_cells) for the points a
    /// submesh's connectivity is renumbered against. A side channel for a reader where a submesh
    /// carries a copy of its points, and the array its `<Geometry>` selects them through where it
    /// does not -- in which case the light data references this and records it nowhere else.
    fn write_submesh_points(&mut self, submesh: usize, points: &Values<'_>) -> Result<DataContent>;

    /// Write one array of attribute data, identified by its position among all the arrays
    /// written in the current time step. Backends name their heavy data by this index
    fn write_data(&mut self, index: usize, data: &Values<'_>) -> Result<DataContent>;

    /// Whether a submesh may reference this storage's arrays instead of being given a copy of
    /// its share of them.
    ///
    /// Only the HDF5 storages can: `ParaView`'s XDMF2 reader honours a `HyperSlab`/`Coordinates`
    /// selection for `Format="HDF"`, while for the ascii storages it reads the source array from
    /// its start instead -- silently, so a selection there would put values in the viewer that
    /// the file does not contain.
    fn supports_selections(&self) -> bool {
        false
    }

    /// Write one array of indices for a submesh's grids to select their share of a field with,
    /// identified by its position among all such arrays.
    ///
    /// Written once, at the step that first needs it, and referenced by every step after --
    /// which is why it belongs with the mesh's arrays rather than with a step's data, and why a
    /// backend that removes a discarded step's data must leave it alone. Only reachable from a
    /// backend that answers [`supports_selections`](Self::supports_selections) with `true`.
    fn write_selection(&mut self, _index: usize, _indices: &Values<'_>) -> Result<DataContent> {
        Err(Error::Internal(
            "this storage cannot be selected out of, so it is never asked to write a selection",
        ))
    }

    fn write_data_initialize(&mut self, _time: &str) -> Result<()> {
        Ok(())
    }

    fn write_data_finalize(&mut self) -> Result<()> {
        Ok(())
    }

    /// End the current time step and remove the heavy data written for it.
    ///
    /// Called instead of [`write_data_finalize`](Self::write_data_finalize) when a time step is
    /// abandoned rather than written -- an attribute was rejected, or the caller's closure
    /// returned an error. Without it the step's heavy data would stay on disk with no `<Grid>`
    /// in the XDMF file referencing it. Backends that write nothing until the step is complete
    /// (e.g. `AsciiInline`) keep the default, which is just to finalize.
    fn write_data_discard(&mut self) -> Result<()> {
        self.write_data_finalize()
    }

    // flush the writer, if applicable
    fn flush(&mut self) -> Result<()> {
        Ok(())
    }
}

// zlib/deflate only accepts levels 0-9; anything else is a caller mistake that should be
// rejected before a writer is constructed, rather than surfacing as a raw HDF5 error later
// (`H5Pset_deflate(): invalid deflate level`) from inside `write_mesh`.
fn validate_deflate_level(deflate_level: Option<u8>) -> Result<()> {
    if let Some(level) = deflate_level
        && level > 9
    {
        return Err(Error::InvalidConfiguration {
            reason: format!("deflate level {level} is out of range, must be between 0 and 9"),
        });
    }
    Ok(())
}

/// Create a writer for the heavy data, based on the chosen data storage.
pub(crate) fn create_writer(
    file_name: &Path,
    data_storage: DataStorage,
) -> Result<Box<dyn DataWriter>> {
    match data_storage {
        DataStorage::Ascii => Ok(Box::new(ascii_writer::AsciiWriter::new(file_name)?)),
        DataStorage::AsciiInline => Ok(Box::new(ascii_writer::AsciiInlineWriter::new())),
        DataStorage::Hdf5SingleFile { deflate_level } => {
            validate_deflate_level(deflate_level)?;
            cfg_select! {
                feature = "hdf5" => Ok(Box::new(hdf5_writer::SingleFileHdf5Writer::new(
                    file_name,
                    deflate_level.unwrap_or(hdf5_writer::DEFAULT_DEFLATE_LEVEL),
                )?)),
                _ => Err(Error::InvalidConfiguration {
                    reason: format!(
                        "using {data_storage:?} DataStorage requires the 'hdf5' feature"
                    ),
                }),
            }
        }
        DataStorage::Hdf5MultipleFiles { deflate_level } => {
            validate_deflate_level(deflate_level)?;
            cfg_select! {
                feature = "hdf5" => Ok(Box::new(hdf5_writer::MultipleFilesHdf5Writer::new(
                    file_name,
                    deflate_level.unwrap_or(hdf5_writer::DEFAULT_DEFLATE_LEVEL),
                )?)),
                _ => Err(Error::InvalidConfiguration {
                    reason: format!(
                        "using {data_storage:?} DataStorage requires the 'hdf5' feature"
                    ),
                }),
            }
        }
        DataStorage::Binary => Ok(Box::new(binary_writer::BinaryWriter::new(file_name)?)),
    }
}

/// Check if the hdf5 feature is enabled.
pub const fn is_hdf5_enabled() -> bool {
    cfg_select! {
        feature = "hdf5" => true,
        _ => false,
    }
}

/// Type of the data (scalar, vector, tensor, etc.)
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum DataAttribute {
    /// single value
    Scalar,
    /// 3D vector (3 components)
    Vector,
    /// 2nd order tensor in 3D (9 components)
    Tensor,
    /// Symmetric 2nd order tensor in 3D (6 components)
    Tensor6,
    /// Matrix with specified number of rows and columns
    Matrix(usize, usize),
    /// Generic data with specified size
    Generic(usize),
}

impl DataAttribute {
    /// Number of components of one entity, [`None`] if [`Matrix`](Self::Matrix)'s two dimensions
    /// do not fit a `usize` -- a caller-supplied product, so it is checked rather than wrapped.
    pub(crate) fn size(&self) -> Option<usize> {
        match self {
            Self::Scalar => Some(1),
            Self::Vector => Some(3),
            Self::Tensor => Some(9),
            Self::Tensor6 => Some(6),
            Self::Matrix(n, m) => n.checked_mul(*m),
            Self::Generic(size) => Some(*size),
        }
    }
}

impl From<DataAttribute> for attribute::AttributeType {
    fn from(data_attr: DataAttribute) -> Self {
        match data_attr {
            DataAttribute::Scalar => Self::Scalar,
            DataAttribute::Vector => Self::Vector,
            DataAttribute::Tensor => Self::Tensor,
            DataAttribute::Tensor6 => Self::Matrix, // written as Matrix to get detected as symmetric tensor
            DataAttribute::Matrix(_, _) => Self::Matrix,
            DataAttribute::Generic(_) => Self::Matrix,
        }
    }
}

/// Create directories in a way that is safe for MPI applications.
///
/// This function will create the directory if it does not exist, and wait for it to appear in the filesystem.
/// This is particularly needed on systems such as clusters with slow filesystems, to ensure that
/// all processes can see the created directory before proceeding.
///
/// For more details check the [reference](https://github.com/KratosMultiphysics/Kratos/pull/9247).
/// Its a battle-tested solution tested with > 1000 processes
pub fn mpi_safe_create_dir_all(path: impl AsRef<Path> + std::fmt::Debug) -> Result<()> {
    if !&path.as_ref().exists() {
        std::fs::create_dir_all(&path)
            .map_err(error::io_ctx("creating directory", path.as_ref()))?;
    }

    if !path.as_ref().exists() {
        // wait for the path to appear in the filesystem
        std::thread::sleep(std::time::Duration::from_millis(50));
    }

    Ok(())
}

/// Remove the files written for an abandoned time step, shared by the file-per-attribute backends
/// (`Ascii`/`Binary`) implementing [`DataWriter::write_data_discard`].
///
/// The paths are drained out of `step_files`, so they are forgotten even if a removal fails -- a
/// second discard attempt on the same paths would only report the same failure again. Every file
/// is attempted before reporting, so one failure does not leave the rest behind; the first error
/// is the one returned.
pub(crate) fn remove_step_files(step_files: &mut Vec<PathBuf>) -> Result<()> {
    let mut first_error = None;

    for path in step_files.drain(..) {
        let result = std::fs::remove_file(&path)
            .map_err(error::io_ctx("removing discarded data file", &path));
        if let Err(error) = result
            && first_error.is_none()
        {
            first_error = Some(error);
        }
    }

    match first_error {
        Some(error) => Err(error),
        None => Ok(()),
    }
}

/// The arrays a mesh is written as, named for the file or `HDF5` dataset each goes into.
pub(crate) const POINTS: &str = "points";
pub(crate) const CELLS: &str = "cells";
pub(crate) const SUBMESH_POINTS: &str = "submesh_points";
pub(crate) const SUBMESH_CELLS: &str = "submesh_cells";
/// Index arrays a submesh's grids select their share of a field with, see
/// [`DataWriter::write_selection`]. Numbered like a submesh's own arrays, but by the order they
/// were needed rather than by submesh -- one submesh needs one per shape of field it carries.
pub(crate) const SELECTIONS: &str = "selections";

/// Name of the `Information` element recording which [`DataStorage`] wrote the document. The
/// reader takes it to reject a file it cannot read when the file is opened, rather than at the
/// first read call that reaches heavy data.
pub(crate) const DATA_STORAGE: &str = "data_storage";

/// Name of the file one of a mesh's arrays goes into, for the backends that write one file per
/// array: `<array>.<extension>` for the mesh's own, `<array>_<index>.<extension>` for the one
/// belonging to the submesh at that position.
pub(crate) fn mesh_file_name(array: &str, submesh: Option<usize>, extension: &str) -> String {
    match submesh {
        Some(index) => format!("{array}_{index}.{extension}"),
        None => format!("{array}.{extension}"),
    }
}

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

    #[test]
    fn test_mpi_safe_create_dir_all() {
        let tmp_dir = temp_dir::TempDir::new().unwrap();
        let dirs_to_create = tmp_dir.path().join("out/xdmf/test/folder/random/testing");

        // Try to create dirs from 100 threads concurrently
        let handles: Vec<_> = (0..100)
            .map(|_| {
                std::thread::spawn({
                    let dir_thread_local = dirs_to_create.clone();
                    move || mpi_safe_create_dir_all(dir_thread_local).unwrap()
                })
            })
            .collect();

        // join threads, will propagate errors if any
        for handle in handles {
            handle.join().unwrap();
        }

        // Check that the directory was created
        assert!(dirs_to_create.exists());
    }

    #[test]
    fn test_data_attribute() {
        let scalar = DataAttribute::Scalar;
        let vector = DataAttribute::Vector;
        let tensor = DataAttribute::Tensor;
        let tensor6 = DataAttribute::Tensor6;
        let matrix = DataAttribute::Matrix(3, 3);
        let generic = DataAttribute::Generic(5);

        assert_eq!(scalar.size(), Some(1));
        assert_eq!(vector.size(), Some(3));
        assert_eq!(tensor.size(), Some(9));
        assert_eq!(tensor6.size(), Some(6));
        assert_eq!(matrix.size(), Some(9));
        assert_eq!(generic.size(), Some(5));
        assert_eq!(DataAttribute::Matrix(usize::MAX, 2).size(), None);

        assert_eq!(attribute::AttributeType::Scalar, scalar.into());
        assert_eq!(attribute::AttributeType::Vector, vector.into());
        assert_eq!(attribute::AttributeType::Tensor, tensor.into());
        assert_eq!(attribute::AttributeType::Matrix, tensor6.into());
        assert_eq!(attribute::AttributeType::Matrix, matrix.into());
        assert_eq!(attribute::AttributeType::Matrix, generic.into());
    }

    #[test]
    fn test_data_storage_from_str() {
        // Test exact case matches
        assert_eq!("ascii".parse::<DataStorage>().unwrap(), DataStorage::Ascii);
        assert_eq!("Ascii".parse::<DataStorage>().unwrap(), DataStorage::Ascii);
        assert_eq!("ASCII".parse::<DataStorage>().unwrap(), DataStorage::Ascii);

        // Test AsciiInline variants
        assert_eq!(
            "asciiinline".parse::<DataStorage>().unwrap(),
            DataStorage::AsciiInline
        );
        assert_eq!(
            "ascii_inline".parse::<DataStorage>().unwrap(),
            DataStorage::AsciiInline
        );
        assert_eq!(
            "ascii-inline".parse::<DataStorage>().unwrap(),
            DataStorage::AsciiInline
        );

        // Test Hdf5SingleFile variants
        assert_eq!(
            "hdf5singlefile".parse::<DataStorage>().unwrap(),
            DataStorage::Hdf5SingleFile {
                deflate_level: None
            }
        );
        assert_eq!(
            "hdf5_single_file".parse::<DataStorage>().unwrap(),
            DataStorage::Hdf5SingleFile {
                deflate_level: None
            }
        );
        assert_eq!(
            "Hdf5-Single-File".parse::<DataStorage>().unwrap(),
            DataStorage::Hdf5SingleFile {
                deflate_level: None
            }
        );

        // Test Hdf5MultipleFiles variants
        assert_eq!(
            "hdf5multiplefiles".parse::<DataStorage>().unwrap(),
            DataStorage::Hdf5MultipleFiles {
                deflate_level: None
            }
        );
        assert_eq!(
            "hdf5_multiple_files".parse::<DataStorage>().unwrap(),
            DataStorage::Hdf5MultipleFiles {
                deflate_level: None
            }
        );
        assert_eq!(
            "HDF5-Multiple-Files".parse::<DataStorage>().unwrap(),
            DataStorage::Hdf5MultipleFiles {
                deflate_level: None
            }
        );

        // Test Binary variant
        assert_eq!(
            "binary".parse::<DataStorage>().unwrap(),
            DataStorage::Binary
        );
        assert_eq!(
            "Binary".parse::<DataStorage>().unwrap(),
            DataStorage::Binary
        );

        // Test invalid input
        let err = "invalid".parse::<DataStorage>().unwrap_err();
        assert_eq!(
            err,
            "Invalid DataStorage variant: 'invalid'. Valid options are: 'Ascii', 'AsciiInline', 'Hdf5SingleFile', 'Hdf5MultipleFiles', 'Binary'"
        );

        let err = "".parse::<DataStorage>().unwrap_err();
        assert_eq!(
            err,
            "Invalid DataStorage variant: ''. Valid options are: 'Ascii', 'AsciiInline', 'Hdf5SingleFile', 'Hdf5MultipleFiles', 'Binary'"
        );
    }

    #[test]
    fn test_validate_deflate_level() {
        validate_deflate_level(None).unwrap();
        validate_deflate_level(Some(0)).unwrap();
        validate_deflate_level(Some(9)).unwrap();

        std::assert_matches!(
            validate_deflate_level(Some(10)).unwrap_err(),
            Error::InvalidConfiguration { reason } if reason.contains("deflate level 10")
        );
        std::assert_matches!(
            validate_deflate_level(Some(255)).unwrap_err(),
            Error::InvalidConfiguration { reason } if reason.contains("deflate level 255")
        );
    }

    #[test]
    fn create_writer_rejects_invalid_deflate_level() {
        let tmp_dir = temp_dir::TempDir::new().unwrap();
        let file_name = tmp_dir.path().join("test.xdmf");

        for storage in [
            DataStorage::Hdf5SingleFile {
                deflate_level: Some(10),
            },
            DataStorage::Hdf5MultipleFiles {
                deflate_level: Some(10),
            },
        ] {
            let Err(err) = create_writer(&file_name, storage) else {
                panic!("expected an error for deflate_level 10");
            };
            std::assert_matches!(err, Error::InvalidConfiguration { reason } if reason.contains("deflate level 10"));
        }
    }
}