reweighting_function Namespace Reference

Classes
class	FluxReweighter

Functions
	tqdm (x, **kwargs)
	load_analysis_config (path)
	resolve_normalization (norm, number_files, target_luminosity=68.6, luminosity_per_file=100.0)
	build_fastsim_curves (fastsim_configs, reference, degree)
	dpm_file_list (data_dir, n_files, pattern, indexed=True)
	build_chain (files, tree_name="cbmsim")
	_incoming_neutrino_energy (tree, pids)
	get_energies_pyroot (files, flavour="mu", tree_name="cbmsim", n_events=None)
	_find_mctrack_branches (tree)
	get_energies_uproot (files, flavour="mu", tree_name="cbmsim", n_events=None)
	read_dpm_energies (ri, reader="auto")
	make_comparison_plot (dpm, curves, normalized_factor, outfile, bins=100, hist_range=(0, 2000), title=None)
	fastsim_entries_from_args (paths)
	main ()

Variables
	ROOT = None
	uproot = None
	ak = None
dict	NU_PID = {"e": 12, "mu": 14, "tau": 16}
list	DEFAULT_COLORS = ["tab:blue", "tab:red", "tab:green", "tab:orange", "tab:purple"]
list	DEFAULT_STYLES = ["--", "-.", ":"]

Detailed Description

Neutrino-flux reweighting and energy-distribution comparison for SND@LHC.

Overview
--------
Hadron-production models disagree about the neutrino flux reaching SND@LHC.
This module derives an energy-dependent weight w(E) that maps one model onto
another, and applies it to a simulated sample.

    w(E) = flux_fastsim(E) / flux_DPMJET(E)

where each flux is the sum of a light-hadron (pi/K) and a heavy-flavour
(charm) contribution:

    flux_DPMJET  = light_DPMJET  + charm_DPMJET      (the reference / denominator)
    flux_fastsim = light_EPOS    + charm_POWHEG      (the target  / numerator)

Multiplying every DPMJET event by w(E) therefore reproduces the fast-sim
prediction, without re-running the simulation:

    DPMJET x w(E)  ==  fast-sim (light + charm)

Two capabilities
----------------
1. ``FluxReweighter`` -- builds the ratio from the flux text files, fits a
   polynomial w(E), and writes a diagnostic plot per polynomial degree.
2. ``main()`` in "compare" mode -- reads neutrino energies straight from
   SND@LHC ROOT files (``cbmsim`` tree), applies w(E), and overlays the
   DPMJET reference with each reweighted fast-sim prediction.

Inputs
------
Flux text files: whitespace-separated, one neutrino per row, no header, with
the columns listed in ``FluxReweighter.COLS``. POWHEG files carry one extra
trailing column (``iEvent``).

ROOT files: standard SND@LHC digitised output holding the ``cbmsim`` tree with
the ``MCTrack`` branch, from which the incoming (primary) neutrino energy is
taken.

Configuration
-------------
Every input is read from YAML; nothing analysis-specific is hard-coded.

    Analysis.yaml      run settings: ROOT input, which fast-sim variants to
                       overlay, normalisation, plot styling
    Config_epos.yaml   one per fast-sim variant: flux file paths, flavour,
                       binning, fit range, cuts, weights, output label

Usage
-----
    # diagnostic ratio-fit plots for one or more reweighter configs
    python3 reweighting_function.py -m reweight -rc Config_epos.yaml

    # full comparison: fit w(E), read ROOT energies, overlay and save
    python3 reweighting_function.py -c Analysis.yaml -rc Config_epos.yaml

    # everything taken from the run YAML, no extra arguments
    python3 reweighting_function.py -c Analysis.yaml

As a library:

    from reweighting_function import FluxReweighter
    poly, chi2red = FluxReweighter("Config_epos.yaml").fit(degree=2)
    weights = poly(energies)      # w(E) evaluated per event

Function Documentation

_find_mctrack_branches()

reweighting_function._find_mctrack_branches (   tree )

protected

Locate the MCTrack member branches whatever the exact prefix is.

Different productions store them as ``MCTrack.fPdgCode``,
``MCTrack/fPdgCode`` and so on, so the names are matched by suffix.
Returns a dict mapping the member name to the branch key found.

Definition at line 732 of file reweighting_function.py.

def _find_mctrack_branches(tree):
    """Locate the MCTrack member branches whatever the exact prefix is.
 
    Different productions store them as ``MCTrack.fPdgCode``,
    ``MCTrack/fPdgCode`` and so on, so the names are matched by suffix.
    Returns a dict mapping the member name to the branch key found.
    """
    want = ["fPdgCode", "fMotherId", "fPx", "fPy", "fPz"]
    keys = tree.keys()
    found = {}
    for w in want:
        cands = ([k for k in keys if "MCTrack" in k and k.endswith("." + w)]
                 or [k for k in keys if "MCTrack" in k and k.endswith(w)]
                 or [k for k in keys if k.endswith(w)])
        if cands:
            found[w] = cands[0]
    return found
 

_incoming_neutrino_energy()

reweighting_function._incoming_neutrino_energy (   tree,   pids  )

protected

Energy of the primary neutrino in the current entry, or None.

The incoming neutrino is the MCTrack whose |pdg| matches the requested
flavour and which has no mother (MotherId == -1).

Definition at line 699 of file reweighting_function.py.

def _incoming_neutrino_energy(tree, pids):
    """Energy of the primary neutrino in the current entry, or None.
 
    The incoming neutrino is the MCTrack whose |pdg| matches the requested
    flavour and which has no mother (MotherId == -1).
    """
    for trk in tree.MCTrack:
        if abs(trk.GetPdgCode()) in pids and trk.GetMotherId() == -1:
            return trk.GetEnergy()
    return None
 
 

build_chain()

reweighting_function.build_chain	(		files,
			tree_name = "cbmsim"
	)

Chain the given ROOT files into one TChain (pyROOT only).

Definition at line 689 of file reweighting_function.py.

def build_chain(files, tree_name="cbmsim"):
    """Chain the given ROOT files into one TChain (pyROOT only)."""
    if ROOT is None:
        raise RuntimeError("ROOT is not available in this environment.")
    chain = ROOT.TChain(tree_name)
    added = sum(bool(chain.Add(f)) for f in files)
    print(f"Chained {added}/{len(files)} files (tree '{tree_name}')")
    return chain
 
 

build_fastsim_curves()

reweighting_function.build_fastsim_curves	(		fastsim_configs,
			reference,
			degree
	)

Fit w(E) for each fast-sim variant and assemble the curves to draw.

One ``FluxReweighter`` is run per config, so the weights come from the fit
and are never hard-coded. The fitted coefficients and chi2/ndf are printed
for the record, and each curve is labelled from its own config.

Returns a list of ``(label, weight_poly | None, color, linestyle)`` with the
unweighted DPMJET reference (``weight_poly`` None) first.

Definition at line 647 of file reweighting_function.py.

def build_fastsim_curves(fastsim_configs, reference, degree):
    """Fit w(E) for each fast-sim variant and assemble the curves to draw.
 
    One ``FluxReweighter`` is run per config, so the weights come from the fit
    and are never hard-coded. The fitted coefficients and chi2/ndf are printed
    for the record, and each curve is labelled from its own config.
 
    Returns a list of ``(label, weight_poly | None, color, linestyle)`` with the
    unweighted DPMJET reference (``weight_poly`` None) first.
    """
    curves = [(reference.get("label", "DPMJET (reference)"),
               None,
               reference.get("color", "black"),
               reference.get("linestyle", "-"))]
    for entry in fastsim_configs:
        rw = FluxReweighter(entry["config"])
        poly, chi2red = rw.fit(degree=degree)
        label = rw.cfg["label"]
        coeffs = np.array2string(poly.coefficients, precision=4, separator=", ")
        print(f"{label}: fast-sim(light+charm)/DPMJET  deg={degree}  "
              f"coeffs(high->low)={coeffs}  chi2/ndf={chi2red:.2f}")
        curves.append((label, poly,
                       entry.get("color", "tab:blue"),
                       entry.get("linestyle", "--")))
    return curves
 
 
# ==========================================================================
# Reading energies from ROOT (SND@LHC cbmsim tree)
# ==========================================================================

dpm_file_list()

reweighting_function.dpm_file_list	(		data_dir,
			n_files,
			pattern,
			indexed = True
	)

Build the list of ROOT files to read.

``indexed=True`` expects the per-job layout ``<data_dir>/<i>/<pattern>``
for i in range(n_files); ``indexed=False`` a single merged file at
``<data_dir>/<pattern>``.

Definition at line 677 of file reweighting_function.py.

def dpm_file_list(data_dir, n_files, pattern, indexed=True):
    """Build the list of ROOT files to read.
 
    ``indexed=True`` expects the per-job layout ``<data_dir>/<i>/<pattern>``
    for i in range(n_files); ``indexed=False`` a single merged file at
    ``<data_dir>/<pattern>``.
    """
    if indexed:
        return [os.path.join(data_dir, str(i), pattern) for i in range(n_files)]
    return [os.path.join(data_dir, pattern)]
 
 

fastsim_entries_from_args()

reweighting_function.fastsim_entries_from_args

(

paths

)

Turn bare config paths from ``-rc`` into styled fast-sim entries.

Definition at line 885 of file reweighting_function.py.

def fastsim_entries_from_args(paths):
    """Turn bare config paths from ``-rc`` into styled fast-sim entries."""
    return [{"config": path,
             "color": DEFAULT_COLORS[i % len(DEFAULT_COLORS)],
             "linestyle": DEFAULT_STYLES[i % len(DEFAULT_STYLES)]}
            for i, path in enumerate(paths)]
 
 

get_energies_pyroot()

reweighting_function.get_energies_pyroot	(		files,
			flavour = "mu",
			tree_name = "cbmsim",
			n_events = None
	)

Read per-event neutrino energies [GeV] with pyROOT.

Loops the chained ``cbmsim`` tree entry by entry. Returns a 1-D array with
one energy per event that contains a matching primary neutrino.

Definition at line 711 of file reweighting_function.py.

def get_energies_pyroot(files, flavour="mu", tree_name="cbmsim", n_events=None):
    """Read per-event neutrino energies [GeV] with pyROOT.
 
    Loops the chained ``cbmsim`` tree entry by entry. Returns a 1-D array with
    one energy per event that contains a matching primary neutrino.
    """
    pids = set(NU_PID.values()) if flavour == "all" else {NU_PID[flavour]}
    chain = build_chain(files, tree_name)
    n_total = chain.GetEntries()
    n = n_total if n_events is None else min(n_events, n_total)
 
    energies = []
    for i in tqdm(range(n), desc="reading neutrino energy (pyROOT)"):
        chain.GetEntry(i)
        E = _incoming_neutrino_energy(chain, pids)
        if E is not None:
            energies.append(E)
    arr = np.asarray(energies, dtype=float)
    print(f"Extracted {arr.size} neutrino energies (flavour={flavour})")
    return arr
 

get_energies_uproot()

reweighting_function.get_energies_uproot	(		files,
			flavour = "mu",
			tree_name = "cbmsim",
			n_events = None
	)

Read per-event neutrino energies [GeV] with uproot, without ROOT.

Selects the incoming neutrino (|pdg| matching the flavour, fMotherId == -1)
and, neutrinos being effectively massless, takes E = |p|. Unreadable files
are skipped with a message.

Definition at line 750 of file reweighting_function.py.

def get_energies_uproot(files, flavour="mu", tree_name="cbmsim", n_events=None):
    """Read per-event neutrino energies [GeV] with uproot, without ROOT.
 
    Selects the incoming neutrino (|pdg| matching the flavour, fMotherId == -1)
    and, neutrinos being effectively massless, takes E = |p|. Unreadable files
    are skipped with a message.
    """
    if uproot is None or ak is None:
        raise RuntimeError("uproot/awkward not available -- pip install uproot awkward")
    pids = set(NU_PID.values()) if flavour == "all" else {NU_PID[flavour]}
 
    out = []
    for path in tqdm(files, desc="reading neutrino energy (uproot)"):
        try:
            tree = uproot.open(f"{path}:{tree_name}")
        except Exception as exc:                      # noqa: BLE001
            print(f"  [skip] {path}: {exc}")
            continue
        br = _find_mctrack_branches(tree)
        if len(br) < 5:
            raise RuntimeError(
                "Could not find MCTrack momentum/pdg branches. Branches seen: "
                f"{tree.keys()[:40]} ... adjust _find_mctrack_branches().")
        a = tree.arrays(list(br.values()), entry_stop=n_events, library="ak")
        pdg, mom = a[br["fPdgCode"]], a[br["fMotherId"]]
        px, py, pz = a[br["fPx"]], a[br["fPy"]], a[br["fPz"]]
 
        is_nu = ak.zeros_like(mom, dtype=bool)
        for pid in pids:
            is_nu = is_nu | (abs(pdg) == pid)
        sel = is_nu & (mom == -1)
 
        E = np.sqrt(px ** 2 + py ** 2 + pz ** 2)
        first = ak.firsts(E[sel])                     # first primary nu per event
        first = first[~ak.is_none(first)]
        out.append(ak.to_numpy(first))
 
    arr = np.concatenate(out) if out else np.array([], dtype=float)
    print(f"Extracted {arr.size} neutrino energies (flavour={flavour})")
    return arr
 
 

load_analysis_config()

reweighting_function.load_analysis_config

(

path

)

Load the top-level run YAML (Analysis.yaml) into a dict.

Definition at line 622 of file reweighting_function.py.

def load_analysis_config(path):
    """Load the top-level run YAML (Analysis.yaml) into a dict."""
    with open(path) as f:
        return yaml.safe_load(f) or {}
 
 

main()

reweighting_function.main

(

)

Parse the command line and run the requested mode.

Modes
-----
reweight : fit each reweighter config and save its diagnostic plots. Needs
           no run YAML when both ``-m reweight`` and ``-rc`` are given.
compare  : fit w(E), read the DPMJET energies from ROOT, and save the
           overlay. Needs the run YAML for the ROOT input and plot settings.

Command-line options override the corresponding run-YAML fields; anything
not given falls back to the YAML.

Definition at line 893 of file reweighting_function.py.

def main():
    """Parse the command line and run the requested mode.
 
    Modes
    -----
    reweight : fit each reweighter config and save its diagnostic plots. Needs
               no run YAML when both ``-m reweight`` and ``-rc`` are given.
    compare  : fit w(E), read the DPMJET energies from ROOT, and save the
               overlay. Needs the run YAML for the ROOT input and plot settings.
 
    Command-line options override the corresponding run-YAML fields; anything
    not given falls back to the YAML.
    """
    p = argparse.ArgumentParser(
        description="Flux reweighting + energy-distribution comparison (YAML driven)")
    p.add_argument("-c", "--config", default="analysis.yaml",
                   help="Top-level analysis YAML with the run settings")
    p.add_argument("-rc", "--reweight-config", nargs="+", default=None,
                   metavar="YAML",
                   help="FluxReweighter config file(s). In compare mode these are "
                        "the fast-sim variants (overrides fastsim_configs in the "
                        "run YAML); in reweight mode each is fit and plotted.")
    p.add_argument("-m", "--mode", choices=["compare", "reweight"], default=None,
                   help="Override the 'mode' set in the run YAML")
    p.add_argument("--reader", choices=["auto", "pyroot", "uproot"], default=None,
                   help="ROOT backend: auto (default), pyroot, or uproot "
                        "(pip install uproot awkward; no ROOT framework needed)")
    args = p.parse_args()
 
    # The run YAML is loaded lazily, so 'reweight' mode with -rc works in a
    # directory that has no Analysis.yaml at all.
    _cfg = {}
 
    def get_cfg():
        nonlocal _cfg
        if not _cfg:
            try:
                _cfg = load_analysis_config(args.config)
            except FileNotFoundError:
                raise SystemExit(
                    f"Run config '{args.config}' not found. It's needed for "
                    "compare mode (and to read 'mode'/'reweight_only' when not "
                    "given on the CLI). Pass -c <file>, or in reweight mode give "
                    "-m reweight with -rc <config(s)> so no run YAML is required.")
        return _cfg
 
    mode = args.mode or get_cfg().get("mode", "compare")
 
    if mode == "reweight":
        paths = args.reweight_config or [get_cfg()["reweight_only"]["config"]]
        for path in paths:
            FluxReweighter(path).run()
        return
 
    # --- compare mode ---
    cfg  = get_cfg()
    ri   = cfg["root_input"]
    rw   = cfg.get("reweight", {})
    plot = cfg.get("plot", {})
    degree = int(rw.get("degree", 2))
 
    fastsim = (fastsim_entries_from_args(args.reweight_config)
               if args.reweight_config else cfg["fastsim_configs"])
 
    # 1) reweighting polynomials, fitted by FluxReweighter
    curves = build_fastsim_curves(fastsim, cfg.get("reference", {}), degree=degree)
 
    # 2) DPMJET neutrino energies from ROOT (pyROOT or uproot)
    reader = args.reader or ri.get("reader", "auto")
    dpm = read_dpm_energies(ri, reader=reader)
    if dpm.size == 0:
        raise SystemExit("No energies extracted -- check files / branch names.")
 
    # 3) overlay DPMJET with each reweighted fast-sim curve
    norm = resolve_normalization(
        plot.get("normalization", "auto"),
        int(ri["number_files"]),
        target_luminosity=float(plot.get("target_luminosity", 68.6)),
        luminosity_per_file=float(plot.get("luminosity_per_file", 100.0)),
    )
    make_comparison_plot(
        dpm, curves, normalized_factor=norm,
        outfile=plot.get("outfile", "energy_comparison.pdf"),
        bins=int(plot.get("bins", 100)),
        hist_range=(0, float(plot.get("energy_max", 2000.0))),
        title=plot.get("title"),
    )
 
 

make_comparison_plot()

reweighting_function.make_comparison_plot	(		dpm,
			curves,
			normalized_factor,
			outfile,
			bins = 100,
			hist_range = (0, 2000),
			title = None
	)

Overlay the DPMJET energy spectrum with each reweighted fast-sim curve.

The same array of DPMJET energies is histogrammed once per curve. The
reference is unweighted; every other curve applies its w(E) per event, so
``DPMJET x w(E)`` is the fast-sim prediction. The binning here is
independent of the binning used to derive w(E): the weight is evaluated
from the polynomial at each event's energy, not looked up by bin.

The integral quoted in each legend entry is the normalised area under the
histogram, i.e. the expected yield at the target luminosity.

Returns a dict of ``{label: integral}``.

Definition at line 820 of file reweighting_function.py.

                         bins=100, hist_range=(0, 2000), title=None):
    """Overlay the DPMJET energy spectrum with each reweighted fast-sim curve.
 
    The same array of DPMJET energies is histogrammed once per curve. The
    reference is unweighted; every other curve applies its w(E) per event, so
    ``DPMJET x w(E)`` is the fast-sim prediction. The binning here is
    independent of the binning used to derive w(E): the weight is evaluated
    from the polynomial at each event's energy, not looked up by bin.
 
    The integral quoted in each legend entry is the normalised area under the
    histogram, i.e. the expected yield at the target luminosity.
 
    Returns a dict of ``{label: integral}``.
    """
    if title is None:
        title = "Energy Distribution: DPMJET vs fast-sim (light + charm) reweighted"
 
    plt.figure(figsize=(12, 7))
    integrals = {}
 
    for label, weight_poly, color, linestyle in curves:
        weights = None if weight_poly is None else weight_poly(dpm)
 
        # Histogram used for the integral only (not drawn).
        counts, edges = np.histogram(
            dpm, bins=bins, range=hist_range, weights=weights
        )
        bin_width = edges[1] - edges[0]
        integral = (normalized_factor * np.sum(counts) * bin_width
                    * (bins / (hist_range[1] - hist_range[0])))
        integrals[label] = integral
 
        plt.hist(dpm, bins=bins, range=hist_range, histtype="step",
                 weights=weights, color=color, linestyle=linestyle,
                 label=f"{label} (\u222bfde \u2248 {integral:.1f})")
 
    plt.xlabel("Energy GeV")
    plt.ylabel(f"Events / {(hist_range[1] - hist_range[0]) / bins} GeV")
    plt.title(title)
    plt.legend()
    plt.grid(alpha=0.3)
    plt.tight_layout()
 
    os.makedirs(os.path.dirname(os.path.abspath(outfile)), exist_ok=True)
    plt.savefig(outfile, dpi=150)
    plt.close()
    print("Saved:", outfile)
 
    print("\nIntegrals (sum(counts)*bin_width, normalized):")
    for label, val in integrals.items():
        print(f"{label:45s}: {val:.3f}")
    return integrals
 
 
# ==========================================================================
# Command line entry point
# ==========================================================================
 
#: Fallback styling when fast-sim configs are given on the command line and so
#: carry no colour/linestyle of their own.

read_dpm_energies()

reweighting_function.read_dpm_energies	(		ri,
			reader = "auto"
	)

Read the DPMJET neutrino energies using the requested ROOT backend.

``ri`` is the ``root_input`` block of the run YAML. ``reader`` is "auto"
(pyROOT when importable, else uproot), "pyroot" or "uproot".

Definition at line 792 of file reweighting_function.py.

def read_dpm_energies(ri, reader="auto"):
    """Read the DPMJET neutrino energies using the requested ROOT backend.
 
    ``ri`` is the ``root_input`` block of the run YAML. ``reader`` is "auto"
    (pyROOT when importable, else uproot), "pyroot" or "uproot".
    """
    files = dpm_file_list(ri["data_dir"], int(ri["number_files"]),
                          ri["pattern"], indexed=ri.get("indexed", True))
    tree_name = ri.get("tree_name", "cbmsim")
    flavour = ri.get("flavour", "mu")
 
    if reader == "auto":
        reader = "pyroot" if ROOT is not None else "uproot"
 
    if reader == "pyroot":
        if ROOT is None:
            raise RuntimeError("reader=pyroot but ROOT is unavailable. Source the "
                               "SND@LHC setup, or use reader: uproot "
                               "(pip install uproot awkward).")
        return get_energies_pyroot(files, flavour, tree_name)
    if reader == "uproot":
        return get_energies_uproot(files, flavour, tree_name)
    raise ValueError(f"unknown reader: {reader!r} (use auto/pyroot/uproot)")
 
 
# ==========================================================================
# Comparison plot
# ==========================================================================

resolve_normalization()

reweighting_function.resolve_normalization	(		norm,
			number_files,
			target_luminosity = 68.6,
			luminosity_per_file = 100.0
	)

Return the factor scaling the simulated sample to the target luminosity.

The automatic value is the luminosity ratio L_data / L_MC,

    target_luminosity / (number_files * luminosity_per_file)

with each input file representing ``luminosity_per_file`` [fb^-1] and
``number_files`` of them merged. Passing a number instead of ``"auto"``
uses it verbatim; ``"auto"`` or a non-positive value uses the formula.

Definition at line 628 of file reweighting_function.py.

                          luminosity_per_file=100.0):
    """Return the factor scaling the simulated sample to the target luminosity.
 
    The automatic value is the luminosity ratio L_data / L_MC,
 
        target_luminosity / (number_files * luminosity_per_file)
 
    with each input file representing ``luminosity_per_file`` [fb^-1] and
    ``number_files`` of them merged. Passing a number instead of ``"auto"``
    uses it verbatim; ``"auto"`` or a non-positive value uses the formula.
    """
    auto = target_luminosity / (number_files * luminosity_per_file)
    if isinstance(norm, str) and norm.strip().lower() == "auto":
        return auto
    norm = float(norm)
    return norm if norm > 0 else auto
 
 

tqdm()

reweighting_function.tqdm	(		x,
		**	kwargs
	)

Definition at line 100 of file reweighting_function.py.

    def tqdm(x, **kwargs):
        return x
 
 

Variable Documentation

ak

reweighting_function.ak = None

Definition at line 95 of file reweighting_function.py.

DEFAULT_COLORS

list reweighting_function.DEFAULT_COLORS = ["tab:blue", "tab:red", "tab:green", "tab:orange", "tab:purple"]

Definition at line 881 of file reweighting_function.py.

DEFAULT_STYLES

list reweighting_function.DEFAULT_STYLES = ["--", "-.", ":"]

Definition at line 882 of file reweighting_function.py.

NU_PID

dict reweighting_function.NU_PID = {"e": 12, "mu": 14, "tau": 16}

Definition at line 619 of file reweighting_function.py.

ROOT

reweighting_function.ROOT = None

Definition at line 84 of file reweighting_function.py.

uproot

reweighting_function.uproot = None

Definition at line 91 of file reweighting_function.py.