fastlem/lem/
generator.rs

1use rand::{rngs::StdRng, Rng, SeedableRng};
2use std::marker::PhantomData;
3use thiserror::Error;
4
5use crate::{
6    core::{
7        parameters::TopographicalParameters,
8        traits::{Model, Site},
9        units::{Length, Step},
10    },
11    lem::drainage_basin::DrainageBasin,
12    lem::stream_tree,
13};
14
15/// The default value of the exponent `m` for calculating stream power.
16const DEFAULT_M_EXP: f64 = 0.5;
17
18#[derive(Error, Debug)]
19pub enum GenerationError {
20    #[error("The number of topographical parameters must be equal to the number of sites")]
21    InvalidNumberOfParameters,
22    #[error("You must set topographical parameters before generating terrain")]
23    ParametersNotSet,
24    #[error("You must set `TerrainModel` before generating terrain")]
25    ModelNotSet,
26}
27
28/// Provides methods for generating terrain.
29///
30/// ### Required properties
31///  - `model` is the vector representation of the terrain network.
32///  - `parameters` is the topographical parameters of sites. Each parameter contains the uplift rates, erodibilities, base elevations and maximum slopes (see [TopographicalParameters] for details).
33/// ### Optional properties
34///  - `max_iteration` is the maximum number of iterations. If not set, the iterations will be repeated until the elevations of all sites are stable.
35///
36#[derive(Clone)]
37pub struct TerrainGenerator<S, M, T>
38where
39    S: Site,
40    M: Model<S, T>,
41{
42    model: Option<M>,
43    parameters: Option<Vec<TopographicalParameters>>,
44    max_iteration: Option<Step>,
45    _phantom: PhantomData<(S, T)>,
46}
47
48impl<S, M, T> Default for TerrainGenerator<S, M, T>
49where
50    S: Site,
51    M: Model<S, T>,
52{
53    fn default() -> Self {
54        Self {
55            model: None,
56            parameters: None,
57            max_iteration: None,
58            _phantom: PhantomData,
59        }
60    }
61}
62
63impl<S, M, T> TerrainGenerator<S, M, T>
64where
65    S: Site,
66    M: Model<S, T>,
67{
68    /// Set the model that contains the set of sites.
69    pub fn set_model(mut self, model: M) -> Self {
70        self.model = Some(model);
71        self
72    }
73
74    /// Set the topographical parameters of sites. See [TopographicalParameters] about the parameters.
75    pub fn set_parameters(mut self, parameters: Vec<TopographicalParameters>) -> Self {
76        self.parameters = Some(parameters);
77        self
78    }
79
80    /// Set the maximum number of iterations.
81    ///
82    /// The iteration(loop) for calculating elevations will be stopped when the number of iterations reaches `max_iteration`.
83    /// If not set, the iterations will be repeated until the elevations of all sites are stable.
84    pub fn set_max_iteration(mut self, max_iteration: Step) -> Self {
85        self.max_iteration = Some(max_iteration);
86        self
87    }
88
89    /// Generate terrain.
90    pub fn generate(self) -> Result<T, GenerationError> {
91        let model = {
92            if let Some(model) = &self.model {
93                model
94            } else {
95                return Err(GenerationError::ModelNotSet);
96            }
97        };
98
99        let (num, sites, areas, graph, default_outlets) = (
100            model.num(),
101            model.sites(),
102            model.areas(),
103            model.graph(),
104            model.default_outlets(),
105        );
106
107        let parameters = {
108            if let Some(parameters) = &self.parameters {
109                if parameters.len() != num {
110                    return Err(GenerationError::InvalidNumberOfParameters);
111                }
112                parameters
113            } else {
114                return Err(GenerationError::ParametersNotSet);
115            }
116        };
117
118        let m_exp = DEFAULT_M_EXP;
119
120        let outlets = {
121            let outlets = parameters
122                .iter()
123                .enumerate()
124                .filter(|(_, param)| param.is_outlet)
125                .map(|(i, _)| i)
126                .collect::<Vec<_>>();
127            if outlets.is_empty() {
128                default_outlets.to_vec()
129            } else {
130                outlets
131            }
132        };
133
134        let mut rng: StdRng = SeedableRng::seed_from_u64(0);
135        let mut elevations = parameters
136            .iter()
137            .map(|a| a.base_elevation + rng.gen::<f64>() * f64::EPSILON)
138            .collect::<Vec<_>>();
139
140        for _ in 0..self.max_iteration.unwrap_or(u32::MAX) {
141            let stream_tree =
142                stream_tree::StreamTree::construct(sites, &elevations, graph, &outlets);
143
144            let mut drainage_areas: Vec<f64> = areas.to_vec();
145            let mut response_times = vec![0.0; num];
146            let mut changed = false;
147
148            // calculate elevations for each drainage basin
149            outlets.iter().for_each(|&outlet| {
150                // construct drainage basin
151                let drainage_basin = DrainageBasin::construct(outlet, &stream_tree, graph);
152
153                // calculate drainage areas
154                drainage_basin.for_each_downstream(|i| {
155                    let j = stream_tree.next[i];
156                    if j != i {
157                        drainage_areas[j] += drainage_areas[i];
158                    }
159                });
160
161                // calculate response times
162                drainage_basin.for_each_upstream(|i| {
163                    let j = stream_tree.next[i];
164                    let distance: Length = {
165                        let (ok, edge) = graph.has_edge(i, j);
166                        if ok {
167                            edge
168                        } else {
169                            1.0
170                        }
171                    };
172                    let celerity = parameters[i].erodibility * drainage_areas[i].powf(m_exp);
173                    response_times[i] += response_times[j] + 1.0 / celerity * distance;
174                });
175
176                // calculate elevations
177                drainage_basin.for_each_upstream(|i| {
178                    let mut new_elevation = elevations[outlet]
179                        + parameters[i].uplift_rate
180                            * (response_times[i] - response_times[outlet]).max(0.0);
181
182                    // check if the slope is too steep
183                    // if max_slope_func is not set, the slope is not checked
184                    if let Some(max_slope) = parameters[i].max_slope {
185                        let j = stream_tree.next[i];
186                        let distance: Length = {
187                            let (ok, edge) = graph.has_edge(i, j);
188                            if ok {
189                                edge
190                            } else {
191                                1.0
192                            }
193                        };
194                        let max_slope = max_slope.tan();
195                        let slope = (new_elevation - elevations[j]) / distance;
196                        if slope > max_slope {
197                            new_elevation = elevations[j] + max_slope * distance;
198                        }
199                    }
200
201                    changed |= new_elevation != elevations[i];
202                    elevations[i] = new_elevation;
203                });
204            });
205
206            // if the elevations of all sites are stable, break
207            if !changed {
208                break;
209            }
210        }
211
212        Ok(model.create_terrain_from_result(&elevations))
213    }
214}