dyna3/app-proto/src/components/diagnostics.rs

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

use crate::{AppState, specified::SpecifiedValue};

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

impl DiagnosticsState {
    fn new(initial_tab: String) -> Self {
        Self { 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(),
                    }
                ))
            }
        }
    }
}

// history step input
#[component]
fn StepInput() -> View {
    // get the assembly
    let state = use_context::<AppState>();
    let assembly = state.assembly;
    
    // the `last_step` signal holds the index of the last step
    let last_step = assembly.descent_history.map(
        |history| match history.config.len() {
            0 => None,
            n => Some(n - 1),
        }
    );
    let input_max = last_step.map(|last| last.unwrap_or(0));
    
    // these signals hold the entered step number
    let value = create_signal(String::new());
    let value_as_number = create_signal(0.0);
    
    create_effect(move || {
        value.set(assembly.step.with(|n| n.spec.clone()));
    });
    
    view! {
        div(id = "step-input") {
            label { "Step" }
            input(
                r#type = "number",
                min = "0",
                max = input_max.with(|max| max.to_string()),
                bind:value = value,
                bind:valueAsNumber = value_as_number,
                on:change = move |_| {
                    if last_step.with(|last| last.is_some()) {
                        // clamp the step within its allowed range. the lower
                        // bound is redundant on browsers that make it
                        // impossible to type negative values into a number
                        // input with a non-negative `min`, but there's no harm
                        // in being careful
                        let step_raw = value.with(
                            |val| SpecifiedValue::try_from(val.clone())
                                .unwrap_or(SpecifiedValue::from_empty_spec()
                            )
                        );
                        let step = SpecifiedValue::from(
                            step_raw.value.map(
                                |val| val.clamp(0.0, input_max.get() as f64)
                            )
                        );
                        
                        // set the input string and the assembly's active step
                        value.set(step.spec.clone());
                        assembly.step.set(step);
                    } else {
                        value.set(String::new());
                    }
                },
            )
        }
    }
}

fn into_log10_time_point((step, value): (usize, f64)) -> Vec<Option<f64>> {
    vec![
        Some(step as f64),
        if value == 0.0 { None } else { Some(value.abs().log10()) },
    ]
}

// 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: Vec<_> = state.assembly.descent_history.with(
                |history| history.scaled_loss
                    .iter()
                    .enumerate()
                    .map(|(step, &loss)| (step, loss))
                    .map(into_log10_time_point)
                    .collect()
            );
            
            // initialize the chart axes
            let step_axis = Axis::new()
                .type_(AxisType::Category)
                .boundary_gap(false);
            let scaled_loss_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 scaled_loss_series = Line::new().data(
                if scaled_loss.len() > 0 {
                    scaled_loss
                } else {
                    vec![vec![Some(0.0), None::<f64>]]
                }
            );
            let chart = Chart::new()
                .animation(false)
                .data_zoom(DataZoom::new().y_axis_index(0).right(40))
                .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, split into its
            // positive, negative, and strictly-zero parts
            let (
                hess_eigvals_zero,
                hess_eigvals_nonzero,
            ): (Vec<_>, Vec<_>) = state.assembly.descent_history.with(
                |history| history.hess_eigvals
                    .iter()
                    .enumerate()
                    .map(
                        |(step, eigvals)| eigvals.iter().map(
                            move |&val| (step, val)
                        )
                    )
                    .flatten()
                    .partition(|&(_, val)| val == 0.0)
            );
            let zero_level = hess_eigvals_nonzero
                .iter()
                .map(|(_, val)| val.abs())
                .reduce(f64::min)
                .map(|val| 0.1 * val)
                .unwrap_or(1.0);
            let (
                hess_eigvals_pos,
                hess_eigvals_neg,
            ): (Vec<_>, Vec<_>) = hess_eigvals_nonzero
                .into_iter()
                .partition(|&(_, val)| val > 0.0);
            
            // initialize the chart axes
            let step_axis = Axis::new()
                .type_(AxisType::Category)
                .boundary_gap(false);
            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 eigval_series_pos = Scatter::new()
                .symbol_size(4.5)
                .data(
                    if hess_eigvals_pos.len() > 0 {
                        hess_eigvals_pos
                            .into_iter()
                            .map(into_log10_time_point)
                            .collect()
                    } else {
                        vec![vec![Some(0.0), None::<f64>]]
                    }
                );
            let eigval_series_neg = Scatter::new()
                .symbol(Symbol::Diamond)
                .symbol_size(6.0)
                .data(
                    if hess_eigvals_neg.len() > 0 {
                        hess_eigvals_neg
                            .into_iter()
                            .map(into_log10_time_point)
                            .collect()
                    } else {
                        vec![vec![Some(0.0), None::<f64>]]
                    }
                );
            let eigval_series_zero = Scatter::new()
                .symbol(Symbol::Triangle)
                .symbol_size(5.0)
                .data(
                    if hess_eigvals_zero.len() > 0 {
                        hess_eigvals_zero
                            .into_iter()
                            .map(|(step, _)| (step, zero_level))
                            .map(into_log10_time_point)
                            .collect()
                    } else {
                        vec![vec![Some(0.0), None::<f64>]]
                    }
                );
            let chart = Chart::new()
                .animation(false)
                .data_zoom(DataZoom::new().y_axis_index(0).right(40))
                .x_axis(step_axis)
                .y_axis(eigval_axis)
                .grid(Grid::new().top(20).right(80).bottom(30).left(60))
                .series(eigval_series_pos)
                .series(eigval_series_neg)
                .series(eigval_series_zero);
            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" }
                }
                StepInput {}
            }
            DiagnosticsPanel(name = "loss") { LossHistory {} }
            DiagnosticsPanel(name = "spectrum") { SpectrumHistory {} }
        }
    }
}