use std::{env, hint::black_box, process, time::Instant};
use laddu::parameter;
use laddu::{
amplitudes::{
kmatrix::{
KopfKMatrixA0, KopfKMatrixA0Channel, KopfKMatrixA2, KopfKMatrixA2Channel,
KopfKMatrixF0, KopfKMatrixF0Channel, KopfKMatrixF2, KopfKMatrixF2Channel,
},
zlm::Zlm,
},
data::DatasetReadOptions,
extensions::NLL,
io,
traits::LikelihoodTerm,
utils::{
enums::{Frame, Sign},
variables::Mass,
},
};
use rayon::{prelude::*, ThreadPoolBuilder};
fn reaction_variables() -> (laddu::Angles, laddu::Polarization, Mass) {
let beam = laddu::Particle::measured("beam", "beam");
let target = laddu::Particle::missing("target");
let kshort1 = laddu::Particle::measured("K_S1", "kshort1");
let kshort2 = laddu::Particle::measured("K_S2", "kshort2");
let kk = laddu::Particle::composite("KK", [&kshort1, &kshort2]).unwrap();
let proton = laddu::Particle::measured("proton", "proton");
let reaction = laddu::Reaction::two_to_two(&beam, &target, &kk, &proton).unwrap();
let angles = reaction
.decay(&kk)
.unwrap()
.angles(&kshort1, Frame::Helicity)
.unwrap();
let polarization = reaction.polarization("pol_magnitude", "pol_angle");
let resonance_mass = reaction.mass(&kk);
(angles, polarization, resonance_mass)
}
fn main() {
let args = Args::parse();
let nll = build_kmatrix_nll(&args.dataset_path);
let pool = ThreadPoolBuilder::new()
.num_threads(args.thread_count.unwrap_or_else(num_cpus::get))
.build()
.expect("failed to build rayon pool");
let nll_ref: &NLL = &nll;
println!(
"Profiling {} iterations ({:?}) on {} threads...",
args.iterations,
args.mode,
pool.current_num_threads()
);
let start = Instant::now();
pool.install(|| {
let nll = nll_ref;
(0..args.iterations)
.into_par_iter()
.for_each(|iteration| run_iteration(nll, iteration, args.mode));
});
let elapsed = start.elapsed();
println!(
"Finished {} iterations in {:.2?} ({:.4?}/iter)",
args.iterations,
elapsed,
elapsed / args.iterations as u32
);
}
#[derive(Clone, Copy, Debug)]
enum Mode {
Value,
Gradient,
ValueAndGradient,
}
impl Mode {
fn parse(value: &str) -> Option<Self> {
match value {
"value" => Some(Self::Value),
"gradient" => Some(Self::Gradient),
"both" => Some(Self::ValueAndGradient),
_ => None,
}
}
}
#[derive(Debug)]
struct Args {
dataset_path: String,
iterations: usize,
thread_count: Option<usize>,
mode: Mode,
}
impl Args {
fn parse() -> Self {
const DEFAULT_DATASET: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/benches/bench.parquet");
let mut dataset_path = DEFAULT_DATASET.to_string();
let mut iterations = 512usize;
let mut thread_count = None;
let mut mode = Mode::ValueAndGradient;
let mut args = env::args().skip(1);
while let Some(arg) = args.next() {
match arg.as_str() {
"--dataset" => {
dataset_path = args
.next()
.unwrap_or_else(|| usage_and_exit("missing dataset path"))
}
_ if arg.starts_with("--dataset=") => {
dataset_path = arg["--dataset=".len()..].to_string();
}
"--iterations" => {
iterations = parse_positive(args.next(), "iterations");
}
_ if arg.starts_with("--iterations=") => {
iterations = parse_positive(
Some(arg["--iterations=".len()..].to_string()),
"iterations",
);
}
"--threads" => {
thread_count = Some(parse_positive(args.next(), "threads"));
}
_ if arg.starts_with("--threads=") => {
thread_count = Some(parse_positive(
Some(arg["--threads=".len()..].to_string()),
"threads",
));
}
"--mode" => {
mode = parse_mode(args.next());
}
_ if arg.starts_with("--mode=") => {
mode = parse_mode(Some(arg["--mode=".len()..].to_string()));
}
"--help" | "-h" => {
print_usage();
process::exit(0);
}
other => {
usage_and_exit(&format!("unknown argument: {other}"));
}
}
}
Args {
dataset_path,
iterations,
thread_count,
mode,
}
}
}
fn parse_positive(value: Option<String>, flag: &str) -> usize {
let parsed = value
.unwrap_or_else(|| usage_and_exit(&format!("missing value for --{flag}")))
.parse::<usize>()
.unwrap_or_else(|_| usage_and_exit(&format!("invalid integer for --{flag}")));
if parsed == 0 {
usage_and_exit(&format!("--{flag} must be greater than zero"));
}
parsed
}
fn parse_mode(value: Option<String>) -> Mode {
let value = value.unwrap_or_else(|| usage_and_exit("missing value for --mode"));
Mode::parse(&value)
.unwrap_or_else(|| usage_and_exit("mode must be one of: value, gradient, both"))
}
fn print_usage() {
eprintln!(
"Usage: cargo run -p laddu --example profile_kmatrix [--dataset PATH] [--iterations N] \
[--threads N] [--mode value|gradient|both]"
);
}
fn usage_and_exit(message: &str) -> ! {
eprintln!("{message}");
print_usage();
process::exit(1);
}
fn build_kmatrix_nll(dataset_path: &str) -> Box<NLL> {
let p4_names = ["beam", "proton", "kshort1", "kshort2"];
let aux_names = ["pol_magnitude", "pol_angle"];
let options = DatasetReadOptions::default()
.p4_names(p4_names)
.aux_names(aux_names);
let ds_data = io::read_parquet(dataset_path, &options).expect("failed to load data sample");
let ds_mc = io::read_parquet(dataset_path, &options).expect("failed to load MC sample");
let (angles, polarization, resonance_mass) = reaction_variables();
let z00p = Zlm::new("Z00+", 0, 0, Sign::Positive, &angles, &polarization).unwrap();
let z00n = Zlm::new("Z00-", 0, 0, Sign::Negative, &angles, &polarization).unwrap();
let z22p = Zlm::new("Z22+", 2, 2, Sign::Positive, &angles, &polarization).unwrap();
let f0p = KopfKMatrixF0::new(
"f0+",
[
[parameter!("f0+ c00 re", 0.0), parameter!("f0+ c00 im", 0.0)],
[
parameter!("f0(980)+ re"),
parameter!("f0(980)+ im_fix", 0.0),
],
[parameter!("f0(1370)+ re"), parameter!("f0(1370)+ im")],
[parameter!("f0(1500)+ re"), parameter!("f0(1500)+ im")],
[parameter!("f0(1710)+ re"), parameter!("f0(1710)+ im")],
],
KopfKMatrixF0Channel::PiPi,
&resonance_mass,
None,
)
.unwrap();
let a0p = KopfKMatrixA0::new(
"a0+",
[
[parameter!("a0(980)+ re"), parameter!("a0(980)+ im")],
[parameter!("a0(1450)+ re"), parameter!("a0(1450)+ im")],
],
KopfKMatrixA0Channel::PiEta,
&resonance_mass,
None,
)
.unwrap();
let f0n = KopfKMatrixF0::new(
"f0-",
[
[parameter!("f0- c00 re", 0.0), parameter!("f0- c00 im", 0.0)],
[
parameter!("f0(980)- re"),
parameter!("f0(980)- im_fix", 0.0),
],
[parameter!("f0(1370)- re"), parameter!("f0(1370)- im")],
[parameter!("f0(1500)- re"), parameter!("f0(1500)- im")],
[parameter!("f0(1710)- re"), parameter!("f0(1710)- im")],
],
KopfKMatrixF0Channel::PiPi,
&resonance_mass,
None,
)
.unwrap();
let a0n = KopfKMatrixA0::new(
"a0-",
[
[parameter!("a0(980)- re"), parameter!("a0(980)- im")],
[parameter!("a0(1450)- re"), parameter!("a0(1450)- im")],
],
KopfKMatrixA0Channel::PiEta,
&resonance_mass,
None,
)
.unwrap();
let f2 = KopfKMatrixF2::new(
"f2",
[
[parameter!("f2(1270) re"), parameter!("f2(1270) im")],
[parameter!("f2(1525) re"), parameter!("f2(1525) im")],
[parameter!("f2(1850) re"), parameter!("f2(1850) im")],
[parameter!("f2(1910) re"), parameter!("f2(1910) im")],
],
KopfKMatrixF2Channel::KKbar,
&resonance_mass,
None,
)
.unwrap();
let a2 = KopfKMatrixA2::new(
"a2",
[
[parameter!("a2(1320) re"), parameter!("a2(1320) im")],
[parameter!("a2(1700) re"), parameter!("a2(1700) im")],
],
KopfKMatrixA2Channel::PiEtaPrime,
&resonance_mass,
None,
)
.unwrap();
let s0p = f0p + a0p;
let s0n = f0n + a0n;
let d2p = f2 + a2;
let pos_re = (&s0p * z00p.real() + &d2p * z22p.real()).norm_sqr();
let pos_im = (&s0p * z00p.imag() + &d2p * z22p.imag()).norm_sqr();
let neg_re = (&s0n * z00n.real()).norm_sqr();
let neg_im = (&s0n * z00n.imag()).norm_sqr();
let expr = pos_re + pos_im + neg_re + neg_im;
NLL::new(&expr, &ds_data, &ds_mc, None).unwrap()
}
fn run_iteration(nll: &NLL, iteration: usize, mode: Mode) {
let params = generate_parameters(iteration, nll.n_free());
match mode {
Mode::Value => {
let _ = black_box(nll.evaluate(¶ms));
}
Mode::Gradient => {
let _ = black_box(nll.evaluate_gradient(¶ms));
}
Mode::ValueAndGradient => {
let value = nll.evaluate(¶ms);
let gradient = nll.evaluate_gradient(¶ms);
let _ = black_box((value, gradient));
}
}
}
fn generate_parameters(iteration: usize, len: usize) -> Vec<f64> {
(0..len)
.map(|idx| {
let seed = iteration as f64 + idx as f64 * 0.5;
(seed.sin() + seed.cos()) * 0.5
})
.collect()
}