From ca5f3cf664bf27ee97e404fce5cdde53a20809b9 Mon Sep 17 00:00:00 2001
From: juacrumar <juacrumar@lairen.eu>
Date: Mon, 15 Apr 2024 10:31:48 +0200
Subject: [PATCH] expose bins and lumi

---
 pineappl_py/pineappl/fk_table.py | 38 ++++++++++++++++++++++++++++++++
 pineappl_py/src/fk_table.rs      | 10 +++++++--
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/pineappl_py/pineappl/fk_table.py b/pineappl_py/pineappl/fk_table.py
index 813f7a7b6..91a7080c4 100644
--- a/pineappl_py/pineappl/fk_table.py
+++ b/pineappl_py/pineappl/fk_table.py
@@ -1,3 +1,5 @@
+import numpy as np
+
 from .pineappl import PyFkTable, PyFkAssumptions
 from .utils import PyWrapper
 
@@ -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):
     """
diff --git a/pineappl_py/src/fk_table.rs b/pineappl_py/src/fk_table.rs
index f03d1c793..17781aba1 100644
--- a/pineappl_py/src/fk_table.rs
+++ b/pineappl_py/src/fk_table.rs
@@ -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;
@@ -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)
     }