eorst/
data_sources.rs

1//! Data source definitions for raster datasets.
2//!
3//! This module provides types for describing the source files and bands
4//! that make up a raster dataset.
5
6use chrono::{DateTime, FixedOffset};
7use gdal::Dataset;
8use std::path::PathBuf;
9
10/// Type for representing dates in raster data.
11#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
12pub enum DateType {
13    /// Specific date.
14    Date(DateTime<FixedOffset>),
15    /// Index into a time dimension.
16    Index(usize),
17}
18
19/// Data source for a raster dataset.
20#[derive(Debug, PartialEq, Clone)]
21pub struct DataSource {
22    /// Path to the source file.
23    pub source: PathBuf,
24    /// Band indices to use.
25    pub bands: Vec<usize>,
26    /// Names of the layers.
27    pub layer_names: Vec<String>,
28}
29
30/// Builder for creating a DataSource.
31pub struct DataSourceBuilder {
32    /// Path to the source file.
33    pub source: PathBuf,
34    /// Band indices to use.
35    pub bands: Vec<usize>,
36    /// Names of the layers.
37    pub layer_names: Vec<String>,
38    /// Date indices for temporal data.
39    pub date_indices: Vec<DateType>,
40}
41
42impl DataSourceBuilder {
43    /// Creates a DataSourceBuilder from a file path.
44    pub fn from_file(file_name: &PathBuf) -> Self {
45        let source = file_name;
46
47        let n_bands = DataSourceBuilder::get_n_bands(file_name);
48        let bands: Vec<usize> = (1..n_bands + 1).collect();
49
50        let layer_names = (0..n_bands)
51            .map(|layer_index| format!("layer_{:?}", layer_index))
52            .collect::<Vec<String>>();
53        let date_indices = vec![DateType::Index(0)];
54
55        DataSourceBuilder {
56            source: source.clone(),
57            bands,
58            layer_names,
59            date_indices,
60        }
61    }
62
63    /// Sets the band dimension (layer or time).
64    pub fn band_dimension(self, dimension: crate::types::Dimension) -> Self {
65        let n_bands = DataSourceBuilder::get_n_bands(&self.source);
66
67        let layer_names = match dimension {
68            crate::types::Dimension::Layer => self.layer_names,
69            crate::types::Dimension::Time => vec!["Layer_0".to_string()],
70        };
71
72        let date_indices = match dimension {
73            crate::types::Dimension::Layer => self.date_indices,
74            crate::types::Dimension::Time => {
75                (0..n_bands).map(DateType::Index).collect::<Vec<DateType>>()
76            }
77        };
78
79        DataSourceBuilder {
80            source: self.source.clone(),
81            bands: self.bands,
82            layer_names,
83            date_indices,
84        }
85    }
86
87    /// Sets the date indices for temporal data.
88    pub fn set_dates(self, dates: Vec<DateType>) -> Self {
89        let n_dates = self.date_indices.len();
90        let n_dates_new = dates.len();
91        assert_eq!(n_dates, n_dates_new);
92
93        DataSourceBuilder {
94            source: self.source.clone(),
95            bands: self.bands,
96            layer_names: self.layer_names,
97            date_indices: dates,
98        }
99    }
100
101    /// Sets the layer names.
102    pub fn set_names(self, names: Vec<&str>) -> Self {
103        let n_names = self.layer_names.len();
104        let n_names_new = names.len();
105        assert_eq!(n_names, n_names_new);
106
107        DataSourceBuilder {
108            source: self.source.clone(),
109            bands: self.bands,
110            layer_names: names.iter().map(|n| n.to_string()).collect(),
111            date_indices: self.date_indices,
112        }
113    }
114
115    /// Builds the DataSource.
116    pub fn build(self) -> DataSource {
117        DataSource {
118            source: self.source,
119            bands: self.bands,
120            layer_names: self.layer_names,
121        }
122    }
123
124    /// Sets the band indices.
125    pub fn bands(mut self, bands: Vec<usize>) -> Self {
126        self.bands = bands.clone();
127        // Update layer_names to match the selected bands
128        self.layer_names = bands
129            .iter()
130            .enumerate()
131            .map(|(i, _)| format!("layer_{}", i))
132            .collect();
133        self
134    }
135
136    fn get_n_bands(source: &PathBuf) -> usize {
137        let ds = Dataset::open(source).unwrap();
138        ds.raster_count()
139    }
140}