Skip to content

Commit 08cc5e0

Browse files
cschwancschwan
committed
Refactor pineappl export code
1 parent ddf48a0 commit 08cc5e0

1 file changed

Lines changed: 73 additions & 131 deletions

File tree

‎pineappl_cli/src/export/applgrid.rs‎

Lines changed: 73 additions & 131 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use anyhow::{Result, bail};
22
use cxx::{UniquePtr, let_cxx_string};
33
use float_cmp::approx_eq;
4+
use itertools::izip;
45
use lhapdf::Pdf;
56
use ndarray::{Axis, s};
67
use pineappl::boc::{Channel, Kinematics, Order};
@@ -261,47 +262,31 @@ pub fn convert_into_applgrid(
261262
grid.channels().len().try_into().unwrap(),
262263
grid.convolutions().len() == 1,
263264
);
264-
let appl_q2: Vec<_> = (0..igrid.Ntau()).map(|i| igrid.getQ2(i)).collect();
265-
let appl_x1: Vec<_> = (0..igrid.Ny1()).map(|i| igrid.getx1(i)).collect();
266-
let appl_x2: Vec<_> = (0..igrid.Ny2()).map(|i| igrid.getx2(i)).collect();
265+
let appl_grids: Vec<Vec<_>> = vec![
266+
(0..igrid.Ntau()).map(|i| igrid.getQ2(i)).collect(),
267+
(0..igrid.Ny1()).map(|i| igrid.getx1(i)).collect(),
268+
(0..igrid.Ny2()).map(|i| igrid.getx2(i)).collect(),
269+
];
267270

268271
for (channel, subgrid) in subgrids
269272
.iter()
270273
.enumerate()
271274
.filter(|(_, subgrid)| !subgrid.is_empty())
272275
{
273-
let appl_q2_idx: Vec<_> = grid.scales().fac.calc(&subgrid.node_values(), grid.kinematics())
274-
.iter()
275-
.map(|&fac| {
276-
appl_q2
277-
.iter()
278-
.position(|&x| subgrid::node_value_eq(x, fac))
279-
.map_or_else(
280-
|| {
281-
if discard_non_matching_values {
282-
Ok(-1)
283-
} else {
284-
bail!(
285-
"factorization scale muf2 = {fac} not found in APPLgrid; try exporting with `--discard-non-matching-values`",
286-
)
287-
}
288-
},
289-
|idx| Ok(idx.try_into().unwrap()),
290-
)
291-
})
292-
.collect::<Result<_>>()?;
293-
294-
let (x1_grid, x2_grid) = if convolutions == 2 {
295-
(
296-
grid.kinematics()
297-
.iter()
298-
.zip(subgrid.node_values())
299-
.find_map(|(kin, node_values)| {
300-
matches!(kin, &Kinematics::X(idx) if idx == 0)
301-
.then_some(node_values)
302-
})
303-
// TODO: convert this into an error
304-
.unwrap(),
276+
let grids = vec![
277+
grid.scales()
278+
.fac
279+
.calc(&subgrid.node_values(), grid.kinematics())
280+
.into_owned(),
281+
grid.kinematics()
282+
.iter()
283+
.zip(subgrid.node_values())
284+
.find_map(|(kin, node_values)| {
285+
matches!(kin, &Kinematics::X(idx) if idx == 0).then_some(node_values)
286+
})
287+
// TODO: convert this into an error
288+
.unwrap(),
289+
if convolutions == 2 {
305290
grid.kinematics()
306291
.iter()
307292
.zip(subgrid.node_values())
@@ -310,118 +295,75 @@ pub fn convert_into_applgrid(
310295
.then_some(node_values)
311296
})
312297
// TODO: convert this into an error
313-
.unwrap(),
314-
)
315-
} else {
316-
(
317-
grid.kinematics()
318-
.iter()
319-
.zip(subgrid.node_values())
320-
.find_map(|(kin, node_values)| {
321-
matches!(kin, &Kinematics::X(idx) if idx == 0)
322-
.then_some(node_values)
323-
})
324-
// TODO: convert this into an error
325-
.unwrap(),
326-
Vec::new(),
327-
)
328-
};
329-
330-
let appl_x1_idx: Vec<_> = x1_grid
331-
.iter()
332-
.map(|&x1| {
333-
appl_x1
334-
.iter()
335-
.position(|&x| subgrid::node_value_eq(x, x1))
336-
.map_or_else(
337-
|| {
338-
if discard_non_matching_values {
339-
Ok(-1)
340-
} else {
341-
bail!("momentum fraction x1 = {x1} not found in APPLgrid; try exporting with `--discard-non-matching-values`")
342-
}
343-
},
344-
|idx| Ok(idx.try_into().unwrap()),
345-
)
346-
})
347-
.collect::<Result<_>>()?;
348-
let appl_x2_idx: Vec<_> = x2_grid
349-
.iter()
350-
.map(|&x2| {
351-
appl_x2
352-
.iter()
353-
.position(|&x| subgrid::node_value_eq(x, x2))
354-
.map_or_else(
355-
|| {
356-
if discard_non_matching_values {
357-
Ok(-1)
358-
} else {
359-
bail!("momentum fraction x2 = {x2} not found in APPLgrid; try exporting with `--discard-non-matching-values`")
360-
}
361-
},
362-
|idx| Ok(idx.try_into().unwrap()),
363-
)
364-
})
365-
.collect::<Result<_>>()?;
298+
.unwrap()
299+
} else {
300+
Vec::new()
301+
},
302+
];
303+
304+
let appl_idx: Vec<Vec<_>> = izip!(&grids, &appl_grids, ["factorization scale muf2", "momentum fraction x1", "momentum fraction x2"])
305+
.map(|(grid, appl_grid, label)| {
306+
grid
307+
.iter()
308+
.map(|&value| {
309+
appl_grid
310+
.iter()
311+
.position(|&appl_value| subgrid::node_value_eq(appl_value, value))
312+
.map_or_else(
313+
|| {
314+
if discard_non_matching_values {
315+
Ok(-1)
316+
} else {
317+
bail!("{label} = {value} not found in APPLgrid; try exporting with `--discard-non-matching-values`")
318+
}
319+
},
320+
|idx| Ok(idx.try_into().unwrap()),
321+
)
322+
})
323+
.collect::<Result<_>>()
324+
}
325+
).collect::<Result<_>>()?;
366326

367327
let mut weightgrid = ffi::igrid_weightgrid(igrid.pin_mut(), channel);
368328

369-
for (indices, value) in subgrid.indexed_iter() {
329+
'looop: for (indices, value) in subgrid.indexed_iter() {
370330
// TODO: here we assume that all X are consecutive starting from the second
371331
// element and are in ascending order
372-
let iq2 = indices[0];
373-
let appl_q2_idx = appl_q2_idx[iq2];
374-
375-
if appl_q2_idx == -1 {
376-
if value != 0.0 {
377-
println!(
378-
"WARNING: discarding non-matching scale muf2 = {} in subgrid {:?}",
379-
grid.scales()
380-
.fac
381-
.calc(&subgrid.node_values(), grid.kinematics())[iq2],
382-
(order, bin, channel)
383-
);
384-
}
385-
386-
continue;
387-
}
388-
389-
let appl_x1_i = appl_x1_idx[indices[1]];
390-
if appl_x1_i == -1 {
391-
if value != 0.0 {
392-
println!(
393-
"WARNING: discarding non-matching momentum fraction x1 = {} in subgrid {:?}",
394-
x1_grid[indices[1]],
395-
(order, bin, channel)
396-
);
397-
}
398332

399-
continue;
400-
}
401-
402-
let appl_x2_i = if convolutions == 2 {
403-
let i = appl_x2_idx[indices[2]];
404-
if i == -1 {
333+
let appl_indices = [
334+
appl_idx[0][indices[0]],
335+
appl_idx[1][indices[1]],
336+
if convolutions == 2 {
337+
appl_idx[2][indices[2]]
338+
} else {
339+
0
340+
},
341+
];
342+
343+
for (&appl_index, grid, &index, label) in izip!(
344+
&appl_indices,
345+
&grids,
346+
&indices,
347+
["scale muf2", "momentum fraction x1", "momentum fraction x2"]
348+
) {
349+
if appl_index == -1 {
405350
if value != 0.0 {
406351
println!(
407-
"WARNING: discarding non-matching momentum fraction x2 = {} in subgrid {:?}",
408-
x2_grid[indices[2]],
352+
"WARNING: discarding non-matching {label} = {} in subgrid {:?}",
353+
grid[index],
409354
(order, bin, channel)
410355
);
411356
}
412357

413-
continue;
358+
continue 'looop;
414359
}
415-
i
416-
} else {
417-
0
418-
};
360+
}
419361

420362
ffi::sparse_matrix_set(
421363
weightgrid.as_mut(),
422-
appl_q2_idx,
423-
appl_x1_i,
424-
appl_x2_i,
364+
appl_indices[0],
365+
appl_indices[1],
366+
appl_indices[2],
425367
factor * value,
426368
);
427369
}

0 commit comments

Comments
 (0)