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
15const 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#[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 pub fn set_model(mut self, model: M) -> Self {
70 self.model = Some(model);
71 self
72 }
73
74 pub fn set_parameters(mut self, parameters: Vec<TopographicalParameters>) -> Self {
76 self.parameters = Some(parameters);
77 self
78 }
79
80 pub fn set_max_iteration(mut self, max_iteration: Step) -> Self {
85 self.max_iteration = Some(max_iteration);
86 self
87 }
88
89 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 outlets.iter().for_each(|&outlet| {
150 let drainage_basin = DrainageBasin::construct(outlet, &stream_tree, graph);
152
153 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 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 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 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 !changed {
208 break;
209 }
210 }
211
212 Ok(model.create_terrain_from_result(&elevations))
213 }
214}