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}