nd_arrow_array/compute/
aggregate.rs

1//! Block and weighted aggregation: `coarsen`, `weighted_sum`, `weighted_mean`, `dot`.
2
3use crate::compute::elementwise::{div, is_not_null_mask, mul, where_};
4use crate::compute::reduce::{Reduction, reduce, sum};
5use crate::compute::shape::{Selector, Slice, isel};
6use crate::dimensions::{Dimension, Dimensions};
7use crate::error::{Result, nd_err};
8use crate::{IntoScalar, NdArrowArray};
9
10/// What [`coarsen`] does when a dimension is not a multiple of its window.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
12pub enum CoarsenBoundary {
13    /// Fail, like xarray `boundary="exact"`.
14    #[default]
15    Exact,
16    /// Drop the last positions, like xarray `boundary="trim"`.
17    Trim,
18}
19
20/// Reduce non-overlapping blocks, like `xarray.DataArray.coarsen`.
21///
22/// Each `(dim, window)` pair splits `dim` into blocks of `window` positions.
23/// The split itself is a reshape without a copy.
24pub fn coarsen(
25    a: &NdArrowArray,
26    windows: &[(&str, usize)],
27    op: Reduction,
28    boundary: CoarsenBoundary,
29) -> Result<NdArrowArray> {
30    let mut trimmed = a.clone();
31    for &(name, window) in windows {
32        if window == 0 {
33            return Err(nd_err!(
34                InvalidArgument,
35                "coarsen window must be at least 1"
36            ));
37        }
38        let size = a
39            .dims()
40            .size_of(name)
41            .ok_or_else(|| nd_err!(InvalidArgument, "dimension '{name}' is not in {}", a.dims()))?;
42        if size % window != 0 {
43            if boundary == CoarsenBoundary::Exact {
44                return Err(nd_err!(
45                    ShapeMismatch,
46                    "dimension '{name}' of size {size} is not a multiple of {window}"
47                ));
48            }
49            let keep = (size - size % window) as isize;
50            trimmed = isel(&trimmed, &[(name, Selector::Slice(Slice::range(0, keep)))])?;
51        }
52    }
53
54    let mut dims: Vec<Dimension> = Vec::new();
55    let mut window_dims: Vec<String> = Vec::new();
56    for dim in trimmed.dims().iter() {
57        match windows.iter().find(|(name, _)| *name == dim.name()) {
58            Some(&(name, window)) => {
59                let inner = format!("{name}__window");
60                if trimmed.dims().contains(&inner) {
61                    return Err(nd_err!(
62                        InvalidArgument,
63                        "dimension '{inner}' already exists"
64                    ));
65                }
66                dims.push(Dimension::new(name, dim.size() / window));
67                dims.push(Dimension::new(inner.as_str(), window));
68                window_dims.push(inner);
69            }
70            None => dims.push(dim.clone()),
71        }
72    }
73    let attrs = trimmed.attributes().clone();
74    let reshaped = trimmed.reshape(Dimensions::try_new(dims)?)?;
75    let names: Vec<&str> = window_dims.iter().map(String::as_str).collect();
76    Ok(reduce(&reshaped, Some(&names), op)?.with_attributes(attrs))
77}
78
79/// Weighted sum over `axes`, like `xarray.DataArray.weighted(w).sum()`.
80/// Null values of `a` count as zero. Result type follows `a * weights`.
81pub fn weighted_sum(
82    a: &NdArrowArray,
83    weights: &NdArrowArray,
84    axes: Option<&[&str]>,
85) -> Result<NdArrowArray> {
86    sum(&mul(a, weights)?, axes)
87}
88
89/// Weighted mean over `axes`, like `xarray.DataArray.weighted(w).mean()`.
90/// Weights of null values do not count. The result is floating point.
91pub fn weighted_mean(
92    a: &NdArrowArray,
93    weights: &NdArrowArray,
94    axes: Option<&[&str]>,
95) -> Result<NdArrowArray> {
96    let numerator = weighted_sum(a, weights, axes)?;
97    let present = is_not_null_mask(a)?;
98    let zero = NdArrowArray::try_new_scalar(0.0f64.into_scalar())?;
99    let used_weights = where_(&present, weights, &zero)?;
100    let denominator = sum(&used_weights, axes)?;
101    let nonzero = crate::compute::ne(&denominator, &zero)?;
102    let null = NdArrowArray::new_null_scalar(&arrow::datatypes::DataType::Float64);
103    where_(&nonzero, &div(&numerator, &denominator)?, &null)
104}
105
106/// Sum of products over the named shared dimensions, like `xarray.dot`.
107///
108/// The product broadcasts first, so memory use follows the broadcast shape.
109/// Nulls count as zero.
110pub fn dot(a: &NdArrowArray, b: &NdArrowArray, dims: &[&str]) -> Result<NdArrowArray> {
111    for name in dims {
112        if !(a.dims().contains(name) && b.dims().contains(name)) {
113            return Err(nd_err!(
114                InvalidArgument,
115                "dot dimension '{name}' must be in both arrays"
116            ));
117        }
118    }
119    sum(&mul(a, b)?, Some(dims))
120}
121
122#[cfg(test)]
123mod tests {
124    use std::sync::Arc;
125
126    use arrow::array::{AsArray, Float64Array, Int32Array};
127    use arrow::datatypes::{Float64Type, Int64Type};
128
129    use super::*;
130
131    fn nd(values: Vec<f64>, dims: &[(&str, usize)]) -> NdArrowArray {
132        NdArrowArray::try_new(
133            Arc::new(Float64Array::from(values)),
134            Dimensions::try_new(dims.iter().copied()).unwrap(),
135        )
136        .unwrap()
137    }
138
139    fn f64s(a: &NdArrowArray) -> Vec<Option<f64>> {
140        a.values().as_primitive::<Float64Type>().iter().collect()
141    }
142
143    #[test]
144    fn coarsen_blocks() {
145        let a = NdArrowArray::try_new(
146            Arc::new(Int32Array::from_iter_values(0..16)),
147            Dimensions::try_new([("y", 4), ("x", 4)]).unwrap(),
148        )
149        .unwrap();
150        let c = coarsen(
151            &a,
152            &[("y", 2), ("x", 2)],
153            Reduction::Sum,
154            CoarsenBoundary::Exact,
155        )
156        .unwrap();
157        assert_eq!(
158            c.dims(),
159            &Dimensions::try_new([("y", 2), ("x", 2)]).unwrap()
160        );
161        let v: Vec<i64> = c.values().as_primitive::<Int64Type>().values().to_vec();
162        assert_eq!(v, vec![10, 18, 42, 50]);
163
164        assert!(coarsen(&a, &[("x", 3)], Reduction::Sum, CoarsenBoundary::Exact).is_err());
165        let t = coarsen(&a, &[("x", 3)], Reduction::Max, CoarsenBoundary::Trim).unwrap();
166        assert_eq!(t.shape(), vec![4, 1]);
167    }
168
169    #[test]
170    fn weighted_mean_ignores_null_weights() {
171        let a = NdArrowArray::try_new(
172            Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])),
173            Dimensions::try_new([("lat", 3)]).unwrap(),
174        )
175        .unwrap();
176        let w = nd(vec![1.0, 5.0, 3.0], &[("lat", 3)]);
177        assert_eq!(f64s(&weighted_sum(&a, &w, None).unwrap()), vec![Some(10.0)]);
178        assert_eq!(f64s(&weighted_mean(&a, &w, None).unwrap()), vec![Some(2.5)]);
179    }
180
181    #[test]
182    fn dot_matches_matrix_product() {
183        let a = nd(vec![1.0, 2.0, 3.0, 4.0], &[("i", 2), ("k", 2)]);
184        let b = nd(vec![5.0, 6.0, 7.0, 8.0], &[("k", 2), ("j", 2)]);
185        let c = dot(&a, &b, &["k"]).unwrap();
186        assert_eq!(c.dims().names(), vec!["i", "j"]);
187        assert_eq!(
188            f64s(&c),
189            vec![Some(19.0), Some(22.0), Some(43.0), Some(50.0)]
190        );
191        assert!(dot(&a, &b, &["i"]).is_err());
192    }
193}
nd_arrow_array/compute/aggregate.rs

nd_arrow_array/compute/
aggregate.rs
