gpqa 0.2.0

Utility to assess Egor GPs quality saved using GP recorder
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

use anyhow::Result;
use clap::Parser;
use egobox_moe::GpMetric;
use egobox_moe::GpMixture;
use egobox_moe::MixtureGpSurrogate;
use rayon::prelude::*;
use std::collections::HashMap;
use std::fs;

#[derive(Parser)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// Binary GP file generated by Egor optimizer
    filename: String,

    /// Use K folding cross validation procedure (default LOO)
    #[arg(short, long, default_value = "0")]
    kfold: usize,
}

#[derive(Debug, Clone)]
struct Metrics {
    pub q2: f64,
    pub pva: f64,
    pub iae_alpha: f64,
}

fn compute_metrics(gp_models: &[Box<dyn MixtureGpSurrogate>], kfold: usize) -> Vec<Metrics> {
    let mut res: Vec<_> = gp_models
        .par_iter()
        .enumerate()
        .map(|(i, gp)| {
            let scores: Vec<_> = [GpMetric::Q2, GpMetric::Pva, GpMetric::IAEAlphaWithPlot]
                .par_iter()
                .map(|m| {
                    let score = gp.as_ref().score(*m, kfold);
                    (m, score)
                })
                .collect();
            let scores: HashMap<_, _> = scores.into_iter().collect();

            if i == 0
                && let Some(data) = &scores.get(&GpMetric::IAEAlphaWithPlot).unwrap().plot_data
            {
                println!("\nIAEα plot data for first GP model:");
                println!("Alpha | Empirical coverage | Target coverage | Delta");
                println!("---------------------------------------------------");
                for i in 0..data.alphas.len() {
                    let alpha = data.alphas[i];
                    let delta = data.deltas[i];

                    println!(
                        "{:5.2}% |       {:5.2}%      |     {:5.2}%    | {:5.2}%",
                        alpha * 100.,
                        delta * 100.,
                        (1. - alpha) * 100.,
                        (delta - (1. - alpha)).abs() * 100.
                    );
                }
                println!();
            }

            (
                i,
                Metrics {
                    q2: scores.get(&GpMetric::Q2).unwrap().value,
                    pva: scores.get(&GpMetric::Pva).unwrap().value,
                    iae_alpha: scores.get(&GpMetric::IAEAlphaWithPlot).unwrap().value,
                },
            )
        })
        .collect();
    res.sort_by_key(|(i, _)| *i);
    res.into_iter().map(|(_, m)| m).collect::<Vec<_>>()
}

fn main() -> Result<()> {
    let args = Args::parse();

    let data: Vec<u8> = fs::read(&args.filename)?;

    // Try to load a vector of GP models (Egor optimizer save)
    let gp_models: Vec<Box<dyn MixtureGpSurrogate>> =
        bincode::serde::decode_from_slice(&data, bincode::config::standard())
            .map(|(res, _)| res)
            .unwrap_or_default();

    let gp_models = if gp_models.is_empty() {
        // Try to load a single GP model (Gpx save)
        let gp: Box<GpMixture> =
            bincode::serde::decode_from_slice(&data, bincode::config::standard())
                .map(|(res, _)| res)?;
        vec![gp as Box<dyn MixtureGpSurrogate>]
    } else {
        gp_models
    };

    println!(
        "Loaded {} GP model(s) from {}",
        gp_models.len(),
        args.filename
    );
    gp_models.iter().for_each(|gp| {
        println!("Loaded GP model: {}", gp);
    });

    let (xt, _yt) = gp_models.first().unwrap().training_data();
    println!("Training data: {} samples ({}-dim)", xt.nrows(), xt.ncols());

    let k = if args.kfold == 0 {
        xt.nrows()
    } else {
        args.kfold
    };
    let res: Vec<_> = compute_metrics(&gp_models, k);

    for (i, m) in res.iter().enumerate() {
        println!(
            "GP({}): Q2 = {:.2}, PVA = {:.2}, IAEα = {:.2}",
            i, m.q2, m.pva, m.iae_alpha
        );
    }

    Ok(())
}