Skip to main content

egui_charts/studies/
custom.rs

1//! Runtime-defined custom indicators.
2//!
3//! The [`CustomIndicator`] type lets users create indicators at runtime by
4//! providing a closure that takes `&[Bar]` and returns `Vec<IndicatorValue>`.
5//! This is useful for prototyping or embedding user-defined logic without
6//! creating a new struct that implements [`Indicator`].
7//!
8//! # Example
9//!
10//! ```rust
11//! use egui_charts::studies::{CustomIndicator, IndicatorValue};
12//!
13//! let momentum = CustomIndicator::new("Momentum(10)", |bars| {
14//!     bars.iter()
15//!         .enumerate()
16//!         .map(|(i, bar)| {
17//!             if i < 10 {
18//!                 IndicatorValue::None
19//!             } else {
20//!                 IndicatorValue::Single(bar.close - bars[i - 10].close)
21//!             }
22//!         })
23//!         .collect()
24//! })
25//! .with_overlay(false)
26//! .with_description("Custom 10-bar price momentum");
27//! ```
28//!
29//! The closure is shared behind an [`Arc`], so a clone computes exactly what
30//! the original does.
31
32use crate::model::Bar;
33use crate::studies::palette;
34use crate::studies::{Indicator, IndicatorValue, IndicatorValues};
35use egui::Color32;
36use std::sync::Arc;
37
38/// Type alias for the custom calculation closure.
39///
40/// The closure receives the full bar series and must return exactly one
41/// [`IndicatorValue`] per bar.
42pub type CalculateFn = Arc<dyn Fn(&[Bar]) -> Vec<IndicatorValue> + Send + Sync>;
43
44/// Calculation callback receiving the validated current inputs.
45pub type InputCalculateFn =
46    Arc<dyn Fn(&[Bar], &[super::StudyInput]) -> Vec<IndicatorValue> + Send + Sync>;
47
48/// Per-bar colour callback, receiving the output line, source bar and calculated columns.
49pub type PlotColorFn = Arc<dyn Fn(usize, usize, &IndicatorValues) -> Option<Color32> + Send + Sync>;
50
51/// An indicator defined at runtime via a calculation closure.
52///
53/// Build one with [`new`](Self::new) and customise it with the builder
54/// methods (`with_description`, `with_color`, `with_overlay`, etc.).
55#[derive(Clone)]
56pub struct CustomIndicator {
57    name: String,
58    desc: String,
59    calculate_fn: CalculateFn,
60    input_calculate_fn: Option<InputCalculateFn>,
61    inputs: Vec<super::StudyInput>,
62    values: IndicatorValues,
63    colors: Vec<Color32>,
64    is_overlay: bool,
65    line_count: usize,
66    line_names: Vec<String>,
67    visible: bool,
68    value_format: Option<crate::scales::StudyFormat>,
69    scale: Option<super::StudyScale>,
70    plot_types: Vec<super::PlotType>,
71    offsets: Vec<isize>,
72    fills: Vec<super::StudyFill>,
73    plot_color: Option<PlotColorFn>,
74}
75
76impl CustomIndicator {
77    /// Create a new custom indicator.
78    ///
79    /// # Arguments
80    /// * `name` -- Display name shown in legends and tooltips.
81    /// * `calculate_fn` -- Closure that computes indicator values from bars.
82    ///
83    /// The indicator defaults to a single-line yellow overlay.
84    pub fn new<F>(name: impl Into<String>, calculate_fn: F) -> Self
85    where
86        F: Fn(&[Bar]) -> Vec<IndicatorValue> + Send + Sync + 'static,
87    {
88        let name = name.into();
89        Self {
90            line_names: vec![name.clone()],
91            name,
92            desc: String::new(),
93            calculate_fn: Arc::new(calculate_fn),
94            input_calculate_fn: None,
95            inputs: Vec::new(),
96            values: IndicatorValues::new(),
97            colors: vec![palette::YELLOW], // Yellow default
98            is_overlay: true,
99            line_count: 1,
100            visible: true,
101            value_format: None,
102            scale: None,
103            plot_types: Vec::new(),
104            offsets: Vec::new(),
105            fills: Vec::new(),
106            plot_color: None,
107        }
108    }
109
110    /// Declare editable inputs and calculate from their current values. Ids must be unique;
111    /// defaults and initial values must satisfy the same schema used by the editor and setters.
112    pub fn with_inputs<F>(
113        mut self,
114        inputs: Vec<super::StudyInput>,
115        calculate: F,
116    ) -> Result<Self, super::StudyError>
117    where
118        F: Fn(&[Bar], &[super::StudyInput]) -> Vec<IndicatorValue> + Send + Sync + 'static,
119    {
120        for (at, input) in inputs.iter().enumerate() {
121            if input.id.is_empty() {
122                return Err(super::StudyError::EmptyInputId { index: at });
123            }
124            let reason = if inputs[..at].iter().any(|other| other.id == input.id) {
125                "is declared twice".into()
126            } else if !input.accepts(&input.value) {
127                format!("has value {:?}, expected {:?}", input.value, input.kind)
128            } else if !input.accepts(&input.default) {
129                format!("has default {:?}, expected {:?}", input.default, input.kind)
130            } else {
131                continue;
132            };
133            return Err(super::StudyError::InvalidInput {
134                id: input.id.clone(),
135                reason,
136            });
137        }
138        self.inputs = inputs;
139        self.input_calculate_fn = Some(Arc::new(calculate));
140        self.values.clear();
141        Ok(self)
142    }
143
144    /// Set a human-readable description (used in tooltips).
145    pub fn with_description(mut self, desc: impl Into<String>) -> Self {
146        self.desc = desc.into();
147        self
148    }
149
150    /// Set the line color (single-line indicators).
151    pub fn with_color(mut self, color: Color32) -> Self {
152        self.colors = vec![color];
153        self
154    }
155
156    /// Set colors for a multi-line indicator.
157    pub fn with_colors(mut self, colors: Vec<Color32>) -> Self {
158        self.colors = colors;
159        self
160    }
161
162    /// Set whether this indicator is drawn on the main price chart (`true`)
163    /// or in a separate sub-pane (`false`).
164    pub fn with_overlay(mut self, is_overlay: bool) -> Self {
165        self.is_overlay = is_overlay;
166        self
167    }
168
169    /// Set the number of output lines (must match the closure's output).
170    pub fn with_line_count(mut self, count: usize) -> Self {
171        self.line_count = count;
172        self
173    }
174
175    /// Set the units used by every readout of this study.
176    pub fn with_value_format(mut self, format: crate::scales::StudyFormat) -> Self {
177        self.value_format = Some(format);
178        self
179    }
180
181    /// Set the value scale independently of overlay or separate-pane placement.
182    pub fn with_scale(mut self, scale: super::StudyScale) -> Self {
183        self.scale = Some(scale);
184        self
185    }
186
187    /// Set geometry per output; omitted outputs remain lines.
188    pub fn with_plot_types(mut self, types: Vec<super::PlotType>) -> Self {
189        self.plot_types = types;
190        self
191    }
192
193    /// Set displacement per output; omitted outputs stay on their source bars.
194    pub fn with_offsets(mut self, offsets: Vec<isize>) -> Self {
195        self.offsets = offsets;
196        self
197    }
198
199    /// Set filled bands between output lines or constant levels.
200    pub fn with_fills(mut self, fills: Vec<super::StudyFill>) -> Self {
201        self.fills = fills;
202        self
203    }
204
205    /// Set a per-bar colour callback. Returning `None` uses the plot colour.
206    pub fn with_plot_color<F>(mut self, color: F) -> Self
207    where
208        F: Fn(usize, usize, &IndicatorValues) -> Option<Color32> + Send + Sync + 'static,
209    {
210        self.plot_color = Some(Arc::new(color));
211        self
212    }
213
214    /// Set display names for each line (used in legends).
215    pub fn with_line_names(mut self, names: Vec<String>) -> Self {
216        self.line_names = names;
217        self
218    }
219}
220
221impl Indicator for CustomIndicator {
222    fn inputs(&self) -> Vec<super::StudyInput> {
223        self.inputs.clone()
224    }
225
226    fn set_input(
227        &mut self,
228        id: &str,
229        value: super::InputValue,
230    ) -> Result<(), crate::studies::StudyError> {
231        let input = self
232            .inputs
233            .iter_mut()
234            .find(|input| input.id == id)
235            .ok_or_else(|| crate::studies::StudyError::UnknownInput(id.into()))?;
236        if !input.accepts(&value) {
237            return Err(crate::studies::StudyError::invalid_input(id));
238        }
239        input.value = value;
240        self.values.clear();
241        Ok(())
242    }
243
244    fn value_format(&self) -> crate::scales::StudyFormat {
245        self.value_format.unwrap_or(if self.is_overlay {
246            crate::scales::StudyFormat::Price
247        } else {
248            crate::scales::StudyFormat::Fixed(2)
249        })
250    }
251
252    fn scale(&self) -> super::StudyScale {
253        self.scale.unwrap_or(if self.is_overlay {
254            super::StudyScale::Main
255        } else {
256            super::StudyScale::Independent
257        })
258    }
259
260    fn plot_style(&self, line: usize) -> super::PlotStyle {
261        super::PlotStyle {
262            color: self.colors.get(line).copied().unwrap_or(Color32::WHITE),
263            width: 1.0,
264            pattern: crate::model::LineStyle::Solid,
265            plot_type: self.plot_types.get(line).copied().unwrap_or_default(),
266            visible: true,
267        }
268    }
269
270    fn plot_offset(&self, line: usize) -> isize {
271        self.offsets.get(line).copied().unwrap_or(0)
272    }
273
274    fn plot_color(&self, line: usize, bar: usize) -> Option<Color32> {
275        self.plot_color
276            .as_ref()
277            .and_then(|color| color(line, bar, &self.values))
278    }
279
280    fn fills(&self) -> Vec<super::StudyFill> {
281        self.fills.clone()
282    }
283
284    fn name(&self) -> &str {
285        &self.name
286    }
287
288    fn desc(&self) -> &str {
289        &self.desc
290    }
291
292    fn calculate(&mut self, data: &[Bar]) {
293        self.values = match &self.input_calculate_fn {
294            Some(calculate) => calculate(data, &self.inputs),
295            None => (self.calculate_fn)(data),
296        }
297        .into();
298    }
299
300    fn values(&self) -> &IndicatorValues {
301        &self.values
302    }
303
304    fn colors(&self) -> Vec<Color32> {
305        self.colors.clone()
306    }
307
308    fn set_colors(&mut self, colors: Vec<Color32>) {
309        if !colors.is_empty() {
310            self.colors = colors;
311        }
312    }
313
314    fn is_overlay(&self) -> bool {
315        self.is_overlay
316    }
317
318    fn line_count(&self) -> usize {
319        self.line_count
320    }
321
322    fn line_names(&self) -> Vec<String> {
323        self.line_names.clone()
324    }
325
326    fn is_visible(&self) -> bool {
327        self.visible
328    }
329
330    fn set_visible(&mut self, visible: bool) {
331        self.visible = visible;
332    }
333
334    fn clone_box(&self) -> Box<dyn Indicator> {
335        Box::new(self.clone())
336    }
337}
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342
343    #[test]
344    fn declared_inputs_validate_edits_and_reach_cloned_calculations() {
345        use crate::model::{PriceSource, Timeframe};
346        use crate::studies::{InputKind, InputValue, StudyInput, StudySource};
347        let schema = vec![
348            StudyInput::new(
349                "length",
350                "Length",
351                InputValue::Length(2),
352                InputKind::Length(1, 10),
353            ),
354            StudyInput::new(
355                "color",
356                "Colour",
357                InputValue::Color(Color32::RED),
358                InputKind::Color,
359            ),
360            StudyInput::new(
361                "source",
362                "Source",
363                InputValue::Source(StudySource::Price(PriceSource::HL2)),
364                InputKind::Source,
365            ),
366            StudyInput::new(
367                "interval",
368                "Interval",
369                InputValue::Timeframe(Some(Timeframe::Hour1)),
370                InputKind::Timeframe,
371            ),
372            StudyInput::new(
373                "symbol",
374                "Symbol",
375                InputValue::Symbol("TEST".into()),
376                InputKind::Symbol,
377            ),
378        ];
379        let mut study = CustomIndicator::new("Inputs", |_| Vec::new())
380            .with_inputs(schema, |bars, inputs| {
381                let InputValue::Length(length) = inputs[0].value else {
382                    unreachable!()
383                };
384                assert!(
385                    bars.iter()
386                        .all(|bar| bar.payload.as_deref() == Some(&serde_json::json!(42)))
387                );
388                bars.iter()
389                    .map(|bar| IndicatorValue::Single(bar.close * length as f64))
390                    .collect()
391            })
392            .unwrap();
393        let mut bar = Bar::new(chrono::DateTime::UNIX_EPOCH, 3.0, 3.0, 3.0, 3.0, 1.0);
394        bar.payload = Some(Arc::new(serde_json::json!(42)));
395        study.calculate(std::slice::from_ref(&bar));
396        for (id, bad) in [
397            ("length", InputValue::Length(0)),
398            ("length", InputValue::Number(2.0)),
399            ("symbol", InputValue::Symbol("\n".into())),
400        ] {
401            assert!(study.set_input(id, bad).is_err());
402            assert_eq!(study.values().value(0, 0), Some(6.0));
403        }
404        study.set_input("length", InputValue::Length(4)).unwrap();
405        assert!(study.values().is_empty());
406        assert_eq!(study.inputs()[0].default, InputValue::Length(2));
407        let saved: Vec<_> = study
408            .inputs()
409            .into_iter()
410            .map(|input| (input.id, input.value))
411            .collect();
412        let restored: Vec<(String, InputValue)> =
413            serde_json::from_str(&serde_json::to_string(&saved).unwrap()).unwrap();
414        let mut cloned = study.clone_box();
415        for (id, value) in restored {
416            cloned.set_input(&id, value).unwrap();
417        }
418        cloned.calculate(std::slice::from_ref(&bar));
419        assert_eq!(cloned.values().value(0, 0), Some(12.0));
420        cloned.update_source(&[bar], 0, &[5.0]);
421        assert_eq!(cloned.values().value(0, 0), Some(20.0));
422        let length = |value| StudyInput::new("length", "Length", value, InputKind::Length(1, 10));
423        for (schema, expected) in [
424            (
425                vec![StudyInput::new(
426                    "",
427                    "Length",
428                    InputValue::Length(2),
429                    InputKind::Length(1, 10),
430                )],
431                "study input 0 has an empty id",
432            ),
433            (
434                vec![length(InputValue::Length(2)), length(InputValue::Length(3))],
435                "study input \"length\" is declared twice",
436            ),
437            (
438                vec![length(InputValue::Length(0))],
439                "study input \"length\" has value Length(0), expected Length(1, 10)",
440            ),
441            (
442                vec![length(InputValue::Length(2)).with_default(InputValue::Length(11))],
443                "study input \"length\" has default Length(11), expected Length(1, 10)",
444            ),
445        ] {
446            let error = CustomIndicator::new("Inputs", |_| Vec::new())
447                .with_inputs(schema, |_, _| Vec::new())
448                .err()
449                .unwrap();
450            assert_eq!(error.to_string(), expected);
451        }
452    }
453
454    #[test]
455    fn clone_keeps_the_calculation() {
456        let indicator = CustomIndicator::new("Constant", |_| vec![IndicatorValue::Single(1.0); 3]);
457
458        let mut cloned = indicator.clone_box();
459        cloned.calculate(&[]);
460        assert_eq!(cloned.values().len(), 3);
461    }
462}