nd_arrow_array/compute/
interp.rs

1//! Linear interpolation along one dimension.
2
3use std::sync::Arc;
4
5use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, Float64Array};
6use arrow::compute::cast as arrow_cast;
7use arrow::datatypes::{DataType, Float64Type};
8
9use crate::NdArrowArray;
10use crate::compute::elementwise::where_;
11use crate::compute::search::partition_point;
12use crate::compute::shape::take;
13use crate::dimensions::{Dimension, Dimensions};
14use crate::error::{Result, nd_err};
15
16fn to_f64(values: &ArrayRef) -> Result<Float64Array> {
17    Ok(arrow_cast(values.as_ref(), &DataType::Float64)?
18        .as_primitive::<Float64Type>()
19        .clone())
20}
21
22/// How [`interp_with`] picks values between coordinates.
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
24pub enum InterpMethod {
25    /// Linear interpolation. The result is floating point.
26    #[default]
27    Linear,
28    /// The value at the nearest coordinate (the lower one on a tie). Keeps the type.
29    Nearest,
30}
31
32/// Interpolate `a` linearly along `dim` from coordinates `x` to `new_x`,
33/// like `xarray.DataArray.interp` with `method="linear"`.
34///
35/// `x` is 1-D, strictly increasing and as long as `dim`. `new_x` is 1-D. The
36/// result keeps the axis order and has `new_x.len()` elements along `dim`.
37/// Points outside `[x[0], x[-1]]`, null points and points next to a null value
38/// give null. The result is `Float64` (`Float32` for `Float32` input).
39pub fn interp(
40    a: &NdArrowArray,
41    dim: &str,
42    x: &NdArrowArray,
43    new_x: &NdArrowArray,
44) -> Result<NdArrowArray> {
45    interp_with(a, dim, x, new_x, InterpMethod::Linear)
46}
47
48/// Interpolate along `dim` with `method`, see [`interp`].
49pub fn interp_with(
50    a: &NdArrowArray,
51    dim: &str,
52    x: &NdArrowArray,
53    new_x: &NdArrowArray,
54    method: InterpMethod,
55) -> Result<NdArrowArray> {
56    let plan = target_plan(a, dim, x, new_x)?;
57    match method {
58        InterpMethod::Linear => interp_linear(a, dim, &plan),
59        InterpMethod::Nearest => interp_nearest(a, dim, &plan),
60    }
61}
62
63/// Per target: (lower index, weight of the upper neighbor), or `None` outside the range.
64fn target_plan(
65    a: &NdArrowArray,
66    dim: &str,
67    x: &NdArrowArray,
68    new_x: &NdArrowArray,
69) -> Result<Vec<Option<(usize, f64)>>> {
70    let axis = a.dims().axis(dim)?;
71    let n = a.dims().as_slice()[axis].size();
72    if x.rank() != 1 || x.len() != n {
73        return Err(nd_err!(
74            InvalidArgument,
75            "interp coordinates must be 1-D with {n} values along '{dim}'"
76        ));
77    }
78    if new_x.rank() != 1 {
79        return Err(nd_err!(InvalidArgument, "interp targets must be 1-D"));
80    }
81    let xs = to_f64(x.values())?;
82    if xs.null_count() > 0 || xs.values().windows(2).any(|w| w[1] <= w[0]) {
83        return Err(nd_err!(
84            InvalidArgument,
85            "interp coordinates must be non-null and strictly increasing"
86        ));
87    }
88    let xs = xs.values();
89    let targets = to_f64(new_x.values())?;
90
91    let plan: Vec<Option<(usize, f64)>> = targets
92        .iter()
93        .map(|t| {
94            let t = t?;
95            if n == 0 || t < xs[0] || t > xs[n - 1] {
96                return None;
97            }
98            let upper = partition_point(n, |i| xs[i] < t);
99            if xs[upper] == t {
100                return Some((upper, 0.0));
101            }
102            let lower = upper - 1;
103            Some((lower, (t - xs[lower]) / (xs[upper] - xs[lower])))
104        })
105        .collect();
106    Ok(plan)
107}
108
109fn interp_nearest(
110    a: &NdArrowArray,
111    dim: &str,
112    plan: &[Option<(usize, f64)>],
113) -> Result<NdArrowArray> {
114    let indices: Vec<isize> = plan
115        .iter()
116        .map(|p| p.map_or(0, |(lower, w)| if w <= 0.5 { lower } else { lower + 1 }) as isize)
117        .collect();
118    if plan.iter().all(Option::is_some) {
119        return take(a, dim, &indices);
120    }
121    let picked = take(a, dim, &indices)?;
122    let valid = NdArrowArray::try_new(
123        Arc::new(BooleanArray::from_iter(
124            plan.iter().map(|p| Some(p.is_some())),
125        )),
126        Dimensions::try_new([(dim, plan.len())])?,
127    )?;
128    where_(
129        &valid,
130        &picked,
131        &NdArrowArray::new_null_scalar(a.data_type()),
132    )
133}
134
135fn interp_linear(
136    a: &NdArrowArray,
137    dim: &str,
138    plan: &[Option<(usize, f64)>],
139) -> Result<NdArrowArray> {
140    let axis = a.dims().axis(dim)?;
141    let n = a.dims().as_slice()[axis].size();
142
143    let values = to_f64(a.values())?;
144    let shape = a.shape();
145    let outer: usize = shape[..axis].iter().product();
146    let inner: usize = shape[axis + 1..].iter().product();
147    let m = plan.len();
148    let mut out: Vec<Option<f64>> = Vec::with_capacity(outer * m * inner);
149    let get = |p: usize| values.is_valid(p).then(|| values.value(p));
150    for o in 0..outer {
151        for step in plan {
152            for j in 0..inner {
153                out.push(step.and_then(|(lower, w)| {
154                    let low = get((o * n + lower) * inner + j)?;
155                    if w == 0.0 {
156                        return Some(low);
157                    }
158                    let high = get((o * n + lower + 1) * inner + j)?;
159                    Some(low + (high - low) * w)
160                }));
161            }
162        }
163    }
164
165    let mut result: ArrayRef = Arc::new(Float64Array::from(out));
166    if matches!(a.data_type(), DataType::Float32 | DataType::Float16) {
167        result = arrow_cast(result.as_ref(), &DataType::Float32)?;
168    }
169    let dims = a.dims().iter().map(|d| {
170        if d.name() == dim {
171            Dimension::new(dim, m)
172        } else {
173            d.clone()
174        }
175    });
176    Ok(NdArrowArray::try_new(result, Dimensions::try_new(dims)?)?
177        .with_attributes(a.attributes().clone()))
178}
179
180#[cfg(test)]
181mod tests {
182    use arrow::array::Int32Array;
183
184    use super::*;
185
186    fn nd(values: ArrayRef, dims: &[(&str, usize)]) -> NdArrowArray {
187        NdArrowArray::try_new(values, Dimensions::try_new(dims.iter().copied()).unwrap()).unwrap()
188    }
189
190    #[test]
191    fn interpolates_each_lane() {
192        let a = nd(
193            Arc::new(Int32Array::from(vec![0, 10, 20, 100, 110, 120])),
194            &[("y", 2), ("x", 3)],
195        );
196        let x = nd(
197            Arc::new(Float64Array::from(vec![0.0, 1.0, 3.0])),
198            &[("x", 3)],
199        );
200        let new_x = nd(
201            Arc::new(Float64Array::from(vec![0.5, 2.0, 3.0, 4.0])),
202            &[("x", 4)],
203        );
204        let out = interp(&a, "x", &x, &new_x).unwrap();
205        assert_eq!(out.shape(), vec![2, 4]);
206        let v: Vec<Option<f64>> = out.values().as_primitive::<Float64Type>().iter().collect();
207        assert_eq!(
208            v,
209            vec![
210                Some(5.0),
211                Some(15.0),
212                Some(20.0),
213                None,
214                Some(105.0),
215                Some(115.0),
216                Some(120.0),
217                None
218            ]
219        );
220    }
221
222    #[test]
223    fn nearest_keeps_type() {
224        let a = nd(Arc::new(Int32Array::from(vec![10, 20, 30])), &[("x", 3)]);
225        let x = nd(
226            Arc::new(Float64Array::from(vec![0.0, 1.0, 2.0])),
227            &[("x", 3)],
228        );
229        let new_x = nd(
230            Arc::new(Float64Array::from(vec![0.4, 0.5, 1.6, 9.0])),
231            &[("x", 4)],
232        );
233        let out = interp_with(&a, "x", &x, &new_x, InterpMethod::Nearest).unwrap();
234        let v: Vec<Option<i32>> = out
235            .values()
236            .as_primitive::<arrow::datatypes::Int32Type>()
237            .iter()
238            .collect();
239        assert_eq!(v, vec![Some(10), Some(10), Some(30), None]);
240    }
241
242    #[test]
243    fn rejects_unsorted_coordinates() {
244        let a = nd(Arc::new(Int32Array::from(vec![1, 2])), &[("x", 2)]);
245        let x = nd(Arc::new(Float64Array::from(vec![1.0, 0.0])), &[("x", 2)]);
246        assert!(interp(&a, "x", &x, &x).is_err());
247    }
248}
nd_arrow_array/compute/interp.rs

nd_arrow_array/compute/
interp.rs
