Skip to main content

chronos_ts/
viz.rs

1// src/viz.rs
2
3use crate::decomposition::ProphetPrediction;
4use crate::errors::{ChronosError, Result};
5use chrono::NaiveDate;
6use serde::{Deserialize, Serialize};
7use std::fs::File;
8use std::io::Write;
9use std::path::Path;
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct ChartDataPoint {
13    pub date: String,
14    pub yhat: f64,
15    pub yhat_lower: Option<f64>,
16    pub yhat_upper: Option<f64>,
17    pub trend: f64,
18    pub seasonal: f64,
19    pub holidays: f64,
20}
21
22#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct DecompositionExport {
24    pub series: Vec<ChartDataPoint>,
25}
26
27pub struct VisualizationExporter;
28
29impl VisualizationExporter {
30    pub fn to_export_data(
31        dates: &[NaiveDate],
32        pred: &ProphetPrediction,
33    ) -> Result<DecompositionExport> {
34        let n = dates.len();
35        if pred.yhat.len() != n {
36            return Err(ChronosError::InvalidParameters(
37                "Dates length does not match prediction length".into(),
38            ));
39        }
40
41        let mut series = Vec::with_capacity(n);
42
43        for i in 0..n {
44            let yhat_lower = pred.yhat_lower.as_ref().map(|v| v[i]);
45            let yhat_upper = pred.yhat_upper.as_ref().map(|v| v[i]);
46
47            series.push(ChartDataPoint {
48                date: dates[i].to_string(),
49                yhat: pred.yhat[i],
50                yhat_lower,
51                yhat_upper,
52                trend: pred.trend[i],
53                seasonal: pred.seasonal[i],
54                holidays: pred.holidays[i],
55            });
56        }
57
58        Ok(DecompositionExport { series })
59    }
60
61    pub fn to_json(dates: &[NaiveDate], pred: &ProphetPrediction) -> Result<String> {
62        let export_data = Self::to_export_data(dates, pred)?;
63        Ok(serde_json::to_string_pretty(&export_data)?)
64    }
65
66    pub fn to_html_string(
67        dates: &[NaiveDate],
68        pred: &ProphetPrediction,
69        title: &str,
70    ) -> Result<String> {
71        let json_data = Self::to_json(dates, pred)?;
72
73        let html = format!(
74            r#"<!DOCTYPE html>
75<html lang="en">
76<head>
77    <meta charset="UTF-8">
78    <meta name="viewport" content="width=device-width, initial-scale=1.0">
79    <title>{title}</title>
80    <script src="https://cdn.plot.ly/plotly-2.32.0.min.js"></script>
81    <style>
82        body {{ font-family: sans-serif; margin: 0; padding: 24px; background: #0f172a; color: #f8fafc; }}
83        .container {{ max-width: 1200px; margin: 0 auto; }}
84        .chart-card {{ background: #1e293b; border-radius: 8px; padding: 16px; margin-bottom: 24px; }}
85        #forecast-chart, #components-chart {{ width: 100%; height: 450px; }}
86    </style>
87</head>
88<body>
89    <div class="container">
90        <h1>{title}</h1>
91        <div class="chart-card"><div id="forecast-chart"></div></div>
92        <div class="chart-card"><div id="components-chart"></div></div>
93    </div>
94    <script>
95        const payload = {json_data};
96        const dates = payload.series.map(d => d.date);
97        const yhat = payload.series.map(d => d.yhat);
98        const trend = payload.series.map(d => d.trend);
99        const seasonal = payload.series.map(d => d.seasonal);
100        const holidays = payload.series.map(d => d.holidays);
101
102        const hasUpper = payload.series[0] && payload.series[0].yhat_upper !== null;
103        const forecastTraces = [];
104
105        if (hasUpper) {{
106            const yhatUpper = payload.series.map(d => d.yhat_upper);
107            const yhatLower = payload.series.map(d => d.yhat_lower);
108            forecastTraces.push({{
109                x: dates.concat(dates.slice().reverse()),
110                y: yhatUpper.concat(yhatLower.slice().reverse()),
111                fill: 'tozerox',
112                fillcolor: 'rgba(59, 130, 246, 0.2)',
113                line: {{ color: 'transparent' }},
114                name: 'Uncertainty Interval',
115                type: 'scatter'
116            }});
117        }}
118
119        forecastTraces.push({{ x: dates, y: yhat, mode: 'lines', name: 'Forecast (yhat)', line: {{ color: '#3b82f6', width: 2.5 }} }});
120        Plotly.newPlot('forecast-chart', forecastTraces, {{ title: 'Forecast', paper_bgcolor: '#1e293b', plot_bgcolor: '#1e293b', font: {{ color: '#f8fafc' }} }});
121
122        const componentTraces = [
123            {{ x: dates, y: trend, mode: 'lines', name: 'Trend', line: {{ color: '#10b981' }} }},
124            {{ x: dates, y: seasonal, mode: 'lines', name: 'Seasonality', line: {{ color: '#f59e0b' }} }},
125            {{ x: dates, y: holidays, mode: 'lines', name: 'Holidays', line: {{ color: '#ec4899' }} }}
126        ];
127        Plotly.newPlot('components-chart', componentTraces, {{ title: 'Components', paper_bgcolor: '#1e293b', plot_bgcolor: '#1e293b', font: {{ color: '#f8fafc' }} }});
128    </script>
129</body>
130</html>"#,
131            title = title,
132            json_data = json_data
133        );
134
135        Ok(html)
136    }
137
138    pub fn save_html<P: AsRef<Path>>(
139        dates: &[NaiveDate],
140        pred: &ProphetPrediction,
141        title: &str,
142        path: P,
143    ) -> Result<()> {
144        let html_content = Self::to_html_string(dates, pred, title)?;
145        let mut file = File::create(path)?;
146        file.write_all(html_content.as_bytes())?;
147        Ok(())
148    }
149}