demes/builder.rs
1use thiserror::Error;
2
3use crate::specification::Graph;
4use crate::specification::GraphDefaults;
5use crate::specification::UnresolvedDeme;
6use crate::specification::UnresolvedDemeHistory;
7use crate::specification::UnresolvedEpoch;
8use crate::specification::UnresolvedGraph;
9use crate::DemesError;
10use crate::InputGenerationTime;
11use crate::InputProportion;
12use crate::InputTime;
13use crate::TimeUnits;
14use crate::UnresolvedMigration;
15
16/// Error type raised by [`GraphBuilder`]
17#[derive(Error, Debug)]
18#[non_exhaustive]
19pub enum BuilderError {
20 /// Error type when defaults fail to serialize
21 #[error(transparent)]
22 SerdeYamlError(#[from] serde_yaml::Error),
23}
24
25/// This type allows building a [`Graph`](crate::Graph) using code
26/// rather then using text input.
27///
28/// # Notes
29///
30/// * A "builder" in rust will never be as convenient
31/// as one in, say, Python or Juilia.
32/// The lack of a rust REPL and the strong type checking
33/// are the primary reasons.
34/// * All error checks are delayed until resolution.
35pub struct GraphBuilder {
36 graph: UnresolvedGraph,
37 metadata: Option<crate::Metadata>,
38}
39
40impl GraphBuilder {
41 /// Constructor
42 ///
43 /// # Returns
44 ///
45 /// This function returns an "builder" containing an unresolved
46 /// [`Graph`](crate::Graph).
47 pub fn new(
48 time_units: TimeUnits,
49 generation_time: Option<InputGenerationTime>,
50 defaults: Option<GraphDefaults>,
51 ) -> Self {
52 Self {
53 graph: UnresolvedGraph::new(time_units, generation_time, defaults),
54 metadata: None,
55 }
56 }
57
58 /// Construct a builder with time units in generations.
59 ///
60 /// This function works by calling [`GraphBuilder::new`](crate::GraphBuilder::new).
61 pub fn new_generations(defaults: Option<GraphDefaults>) -> Self {
62 Self {
63 graph: UnresolvedGraph::new(TimeUnits::Generations, None, defaults),
64 metadata: None,
65 }
66 }
67
68 /// Add a [`Deme`](crate::Deme) to the graph.
69 ///
70 /// # Examples
71 ///
72 /// ```
73 /// let start_size = demes::InputDemeSize::from(100.);
74 /// let epoch = demes::UnresolvedEpoch{start_size: Some(start_size), ..Default::default()};
75 /// let history = demes::UnresolvedDemeHistory::default();
76 /// let mut b = demes::GraphBuilder::new_generations(None);
77 /// b.add_deme("A", vec![epoch], history, Some("this is deme A"));
78 /// b.resolve().unwrap();
79 /// ```
80 ///
81 /// # Notes
82 pub fn add_deme<I: IntoIterator<Item = UnresolvedEpoch>>(
83 &mut self,
84 name: &str,
85 epochs: I,
86 history: UnresolvedDemeHistory,
87 description: Option<&str>,
88 ) {
89 let epochs = epochs.into_iter().collect::<Vec<_>>();
90 let ptr = UnresolvedDeme::new_via_builder(name, epochs, history, description);
91 self.graph.add_deme(ptr);
92 }
93
94 /// Add a migration to the graph.
95 ///
96 /// # Examples
97 ///
98 /// ## Adding an asymmetric migration
99 ///
100 /// See [`UnresolvedMigration`].
101 ///
102 /// ```
103 /// let start_size = demes::InputDemeSize::from(100.);
104 /// let epoch = demes::UnresolvedEpoch{start_size: Some(start_size), ..Default::default()};
105 /// let history = demes::UnresolvedDemeHistory::default();
106 /// let mut b = demes::GraphBuilder::new_generations(None);
107 /// b.add_deme("A", vec![epoch], history.clone(), Some("this is deme A"));
108 /// b.add_deme("B", vec![epoch], history, Some("this is deme B"));
109 /// let migration = demes::UnresolvedMigration::default().set_source("A").set_dest("B").set_rate(1e-4);
110 /// b.add_migration(migration);
111 /// b.resolve().unwrap();
112 /// ```
113 ///
114 /// ## Adding a symmetric migration
115 ///
116 /// See [`UnresolvedMigration`].
117 ///
118 /// ```
119 /// let start_size = demes::InputDemeSize::from(100.);
120 /// let epoch = demes::UnresolvedEpoch{start_size: Some(start_size), ..Default::default()};
121 /// let history = demes::UnresolvedDemeHistory::default();
122 /// let mut b = demes::GraphBuilder::new_generations(None);
123 /// b.add_deme("A", vec![epoch], history.clone(), Some("this is deme A"));
124 /// b.add_deme("B", vec![epoch], history, Some("this is deme B"));
125 /// let migration = demes::UnresolvedMigration::default().set_demes(["A","B"].as_slice()).set_rate(1e-4);
126 /// b.add_migration(migration);
127 /// b.resolve().unwrap();
128 /// ```
129 ///
130 /// We can also use a `Vec` instead of an array:
131 ///
132 /// ```
133 /// # let start_size = demes::InputDemeSize::from(100.);
134 /// # let epoch = demes::UnresolvedEpoch{start_size: Some(start_size), ..Default::default()};
135 /// # let history = demes::UnresolvedDemeHistory::default();
136 /// # let mut b = demes::GraphBuilder::new_generations(None);
137 /// # b.add_deme("A", vec![epoch], history.clone(), Some("this is deme A"));
138 /// # b.add_deme("B", vec![epoch], history, Some("this is deme B"));
139 /// let migration = demes::UnresolvedMigration::default().set_demes(vec!["A","B"]).set_rate(1e-4);
140 /// b.add_migration(migration);
141 /// # b.resolve().unwrap();
142 /// ```
143 pub fn add_migration(&mut self, migration: UnresolvedMigration) {
144 self.graph.add_migration(migration);
145 }
146
147 /// Add an [`UnresolvedPulse`](crate::UnresolvedPulse) to the graph.
148 ///
149 /// # Examples
150 ///
151 /// ```
152 /// let start_size = demes::InputDemeSize::from(100.);
153 /// let epoch = demes::UnresolvedEpoch{start_size: Some(start_size), ..Default::default()};
154 /// let history = demes::UnresolvedDemeHistory::default();
155 /// let mut b = demes::GraphBuilder::new_generations(None);
156 /// b.add_deme("A", vec![epoch], history.clone(), Some("this is deme A"));
157 /// b.add_deme("B", vec![epoch], history, Some("this is deme B"));
158 /// b.add_pulse(Some(&["A"]),
159 /// Some("B"),
160 /// Some(50.0),
161 /// Some([0.5].as_slice()));
162 /// b.resolve().unwrap();
163 /// ```
164 pub fn add_pulse<T: Into<InputTime>, P: Into<InputProportion>, I: IntoIterator<Item = P>>(
165 &mut self,
166 sources: Option<&[&str]>,
167 dest: Option<&str>,
168 time: Option<T>,
169 proportions: Option<I>,
170 ) {
171 let sources = sources.map(|value| value.iter().map(|v| v.to_string()).collect::<Vec<_>>());
172 let dest = dest.map(|value| value.to_string());
173 let time = time.map(|t| t.into());
174 let proportions = proportions.map(|s| {
175 s.into_iter()
176 .map(|p| p.into())
177 .collect::<Vec<InputProportion>>()
178 });
179 self.graph.add_pulse(sources, dest, time, proportions);
180 }
181
182 /// Generate and return a resolved [`Graph`](crate::Graph).
183 ///
184 /// # Errors
185 ///
186 /// Returns [`DemesError'](crate::DemesError) if any
187 /// of the data are invalid.
188 pub fn resolve(self) -> Result<Graph, DemesError> {
189 let mut builder = self;
190 match builder.metadata {
191 None => (),
192 Some(m) => builder.graph.set_metadata(m),
193 }
194 builder.graph.resolve()?;
195 builder.graph.try_into()
196 }
197
198 /// Set top-level metadata
199 ///
200 /// # Parameters
201 ///
202 /// * `metadata`: the metadata type
203 ///
204 /// # Note
205 ///
206 /// Repeated calls will overwrite existing metadata.
207 ///
208 /// # Errors
209 ///
210 /// * [`BuilderError`] if serialization to YAML fails.
211 ///
212 /// # Example
213 ///
214 /// ```
215 /// #[derive(serde::Serialize, serde::Deserialize)]
216 /// struct MyMetaData {
217 /// foo: i32,
218 /// bar: String
219 /// }
220 /// # let mut builder = demes::GraphBuilder::new_generations(None);
221 /// builder.set_toplevel_metadata(&MyMetaData{foo: 3, bar: "string".to_owned()}).unwrap();
222 /// ```
223 pub fn set_toplevel_metadata<T: serde::Serialize>(
224 &mut self,
225 metadata: &T,
226 ) -> Result<(), BuilderError> {
227 let yaml = serde_yaml::to_string(metadata)?;
228 let metadata: crate::Metadata = serde_yaml::from_str(&yaml)?;
229 self.metadata = Some(metadata);
230 Ok(())
231 }
232}
233
234#[cfg(test)]
235mod tests {
236 use super::*;
237 use crate::specification::DemeDefaults;
238 use crate::InputDemeSize;
239
240 #[test]
241 #[should_panic]
242 fn new_builder() {
243 let b = GraphBuilder::new(TimeUnits::Generations, None, None);
244 b.resolve().unwrap();
245 }
246
247 #[test]
248 fn add_deme_with_epochs() {
249 let mut b = GraphBuilder::new_generations(Some(GraphDefaults::default()));
250 let edata = UnresolvedEpoch {
251 start_size: Some(InputDemeSize::from(100.0)),
252 ..Default::default()
253 };
254 b.add_deme("CEU", vec![edata], UnresolvedDemeHistory::default(), None);
255 let _graph = b.resolve().unwrap();
256 }
257
258 #[test]
259 fn use_proportion_for_proportions() {
260 let p = InputProportion::from(0.5);
261 let _ = UnresolvedDemeHistory {
262 proportions: Some(vec![p, p]),
263 ..Default::default()
264 };
265 }
266
267 #[test]
268 fn builder_deme_defaults() {
269 let defaults = DemeDefaults {
270 epoch: UnresolvedEpoch {
271 end_size: Some(InputDemeSize::from(100.)),
272 ..Default::default()
273 },
274 };
275 let history = UnresolvedDemeHistory {
276 defaults,
277 ..Default::default()
278 };
279 let mut b = GraphBuilder::new_generations(None);
280 b.add_deme("YRB", vec![], history, None);
281 b.resolve().unwrap();
282 }
283}