Skip to main content

conspire/geometry/grid/read/
mod.rs

1mod spn;
2mod vti;
3
4use crate::{
5    geometry::grid::Grid,
6    io::{Npy, NpyType},
7};
8use std::{
9    fmt::Display,
10    io::{Error as ErrorIO, ErrorKind},
11    path::Path,
12    str::FromStr,
13};
14
15pub enum Input<P>
16where
17    P: AsRef<Path>,
18{
19    Npy(P),
20    Spn(P, Vec<usize>),
21    Vti(P),
22}
23
24impl<P> AsRef<Path> for Input<P>
25where
26    P: AsRef<Path>,
27{
28    fn as_ref(&self) -> &Path {
29        match self {
30            Input::Npy(path) => path.as_ref(),
31            Input::Spn(path, _) => path.as_ref(),
32            Input::Vti(path) => path.as_ref(),
33        }
34    }
35}
36
37impl<const D: usize, T, P> TryFrom<Input<P>> for Grid<D, T>
38where
39    P: AsRef<Path>,
40    T: NpyType + FromStr,
41    <T as FromStr>::Err: Display,
42{
43    type Error = ErrorIO;
44    fn try_from(input: Input<P>) -> Result<Self, Self::Error> {
45        match input {
46            Input::Spn(path, nel) => spn::read(path, nel),
47            Input::Npy(path) => {
48                let npy = Npy::<T>::read(path)?;
49                let nel: [usize; D] = npy.shape.try_into().map_err(|shape: Vec<usize>| {
50                    ErrorIO::new(
51                        ErrorKind::InvalidData,
52                        format!("npy has {} axes but Grid was asked for D={D}", shape.len()),
53                    )
54                })?;
55                if npy.fortran_order {
56                    Ok(Grid::new(npy.data, nel))
57                } else {
58                    Ok(Grid::new_row_major(npy.data, nel))
59                }
60            }
61            Input::Vti(path) => vti::read(path),
62        }
63    }
64}
65
66pub(super) fn transpose<const D: usize, T: Copy>(c_order: Vec<T>, nel: [usize; D]) -> Vec<T> {
67    let total: usize = nel.iter().product();
68    let mut f_stride = [1usize; D];
69    for axis in 1..D {
70        f_stride[axis] = f_stride[axis - 1] * nel[axis - 1];
71    }
72    let mut c_stride = [1usize; D];
73    for axis in (0..D.saturating_sub(1)).rev() {
74        c_stride[axis] = c_stride[axis + 1] * nel[axis + 1];
75    }
76    let mut out: Vec<T> = Vec::with_capacity(total);
77    if total > 0 {
78        unsafe {
79            transpose_block(&c_order, out.as_mut_ptr(), 0, 0, nel, &c_stride, &f_stride);
80            out.set_len(total);
81        }
82    }
83    out
84}
85
86unsafe fn transpose_block<const D: usize, T: Copy>(
87    src: &[T],
88    dst: *mut T,
89    c_off: usize,
90    f_off: usize,
91    len: [usize; D],
92    c_stride: &[usize; D],
93    f_stride: &[usize; D],
94) {
95    const TILE: usize = 16;
96    if let Some(axis) = (0..D)
97        .filter(|&axis| len[axis] > TILE)
98        .max_by_key(|&axis| len[axis])
99    {
100        let half = len[axis] / 2;
101        let mut lo = len;
102        lo[axis] = half;
103        let mut hi = len;
104        hi[axis] = len[axis] - half;
105        unsafe {
106            transpose_block(src, dst, c_off, f_off, lo, c_stride, f_stride);
107            transpose_block(
108                src,
109                dst,
110                c_off + half * c_stride[axis],
111                f_off + half * f_stride[axis],
112                hi,
113                c_stride,
114                f_stride,
115            );
116        }
117        return;
118    }
119    let volume: usize = len.iter().product();
120    let mut idx = [0usize; D];
121    let (mut c, mut f) = (c_off, f_off);
122    for _ in 0..volume {
123        unsafe {
124            *dst.add(f) = src[c];
125        }
126        for axis in 0..D {
127            idx[axis] += 1;
128            c += c_stride[axis];
129            f += f_stride[axis];
130            if idx[axis] < len[axis] {
131                break;
132            }
133            idx[axis] = 0;
134            c -= len[axis] * c_stride[axis];
135            f -= len[axis] * f_stride[axis];
136        }
137    }
138}