Python functions that would be useful for pineako #183

Author: Radonirinaunimi

pineappl_py/src/fk_table.rs

@@ -2,10 +2,12 @@
 
 use super::convolutions::PyConv;
 use super::grid::PyGrid;
+use super::pids::PyPidBasis;
 use numpy::{IntoPyArray, PyArray1, PyArrayDyn};
 use pineappl::convolutions::ConvolutionCache;
 use pineappl::fk_table::{FkAssumptions, FkTable};
 use pineappl::grid::Grid;
+use pineappl::pids::PidBasis;
 use pyo3::prelude::*;
 use std::collections::BTreeMap;
 use std::fs::File;
@@ -206,6 +208,30 @@ impl PyFkTable {
         self.fk_table.x_grid().into_pyarray(py)
     }
 
+    /// Return the convention by which the channels' PIDS are encoded.
+    #[getter]
+    #[must_use]
+    pub const fn pid_basis(&self) -> PyPidBasis {
+        match self.fk_table.grid().pid_basis() {
+            PidBasis::Pdg => PyPidBasis::Pdg,
+            PidBasis::Evol => PyPidBasis::Evol,
+        }
+    }
+
+    /// Rotate the F Table into the specified basis.
+    ///
+    /// Parameters
+    /// ----------
+    /// pid_basis: PyPidBasis
+    ///     PID basis of the resulting FK Table
+    pub fn rotate_pid_basis(&mut self, pid_basis: PyPidBasis) -> PyGrid {
+        let mut grid_mut = self.fk_table.grid().clone();
+        grid_mut.rotate_pid_basis(pid_basis.into());
+        PyGrid {
+            grid: grid_mut.clone(),
+        }
+    }
+
     /// Write to file.
     ///
     /// # Panics

pineappl_py/src/grid.rs

@@ -696,7 +696,7 @@ impl PyGrid {
         self.grid.dedup_channels(ulps);
     }
 
-    /// Rotate the Grid into the specified basis
+    /// Rotate the Grid into the specified basis.
     ///
     /// Parameters
     /// ----------

pineappl_py/tests/test_fk_table.py

@@ -11,6 +11,7 @@
 from pineappl.convolutions import Conv, ConvType
 from pineappl.fk_table import FkAssumptions, FkTable
 from pineappl.subgrid import ImportSubgridV1
+from pineappl.pids import PidBasis
 
 
 class TestFkTable:
@@ -78,7 +79,7 @@ def test_fktable(
 
         assert fk.table().shape == (1, 51, 34, 34)
         np.testing.assert_allclose(fk.fac0(), 2.7224999999999997)
-        assert fk.frg0() == None
+        assert fk.frg0() is None
 
         # Check the various aspects of the Bins
         assert fk.bins() == 1
@@ -110,6 +111,11 @@ def test_fktable(
         assumption = FkAssumptions("Nf6Sym")
         fk.optimize(assumption)
 
+        # Check that FK table is in the Evolution basis and rotate into PDG
+        assert fk.pid_basis == PidBasis.Evol
+        new_fk = fk.rotate_pid_basis(PidBasis.Pdg)
+        assert new_fk.pid_basis == PidBasis.Pdg
+
     def test_unpolarized_convolution(
         self,
         pdf,