malevich-cli 0.3.0

kaz: pipe data to an honest terminal plot — a stdin-first CLI over malevich
//! `--emit-code`: translate a normalized recipe into a malevich Rust program.
//!
//! The exploratory-to-production bridge: pipe data through kaz until the chart
//! is right, then take the program and keep going in Rust. Input semantics live
//! in [`crate::recipe`]; this module only writes the equivalent public grammar
//! with normalized values baked in as literals. Every shape is compile-tested.

use std::fmt::Write;

use malevich::scale::Colormap;

use crate::recipe::{Chart, DistributionKind, Furniture, GroupedKind, Recipe, ValueMark};
use crate::series::Dataset;

/// The complete standalone program for one prepared chart.
pub fn program(recipe: &Recipe) -> String {
    let mut body = String::new();
    let plot = match &recipe.chart {
        Chart::Value { mark, data } => value(&mut body, *mark, data),
        Chart::ScatterBy { x, y, groups } => scatter_by(&mut body, x, y, groups),
        Chart::Histogram {
            start,
            width,
            counts,
        } => histogram(&mut body, *start, *width, counts),
        Chart::Bars { labels, values } => bars(&mut body, labels, values),
        Chart::Distribution { kind, values } => distribution(&mut body, *kind, values),
        Chart::Grouped {
            kind,
            categories,
            groups,
        } => grouped(&mut body, *kind, categories, groups),
        Chart::Grid {
            columns,
            values,
            extents,
            colormap,
            labels_x,
            labels_y,
            reduce,
        } => grid(
            &mut body,
            *columns,
            values,
            *extents,
            colormap,
            (labels_x.as_deref(), labels_y.as_deref()),
            *reduce,
        ),
        Chart::Empty => "malevich::Plot::new()".to_string(),
    };
    let chart = furniture(plot, &recipe.furniture);

    let mut program = String::new();
    let _ = writeln!(
        program,
        "//! Generated by `kaz {} --emit-code`: the equivalent malevich program,\n\
         //! with the piped data inlined. A starting point — it renders for the\n\
         //! terminal it runs in.\n\n\
         use malevich::Frame;\n\n\
         fn main() {{",
        recipe.command.name()
    );
    program.push_str(&body);
    let _ = writeln!(program, "    let plot = {chart};");
    let sized = recipe.frame.width.is_some() || recipe.frame.height.is_some();
    let binding = if sized { "let mut frame" } else { "let frame" };
    let _ = writeln!(program, "    {binding} = Frame::detect();");
    if let Some(width) = recipe.frame.width {
        let _ = writeln!(program, "    frame.width = {width};");
    }
    if let Some(height) = recipe.frame.height {
        let _ = writeln!(program, "    frame.height = {height};");
    }
    program.push_str("    println!(\"{}\", plot.render(&frame));\n}\n");
    program
}

/// One layer per already-normalized value series.
fn value(body: &mut String, mark: ValueMark, data: &Dataset) -> String {
    let mark = match mark {
        ValueMark::Line => "Line",
        ValueMark::Scatter => "Points",
    };
    for (index, channel) in data.channels().iter().enumerate() {
        let _ = writeln!(
            body,
            "    let channel{index}: Vec<f64> = {};",
            floats(channel)
        );
    }
    let mut chart = String::from("malevich::Plot::new()");
    for series in &data.series {
        let y = series.y.index();
        let constructor = match series.x {
            Some(x) => {
                let x = x.index();
                format!("malevich::{mark}::xy(&channel{x}, &channel{y})")
            }
            None => format!("malevich::{mark}::y(&channel{y})"),
        };
        let labeled = match &series.label {
            Some(name) => format!("{constructor}.label({name:?})"),
            None => constructor,
        };
        let _ = write!(chart, "\n        .layer({labeled})");
    }
    chart
}

fn scatter_by(body: &mut String, x: &[f64], y: &[f64], groups: &[String]) -> String {
    let _ = writeln!(body, "    let x: Vec<f64> = {};", floats(x));
    let _ = writeln!(body, "    let y: Vec<f64> = {};", floats(y));
    let _ = writeln!(body, "    let groups: Vec<&str> = {};", strings(groups));
    "malevich::Plot::new()\n        .layer(malevich::Points::xy(x, y).color_by(groups))".to_string()
}

fn histogram(body: &mut String, start: f64, width: f64, counts: &[f64]) -> String {
    let _ = writeln!(body, "    let counts: Vec<f64> = {};", floats(counts));
    format!(
        "malevich::Plot::new()\n        .layer(malevich::Bars::spans({}, {}, counts))",
        float(start),
        float(width)
    )
}

fn bars(body: &mut String, labels: &[String], values: &[f64]) -> String {
    let _ = writeln!(body, "    let labels: Vec<&str> = {};", strings(labels));
    let _ = writeln!(body, "    let values: Vec<f64> = {};", floats(values));
    "malevich::bar(labels, values)".to_string()
}

fn distribution(body: &mut String, kind: DistributionKind, values: &[f64]) -> String {
    let preset = match kind {
        DistributionKind::Density => "density",
        DistributionKind::Ecdf => "ecdf",
    };
    let _ = writeln!(body, "    let values: Vec<f64> = {};", floats(values));
    format!("malevich::{preset}(values)")
}

fn grouped(
    body: &mut String,
    kind: GroupedKind,
    categories: &[String],
    groups: &[Vec<f64>],
) -> String {
    let preset = match kind {
        GroupedKind::Box => "box_plot",
        GroupedKind::Violin => "violin",
    };
    let _ = writeln!(body, "    let groups: Vec<Vec<f64>> = vec![");
    for group in groups {
        let _ = writeln!(body, "        {},", floats(group));
    }
    let _ = writeln!(body, "    ];");
    let _ = writeln!(
        body,
        "    let categories: Vec<&str> = {};",
        strings(categories)
    );
    format!("malevich::{preset}(categories, groups)")
}

fn grid(
    body: &mut String,
    columns: usize,
    values: &[f64],
    extents: Option<((f64, f64), (f64, f64))>,
    map: &Colormap,
    labels: (Option<&[String]>, Option<&[String]>),
    reduce: Option<malevich::stat::Reducer>,
) -> String {
    let _ = writeln!(body, "    let values: Vec<f64> = {};", floats(values));
    let mut cells = format!(
        "malevich::Cells::matrix({columns}, values)\n            .colormap({})",
        colormap(map)
    );
    if let Some(reducer) = reduce {
        let _ = write!(cells, "\n            .reduce({})", reducer_name(reducer));
    }
    if let Some((x, y)) = extents {
        let _ = write!(
            cells,
            "\n            .extents(({}, {}), ({}, {}))",
            float(x.0),
            float(x.1),
            float(y.0),
            float(y.1)
        );
    }
    let mut chart = format!("malevich::Plot::new()\n        .layer({cells})\n        .colorbar()");
    if let Some(names) = labels.0 {
        let _ = write!(
            chart,
            "\n        .x_scale(malevich::Scale::bands({}))",
            strings(names)
        );
    }
    if let Some(names) = labels.1 {
        let _ = write!(
            chart,
            "\n        .y_scale(malevich::Scale::bands({}))",
            strings(names)
        );
    }
    chart
}

/// The path expression for a CLI-reachable reducer.
fn reducer_name(reducer: malevich::stat::Reducer) -> &'static str {
    use malevich::stat::Reducer;
    match reducer {
        Reducer::Max => "malevich::stat::Reducer::Max",
        Reducer::Min => "malevich::stat::Reducer::Min",
        Reducer::Median => "malevich::stat::Reducer::Median",
        _ => "malevich::stat::Reducer::Mean",
    }
}

/// Shared plot furniture as chained builder calls.
fn furniture(mut chart: String, furniture: &Furniture) -> String {
    let mut push = |call: String| {
        chart.push_str("\n        ");
        chart.push_str(&call);
    };
    if let Some(title) = &furniture.title {
        push(format!(".title({title:?})"));
    }
    if let Some(xlabel) = &furniture.xlabel {
        push(format!(".x_label({xlabel:?})"));
    }
    if let Some(ylabel) = &furniture.ylabel {
        push(format!(".y_label({ylabel:?})"));
    }
    if let Some((lo, hi)) = furniture.xlim {
        push(format!(".x_domain({}, {})", float(lo), float(hi)));
    }
    if let Some((lo, hi)) = furniture.ylim {
        push(format!(".y_domain({}, {})", float(lo), float(hi)));
    }
    if furniture.time_x {
        push(".time_x()".to_string());
    }
    if furniture.log_x {
        push(".log_x()".to_string());
    }
    if furniture.log_y {
        push(".log_y()".to_string());
    }
    chart
}

/// The named-constant expression for a parsed colormap, re-centered as needed.
fn colormap(map: &Colormap) -> String {
    let named = [
        ("VIRIDIS", Colormap::VIRIDIS),
        ("MAGMA", Colormap::MAGMA),
        ("CIVIDIS", Colormap::CIVIDIS),
        ("GREYS", Colormap::GREYS),
        ("RED_BLUE", Colormap::RED_BLUE),
        ("PURPLE_ORANGE", Colormap::PURPLE_ORANGE),
    ];
    let base = named
        .iter()
        .find(|(_, candidate)| candidate.stops() == map.stops())
        .map_or(
            "malevich::scale::Colormap::DEFAULT".to_string(),
            |(name, _)| format!("malevich::scale::Colormap::{name}"),
        );
    match (map.midpoint(), map.is_log()) {
        (Some(midpoint), _) => format!("{base}.centered_at({})", float(midpoint)),
        (None, true) => format!("{base}.log()"),
        (None, false) => base,
    }
}

/// One `f64` as a valid Rust expression (`NaN` and infinities have no literal).
fn float(value: f64) -> String {
    if value.is_nan() {
        "f64::NAN".to_string()
    } else if value == f64::INFINITY {
        "f64::INFINITY".to_string()
    } else if value == f64::NEG_INFINITY {
        "f64::NEG_INFINITY".to_string()
    } else {
        format!("{value:?}")
    }
}

fn floats(values: &[f64]) -> String {
    let items: Vec<String> = values.iter().copied().map(float).collect();
    format!("vec![{}]", items.join(", "))
}

fn strings(values: &[String]) -> String {
    let items: Vec<String> = values.iter().map(|value| format!("{value:?}")).collect();
    format!("vec![{}]", items.join(", "))
}

#[cfg(test)]
#[path = "tests/emit_tests.rs"]
mod tests;