wxy_domain/
mapping.rs

1//! Mappings: positions for some of a domain's dims, in a named frame.
2
3use crate::frame::Frame;
4use crate::positions::{Positions, Rung};
5
6/// Where stored coordinates sit relative to the elements they describe.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
8pub enum CellConvention {
9    /// Coordinates are element centres (CF default, GeoTIFF PixelIsPoint
10    /// read as centres). `n` values per axis.
11    Centre,
12    /// Coordinates are element edges. `n + 1` values per axis.
13    Edge,
14}
15
16/// Positions for a subset of a domain's dims, and the frame they are in.
17#[derive(Debug, Clone)]
18pub struct Mapping {
19    /// Indices into the domain's dims, in the order `positions` expects.
20    /// For [`Positions::Affine`] this is `[row_dim, col_dim]`.
21    pub dims: Vec<usize>,
22    pub positions: Positions,
23    pub frame: Frame,
24    pub cell: CellConvention,
25    /// Where this mapping came from: `"stored"` for what a reader found, or
26    /// the name of the step that derived it (such as `"normalise"`).
27    pub provenance: String,
28}
29
30impl Mapping {
31    pub fn new(dims: Vec<usize>, positions: Positions, frame: Frame) -> Self {
32        Self {
33            dims,
34            positions,
35            frame,
36            cell: CellConvention::Centre,
37            provenance: "stored".into(),
38        }
39    }
40
41    pub fn with_cell(mut self, cell: CellConvention) -> Self {
42        self.cell = cell;
43        self
44    }
45
46    pub fn with_provenance(mut self, provenance: impl Into<String>) -> Self {
47        self.provenance = provenance.into();
48        self
49    }
50
51    pub fn rung(&self) -> Rung {
52        self.positions.rung()
53    }
54
55    /// The position at a continuous index, one value per covered dim in
56    /// `self.dims` order (see the crate docs for the index convention).
57    /// Returns one value per output coordinate, or `None` where the mapping
58    /// cannot place the index (R0, outside a curvilinear grid, or a rung
59    /// not evaluated in S1).
60    pub fn position(&self, index: &[f64]) -> Option<Vec<f64>> {
61        match &self.positions {
62            Positions::IndexOnly => None,
63            Positions::Affine(a) => {
64                let (x, y) = a.forward(index[1], index[0]);
65                Some(vec![x, y])
66            }
67            Positions::Regular1 { origin, step } => Some(vec![origin + index[0] * step]),
68            Positions::Rectilinear(axes) => axes
69                .iter()
70                .zip(index)
71                .map(|(c, &i)| interp_axis(c, i, self.cell))
72                .collect(),
73            Positions::Curvilinear(arrays) => {
74                let shape = &arrays.first()?.shape;
75                arrays
76                    .iter()
77                    .map(|a| interp_nd(&a.values, shape, index))
78                    .collect()
79            }
80            Positions::Fitted(m) => m.forward(index),
81            Positions::Scheme(_) | Positions::Explicit(_) => None,
82        }
83    }
84}
85
86/// Linear interpolation along one coordinate vector at continuous index
87/// `i`, extrapolating from the end steps.
88fn interp_axis(c: &[f64], i: f64, cell: CellConvention) -> Option<f64> {
89    if c.is_empty() || !i.is_finite() {
90        return None;
91    }
92    // Position of index `i` in "sample space", where sample k is c[k].
93    let s = match cell {
94        CellConvention::Centre => i - 0.5,
95        CellConvention::Edge => i,
96    };
97    if c.len() == 1 {
98        return Some(c[0]);
99    }
100    let k = (s.floor() as isize).clamp(0, c.len() as isize - 2) as usize;
101    let t = s - k as f64;
102    Some(c[k] + t * (c[k + 1] - c[k]))
103}
104
105/// Multilinear interpolation in a centre-registered array at a continuous
106/// index. Returns `None` outside the array's sample hull.
107fn interp_nd(values: &[f64], shape: &[usize], index: &[f64]) -> Option<f64> {
108    let nd = shape.len();
109    if index.len() != nd {
110        return None;
111    }
112    let mut base = vec![0usize; nd];
113    let mut frac = vec![0f64; nd];
114    for d in 0..nd {
115        let s = index[d] - 0.5;
116        if !(s >= 0.0 && s <= (shape[d] - 1) as f64) {
117            return None;
118        }
119        let k = (s.floor() as usize).min(shape[d].saturating_sub(2));
120        base[d] = k;
121        frac[d] = s - k as f64;
122    }
123    let mut strides = vec![1usize; nd];
124    for d in (0..nd.saturating_sub(1)).rev() {
125        strides[d] = strides[d + 1] * shape[d + 1];
126    }
127    let mut acc = 0.0;
128    for corner in 0..(1usize << nd) {
129        let mut w = 1.0;
130        let mut off = 0;
131        for d in 0..nd {
132            let up = (corner >> d) & 1 == 1;
133            let k = if up && shape[d] > 1 {
134                base[d] + 1
135            } else {
136                base[d]
137            };
138            w *= if up { frac[d] } else { 1.0 - frac[d] };
139            off += k * strides[d];
140        }
141        if w != 0.0 {
142            acc += w * values[off];
143        }
144    }
145    Some(acc)
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151    use crate::affine::Affine;
152    use crate::positions::CoordArray;
153
154    #[test]
155    fn affine_centre_of_first_cell() {
156        let m = Mapping::new(
157            vec![0, 1],
158            Positions::Affine(Affine::from_extent([0.0, 360.0, -90.0, 90.0], [360, 180])),
159            Frame::crs("EPSG:4326"),
160        );
161        assert_eq!(m.position(&[0.5, 0.5]).unwrap(), vec![0.5, 89.5]);
162    }
163
164    #[test]
165    fn rectilinear_centres_interpolate_and_extrapolate() {
166        let m = Mapping::new(
167            vec![0],
168            Positions::Rectilinear(vec![vec![0.5, 1.5, 3.5]]),
169            Frame::Index,
170        );
171        assert_eq!(m.position(&[0.5]).unwrap(), vec![0.5]);
172        assert_eq!(m.position(&[2.5]).unwrap(), vec![3.5]);
173        // outer edge of the first element extrapolates by the first step
174        assert_eq!(m.position(&[0.0]).unwrap(), vec![0.0]);
175    }
176
177    #[test]
178    fn curvilinear_bilinear() {
179        // 2 x 2 centres, x = col, y = 10 * row
180        let x = CoordArray::new(vec![2, 2], vec![0.0, 1.0, 0.0, 1.0]);
181        let y = CoordArray::new(vec![2, 2], vec![0.0, 0.0, 10.0, 10.0]);
182        let m = Mapping::new(vec![0, 1], Positions::Curvilinear(vec![x, y]), Frame::Index);
183        assert_eq!(m.position(&[1.0, 1.0]).unwrap(), vec![0.5, 5.0]);
184        assert!(m.position(&[0.0, 0.0]).is_none());
185    }
186}
wxy_domain/mapping.rs

wxy_domain/
mapping.rs
