Skip to content

Commit 6f98239

Browse files
authored
Merge pull request #308 from NNPDF/python-list-args
Fix Python list arguments
2 parents a30d3b5 + 1b5c6e5 commit 6f98239

File tree

5 files changed

+101
-77
lines changed

5 files changed

+101
-77
lines changed

pineappl_py/src/bin.rs

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
//! Binnning interface.
22
3-
use numpy::{PyArrayMethods, PyReadonlyArray1};
43
use pineappl::bin::BinRemapper;
54
use pyo3::prelude::*;
65

@@ -23,9 +22,9 @@ impl PyBinRemapper {
2322
/// limits : list(tuple(float, float))
2423
/// bin limits
2524
#[new]
26-
pub fn new(normalizations: PyReadonlyArray1<f64>, limits: Vec<(f64, f64)>) -> Self {
25+
pub fn new(normalizations: Vec<f64>, limits: Vec<(f64, f64)>) -> Self {
2726
Self {
28-
bin_remapper: BinRemapper::new(normalizations.to_vec().unwrap(), limits).unwrap(),
27+
bin_remapper: BinRemapper::new(normalizations, limits).unwrap(),
2928
}
3029
}
3130
}

pineappl_py/src/fk_table.rs

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
//! FK table interface.
22
33
use super::grid::PyGrid;
4-
use numpy::{IntoPyArray, PyArray1, PyArray4, PyArrayMethods, PyReadonlyArray1};
4+
use numpy::{IntoPyArray, PyArray1, PyArray4};
55
use pineappl::convolutions::LumiCache;
66
use pineappl::fk_table::{FkAssumptions, FkTable};
77
use pineappl::grid::Grid;
@@ -238,8 +238,8 @@ impl PyFkTable {
238238
&self,
239239
pdg_id: i32,
240240
xfx: &Bound<'py, PyAny>,
241-
bin_indices: Option<PyReadonlyArray1<usize>>,
242-
channel_mask: Option<PyReadonlyArray1<bool>>,
241+
bin_indices: Option<Vec<usize>>,
242+
channel_mask: Option<Vec<bool>>,
243243
py: Python<'py>,
244244
) -> Bound<'py, PyArray1<f64>> {
245245
let mut xfx = |id, x, q2| xfx.call1((id, x, q2)).unwrap().extract().unwrap();
@@ -248,8 +248,8 @@ impl PyFkTable {
248248
self.fk_table
249249
.convolve(
250250
&mut lumi_cache,
251-
&bin_indices.map_or(vec![], |b| b.to_vec().unwrap()),
252-
&channel_mask.map_or(vec![], |l| l.to_vec().unwrap()),
251+
&bin_indices.unwrap_or_default(),
252+
&channel_mask.unwrap_or_default(),
253253
)
254254
.into_pyarray_bound(py)
255255
}
@@ -278,8 +278,8 @@ impl PyFkTable {
278278
xfx1: &Bound<'py, PyAny>,
279279
pdg_id2: i32,
280280
xfx2: &Bound<'py, PyAny>,
281-
bin_indices: Option<PyReadonlyArray1<usize>>,
282-
channel_mask: Option<PyReadonlyArray1<bool>>,
281+
bin_indices: Option<Vec<usize>>,
282+
channel_mask: Option<Vec<bool>>,
283283
py: Python<'py>,
284284
) -> Bound<'py, PyArray1<f64>> {
285285
let mut xfx1 = |id, x, q2| xfx1.call1((id, x, q2)).unwrap().extract().unwrap();
@@ -290,8 +290,8 @@ impl PyFkTable {
290290
self.fk_table
291291
.convolve(
292292
&mut lumi_cache,
293-
&bin_indices.map_or(vec![], |b| b.to_vec().unwrap()),
294-
&channel_mask.map_or(vec![], |l| l.to_vec().unwrap()),
293+
&bin_indices.unwrap_or_default(),
294+
&channel_mask.unwrap_or_default(),
295295
)
296296
.into_pyarray_bound(py)
297297
}

pineappl_py/src/grid.rs

Lines changed: 40 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ use super::fk_table::PyFkTable;
77
use super::subgrid::{PySubgridEnum, PySubgridParams};
88
use itertools::izip;
99
use ndarray::CowArray;
10-
use numpy::{IntoPyArray, PyArray1, PyArrayMethods, PyReadonlyArray1, PyReadonlyArray4};
10+
use numpy::{IntoPyArray, PyArray1, PyReadonlyArray4};
1111
use pineappl::convolutions::LumiCache;
1212
use pineappl::evolution::AlphasTable;
1313
use pineappl::grid::{Grid, Ntuple};
@@ -45,14 +45,14 @@ impl PyGrid {
4545
pub fn new_grid(
4646
channels: Vec<PyRef<PyChannel>>,
4747
orders: Vec<PyRef<PyOrder>>,
48-
bin_limits: PyReadonlyArray1<f64>,
48+
bin_limits: Vec<f64>,
4949
subgrid_params: PySubgridParams,
5050
) -> Self {
5151
Self {
5252
grid: Grid::new(
5353
channels.iter().map(|pyc| pyc.entry.clone()).collect(),
5454
orders.iter().map(|pyo| pyo.order.clone()).collect(),
55-
bin_limits.to_vec().unwrap(),
55+
bin_limits,
5656
subgrid_params.subgrid_params,
5757
),
5858
}
@@ -116,20 +116,20 @@ impl PyGrid {
116116
/// cross section weight for all events
117117
pub fn fill_array(
118118
&mut self,
119-
x1s: PyReadonlyArray1<f64>,
120-
x2s: PyReadonlyArray1<f64>,
121-
q2s: PyReadonlyArray1<f64>,
119+
x1s: Vec<f64>,
120+
x2s: Vec<f64>,
121+
q2s: Vec<f64>,
122122
order: usize,
123-
observables: PyReadonlyArray1<f64>,
123+
observables: Vec<f64>,
124124
channel: usize,
125-
weights: PyReadonlyArray1<f64>,
125+
weights: Vec<f64>,
126126
) {
127127
for (&x1, &x2, &q2, &observable, &weight) in izip!(
128-
x1s.as_array().iter(),
129-
x2s.as_array().iter(),
130-
q2s.as_array().iter(),
131-
observables.as_array().iter(),
132-
weights.as_array().iter(),
128+
x1s.iter(),
129+
x2s.iter(),
130+
q2s.iter(),
131+
observables.iter(),
132+
weights.iter(),
133133
) {
134134
self.grid.fill(
135135
order,
@@ -163,7 +163,7 @@ impl PyGrid {
163163
q2: f64,
164164
order: usize,
165165
observable: f64,
166-
weights: PyReadonlyArray1<f64>,
166+
weights: Vec<f64>,
167167
) {
168168
self.grid.fill_all(
169169
order,
@@ -174,7 +174,7 @@ impl PyGrid {
174174
q2,
175175
weight: (),
176176
},
177-
&weights.to_vec().unwrap(),
177+
&weights,
178178
);
179179
}
180180

@@ -267,9 +267,9 @@ impl PyGrid {
267267
pdg_id: i32,
268268
xfx: &Bound<'py, PyAny>,
269269
alphas: &Bound<'py, PyAny>,
270-
order_mask: Option<PyReadonlyArray1<bool>>,
271-
bin_indices: Option<PyReadonlyArray1<usize>>,
272-
channel_mask: Option<PyReadonlyArray1<bool>>,
270+
order_mask: Option<Vec<bool>>,
271+
bin_indices: Option<Vec<usize>>,
272+
channel_mask: Option<Vec<bool>>,
273273
xi: Option<Vec<(f64, f64)>>,
274274
py: Python<'py>,
275275
) -> Bound<'py, PyArray1<f64>> {
@@ -280,10 +280,10 @@ impl PyGrid {
280280
self.grid
281281
.convolve(
282282
&mut lumi_cache,
283-
&order_mask.map_or(vec![], |b| b.to_vec().unwrap()),
284-
&bin_indices.map_or(vec![], |c| c.to_vec().unwrap()),
285-
&channel_mask.map_or(vec![], |d| d.to_vec().unwrap()),
286-
&xi.map_or(vec![(1.0, 1.0)], |m| m),
283+
&order_mask.unwrap_or_default(),
284+
&bin_indices.unwrap_or_default(),
285+
&channel_mask.unwrap_or_default(),
286+
&xi.unwrap_or(vec![(1.0, 1.0)]),
287287
)
288288
.into_pyarray_bound(py)
289289
}
@@ -332,9 +332,9 @@ impl PyGrid {
332332
pdg_id2: i32,
333333
xfx2: &Bound<'py, PyAny>,
334334
alphas: &Bound<'py, PyAny>,
335-
order_mask: Option<PyReadonlyArray1<bool>>,
336-
bin_indices: Option<PyReadonlyArray1<usize>>,
337-
channel_mask: Option<PyReadonlyArray1<bool>>,
335+
order_mask: Option<Vec<bool>>,
336+
bin_indices: Option<Vec<usize>>,
337+
channel_mask: Option<Vec<bool>>,
338338
xi: Option<Vec<(f64, f64)>>,
339339
py: Python<'py>,
340340
) -> Bound<'py, PyArray1<f64>> {
@@ -347,10 +347,10 @@ impl PyGrid {
347347
self.grid
348348
.convolve(
349349
&mut lumi_cache,
350-
&order_mask.map_or(vec![], |b| b.to_vec().unwrap()),
351-
&bin_indices.map_or(vec![], |c| c.to_vec().unwrap()),
352-
&channel_mask.map_or(vec![], |d| d.to_vec().unwrap()),
353-
&xi.map_or(vec![(1.0, 1.0)], |m| m),
350+
&order_mask.unwrap_or_default(),
351+
&bin_indices.unwrap_or_default(),
352+
&channel_mask.unwrap_or_default(),
353+
&xi.unwrap_or(vec![(1.0, 1.0)]),
354354
)
355355
.into_pyarray_bound(py)
356356
}
@@ -365,10 +365,10 @@ impl PyGrid {
365365
/// Returns
366366
/// -------
367367
/// PyEvolveInfo :
368-
/// evolution informations
369-
pub fn evolve_info(&self, order_mask: PyReadonlyArray1<bool>) -> PyEvolveInfo {
368+
/// evolution information
369+
pub fn evolve_info(&self, order_mask: Vec<bool>) -> PyEvolveInfo {
370370
PyEvolveInfo {
371-
evolve_info: self.grid.evolve_info(order_mask.as_slice().unwrap()),
371+
evolve_info: self.grid.evolve_info(order_mask.as_slice()),
372372
}
373373
}
374374

@@ -394,7 +394,7 @@ impl PyGrid {
394394
pub fn evolve_with_slice_iter<'py>(
395395
&self,
396396
slices: &Bound<'py, PyIterator>,
397-
order_mask: PyReadonlyArray1<bool>,
397+
order_mask: Vec<bool>,
398398
xi: (f64, f64),
399399
ren1: Vec<f64>,
400400
alphas: Vec<f64>,
@@ -413,8 +413,7 @@ impl PyGrid {
413413
CowArray::from(op.as_array().to_owned()),
414414
))
415415
}),
416-
// TODO: make `order_mask` a `Vec<f64>`
417-
&order_mask.to_vec().unwrap(),
416+
&order_mask,
418417
xi,
419418
&AlphasTable { ren1, alphas },
420419
)
@@ -448,7 +447,7 @@ impl PyGrid {
448447
&self,
449448
slices_a: &Bound<'py, PyIterator>,
450449
slices_b: &Bound<'py, PyIterator>,
451-
order_mask: PyReadonlyArray1<bool>,
450+
order_mask: Vec<bool>,
452451
xi: (f64, f64),
453452
ren1: Vec<f64>,
454453
alphas: Vec<f64>,
@@ -478,8 +477,7 @@ impl PyGrid {
478477
CowArray::from(op.as_array().to_owned()),
479478
))
480479
}),
481-
// TODO: make `order_mask` a `Vec<f64>`
482-
&order_mask.to_vec().unwrap(),
480+
&order_mask,
483481
xi,
484482
&AlphasTable { ren1, alphas },
485483
)
@@ -648,8 +646,8 @@ impl PyGrid {
648646
/// ----------
649647
/// factors : numpy.ndarray[float]
650648
/// bin-dependent factors by which to scale
651-
pub fn scale_by_bin(&mut self, factors: PyReadonlyArray1<f64>) {
652-
self.grid.scale_by_bin(&factors.to_vec().unwrap());
649+
pub fn scale_by_bin(&mut self, factors: Vec<f64>) {
650+
self.grid.scale_by_bin(&factors);
653651
}
654652

655653
/// Delete bins.
@@ -660,8 +658,8 @@ impl PyGrid {
660658
/// ----------
661659
/// bin_indices : numpy.ndarray[int]
662660
/// list of indices of bins to removed
663-
pub fn delete_bins(&mut self, bin_indices: PyReadonlyArray1<usize>) {
664-
self.grid.delete_bins(&bin_indices.to_vec().unwrap())
661+
pub fn delete_bins(&mut self, bin_indices: Vec<usize>) {
662+
self.grid.delete_bins(&bin_indices)
665663
}
666664
}
667665

pineappl_py/src/import_only_subgrid.rs

Lines changed: 13 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
//! PyImportOnlySubgrid* interface.
22
33
use super::subgrid::PySubgridEnum;
4-
use numpy::{PyArrayMethods, PyReadonlyArray1, PyReadonlyArray3};
4+
use numpy::PyReadonlyArray3;
55
use pineappl::import_only_subgrid::ImportOnlySubgridV1;
66
use pineappl::import_only_subgrid::ImportOnlySubgridV2;
77
use pineappl::sparse_array3::SparseArray3;
@@ -34,14 +34,10 @@ impl PyImportOnlySubgridV2 {
3434
pub fn new(
3535
array: PyReadonlyArray3<f64>,
3636
mu2_grid: Vec<(f64, f64)>,
37-
x1_grid: PyReadonlyArray1<f64>,
38-
x2_grid: PyReadonlyArray1<f64>,
37+
x1_grid: Vec<f64>,
38+
x2_grid: Vec<f64>,
3939
) -> Self {
40-
let mut sparse_array = SparseArray3::new(
41-
mu2_grid.len(),
42-
x1_grid.as_slice().unwrap().len(),
43-
x2_grid.as_slice().unwrap().len(),
44-
);
40+
let mut sparse_array = SparseArray3::new(mu2_grid.len(), x1_grid.len(), x2_grid.len());
4541

4642
for ((imu2, ix1, ix2), value) in array
4743
.as_array()
@@ -60,8 +56,8 @@ impl PyImportOnlySubgridV2 {
6056
fac: *fac,
6157
})
6258
.collect(),
63-
x1_grid.to_vec().unwrap(),
64-
x2_grid.to_vec().unwrap(),
59+
x1_grid,
60+
x2_grid,
6561
),
6662
}
6763
}
@@ -112,15 +108,11 @@ impl PyImportOnlySubgridV1 {
112108
#[new]
113109
pub fn new_import_only_subgrid(
114110
array: PyReadonlyArray3<f64>,
115-
q2_grid: PyReadonlyArray1<f64>,
116-
x1_grid: PyReadonlyArray1<f64>,
117-
x2_grid: PyReadonlyArray1<f64>,
111+
q2_grid: Vec<f64>,
112+
x1_grid: Vec<f64>,
113+
x2_grid: Vec<f64>,
118114
) -> Self {
119-
let mut sparse_array = SparseArray3::new(
120-
q2_grid.as_slice().unwrap().len(),
121-
x1_grid.as_slice().unwrap().len(),
122-
x2_grid.as_slice().unwrap().len(),
123-
);
115+
let mut sparse_array = SparseArray3::new(q2_grid.len(), x1_grid.len(), x2_grid.len());
124116

125117
for ((iq2, ix1, ix2), value) in array
126118
.as_array()
@@ -132,9 +124,9 @@ impl PyImportOnlySubgridV1 {
132124

133125
Self::new(ImportOnlySubgridV1::new(
134126
sparse_array,
135-
q2_grid.to_vec().unwrap(),
136-
x1_grid.to_vec().unwrap(),
137-
x2_grid.to_vec().unwrap(),
127+
q2_grid,
128+
x1_grid,
129+
x2_grid,
138130
))
139131
}
140132

pineappl_py/tests/test_grid.py

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -96,13 +96,48 @@ def test_convolve_with_one(self):
9696
g.convolve_with_one(2212, lambda pid, x, q2: 0.0, lambda q2: 0.0),
9797
[0.0] * 2,
9898
)
99+
v = 5e6 / 9999
99100
np.testing.assert_allclose(
100101
g.convolve_with_one(2212, lambda pid, x, q2: 1, lambda q2: 1.0),
101-
[5e6 / 9999, 0.0],
102+
[v, 0.0],
103+
)
104+
np.testing.assert_allclose(
105+
g.convolve_with_one(
106+
2212, lambda pid, x, q2: 1, lambda q2: 1.0, bin_indices=[0]
107+
),
108+
[v],
109+
)
110+
# block first bins with additional args
111+
np.testing.assert_allclose(
112+
g.convolve_with_one(
113+
2212,
114+
lambda pid, x, q2: 1,
115+
lambda q2: 1.0,
116+
bin_indices=[0],
117+
order_mask=[False],
118+
),
119+
[0.0],
120+
)
121+
np.testing.assert_allclose(
122+
g.convolve_with_one(
123+
2212,
124+
lambda pid, x, q2: 1,
125+
lambda q2: 1.0,
126+
bin_indices=[0],
127+
channel_mask=[False],
128+
),
129+
[0.0],
130+
)
131+
# second bin is empty
132+
np.testing.assert_allclose(
133+
g.convolve_with_one(
134+
2212, lambda pid, x, q2: 1, lambda q2: 1.0, bin_indices=[1]
135+
),
136+
[0.0],
102137
)
103138
np.testing.assert_allclose(
104139
g.convolve_with_one(2212, lambda pid, x, q2: 1, lambda q2: 2.0),
105-
[2**3 * 5e6 / 9999, 0.0],
140+
[2**3 * v, 0.0],
106141
)
107142

108143
def test_io(self, tmp_path):

0 commit comments

Comments
 (0)