eorst/
blocks.rs

1//! Block types for raster dataset processing.
2//!
3//! This module provides types for representing blocks of raster data.
4
5use crate::core_types::RasterType;
6use crate::gdal_utils::{open_for_update, raster_from_size, write_bands_to_file};
7use crate::metadata::RasterMetadata;
8use crate::types::{BlockSize, GeoTransform, Overlap, RasterDataShape, ReadWindow};
9
10use gdal::raster::Buffer;
11use log::debug;
12use ndarray::{s, Array2, Array3, Array4, Axis};
13use std::path::{Path, PathBuf};
14
15impl<U> RasterBlock<U>
16where
17    U: RasterType,
18{
19    /// Converts a 2D array into a 3D feature array.
20    pub fn into_frc(&self, data: &Array2<U>) -> Array3<U> {
21        let n_features = data.shape()[1];
22        let n_rows = self.read_window.size.rows as usize - self.overlap.top - self.overlap.bottom;
23        let n_cols = self.read_window.size.cols as usize; // - self.overlap.left - self.overlap.right;
24        let mut result: Array3<U> = Array3::zeros((n_features, n_rows, n_cols));
25        for n_feature in 0..n_features {
26            let feature = data.slice(s![.., n_feature]);
27            let feature = feature.to_shape((n_rows, n_cols)).unwrap();
28            feature.assign_to(result.slice_mut(s![n_feature, .., ..]))
29        }
30
31        result
32    }
33
34    /// Writes sample data to a vector file.
35    pub fn write_samples_feature<T>(&self, data: &Array2<T>, file_name: &PathBuf, na: T)
36    where
37        T: RasterType,
38    {
39        let n_bands = data.shape()[1] as isize;
40        let block_geotransform = self.geo_transform;
41        let epsg_code = self.epsg_code;
42        let size_x = self.read_window.size.cols as usize - self.overlap.left - self.overlap.right;
43        let size_y = self.read_window.size.rows as usize - self.overlap.top - self.overlap.bottom;
44        let block_size: BlockSize = BlockSize {
45            rows: size_y,
46            cols: size_x,
47        };
48        let na = T::from(na).unwrap();
49        raster_from_size::<T>(
50            file_name,
51            block_geotransform,
52            epsg_code as u32,
53            block_size,
54            n_bands,
55            na,
56        );
57        let out_ds = open_for_update(Path::new(&file_name));
58
59        // each feature goes to a band
60        for feature in 1..n_bands + 1 {
61            let mut out_band = out_ds.rasterband(feature as usize).unwrap();
62            // out_band.set_no_data_value_u64(Some(na)).unwrap();
63            let out_data = data.slice(s![.., feature - 1]);
64            let out_data_u8: Vec<T> = out_data.into_iter().map(|v| *v as T).collect();
65            let mut data_buffer = Buffer::new((block_size.cols, block_size.rows), out_data_u8);
66
67            out_band
68                .write((0, 0), (size_x, size_y), &mut data_buffer)
69                .unwrap();
70        }
71    }
72
73    /// Writes 3D data to a file.
74    pub fn write3<T>(&self, data: Array3<T>, out_fn: &PathBuf)
75    where
76        T: RasterType,
77    {
78        // create and empty raster with the right size
79        debug!("Writing block {:?}", out_fn);
80        let block_geotransform = self.geo_transform;
81        let block_size: BlockSize = BlockSize {
82            rows: self.read_window.size.rows as usize,
83            cols: self.read_window.size.cols as usize,
84        };
85        let n_bands = data.shape()[0] as isize;
86        let epsg_code = self.epsg_code;
87        let na_value = T::zero();
88        raster_from_size::<T>(
89            out_fn,
90            block_geotransform,
91            epsg_code.try_into().unwrap(),
92            block_size,
93            n_bands,
94            na_value,
95        );
96        let out_ds = open_for_update(out_fn);
97        write_bands_to_file(&out_ds, data.view(), (0, 0), (block_size.cols, block_size.rows));
98    }
99
100    /// Writes each time step (axis 0) of a 4D array as a separate block file.
101    ///
102    /// Replaces the duplicated loop pattern found in `apply()`, `apply_with_mask()`
103    /// and `mosaic()` in processing.rs.
104    pub fn write_time_step_blocks<T>(
105        &self,
106        data: &Array4<T>,
107        tmp_file: &Path,
108        file_stem: &str,
109        bid: usize,
110    ) -> Vec<PathBuf>
111    where
112        T: RasterType,
113    {
114        let mut block_fns: Vec<PathBuf> = Vec::new();
115        for (tid, result3) in data.axis_iter(Axis(0)).enumerate() {
116            let block_fn = tmp_file.with_file_name(format!("{}_{}_{}.tif", file_stem, tid, bid));
117            self.write3(result3.to_owned(), &block_fn);
118            block_fns.push(block_fn);
119        }
120        block_fns
121    }
122}
123
124/// A block representing a partition of a raster dataset for parallel processing.
125#[derive(Debug, Clone, PartialEq)]
126pub struct RasterBlock<U>
127where
128    U: RasterType,
129{
130    /// Index of this block in the dataset
131    pub block_index: usize,
132    /// Window describing which portion of the dataset this block covers
133    pub read_window: ReadWindow,
134    /// Size of overlap with neighboring blocks
135    pub overlap_size: usize,
136    /// Geographic transformation parameters
137    pub geo_transform: GeoTransform,
138    /// Overlap data with neighboring blocks
139    pub overlap: Overlap,
140    /// Shape of the raster data
141    pub shape: RasterDataShape,
142    /// EPSG coordinate reference system code
143    pub epsg_code: usize,
144    /// Full raster metadata
145    pub raster_metadata: RasterMetadata<U>,
146}
eorst/blocks.rs

eorst/
blocks.rs
