dyna3/app-proto/src/diagnostics.rs

use charming::{
    Chart,
    WasmRenderer,
    component::{Axis, DataZoom, Grid},
    element::AxisType,
    series::{Line, Scatter},
};
use sycamore::prelude::*;

use crate::AppState;

#[derive(Clone)]
struct DiagnosticsState {
    active_tab: Signal<String>
}

impl DiagnosticsState {
    fn new(initial_tab: String) -> DiagnosticsState {
        DiagnosticsState {
            active_tab: create_signal(initial_tab)
        }
    }
}

// a realization status indicator
#[component]
fn RealizationStatus() -> View {
    let state = use_context::<AppState>();
    let realization_status = state.assembly.realization_status;
    view! {
        div(
            id="realization-status",
            class=realization_status.with(
                |status| match status {
                    Ok(_) => "",
                    Err(_) => "invalid"
                }
            )
        ) {
            div(class="status")
            div {
                (realization_status.with(
                    |status| match status {
                        Ok(_) => "Target accuracy achieved".to_string(),
                        Err(message) => message.clone()
                    }
                ))
            }
        }
    }
}

// the loss history from the last realization
#[component]
fn LossHistory() -> View {
    const CONTAINER_ID: &str = "loss-history";
    let state = use_context::<AppState>();
    let renderer = WasmRenderer::new_opt(None, Some(178));
    
    on_mount(move || {
        create_effect(move || {
            // get the loss history
            let scaled_loss = state.assembly.descent_history.with(
                |history| history.scaled_loss.clone()
            );
            let step_cnt = scaled_loss.len();
            
            // initialize the chart axes and series
            const MIN_INTERVAL: f64 = 0.01;
            let mut step_axis = Axis::new()
                .type_(AxisType::Category)
                .boundary_gap(false);
            let scaled_loss_axis = Axis::new()
                .type_(AxisType::Value)
                .min(0)
                .min_interval(MIN_INTERVAL);
            
            // load the chart data. when there's no history, we load the data
            // point (0, None) to clear the chart. it would feel more natural to
            // load empty data vectors, but that turns out not to clear the
            // chart: it instead leads to previous data being re-used
            let mut scaled_loss_series = Line::new();
            if step_cnt > 0 {
                step_axis = step_axis.data(
                    (0..step_cnt).map(|step| step.to_string()).collect()
                );
                scaled_loss_series = scaled_loss_series.data(scaled_loss);
            } else {
                step_axis = step_axis.data(vec![0.to_string()]);
                scaled_loss_series = scaled_loss_series.data(vec![None::<f64>]);
            }
            let chart = Chart::new()
                .animation(false)
                .data_zoom(DataZoom::new().y_axis_index(0).right(40).start(0).end(100))
                .x_axis(step_axis)
                .y_axis(scaled_loss_axis)
                .grid(Grid::new().top(20).right(80).bottom(30).left(60))
                .series(scaled_loss_series);
            renderer.render(CONTAINER_ID, &chart).unwrap();
        });
    });
    
    view! {
        div(id=CONTAINER_ID, class="diagnostics-chart")
    }
}

// the spectrum of the Hessian during the last realization
#[component]
fn SpectrumHistory() -> View {
    const CONTAINER_ID: &str = "spectrum-history";
    let state = use_context::<AppState>();
    let renderer = WasmRenderer::new(478, 178);
    
    on_mount(move || {
        create_effect(move || {
            // get the spectrum of the Hessian at each step
            let hess_eigvals = state.assembly.descent_history.with(
                |history| history.hess_eigvals
                    .iter()
                    .enumerate()
                    .map(
                        |(step, eigvals)| eigvals.iter().map(
                            move |val| vec![step as f64, *val]
                        )
                    )
                    .flatten()
                    .collect::<Vec<_>>()
            );
            
            // initialize the chart axes and series
            let step_axis = Axis::new();
            let eigval_axis = Axis::new();
            
            // load the chart data. when there's no history, we load the data
            // point (0, None) to clear the chart. it would feel more natural to
            // load empty data vectors, but that turns out not to clear the
            // chart: it instead leads to previous data being re-used
            let mut eigval_series = Scatter::new().symbol_size(7);
            if hess_eigvals.len() > 0 {
                eigval_series = eigval_series.data(hess_eigvals);
            } else {
                eigval_series = eigval_series.data(vec![None::<f64>, None::<f64>]);
            }
            let chart = Chart::new()
                .animation(false)
                .data_zoom(DataZoom::new().y_axis_index(0).right(40).start(0).end(100))
                .x_axis(step_axis)
                .y_axis(eigval_axis)
                .grid(Grid::new().top(20).right(80).bottom(30).left(60))
                .series(eigval_series);
            renderer.render(CONTAINER_ID, &chart).unwrap();
        });
    });
    
    view! {
        div(id=CONTAINER_ID, class="diagnostics-chart")
    }
}

#[component(inline_props)]
fn DiagnosticsPanel(name: &'static str, children: Children) -> View {
    let diagnostics_state = use_context::<DiagnosticsState>();
    view! {
        div(
            class="diagnostics-panel",
            "hidden"=diagnostics_state.active_tab.with(
                |active_tab| {
                    if active_tab == name {
                        None
                    } else {
                        Some("")
                    }
                }
            )
        ) {
            (children)
        }
    }
}

#[component]
pub fn Diagnostics() -> View {
    let diagnostics_state = DiagnosticsState::new("loss".to_string());
    let active_tab = diagnostics_state.active_tab.clone();
    provide_context(diagnostics_state);
    
    view! {
        div(id="diagnostics") {
            div(id="diagnostics-bar") {
                RealizationStatus {}
                select(bind:value=active_tab) {
                    option(value="loss") { "Loss" }
                    option(value="spectrum") { "Spectrum" }
                }
            }
            DiagnosticsPanel(name="loss") { LossHistory {} }
            DiagnosticsPanel(name="spectrum") { SpectrumHistory {} }
        }
    }
}