RDF define new column via Numba

[riccardomanzoni]

Hi,

I want to mingle Numba, RDF and numpy to add a column to my RDataFrame.

This is a minimal snippet, where I define the function to be jitted and test it via the command line.

import ROOT
import numpy as np

bins = np.linspace(1, 10, 10)
weights = bins

@ROOT.Numba.Declare(['float'], 'float')
def GetWeight(pt):
    idx = np.sum(pt>bins)
    return weights[idx-1] if idx>0 else 1.

ROOT.gInterpreter.ProcessLine('Numba::GetWeight(1.5)')

so far so good

(float) 1.00000f
Out[1]: 1

But when I try to use this for real and add it to my RDF,

tree = ROOT.RDataFrame(4).Define('x', '(float)rdfentry_')
tree.Define('weight', 'Numba::GetWeight(x)')

histo = tree.Histo1D(ROOT.RDF.TH1DModel('pt', '', len(bins)-1, bins), 'x', 'weight')

it fails

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-1-11f3f0f38647> in <module>
     15 tree.Define('weight', 'Numba::GetWeight(x)')
     16
---> 17 histo = tree.Histo1D(ROOT.RDF.TH1DModel('pt', '', len(bins)-1, bins), 'x', 'weight')
     18

/work/manzoni/miniconda3/envs/hammer3p8/lib/python3.8/site-packages/ROOT/pythonization/_rdataframe.py in _histo_profile(self, fixed_args, *args)
    120     # the original implementation
    121     else:
--> 122         res = original_method(*args)
    123
    124     return res

TypeError: Template method resolution failed:
  none of the 4 overloaded methods succeeded. Full details:
  ROOT::RDF::RResultPtr<TH1D> ROOT::RDF::RInterface<ROOT::Detail::RDF::RLoopManager,void>::Histo1D(basic_string_view<char,char_traits<char> > vName) =>
    TypeError: takes at most 1 arguments (3 given)
  ROOT::RDF::RResultPtr<TH1D> ROOT::RDF::RInterface<ROOT::Detail::RDF::RLoopManager,void>::Histo1D(const ROOT::RDF::TH1DModel& model, basic_string_view<char,char_traits<char> > vName, basic_string_view<char,char_traits<char> > wName) =>
    runtime_error: Unknown column: weight
  ROOT::RDF::RResultPtr<TH1D> ROOT::RDF::RInterface<ROOT::Detail::RDF::RLoopManager,void>::Histo1D(basic_string_view<char,char_traits<char> > vName, basic_string_view<char,char_traits<char> > wName) =>
    TypeError: takes at most 2 arguments (3 given)
  ROOT::RDF::RResultPtr<TH1D> ROOT::RDF::RInterface<ROOT::Detail::RDF::RLoopManager,void>::Histo1D(const ROOT::RDF::TH1DModel& model = {"", "", 128U, 0., 0.}, basic_string_view<char,char_traits<char> > vName = "") =>
    TypeError: takes at most 2 arguments (3 given)
  Failed to instantiate "Histo1D(ROOT::RDF::TH1DModel*,std::string,std::string)"

Any suggestions?

Thanks!
Riccardo

---

ROOT Version: 6.22/06
Platform: linuxx8664gcc
Compiler: x86_64-conda-linux-gnu-c++

---

[riccardomanzoni]

ooops, it’s totally my fault.

This

tree.Define('weight', 'Numba::GetWeight(x)')

should’ve been

tree = tree.Define('weight', 'Numba::GetWeight(x)')

sorry, I guess the thread can be closed.

[eguiraud]

Hi Riccardo,
the problem (at least this problem ) does not seem to be related to numba but to the construction of the TH1DModel object.

You should see that if you put ROOT.RDF.TH1DModel('pt', '', len(bins)-1, bins) on its on line it fails. In fact it doesn’t have a constructor with that signature: ROOT: ROOT::RDF::TH1DModel Class Reference .

Does fixing that fix everything?
Cheers,
Enrico

[riccardomanzoni]

Hi Enrico,
the problem was cause by the dumb mistake I reported above.
With that fix everything works (the construction of TH1DModel seems ok, or at least it does what it should without complaining )
Cheers,
Riccardo