use std::path::PathBuf;
use clap::Parser;
use lace::cc::alg::ColAssignAlg;
use lace::cc::alg::RowAssignAlg;
use lace::cc::feature::FType;
use lace::cc::state::State;
use lace::cc::state::StateGewekeSettings;
use lace::cc::transition::StateTransition;
use lace::geweke::GewekeTester;
use lace::stats::prior_process::PriorProcessType;
use plotly::common::Mode;
use plotly::layout::Layout;
use plotly::Plot;
use plotly::Scatter;
use rand::SeedableRng;
use rand_xoshiro::Xoshiro256Plus;
#[derive(Parser, Debug)]
#[clap(rename_all = "kebab")]
struct Opt {
#[clap(long, default_value = "gibbs")]
pub row_alg: RowAssignAlg,
#[clap(long, default_value = "gibbs")]
pub col_alg: ColAssignAlg,
#[clap(long, default_value = "50")]
pub nrows: usize,
#[clap(long, default_value = "5")]
pub ncols: usize,
#[clap(long)]
pub no_row_reassign: bool,
#[clap(long)]
pub no_col_reassign: bool,
#[clap(long)]
pub no_view_alpha: bool,
#[clap(long)]
pub no_state_alpha: bool,
#[clap(long)]
pub no_priors: bool,
#[clap(long)]
pub plot_var: Option<String>,
#[clap(long, short, default_value = "10000")]
pub niters: usize,
#[clap(long)]
dst: Option<PathBuf>,
}
fn main() {
let opt = Opt::parse();
let mut rng = Xoshiro256Plus::from_os_rng();
let ftypes: Vec<FType> = (0..opt.ncols)
.map(|i| match i % 3 {
0 => FType::Continuous,
1 => FType::Categorical,
2 => FType::Count,
_ => unreachable!(),
})
.collect();
let mut settings = StateGewekeSettings::new(
opt.nrows,
ftypes,
PriorProcessType::Dirichlet,
PriorProcessType::Dirichlet,
);
let mut transitions: Vec<StateTransition> = Vec::new();
if !opt.no_col_reassign {
transitions.push(StateTransition::ColumnAssignment(opt.col_alg));
}
if !opt.no_state_alpha {
transitions.push(StateTransition::StatePriorProcessParams);
}
if !opt.no_row_reassign {
transitions.push(StateTransition::RowAssignment(opt.row_alg));
}
if !opt.no_view_alpha {
transitions.push(StateTransition::ViewPriorProcessParams);
}
if !opt.no_priors {
transitions.push(StateTransition::FeaturePriors);
}
transitions.push(StateTransition::ComponentParams);
settings.transitions = transitions;
let mut geweke: GewekeTester<State> = GewekeTester::new(settings);
geweke.run(opt.niters, Some(5), &mut rng);
let res = geweke.result();
res.report();
if let Some(ref key) = opt.plot_var {
use lace::stats::EmpiricalCdf;
use lace::utils::minmax;
let (min_f, max_f) = minmax(res.forward.get(key).unwrap());
let (min_p, max_p) = minmax(res.posterior.get(key).unwrap());
let x_min = min_f.min(min_p);
let x_max = max_f.max(max_p);
let step = (x_max - x_min) / 200.0;
let xs: Vec<f64> =
(0..200).map(|i| x_min + (i as f64) * step).collect();
let cdf_f = EmpiricalCdf::new(res.forward.get(key).unwrap());
let cdf_p = EmpiricalCdf::new(res.posterior.get(key).unwrap());
let x_f: Vec<f64> = cdf_f.f(&xs);
let x_p: Vec<f64> = cdf_p.f(&xs);
let trace = Scatter::new(x_f, x_p).name("P-P").mode(Mode::Lines);
let ideal = Scatter::new(vec![0_f64, 1_f64], vec![0_f64, 1_f64])
.name("Reference")
.mode(Mode::Lines);
let mut plot = Plot::new();
plot.add_trace(ideal);
plot.add_trace(trace);
plot.set_layout(Layout::new());
plot.show();
}
if let Some(dst) = opt.dst {
let f = std::fs::OpenOptions::new()
.create(true)
.write(true)
.truncate(true)
.open(dst)
.unwrap();
serde_json::to_writer(f, &res).unwrap();
}
}