Skip to main content

egui_charts/widget/
study_templates.rs

1//! A chart's studies saved and applied as a set, with their links and shared panes.
2
3use super::{Chart, PaneId, SeriesId};
4use crate::studies::{IndicatorFactory, InputValue, StudySource, StudyTemplate, StudyTemplateSet};
5
6/// The position saved for a link to a study the chart no longer has: past the end of any set,
7/// and still a `usize` where that is 32 bits wide.
8const NO_STUDY: usize = u32::MAX as usize;
9
10impl Chart {
11    /// The chart's studies as a [`StudyTemplateSet`], in the order they were added. A link to
12    /// another study's plot is saved as that study's position in the set, or past the end of
13    /// the set when the chart no longer has that study, and the panes record which studies
14    /// share a pane. Apply the set to this or another chart with
15    /// [`apply_study_template_set`](Self::apply_study_template_set).
16    pub fn study_template_set(&self) -> StudyTemplateSet {
17        let ids: Vec<_> = self.indicators().map(|(id, _)| id).collect();
18        let position = |id: SeriesId| ids.iter().position(|study| *study == id);
19        let studies = self
20            .indicators()
21            .map(|(_, study)| {
22                let mut template = StudyTemplate::from_study(study);
23                for (_, value) in &mut template.inputs {
24                    if let InputValue::Source(StudySource::Study { id, .. }) = value {
25                        *id = position(SeriesId(*id)).unwrap_or(NO_STUDY);
26                    }
27                }
28                template
29            })
30            .collect();
31        let price_pane: Vec<_> = self
32            .indicators()
33            .enumerate()
34            .filter(|(_, (_, study))| study.is_overlay())
35            .map(|(at, _)| at)
36            .collect();
37        let mut panes: Vec<Vec<_>> = self
38            .panes()
39            .into_iter()
40            .filter(|pane| *pane != PaneId::MAIN)
41            .map(|pane| {
42                self.pane_series(pane)
43                    .into_iter()
44                    .filter_map(position)
45                    .collect::<Vec<_>>()
46            })
47            .filter(|pane| !pane.is_empty())
48            .collect();
49        // A study taken off the price pane since the last frame gets a pane of its own at the
50        // bottom in the next one.
51        let placed: Vec<_> = price_pane
52            .iter()
53            .chain(panes.iter().flatten())
54            .copied()
55            .collect();
56        panes.extend(
57            (0..ids.len())
58                .filter(|at| !placed.contains(at))
59                .map(|at| vec![at]),
60        );
61        StudyTemplateSet {
62            studies,
63            price_pane,
64            panes,
65        }
66    }
67
68    /// Replaces the chart's studies with those of `set`, built from the built-in studies, as
69    /// [`apply_study_template_set_with`](Self::apply_study_template_set_with) does with its
70    /// factory.
71    pub fn apply_study_template_set(
72        &mut self,
73        set: &StudyTemplateSet,
74    ) -> Result<Vec<SeriesId>, String> {
75        self.apply_study_template_set_with(set, &IndicatorFactory::new())
76    }
77
78    /// Replaces the chart's studies with those of `set`, built through `factory`, which
79    /// resolves the names of built-in and host-registered studies, and returns their ids in the
80    /// set's order. Each link to a study of the set reads the new study, and a link past the end
81    /// of the set reads nothing. The studies listed in
82    /// [`price_pane`](StudyTemplateSet::price_pane) are drawn on the price pane, each list in
83    /// [`panes`](StudyTemplateSet::panes) becomes a new pane at the bottom, in order, and the
84    /// other studies are placed as [`add_indicator`](Self::add_indicator) places them. The
85    /// studies are computed at once from the bars. The chart's volume setting follows the set:
86    /// a volume study in it becomes the chart's volume, and without one the setting is off, as
87    /// when the volume study is removed.
88    ///
89    /// Fails without changing the chart when the set holds more studies than the
90    /// [study cap](Self::set_study_limit) allows, when its panes name a position twice or past
91    /// its studies, or when a study cannot be built, with the error of
92    /// [`StudyTemplate::create_with`]. Wrap the call in [`user_action`](Self::user_action) when
93    /// it applies a user's choice, so that one undo restores the previous studies.
94    pub fn apply_study_template_set_with(
95        &mut self,
96        set: &StudyTemplateSet,
97        factory: &IndicatorFactory,
98    ) -> Result<Vec<SeriesId>, String> {
99        if self
100            .study_limit
101            .is_some_and(|limit| set.studies.len() > limit)
102        {
103            return Err("Study limit reached".into());
104        }
105        let mut placed = vec![false; set.studies.len()];
106        for &at in set.price_pane.iter().chain(set.panes.iter().flatten()) {
107            match placed.get_mut(at) {
108                Some(seen) if !*seen => *seen = true,
109                _ => return Err("Study template panes name a missing or repeated study".into()),
110            }
111        }
112        let studies = set
113            .studies
114            .iter()
115            .map(|template| {
116                template
117                    .create_with(factory)
118                    .map_err(|error| error.to_string())
119            })
120            .collect::<Result<Vec<_>, _>>()?;
121        let replaced: Vec<_> = self.indicators().map(|(id, _)| id).collect();
122        for id in replaced {
123            self.remove_series(id);
124        }
125        let ids: Vec<_> = studies
126            .into_iter()
127            .map(|study| self.add_indicator(study))
128            .collect();
129        for (id, template) in ids.iter().zip(&set.studies) {
130            for (input, value) in &template.inputs {
131                if let InputValue::Source(StudySource::Study { id: parent, plot }) = value {
132                    let source = StudySource::Study {
133                        id: ids.get(*parent).map_or(usize::MAX, |parent| parent.0),
134                        plot: *plot,
135                    };
136                    self.indicator_mut(*id)
137                        .expect("added study")
138                        .set_input(input, InputValue::Source(source))
139                        .expect("validated study source");
140                }
141            }
142        }
143        for &at in &set.price_pane {
144            self.move_series_to_pane(ids[at], PaneId::MAIN);
145        }
146        for pane in &set.panes {
147            if let Some((&first, rest)) = pane.split_first() {
148                self.move_indicator_to_pane(ids[first], ids[first]);
149                for &at in rest {
150                    self.move_indicator_to_pane(ids[at], ids[first]);
151                }
152            }
153        }
154        // Removing the volume study turns the volume setting off; applying one turns it on.
155        let volume = self
156            .indicators()
157            .find(|(_, study)| study.is_volume())
158            .map(|(id, _)| id);
159        self.volume_study = volume;
160        self.volume_setting = self.volume_study.is_some();
161        self.config.show_volume = self.volume_setting;
162        self.recalculate_studies();
163        Ok(ids)
164    }
165}
166
167#[cfg(test)]
168mod tests {
169    use super::super::test_frames;
170    use crate::studies::{CCI, EMA, MACD, RSI, SMA, StudySource, StudyTemplateSet, Volume};
171    use crate::widget::{Chart, SeriesId};
172
173    #[test]
174    fn study_template_sets_keep_links_shared_panes_and_volume_on_another_chart() {
175        let closes: Vec<f64> = (0..100).map(|i| 100.0 + f64::from(i)).collect();
176        let opens: Vec<f64> = closes.iter().map(|close| close - 1.0).collect();
177        let mut source = test_frames::chart(&opens, &closes);
178        let linked = source.add_indicator(Box::new(SMA::new(3)));
179        let rsi = source.add_indicator(Box::new(RSI::new(14)));
180        source.add_indicator(Box::new(MACD::new(12, 26, 9)));
181        source.add_indicator(Box::new(SMA::new(5)));
182        source.add_indicator(Box::new(Volume::default()));
183        source
184            .set_indicator_source(linked, StudySource::Study { id: rsi.0, plot: 0 })
185            .unwrap();
186        assert!(source.move_indicator_to_pane(linked, rsi));
187        let saved = serde_json::to_string(&source.study_template_set()).unwrap();
188
189        let mut chart = test_frames::chart(&opens, &closes);
190        chart.add_indicator(Box::new(EMA::new(5)));
191        chart.add_indicator(Box::new(CCI::new(20)));
192        chart.config_mut().show_volume = false;
193        let set: StudyTemplateSet = serde_json::from_str(&saved).unwrap();
194        let ids = chart.apply_study_template_set(&set).unwrap();
195
196        let studies: Vec<_> = chart
197            .indicators()
198            .map(|(id, study)| (id, study.name().to_owned()))
199            .collect();
200        let names = ["SMA", "RSI", "MACD", "SMA", "Volume"].map(String::from);
201        assert_eq!(studies, ids.iter().copied().zip(names).collect::<Vec<_>>());
202        let panes: Vec<_> = chart
203            .panes()
204            .into_iter()
205            .map(|pane| chart.pane_series(pane))
206            .collect();
207        assert_eq!(
208            panes,
209            [
210                vec![SeriesId::MAIN, ids[3], ids[4]],
211                vec![ids[1], ids[0]],
212                vec![ids[2]],
213            ]
214        );
215        let linked = chart.indicators().next().unwrap().1;
216        assert_eq!(
217            linked.source(),
218            Some(StudySource::Study {
219                id: ids[1].0,
220                plot: 0
221            })
222        );
223        // Closes that only rise have an RSI of 100, so its 3-bar mean is 100; the mean of the
224        // last three closes is 198.
225        assert_eq!(linked.values().value(0, 99), Some(100.0));
226        assert!(chart.config().show_volume);
227    }
228
229    #[test]
230    fn study_template_sets_that_cannot_be_applied_leave_the_chart_unchanged() {
231        let mut chart = test_frames::chart(&[1.0], &[2.0]);
232        chart.add_indicator(Box::new(EMA::new(5)));
233        chart.add_indicator(Box::new(CCI::new(20)));
234        let saved = chart.study_template_set();
235        let layout = |chart: &Chart| {
236            (
237                chart.indicators().map(|(id, _)| id).collect::<Vec<_>>(),
238                chart
239                    .panes()
240                    .into_iter()
241                    .map(|pane| chart.pane_series(pane))
242                    .collect::<Vec<_>>(),
243            )
244        };
245        let before = layout(&chart);
246        let mut unknown = saved.clone();
247        unknown.studies[1].name = "Missing".into();
248        let mut past_end = saved.clone();
249        past_end.panes = vec![vec![1, 2]];
250        let mut twice = saved.clone();
251        twice.price_pane = vec![0, 1];
252        let refused = "Study template panes name a missing or repeated study";
253        for (set, limit, error) in [
254            (unknown, None, "unknown study \"Missing\""),
255            (past_end, None, refused),
256            (twice, None, refused),
257            (saved, Some(1), "Study limit reached"),
258        ] {
259            chart.set_study_limit(limit);
260            assert_eq!(chart.apply_study_template_set(&set), Err(error.to_owned()));
261            assert_eq!(layout(&chart), before);
262        }
263    }
264
265    #[test]
266    fn stored_template_lists_apply_with_their_links_and_default_places() {
267        let mut source = test_frames::chart(&[1.0], &[2.0]);
268        let rsi = source.add_indicator(Box::new(RSI::new(14)));
269        let mean = source.add_indicator(Box::new(SMA::new(3)));
270        source
271            .set_indicator_source(mean, StudySource::Study { id: rsi.0, plot: 0 })
272            .unwrap();
273        let list = serde_json::to_value(source.study_template_set().studies).unwrap();
274
275        let mut chart = test_frames::chart(&[1.0], &[2.0]);
276        let set: StudyTemplateSet = serde_json::from_value(list).unwrap();
277        let ids = chart.apply_study_template_set(&set).unwrap();
278        assert_eq!(
279            chart.indicators().nth(1).unwrap().1.source(),
280            Some(StudySource::Study {
281                id: ids[0].0,
282                plot: 0
283            })
284        );
285        let panes: Vec<_> = chart
286            .panes()
287            .into_iter()
288            .map(|pane| chart.pane_series(pane))
289            .collect();
290        assert_eq!(panes, [vec![SeriesId::MAIN, ids[1]], vec![ids[0]]]);
291    }
292}