Skip to content

Commit

Permalink
expose bins and lumi
Browse files Browse the repository at this point in the history
  • Loading branch information
scarlehoff committed Apr 15, 2024
1 parent 0189f87 commit ca5f3cf
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 2 deletions.
38 changes: 38 additions & 0 deletions pineappl_py/pineappl/fk_table.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import numpy as np

from .pineappl import PyFkTable, PyFkAssumptions
from .utils import PyWrapper

Expand Down Expand Up @@ -50,6 +52,42 @@ def optimize(self, assumptions = "Nf6Ind"):
assumptions = FkAssumptions(assumptions)
return self._raw.optimize(assumptions._raw)

def convolute_with_one(
self,
pdg_id,
xfx,
bin_indices=np.array([], dtype=np.uint64),
lumi_mask=np.array([], dtype=bool),
):
r"""Convolute FkTable with a pdf.
Parameters
----------
pdg_id : int
PDG Monte Carlo ID of the hadronic particle `xfx` is the PDF for
xfx : callable
lhapdf like callable with arguments `pid, x, Q2` returning x*pdf for :math:`x`-grid
bin_indices : sequence(int)
A list with the indices of the corresponding bins that should be calculated. An
empty list means that all orders should be calculated.
lumi_mask : sequence(bool)
Mask for selecting specific luminosity channels. The value `True` means the
corresponding channel is included. An empty list corresponds to all channels being
enabled.
Returns
-------
list(float) :
cross sections for all bins, for each scale-variation tuple (first all bins, then
the scale variation)
"""
return self.raw.convolute_with_one(
pdg_id,
xfx,
np.array(bin_indices),
np.array(lumi_mask),
)


class FkAssumptions(PyWrapper):
"""
Expand Down
10 changes: 8 additions & 2 deletions pineappl_py/src/fk_table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use pineappl::fk_table::{FkAssumptions, FkTable};
use pineappl::grid::Grid;
use pineappl::lumi::LumiCache;

use numpy::{IntoPyArray, PyArray1, PyArray4};
use numpy::{IntoPyArray, PyArray1, PyArray4, PyReadonlyArray1};
use pyo3::prelude::*;

use std::collections::HashMap;
Expand Down Expand Up @@ -218,13 +218,19 @@ impl PyFkTable {
&self,
pdg_id: i32,
xfx: &PyAny,
bin_indices: PyReadonlyArray1<usize>,
lumi_mask: PyReadonlyArray1<bool>,
py: Python<'py>,
) -> &'py PyArray1<f64> {
let mut xfx = |id, x, q2| f64::extract(xfx.call1((id, x, q2)).unwrap()).unwrap();
let mut alphas = |_| 1.0;
let mut lumi_cache = LumiCache::with_one(pdg_id, &mut xfx, &mut alphas);
self.fk_table
.convolute(&mut lumi_cache, &[], &[])
.convolute(
&mut lumi_cache,
&bin_indices.to_vec().unwrap(),
&lumi_mask.to_vec().unwrap(),
)
.into_pyarray(py)
}

Expand Down

0 comments on commit ca5f3cf

Please sign in to comment.