eorst/rasterdataset/
composition.rs

1//! Composition methods for RasterDataset.
2//!
3//! This module contains methods for stacking, extending, and generating column names.
4
5use crate::core_types::RasterType;
6use crate::types::Dimension as UtilsDimension;
7use crate::rasterdataset::RasterDataset;
8use crate::selection::Stack;
9
10use crate::data_sources::DateType;
11#[cfg(feature = "use_polars")]
12use crate::gdal_utils::create_rayon_pool;
13#[cfg(feature = "use_polars")]
14use rayon::iter::{IntoParallelIterator, ParallelIterator};
15
16impl<R> RasterDataset<R>
17where
18    R: RasterType,
19{
20    /// Stack RasterDataset along a dimension.
21    pub fn stack(
22        &mut self,
23        other: &RasterDataset<R>,
24        dimension_to_stack: UtilsDimension,
25    ) -> &mut RasterDataset<R> {
26        match dimension_to_stack {
27            UtilsDimension::Layer => {
28                if self.metadata.date_indices != other.metadata.date_indices {
29                    panic!(
30                        "To stack datasets over Layers, both datasets have to the same time indics"
31                    )
32                }
33                self.metadata
34                    .layer_indices
35                    .extend_from_slice(&other.metadata.layer_indices)
36            }
37            UtilsDimension::Time => {
38                if self.metadata.layer_indices != other.metadata.layer_indices {
39                    panic!("To stack datasets over Time, both have to have the same layers.")
40                }
41                self.metadata
42                    .date_indices
43                    .extend_from_slice(&other.metadata.date_indices);
44            }
45        }
46
47        let axis = dimension_to_stack.get_axis();
48
49        let max_dim_self = self.metadata.shape[axis];
50        self.metadata
51            .shape
52            .stack(other.metadata.shape, dimension_to_stack);
53
54        for other_layer in &other.metadata.layers {
55            let mut layer = other_layer.clone();
56            layer.stack_position(other_layer.to_owned(), dimension_to_stack, max_dim_self);
57            self.metadata.layers.push(layer);
58        }
59
60        self
61    }
62
63    /// Extend a RasterDataset.
64    pub fn extend<V>(&mut self, other: &RasterDataset<V>) -> &mut RasterDataset<R>
65    where
66        V: RasterType,
67    {
68        self.metadata.shape.extend(other.metadata.shape);
69
70        for layer in &other.metadata.layers {
71            self.metadata.layers.push(layer.to_owned());
72        }
73        self
74    }
75
76    /// Create names like time0_layer0, time0_layer1, time1_layer0, etc.
77    pub fn column_names(&self) -> Vec<String> {
78        let mut col_names = Vec::new();
79        for time in &self.metadata.date_indices {
80            for layer in &self.metadata.layer_indices {
81                let time_str = match time {
82                    DateType::Date(date) => format!("{}", date.format("%Y-%m-%d %H:%M:%S UTC%Z")),
83
84                    DateType::Index(index) => format!("time_{}", index),
85                };
86                let c = format!("{time_str}--{layer}");
87                col_names.push(c);
88            }
89        }
90        col_names
91    }
92
93    /// Gets unique values across all blocks.
94    #[cfg(feature = "use_polars")]
95    pub(crate) fn get_unique_values(&self) -> Vec<i32> {
96        let block_to_process: Vec<usize> = (0..self.blocks.len()).collect();
97        let pool = create_rayon_pool(1);
98        let handle = pool.install(|| {
99            block_to_process.into_par_iter().map(|id| {
100                let data = self.read_block::<i32>(id);
101
102                let unique_values: std::collections::HashSet<_> = data.into_iter().collect();
103
104                unique_values
105            })
106        });
107        let collected: Vec<_> = handle.collect();
108        let mut unique: std::collections::HashSet<_> = std::collections::HashSet::new();
109        for set in collected {
110            unique.extend(&set);
111        }
112        let mut unique: Vec<_> = unique.iter().cloned().collect();
113        unique.sort();
114        unique
115    }
116
117    /// An iterator over the [`crate::blocks::RasterBlock`]s of a [`RasterDataset`].
118    pub fn iter(&self) -> RasterDatasetIter<'_, R> {
119        let iter_index = 0;
120        RasterDatasetIter {
121            parent: self,
122            block: &self.blocks[iter_index],
123            iter_index,
124        }
125    }
126}
127
128/// Iterator over RasterDataset blocks.
129pub struct RasterDatasetIter<'a, R>
130where
131    R: RasterType,
132{
133    /// Reference to the parent dataset
134    pub parent: &'a RasterDataset<R>,
135    /// Current block
136    pub block: &'a crate::blocks::RasterBlock<R>,
137    /// Current iteration index
138    pub iter_index: usize,
139}
140
141impl<'a, R> Iterator for RasterDatasetIter<'a, R>
142where
143    R: RasterType,
144{
145    type Item = RasterDatasetIter<'a, R>;
146
147    fn next(&mut self) -> Option<Self::Item> {
148        if self.iter_index < self.parent.blocks.len() {
149            let current = RasterDatasetIter {
150                parent: self.parent,
151                block: self.block,
152                iter_index: self.iter_index,
153            };
154            self.iter_index += 1;
155            if self.iter_index < self.parent.blocks.len() {
156                self.block = &self.parent.blocks[self.iter_index];
157            }
158            Some(current)
159        } else {
160            None
161        }
162    }
163}
eorst/rasterdataset/composition.rs

eorst/rasterdataset/
composition.rs
