diff --git a/README.md b/README.md index 9ae71d0..6eda544 100644 --- a/README.md +++ b/README.md @@ -7,11 +7,11 @@ Meson Structure Analysis ## Λ-reconstruction and Kaon Structure Function Studies -The branch ***analysis/multi-calo-lambda*** of the **meson-structure** repository provides a lightweight and user-friendly Python framework to: +The folder ***analysis/multi-calo-lambda*** of the **meson-structure** repository provides a lightweight and user-friendly Python framework to: - Read reconstructed Λ candidates from EICrecon outputs, - Compare reconstructed vs Geant4 and afterburner angles and spectra, -- Produce kinematic maps, +- Produce kinematic maps (trhuth, electron method, JB method), - Propagate statistical uncertainties to the kaon structure function*. \* The error propagation is based on the assumption that the kaon structure function is proportional to the number of $\Lambda^0$ hyperons produced by the Sullivan process, $F_k\propto N_\Lambda$. This number is related to the reconstruction efficiency $\varepsilon$, the cross-section $\sigma$, and the EIC luminosity $L$, as $N_\Lambda=\varepsilon\times \sigma \times L $. Assuming Poisson statistics, the relative error is estimated as @@ -82,7 +82,6 @@ Examples: ### TO DO LIST -- [physics] Include the experimental kinematics reconstruction (e-, JB) from EICrecon. -- [physics] Attach the uncertainties to Kaon SF models to see if discriminating. -- [physics] Include other sources of error than purely statistical ones (e.g. $\Delta\sigma$ from event generator) -- [ESR] Include Zenobo figures plot format (https://zenodo.org/records/16615455) \ No newline at end of file +- [physics] Event generator sensitivity studies +- [physics] Background contamination studies +- [ESR] Include Zenodo figures plot format (https://zenodo.org/records/16615455) \ No newline at end of file diff --git a/analysis/multicalo-lambda/config.py b/analysis/multicalo-lambda/config.py index e306cec..35fe453 100644 --- a/analysis/multicalo-lambda/config.py +++ b/analysis/multicalo-lambda/config.py @@ -1,8 +1,14 @@ from __future__ import annotations + +import os from dataclasses import dataclass from pathlib import Path from typing import Sequence +import math +# ============================================================================= +# Paths +# ============================================================================= @dataclass(frozen=True) class Paths: @@ -13,8 +19,14 @@ class Paths: "/work/eic3/users/romanov/meson-structure-2026-02/afterburner/{beam}-priority/" "k_lambda_{beam}_5000evt_{idx:04d}.afterburner.hepmc" ) - gen_base_dir: Path = Path("/work/eic3/users/romanov/meson-structure-2026-02/eg-orig-kaon-lambda/") + gen_base_dir: Path = Path("/work/eic3/users/romanov/eg-orig-kaon-lambda-2025-08/") + +PATHS = Paths() +PATHS.outputs.mkdir(parents=True, exist_ok=True) +# ============================================================================= +# Physics constants (global run settings) +# ============================================================================= @dataclass(frozen=True) class PhysicsConstants: @@ -27,6 +39,40 @@ class PhysicsConstants: theta_proton_rad: float = 25e-3 mp_gev: float = 0.9382720813 +CONST = PhysicsConstants() + +# ============================================================================= +# LHAPDF setup +# ============================================================================= + +DEFAULT_LHAPDF_DATA_PATH = "/work/eic3/users/fraisse/lhapdf" + +def setup_lhapdf_env() -> str: + """ + Ensure LHAPDF_DATA_PATH is set before importing python lhapdf. + Returns the resolved LHAPDF path. + """ + lhapdf_path = os.environ.get("LHAPDF_DATA_PATH", DEFAULT_LHAPDF_DATA_PATH) + os.environ["LHAPDF_DATA_PATH"] = lhapdf_path + return lhapdf_path + +PION_PDFSET_DEFAULT = "JAM21PionPDFnlo" +PION_PDFMEMBER_DEFAULT = 0 + +def get_pion_pdfset() -> str: + return os.environ.get("PION_PDFSET", PION_PDFSET_DEFAULT) + +def get_pion_pdfmember() -> int: + return int(os.environ.get("PION_PDFMEMBER", str(PION_PDFMEMBER_DEFAULT))) + +# ============================================================================= +# Default grid settings for structure function evaluation +# ============================================================================= + +GRID_DEFAULTS = dict( + x_min=0.0, x_max=1.0, + q2_min=1.0, q2_max=500.0, + nx=20, nq2=20, +) + -DEFAULT_PATHS = Paths() -DEFAULT_CONST = PhysicsConstants() \ No newline at end of file diff --git a/analysis/multicalo-lambda/kaon_models.py b/analysis/multicalo-lambda/kaon_models.py new file mode 100644 index 0000000..aea2207 --- /dev/null +++ b/analysis/multicalo-lambda/kaon_models.py @@ -0,0 +1,640 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable, Literal +import math +import numpy as np + +from config import ( + setup_lhapdf_env, + get_pion_pdfset, + get_pion_pdfmember, + GRID_DEFAULTS, +) +from plotting import ( + apply_mpl_style, + plot_F2_map +) +from physics import ( + ToyParams, + KaonModelParams, + F2_from_pdf_func, + toy_norm, + toy_shape_x +) +from utils import ( + Grid, + is_lhapdf_available, + make_pdf_kaon_toy_times_pion_evolution, + make_lin_grid, + make_log_grid, + compute_F2_grid, +) + +# global + +GridKind = Literal["lin", "log"] + +ModelKind = Literal[ + # toy models + "toy_baseline", + "toy_soft_valence", + "toy_hard_valence", + "toy_sea_enhanced", + "toy_su3_breaking", + # JAM models + "jam25_rep165", + "jam25_rep390", + "jam25_rep400", + "jam25_mean", + # theoretical models + "cosmao22", + ] + +JAM_REPLICA_MODELS: dict[str, int] = { + "jam25_rep165": 165, + "jam25_rep390": 390, + "jam25_rep400": 400, +} + +# configuration call + +@dataclass(frozen=True) +class EvalConfig: + grid_kind: GridKind = "lin" + grid_overrides: dict[str, float | int] | None = None + + model: ModelKind = "toy_baseline" + toy: ToyParams = ToyParams(a=0.5, b=1.5, Q0=0.5) + kparams: KaonModelParams = KaonModelParams() + + prefix: str | None = None + + +def build_grid(kind: GridKind = "lin", overrides: dict[str, float | int] | None = None) -> Grid: + params = dict(GRID_DEFAULTS) + if overrides: + params.update(overrides) + + if kind == "log": + return make_log_grid(**params) + if kind == "lin": + return make_lin_grid(**params) + raise ValueError(f"Unknown grid kind: {kind}") + + +# model builders + +PdfFunc = Callable[[int, float, float], float] + +def _make_toy_valence_pdf(toy: ToyParams) -> PdfFunc: + """ + Simplistic K+ valence-only PDF at Q0: + u(x) ~ sbar(x) ~ x^a (1-x)^b normalized to 1 + """ + a, b = toy.a, toy.b + norm = toy_norm(a, b) + + def pdf(pid: int, x: float, Q2: float) -> float: + if pid in (2, -3): # u, sbar + return toy_shape_x(x, a, b, norm) + return 0.0 + + return pdf + +def _make_pdf_from_lhapdf(setname: str, member: int = 0) -> PdfFunc: + import lhapdf + pdfset = lhapdf.mkPDF(setname, member) + + def pdf(pid: int, x: float, Q2: float) -> float: + if x <= 0.0 or x >= 1.0 or Q2 <= 0.0: + return 0.0 + return pdfset.xfxQ2(pid, x, Q2) / x + + return pdf + +# theoretical model (DSE) + +def make_pdf_kaon_cosmao22( + setname: str = "CoSMAO22Kaon", + member: int = 0, +) -> PdfFunc: + """ + CoSMAO22 K+ PDFs from LHAPDF. + + LHAPDF set: + CoSMAO22Kaon, ID 9950, Particle 321 = K+ + + Returns f(x,Q2), not x*f(x,Q2). + """ + return _make_pdf_from_lhapdf(setname, member) + + +def _get_pion_uval(pdf_pi: PdfFunc, pion_set: str, x: float, Q2: float) -> float: + u_pi = pdf_pi(2, x, Q2) + ubar_pi = pdf_pi(-2, x, Q2) + + if pion_set == "GRVPI0": + return max(u_pi, 0.0) + else: + return max(u_pi - ubar_pi, 0.0) + + +# JAM model + +F2KComponent = Literal["total", "valence", "sea", "gluon"] +JAM_F2K_FLAVS = [21, -5, -4, -3, -2, -1, 1, 2, 3, 4, 5] +_JAM_PDFS_CACHE: dict[str, list] = {} + +def _get_jam_pdfset(setname: str = "JAM25kaon_nlonll_F2K"): + setup_lhapdf_env() + import lhapdf + + if setname not in _JAM_PDFS_CACHE: + _JAM_PDFS_CACHE[setname] = lhapdf.mkPDFs(setname) + + return _JAM_PDFS_CACHE[setname] + + +def get_n_jam_replicas(setname: str = "JAM25kaon_nlonll_F2K") -> int: + return len(_get_jam_pdfset(setname)) + + +def make_f2k_jam_replica(replica: int, setname: str = "JAM25kaon_nlonll_F2K"): + JAM_F2K_FLAVS = [21, -5, -4, -3, -2, -1, 1, 2, 3, 4, 5] + pdfs = _get_jam_pdfset(setname) + + if replica < 0 or replica >= len(pdfs): + raise ValueError( + f"Replica index {replica} out of range for set {setname} " + f"(available: 0..{len(pdfs)-1})" + ) + + grid = pdfs[replica] + + def f2k(x: float, Q2: float) -> float: + if x <= 0.0 or x >= 1.0 or Q2 <= 0.0: + return np.nan + + vals = np.array([grid.xfxQ2(f, x, Q2) for f in JAM_F2K_FLAVS], dtype=float) + if not np.all(np.isfinite(vals)): + return np.nan + + return float(np.sum(vals)) + + return f2k + + +def get_all_jam_replica_ids(setname: str = "JAM25kaon_nlonll_F2K") -> list[int]: + return list(range(get_n_jam_replicas(setname))) + + +def evaluate_f2k_for_jam_replicas( + replica_ids: list[int], + x_arr: np.ndarray, + q2_arr: np.ndarray, + setname: str = "JAM25kaon_nlonll_F2K", +) -> np.ndarray: + pdfs = _get_jam_pdfset(setname) + JAM_F2K_FLAVS = [21, -5, -4, -3, -2, -1, 1, 2, 3, 4, 5] + + x_arr = np.asarray(x_arr, dtype=float) + q2_arr = np.asarray(q2_arr, dtype=float) + + out = np.full((len(replica_ids), len(x_arr), len(q2_arr)), np.nan, dtype=float) + + for irep, rep in enumerate(replica_ids): + if rep < 0 or rep >= len(pdfs): + raise ValueError( + f"Replica index {rep} out of range for set {setname} " + f"(available: 0..{len(pdfs)-1})" + ) + + grid = pdfs[rep] + + for j, q2 in enumerate(q2_arr): + for i, x in enumerate(x_arr): + if x <= 0.0 or x >= 1.0 or q2 <= 0.0: + continue + + vals = np.array([grid.xfxQ2(f, x, q2) for f in JAM_F2K_FLAVS], dtype=float) + if np.all(np.isfinite(vals)): + out[irep, i, j] = float(np.sum(vals)) + + return out + + +def is_jam_replica_model(model: str) -> bool: + return model in JAM_REPLICA_MODELS or model == "jam25_mean" + + +def get_jam_f2k_flavor_values(grid, x: float, Q2: float) -> dict[str, float]: + """ + Return LHAPDF-style F2K flavor contributions for a JAM replica grid. + + Important: + These are not PDFs q(x,Q2), but flavor contributions to F2K. + """ + pid_to_name = { + 21: "g", + -5: "bb", + -4: "cb", + -3: "sb", + -2: "ub", + -1: "db", + 1: "d", + 2: "u", + 3: "s", + 4: "c", + 5: "b", + } + + vals = {} + for pid, name in pid_to_name.items(): + vals[name] = float(grid.xfxQ2(pid, x, Q2)) + + return vals + +def compute_jam_f2k_component( + flav: dict[str, float], + component: F2KComponent = "total", +) -> float: + """ + Compute total, valence, sea, or gluon contribution to F2K + from JAM LHAPDF-style flavor contributions. + + For K-: + valence = (s - sb) + (ub - u) + sea = 2*u + d + db + 2*sb + (c + cb + b + bb) + gluon = g + """ + + if component == "total": + return sum(flav.values()) + + if component == "valence": + return (flav["s"] - flav["sb"]) + (flav["ub"] - flav["u"]) + + if component == "sea": + return ( + 2.0 * flav["u"] + + flav["d"] + + flav["db"] + + 2.0 * flav["sb"] + + flav["c"] + + flav["cb"] + + flav["b"] + + flav["bb"] + ) + + if component == "gluon": + return flav["g"] + + raise ValueError( + f"Unknown component '{component}'. " + "Use 'total', 'valence', 'sea', or 'gluon'." + ) + + +def make_f2k_jam_replica_component( + replica: int, + component: F2KComponent = "total", + setname: str = "JAM25kaon_nlonll_F2K", +): + pdfs = _get_jam_pdfset(setname) + + if replica < 0 or replica >= len(pdfs): + raise ValueError( + f"Replica index {replica} out of range for set {setname} " + f"(available: 0..{len(pdfs)-1})" + ) + + grid = pdfs[replica] + + def f2k_component(x: float, Q2: float) -> float: + if x <= 0.0 or x >= 1.0 or Q2 <= 0.0: + return np.nan + + flav = get_jam_f2k_flavor_values(grid, x, Q2) + + if not np.all(np.isfinite(list(flav.values()))): + return np.nan + + return float(compute_jam_f2k_component(flav, component)) + + return f2k_component + + +def evaluate_f2k_component_for_jam_replicas( + replica_ids: list[int], + x_arr: np.ndarray, + q2_arr: np.ndarray, + component: F2KComponent = "total", + setname: str = "JAM25kaon_nlonll_F2K", +) -> np.ndarray: + pdfs = _get_jam_pdfset(setname) + + x_arr = np.asarray(x_arr, dtype=float) + q2_arr = np.asarray(q2_arr, dtype=float) + + out = np.full((len(replica_ids), len(x_arr), len(q2_arr)), np.nan, dtype=float) + + for irep, rep in enumerate(replica_ids): + if rep < 0 or rep >= len(pdfs): + raise ValueError( + f"Replica index {rep} out of range for set {setname} " + f"(available: 0..{len(pdfs)-1})" + ) + + grid = pdfs[rep] + + for j, q2 in enumerate(q2_arr): + for i, x in enumerate(x_arr): + if x <= 0.0 or x >= 1.0 or q2 <= 0.0: + continue + + flav = get_jam_f2k_flavor_values(grid, x, q2) + + if np.all(np.isfinite(list(flav.values()))): + out[irep, i, j] = compute_jam_f2k_component( + flav, + component=component, + ) + + return out + + +def make_f2k_jam_mean(setname: str = "JAM25kaon_nlonll_F2K"): + pdfs = _get_jam_pdfset(setname) + + def f2k_mean(x: float, Q2: float) -> float: + if x <= 0.0 or x >= 1.0 or Q2 <= 0.0: + return np.nan + + vals = [] + + for grid in pdfs: + flavs = np.array([grid.xfxQ2(f, x, Q2) for f in JAM_F2K_FLAVS], dtype=float) + if np.all(np.isfinite(flavs)): + vals.append(np.sum(flavs)) + + if len(vals) == 0: + return np.nan + + return float(np.mean(vals)) + + return f2k_mean + +# make structure function + +def make_kaon_pdf_func( + model: ModelKind, + toy: ToyParams, + kparams: KaonModelParams, + *, + pion_set: str, + pion_member: int, + require_lhapdf: bool = True, +) -> PdfFunc: + """ + Build a kaon PDF model as a function pdf(pid, x, Q2) -> f(x,Q2). + """ + + # theorerical models + + if model == "cosmao22": + return make_pdf_kaon_cosmao22( + setname="CoSMAO22Kaon", + member=0, + ) + + # toy models + + a = toy.a + b = toy.b + + if model == "toy_soft_valence": + b = toy.b + float(kparams.db_soft_valence) + + elif model == "toy_hard_valence": + b = toy.b + float(kparams.db_hard_valence) + b = max(b, 0.05) + + elif model in ("toy_baseline", "toy_su3_breaking", "toy_sea_enhanced"): + pass + + else: + raise ValueError( + f"Unknown model: {model}. " + "Use: toy_baseline, toy_soft_valence, toy_hard_valence, " + "toy_sea_enhanced or toy_su3_breaking." + ) + + toy_eff = ToyParams(a=a, b=b, Q0=toy.Q0) + + # pdf + + if require_lhapdf: + pdf_base = make_pdf_kaon_toy_times_pion_evolution( + toy_eff, + pion_set=pion_set, + member=pion_member, + ) + else: + pdf_base = _make_toy_valence_pdf(toy_eff) + + # pdf toy models + + if model in ("toy_baseline", "toy_soft_valence", "toy_hard_valence"): + return pdf_base + + if model == "toy_su3_breaking": + a_u, b_u = toy_eff.a, toy_eff.b + norm_u = toy_norm(a_u, b_u) + + a_s = a_u + float(kparams.da_sbar_su3_break) + b_s = b_u + float(kparams.db_sbar_su3_break) + b_s = max(b_s, 0.05) + norm_s = toy_norm(a_s, b_s) + + def pdf_su3(pid: int, x: float, Q2: float) -> float: + base = pdf_base(pid, x, Q2) + + if pid != -3: + return base + + u0 = toy_shape_x(x, a_u, b_u, norm_u) + s0 = toy_shape_x(x, a_s, b_s, norm_s) + + if u0 <= 0.0: + return 0.0 + + return base * (s0 / u0) + + return pdf_su3 + + if model == "toy_sea_enhanced": + sea_norm = toy_norm(kparams.sea_a, kparams.sea_b) + + def pdf_with_sea(pid: int, x: float, Q2: float) -> float: + base = pdf_base(pid, x, Q2) + sea = float(kparams.sea_amp) * toy_shape_x( + x, kparams.sea_a, kparams.sea_b, sea_norm + ) + + if pid in (1, -1, 2, -2, 3, -3): + return base + sea + + return base + + return pdf_with_sea + + raise ValueError(f"Unhandled model: {model}") + + +# compute structure function + +def compute_F2_kaon_model( + grid: Grid, + model: ModelKind, + toy: ToyParams, + kparams: KaonModelParams, + pion_set: str, + pion_member: int, +) -> np.ndarray: + + # Direct F2K LHAPDF grids (JAM replicas) + + if model == "jam25_mean": + f2k = make_f2k_jam_mean(setname="JAM25kaon_nlonll_F2K") + + return compute_F2_grid( + grid, + f2_point=lambda x, Q2: f2k(x, Q2), + ) + + if is_jam_replica_model(model): + replica = JAM_REPLICA_MODELS[model] + f2k = make_f2k_jam_replica(replica, setname="JAM25kaon_nlonll_F2K") + + return compute_F2_grid( + grid, + f2_point=lambda x, Q2: f2k(x, Q2), + ) + + # Standard PDF-based models + + pdf = make_kaon_pdf_func( + model=model, + toy=toy, + kparams=kparams, + pion_set=pion_set, + pion_member=pion_member, + require_lhapdf=True, + ) + + return compute_F2_grid( + grid, + f2_point=lambda x, Q2: F2_from_pdf_func(pdf, x, Q2), + ) + +# high level structure function + +def evaluate_and_plot_F2K(cfg: EvalConfig = EvalConfig()) -> str | None: + """ + High-level entry point called by run.py. + Returns output path as string, or None if skipped/failure. + """ + setup_lhapdf_env() + apply_mpl_style() + + if not is_lhapdf_available(): + print("WARNING: python lhapdf not available -> skipping plot.") + return None + + pion_set = get_pion_pdfset() + pion_member = get_pion_pdfmember() + + grid = build_grid(cfg.grid_kind, cfg.grid_overrides) + + try: + F2 = compute_F2_kaon_model( + grid=grid, + model=cfg.model, + toy=cfg.toy, + kparams=cfg.kparams, + pion_set=pion_set, + pion_member=pion_member, + ) + except Exception as e: + print(f"WARNING: Cannot build/compute kaon model '{cfg.model}': {e}") + return None + + prefix = cfg.prefix or f"F2K_{cfg.model}_times_{pion_set}" + + toy = cfg.toy + k = cfg.kparams + + # legends + + title = { + "toy_baseline": + fr"Model $\mathbf{{toy\_baseline}}$: $u=s \propto x^a(1-x)^b$" + "\n" + fr"$a={toy.a:.2f},\; b={toy.b:.2f}$", + + "toy_soft_valence": + fr"Model $\mathbf{{toy\_soft\_valence}}$: $u=s \propto x^a(1-x)^{{b+\Delta b}}$" + "\n" + fr"$a={toy.a:.2f},\; b={toy.b:.2f},\; \Delta b={k.db_soft_valence:+.2f}$", + + "toy_hard_valence": + fr"Model $\mathbf{{toy\_hard\_valence}}$: $u=s \propto x^a(1-x)^{{b+\Delta b}}$" + "\n" + fr"$a={toy.a:.2f},\; b={toy.b:.2f},\; \Delta b={k.db_hard_valence:+.2f}$", + + "toy_sea_enhanced": + fr"Model $\mathbf{{toy\_sea\_enhanced}}$: $u=s \propto x^a(1-x)^b + A_s x^{{a_s}}(1-x)^{{b_s}}$" + "\n" + fr"$a={toy.a:.2f},\; b={toy.b:.2f},\; A_{{sea}}={k.sea_amp:.3f},\; a_s={k.sea_a:.2f},\; b_s={k.sea_b:.2f}$", + + "toy_su3_breaking": + fr"Model $\mathbf{{toy\_su3\_breaking}}$: $u \propto x^a(1-x)^b \neq s \propto x^{{a'}}(1-x)^{{b'}} $" + "\n" + fr"$a={toy.a:.2f},\; b={toy.b:.2f},\; a'=a+{k.da_sbar_su3_break:.2f},\; b'=b{k.db_sbar_su3_break:.2f}$", + + "jam25_rep165": + r"Replica 165 (JAM)", + + "jam25_rep390": + r"Replica 390 (JAM)", + + "jam25_rep400": + r"Replica 400 (JAM)", + + "jam25_mean": + r"Replica mean (JAM)", + + "cosmao22": + fr"CoSMAO22 (Dyson-Schwinger)", + + }[cfg.model] + + # colorbar title + + if is_jam_replica_model(cfg.model): + cbarlabel = r"$F_2^{K}(x,Q^2)$" + else: + cbarlabel = r"$F_2^K(x,Q^2) \approx x \left[ \frac{4}{9} u(x) + \frac{1}{9} s(x) \right] \times R_{\pi,\mathrm{LHAPDF}}(Q^2)$" + + out = plot_F2_map( + prefix=prefix, + grid=grid, + F2=F2, + title=title, + cbar_label=cbarlabel, + xscale='lin', + q2scale='lin' + ) + + return str(out) diff --git a/analysis/multicalo-lambda/kaon_studies.py b/analysis/multicalo-lambda/kaon_studies.py new file mode 100644 index 0000000..ac47e68 --- /dev/null +++ b/analysis/multicalo-lambda/kaon_studies.py @@ -0,0 +1,1762 @@ +from __future__ import annotations + +from pathlib import Path +import numpy as np +import awkward as ak +import uproot +import matplotlib.pyplot as plt +from matplotlib.colors import LogNorm +from matplotlib.ticker import LogLocator, NullFormatter, FixedLocator, LogFormatterMathtext, FixedFormatter + +import uproot + +from plotting import ( + apply_mpl_style, + ensure_outdir, savefig, + get_color, get_style, + plot_relerr_map +) + +from physics import ( + KaonModelParams, + proton_kinematics_for_beam, + xk_from_xb_xl, + F2_from_pdf_func, + ToyParams, + KaonModelParams, + toy_norm, + toy_shape_x +) + +from uq import ( + efficiency_from_counts, + expected_yields_nb, + relerr_combined, + migration_probability_xk_given_xb_q2, + project_xbq2_to_xkq2, + projected_efficiency_from_migration, +) + +from kaon_models import ( + EvalConfig, + make_kaon_pdf_func, + make_f2k_jam_replica, + make_f2k_jam_mean, + get_all_jam_replica_ids, + evaluate_f2k_for_jam_replicas, + evaluate_f2k_component_for_jam_replicas, +) + +from config import ( + PhysicsConstants, + CONST, + setup_lhapdf_env, + get_pion_pdfset, + get_pion_pdfmember, +) +from utils import ( + is_lhapdf_available, + make_pdf_kaon_toy_times_pion_evolution, +) + +# event generator cross-section + +def build_sigma_map_xb_q2( + beam: str, + base_dir: str | Path, + xb_range=(0.0, 1.0), + xb_bins=10, + q2_range=(1.0, 500.0), + q2_bins=10, + chunk=2_000_000, + do_savefig: bool = True, + fig_dir: str | Path = "./outputs", + plot_density: bool = True, + log_Q2: bool = True, + log_x: bool = False, +): + Ee_str, Ep_str = beam.lower().split("x") + Ee, Ep = float(Ee_str), float(Ep_str) + + fname = f"k_lambda_crossing_0.000-{Ee:.1f}on{Ep:.1f}_x0.0001-1.0000_q1.0-500.0.root" + fpath = Path(base_dir) / fname + if not fpath.exists(): + raise FileNotFoundError(fpath) + + if log_x: + if xb_range[0] <= 0 or xb_range[1] <= 0: + raise ValueError("xb_range must be strictly positive when log_x=True") + if log_Q2: + if q2_range[0] <= 0 or q2_range[1] <= 0: + raise ValueError("q2_range must be strictly positive when log_Q2=True") + + if log_x: + xb_edges = np.logspace(np.log10(xb_range[0]), np.log10(xb_range[1]), xb_bins + 1) + else: + xb_edges = np.linspace(xb_range[0], xb_range[1], xb_bins + 1) + + if log_Q2: + q2_edges = np.logspace(np.log10(q2_range[0]), np.log10(q2_range[1]), q2_bins + 1) + else: + q2_edges = np.linspace(q2_range[0], q2_range[1], q2_bins + 1) + + sumw_map = np.zeros((xb_bins, q2_bins), dtype=np.float64) + Ngen_map = np.zeros((xb_bins, q2_bins), dtype=np.float64) + + with uproot.open(fpath) as f: + meta = f["Meta"] + evnts = f["Evnts"] + process = f["Process"] + n_use = min(meta.num_entries, evnts.num_entries) + + mc0 = meta["MC"].array(entry_start=0, entry_stop=1, library="np") + Ngen_declared = int(mc0["nEvts"][0]) + + for start in range(0, n_use, chunk): + stop = min(start + chunk, n_use) + + jac = meta["Jacob"].array(entry_start=start, entry_stop=stop, library="np") + mc = meta["MC"].array(entry_start=start, entry_stop=stop, library="np") + inv = evnts["invts"].array(entry_start=start, entry_stop=stop, library="np") + xdis = process["xsec_tagged_dis"].array(entry_start=start, entry_stop=stop, library="np") + kint = process["k_int"].array(entry_start=start, entry_stop=stop, library="np") + + w = mc["PhSpFct"] * jac * xdis * kint # nb + xb = inv["xBj"] + q2 = inv["Q2"] + + ok = np.isfinite(w) & np.isfinite(xb) & np.isfinite(q2) & (q2 > 0) + if not np.any(ok): + continue + + w = w[ok] + xb = xb[ok] + q2 = q2[ok] + + sumw_map += np.histogram2d(xb, q2, bins=(xb_edges, q2_edges), weights=w)[0] + Ngen_map += np.histogram2d(xb, q2, bins=(xb_edges, q2_edges))[0] + + sigma_map_nb = sumw_map / float(Ngen_declared) + sigma_tot_nb = float(sumw_map.sum()) / float(Ngen_declared) + + if do_savefig: + apply_mpl_style() + outdir = ensure_outdir(fig_dir) + outpath = outdir / f"sigma_xb_q2_{beam}.png" + + if plot_density: + dx = np.diff(xb_edges)[:, None] + dQ2 = np.diff(q2_edges)[None, :] + z = sigma_map_nb / (dx * dQ2) + cbar_label = r"$d\sigma/(dx_B\,dQ^2)$ [nb/GeV$^2$]" + title = rf"Sullivan $d\sigma/(dx_B dQ^2)$ — {beam}" + else: + z = sigma_map_nb + cbar_label = r"$\sigma$ [nb]" + title = rf"Sullivan $\sigma(x_B,Q^2)$ — {beam}" + + pos = z[z > 0] + norm = LogNorm(vmin=float(pos.min()), vmax=float(pos.max())) if pos.size else None + + fig, ax = plt.subplots(figsize=(6, 5)) + m = ax.pcolormesh(xb_edges, q2_edges, z.T, shading="auto", norm=norm, cmap="coolwarm") + cb = fig.colorbar(m, ax=ax) + cb.set_label(cbar_label) + ax.set_xlabel(r"$x_B$") + ax.set_ylabel(r"$Q^2$ [GeV$^2$]") + if log_x: + ax.set_xscale("log") + if log_Q2: + ax.set_yscale("log") + ax.set_title(title) + fig.tight_layout() + savefig(fig, outpath, dpi=200) + + print('CHECK. Sigma tot =', sigma_tot_nb, 'nb') + + return sigma_map_nb, Ngen_map, sigma_tot_nb + +# relative errors computation + +def plot_relerr_kaon_sf_xK_Q2( + beam: str, + nfiles: int, + root_base_dir: Path, + suffix: str, + gen_base_dir: Path, + outdir: Path, + tag: str, + xB_bins: int = 20, + xB_range: tuple[float, float] = (0.0, 1.0), + xK_bins: int = 20, + xK_range: tuple[float, float] = (0.0, 2.0), + Q2_bins: int = 20, + Q2_range: tuple[float, float] = (1.0, 500.0), + log_Q2: bool = True, + log_x: bool = False, + tmax: float | None = None, + lumin_fb: float = 1.0, + vmax_percent: float | None = None, + c: PhysicsConstants = CONST, + kinmethod: str = "Truth", +): + """ + Fully consistent version. + + Strategy + -------- + 1) Build sigma_gen(xB_truth, Q2_truth) from generator. + 2) Build eps(xB_truth, Q2_truth) from reco sample, binned in truth kinematics: + eps = Nreco_lambda(truth bins) / Ngen_evt(truth bins) + 3) Build expected reconstructed yields in truth bins: + Nexp_truth = sigma_gen * eps * L + 4) Build a 4D response: + P(xK_reco, Q2_reco | xB_truth, Q2_truth) + 5) Project Nexp_truth to a fully reconstructed map: + Nexp_reco(xK_reco, Q2_reco) + 6) Compute Poisson relative uncertainty: + relerr = 100 / sqrt(Nexp_reco) + + Notes + ----- + - The output xK-Q2 map is now fully reconstructed: + (xK_reco, Q2_reco) + and no longer mixes xK_reco with Q2_truth. + - The MC uncertainty on projected efficiency is intentionally NOT used here, + because the previous projected_efficiency_from_migration construction was + not formally consistent in reco space. + """ + + def _response_probability_xkq2reco_given_xbq2truth(R: np.ndarray) -> np.ndarray: + """ + Build P(xK_reco, Q2_reco | xB_truth, Q2_truth) + from raw counts R with shape: + (xK_bins, Q2_bins, xB_bins, Q2_bins) + """ + denom = np.sum(R, axis=(0, 1)) # shape: (xB_bins, Q2_bins) + P = np.zeros_like(R, dtype=np.float64) + for ixB in range(R.shape[2]): + for iQ in range(R.shape[3]): + if denom[ixB, iQ] > 0: + P[:, :, ixB, iQ] = R[:, :, ixB, iQ] / denom[ixB, iQ] + return P + + def _project_xbq2truth_to_xkq2reco( + P: np.ndarray, + A_truth: np.ndarray, + ) -> np.ndarray: + """ + Project A_truth(xB_truth, Q2_truth) to reco space using + P(xK_reco, Q2_reco | xB_truth, Q2_truth). + + P shape : (xK_bins, Q2_bins, xB_bins, Q2_bins) + A_truth shape: (xB_bins, Q2_bins) + Returns : (xK_bins, Q2_bins) + """ + return np.einsum("kqij,ij->kq", P, A_truth) + + apply_mpl_style() + outdir = ensure_outdir(outdir) + + if log_x: + if xB_range[0] <= 0 or xK_range[0] <= 0: + raise ValueError("xB_range and xK_range must be strictly positive when log_x=True") + xB_edges = np.logspace(np.log10(xB_range[0]), np.log10(xB_range[1]), xB_bins + 1) + xK_edges = np.logspace(np.log10(xK_range[0]), np.log10(xK_range[1]), xK_bins + 1) + else: + xB_edges = np.linspace(xB_range[0], xB_range[1], xB_bins + 1) + xK_edges = np.linspace(xK_range[0], xK_range[1], xK_bins + 1) + + if log_Q2: + if Q2_range[0] <= 0: + raise ValueError("Q2_range must be strictly positive when log_Q2=True") + Q2_edges = np.logspace(np.log10(Q2_range[0]), np.log10(Q2_range[1]), Q2_bins + 1) + else: + Q2_edges = np.linspace(Q2_range[0], Q2_range[1], Q2_bins + 1) + + sigma_gen_nb, _, sigma_tot_nb = build_sigma_map_xb_q2( + beam=beam, + base_dir=gen_base_dir, + xb_range=xB_range, + xb_bins=xB_bins, + q2_range=Q2_range, + q2_bins=Q2_bins, + chunk=2_000_000, + do_savefig=True, + log_Q2=log_Q2, + log_x=log_x, + ) + + if kinmethod not in ("Truth", "Electron", "JB", "Sigma", "DA", "ML"): + raise ValueError( + f"Invalid kinmethod: {kinmethod}. Choose Truth, Electron, JB, Sigma, DA, or ML." + ) + + # truth kinematics + + xB_truth_branch = "InclusiveKinematicsTruth.x" + Q2_truth_branch = "InclusiveKinematicsTruth.Q2" + + # reconstructed kinematics + + xB_reco_branch = f"InclusiveKinematics{kinmethod}.x" + Q2_reco_branch = f"InclusiveKinematics{kinmethod}.Q2" + + # reconstructed lambdas + + lamE_branch = "ReconstructedFarForwardLambdas.energy" + lampx_branch = "ReconstructedFarForwardLambdas.momentum.x" + lampy_branch = "ReconstructedFarForwardLambdas.momentum.y" + lampz_branch = "ReconstructedFarForwardLambdas.momentum.z" + + pk = proton_kinematics_for_beam(beam, c=c) + nhat = pk["nhat"] + pplus_p = float(pk["pplus_p"]) + ppx, ppy, ppz, ppE = float(pk["ppx"]), float(pk["ppy"]), float(pk["ppz"]), float(pk["ppE"]) + + # truth-space quantities + + Ngen_evt = np.zeros((xB_bins, Q2_bins), dtype=np.float64) + Nreco_lam = np.zeros((xB_bins, Q2_bins), dtype=np.float64) + + # phase-space matrix + + R = np.zeros((xK_bins, Q2_bins, xB_bins, Q2_bins), dtype=np.float64) + + for i in range(1, nfiles + 1): + fpath = root_base_dir / f"k_lambda_{beam}_5000evt_{i:03d}_{suffix}.root" + if not fpath.exists(): + continue + + try: + with uproot.open(fpath) as f: + tree = f["events"] + + # event-level truth kinematics + + xB_evt_truth = ak.to_numpy( + ak.flatten(tree[xB_truth_branch].array(library="ak")) + ) + Q2_evt_truth = ak.to_numpy( + ak.flatten(tree[Q2_truth_branch].array(library="ak")) + ) + + # event-level reconstructed kinematics + + xB_evt_reco = ak.to_numpy( + ak.flatten(tree[xB_reco_branch].array(library="ak")) + ) + Q2_evt_reco = ak.to_numpy( + ak.flatten(tree[Q2_reco_branch].array(library="ak")) + ) + + n_evt = min( + len(xB_evt_truth), len(Q2_evt_truth), + len(xB_evt_reco), len(Q2_evt_reco) + ) + if n_evt == 0: + continue + + xB_evt_truth = xB_evt_truth[:n_evt] + Q2_evt_truth = Q2_evt_truth[:n_evt] + xB_evt_reco = xB_evt_reco[:n_evt] + Q2_evt_reco = Q2_evt_reco[:n_evt] + + # truth event count: denominator of efficiency + + Ngen_evt += np.histogram2d( + xB_evt_truth, Q2_evt_truth, bins=(xB_edges, Q2_edges) + )[0] + + # reconstructed lambdas + + lamE_j = tree[lamE_branch].array(library="ak")[:n_evt] + px_j = tree[lampx_branch].array(library="ak")[:n_evt] + py_j = tree[lampy_branch].array(library="ak")[:n_evt] + pz_j = tree[lampz_branch].array(library="ak")[:n_evt] + + # Broadcast truth event kinematics to lambda candidates + + xB_truth_bc, _ = ak.broadcast_arrays(ak.Array(xB_evt_truth), lamE_j) + Q2_truth_bc, _ = ak.broadcast_arrays(ak.Array(Q2_evt_truth), lamE_j) + + # Broadcast reco event kinematics to lambda candidates + + xB_reco_bc, _ = ak.broadcast_arrays(ak.Array(xB_evt_reco), lamE_j) + Q2_reco_bc, _ = ak.broadcast_arrays(ak.Array(Q2_evt_reco), lamE_j) + + xB_truth_c = ak.to_numpy(ak.flatten(xB_truth_bc)) + Q2_truth_c = ak.to_numpy(ak.flatten(Q2_truth_bc)) + xB_reco_c = ak.to_numpy(ak.flatten(xB_reco_bc)) + Q2_reco_c = ak.to_numpy(ak.flatten(Q2_reco_bc)) + + E_c = ak.to_numpy(ak.flatten(lamE_j)) + px_c = ak.to_numpy(ak.flatten(px_j)) + py_c = ak.to_numpy(ak.flatten(py_j)) + pz_c = ak.to_numpy(ak.flatten(pz_j)) + + if xB_truth_c.size == 0: + continue + + # Lambda-derived xL from reconstructed 4-vector + + pL_par = px_c * nhat[0] + py_c * nhat[1] + pz_c * nhat[2] + xL = (E_c + pL_par) / pplus_p + + # Reconstructed xK using selected reconstructed inclusive kinematics + + xK_reco_c = xk_from_xb_xl(xB_reco_c, xL) + + ok = ( + np.isfinite(xB_truth_c) & np.isfinite(Q2_truth_c) & + np.isfinite(xB_reco_c) & np.isfinite(Q2_reco_c) & + np.isfinite(xK_reco_c) & + (Q2_truth_c > 0.0) & (Q2_reco_c > 0.0) + ) + + if tmax is not None: + dE = ppE - E_c + dpx = ppx - px_c + dpy = ppy - py_c + dpz = ppz - pz_c + t = dE * dE - (dpx * dpx + dpy * dpy + dpz * dpz) + tneg = -t + ok &= np.isfinite(tneg) & (tneg < tmax) + + if not np.any(ok): + continue + + # efficiency numerator in truth bins + + xB_truth_sel = xB_truth_c[ok] + Q2_truth_sel = Q2_truth_c[ok] + + Nreco_lam += np.histogram2d( + xB_truth_sel, + Q2_truth_sel, + bins=(xB_edges, Q2_edges), + )[0] + + # (xB_truth, Q2_truth) -> (xK_reco, Q2_reco) + + xK_sel = xK_reco_c[ok] + Q2_reco_sel = Q2_reco_c[ok] + + ixB_truth = np.searchsorted(xB_edges, xB_truth_sel, side="right") - 1 + iQ_truth = np.searchsorted(Q2_edges, Q2_truth_sel, side="right") - 1 + ixK_reco = np.searchsorted(xK_edges, xK_sel, side="right") - 1 + iQ_reco = np.searchsorted(Q2_edges, Q2_reco_sel, side="right") - 1 + + good = ( + (ixB_truth >= 0) & (ixB_truth < xB_bins) & + (iQ_truth >= 0) & (iQ_truth < Q2_bins) & + (ixK_reco >= 0) & (ixK_reco < xK_bins) & + (iQ_reco >= 0) & (iQ_reco < Q2_bins) + ) + + np.add.at( + R, + ( + ixK_reco[good], + iQ_reco[good], + ixB_truth[good], + iQ_truth[good], + ), + 1.0, + ) + + except Exception as e: + print(f"[plot_relerr_kaon_sf_xK_Q2] Warning {fpath.name}: {e}") + + # Truth-space efficiency and expected yields + + eps = efficiency_from_counts(Nreco_lam, Ngen_evt) + print( + "CHECK. Nreco_lam =", Nreco_lam.sum(), + "Ngen_evt =", Ngen_evt.sum(), + "eps_mean =", np.nanmean(eps), + ) + + Nexp_xB_Q2 = expected_yields_nb(sigma_gen_nb, eps, lumin_fb) + + # truth-space relative uncertainty (for diagnostic only) + + relerr_xB_Q2 = np.full_like(Nexp_xB_Q2, np.nan, dtype=np.float64) + mask_truth = Nexp_xB_Q2 > 0 + relerr_xB_Q2[mask_truth] = 100.0 / np.sqrt(Nexp_xB_Q2[mask_truth]) + + # reconstructed projections + + P = _response_probability_xkq2reco_given_xbq2truth(R) + denom_R = np.sum(R, axis=(0, 1)) + populated = denom_R > 0 + if np.any(populated): + check = np.sum(P, axis=(0, 1)) + max_dev = np.max(np.abs(check[populated] - 1.0)) + print(f"CHECK. max deviation of response normalization = {max_dev:.3e}") + + Nexp_xK_Q2 = _project_xbq2truth_to_xkq2reco(P, Nexp_xB_Q2) + + # final uncertainty map + + relerr_xK_Q2 = np.full_like(Nexp_xK_Q2, np.nan, dtype=np.float64) + mask_reco = Nexp_xK_Q2 > 0 + relerr_xK_Q2[mask_reco] = 100.0 / np.sqrt(Nexp_xK_Q2[mask_reco]) + + q2_tag = "logQ2" if log_Q2 else "linQ2" + + # plotting + + plot_relerr_map( + relerr_xB_Q2, + xB_edges, + Q2_edges, + xlabel=r"$x_B^{\rm truth}$", + ylabel=r"$Q^{2,\,\rm truth}\ (\mathrm{GeV}^2)$", + cbarlabel=r"Relative uncertainty on $F_2^K$ (%)", + title=rf"{beam}, $\mathcal{{L}}$={lumin_fb:g} fb$^{{-1}}$, truth normalization", + outpath=outdir / f"relerr_xBtruth_Q2truth_{beam}_{tag}_L{lumin_fb:g}fb_{q2_tag}_kin{kinmethod}.png", + log_y=log_Q2, + log_x=log_x, + log_z=False, + vmax_percent=vmax_percent, + ) + + plot_relerr_map( + relerr_xK_Q2, + xK_edges, + Q2_edges, + xlabel=rf"$x_K^{{\rm reco}} \simeq x_B^{{{kinmethod}}}/(1-x_\Lambda)$", + ylabel=rf"$Q^{{2,\,\rm reco}}_{{{kinmethod}}}\ (\mathrm{{GeV}}^2)$", + cbarlabel=r"Relative uncertainty on $F_2^K$ (%)", + title=rf"{beam}, $\mathcal{{L}}$={lumin_fb:g} fb$^{{-1}}$, {kinmethod} kinematics", + outpath=outdir / f"relerr_xKreco_Q2reco_{beam}_{tag}_L{lumin_fb:g}fb_{q2_tag}_kin{kinmethod}.png", + log_y=log_Q2, + log_x=log_x, + log_z=False, + vmax_percent=vmax_percent, + ) + + return relerr_xB_Q2, relerr_xK_Q2, dict( + eps_truth=eps, + Ngen_evt_truth=Ngen_evt, + Nreco_lam_truth=Nreco_lam, + Nexp_xB_Q2_truth=Nexp_xB_Q2, + Nexp_xK_Q2_reco=Nexp_xK_Q2, + sigma_gen_nb=sigma_gen_nb, + sigma_tot_nb=sigma_tot_nb, + response_counts=R, + response_prob=P, + kinmethod=kinmethod, + ) + +# high-level plotting : models with attached error + +def plot_F2K_xK_slices_with_attached_errors( + relerr_xK_Q2: np.ndarray, + Nexp_xK_Q2: np.ndarray, + xK_edges: np.ndarray, + Q2_edges: np.ndarray, + q2_slices: list[tuple[float, float]] | None = None, + eval_cfg=None, + outdir: str | Path = "./outputs", + outname: str = "F2K_xK_slices_with_errors.png", + title: str | None = None, + logx: bool = False, + logy: bool = False, + ymin: float | None = None, + ymax: float | None = None, + xK_range: tuple[float, float] = (0, 1), + Q2_range: tuple[float, float] = (1, 500), +): + """ + Build F2K(xK) curves for Q2 slices and attach "artificial" uncertainties: + dF2 = F2 * relerr(xK,Q2_slice) + + - F2K comes from eval.py model (toy x-shape × pion Q2 evolution). + - relerr_xK_Q2 is your propagated relative uncertainty map in percent. + - In each Q2 slice, we compute a yield-weighted average: + F2_slice(xK) = sum_i w_i * F2(xK,Q2_i) / sum_i w_i + (relerr_slice)^2 = sum_i w_i * (relerr_i/100)^2 / sum_i w_i + with w_i = Nexp_xK_Q2(xK,Q2_i). + """ + + apply_mpl_style() + outdir = ensure_outdir(outdir) + + if q2_slices is None: + q2_slices = [(0.0, 100.0), (100.0, 200.0), (200.0, 500.0)] + + if eval_cfg is None: + eval_cfg = EvalConfig() + + print("[DEBUG] Nexp_xK_Q2: finite frac =", np.isfinite(Nexp_xK_Q2).mean(), + "min =", np.nanmin(Nexp_xK_Q2), "max =", np.nanmax(Nexp_xK_Q2)) + print("[DEBUG] relerr_xK_Q2: finite frac =", np.isfinite(relerr_xK_Q2).mean(), + "min% =", np.nanmin(relerr_xK_Q2), "max% =", np.nanmax(relerr_xK_Q2)) + + setup_lhapdf_env() + if not is_lhapdf_available(): + print("WARNING: LHAPDF not available -> cannot compute F2K model for slice plot.") + return None + + if eval_cfg is None: + eval_cfg = EvalConfig() + + pion_set = get_pion_pdfset() + pion_member = get_pion_pdfmember() + + pdf = make_kaon_pdf_func( + model="toy_baseline", + toy=eval_cfg.toy, + kparams=KaonModelParams(), + pion_set=pion_set, + pion_member=pion_member, + require_lhapdf=True, + ) + + def F2_point(x: float, Q2: float) -> float: + return float(F2_from_pdf_func(pdf, float(x), float(Q2))) + + def F2_on_grid(x_arr: np.ndarray, Q2_arr: np.ndarray) -> np.ndarray: + """ + Return F2(x_i, Q2_j) as array shape (len(x_arr), len(Q2_arr)). + """ + x_arr = np.asarray(x_arr, dtype=np.float64) + Q2_arr = np.asarray(Q2_arr, dtype=np.float64) + out = np.empty((x_arr.size, Q2_arr.size), dtype=np.float64) + for j, q2 in enumerate(Q2_arr): + # evaluate all x for this q2 + out[:, j] = np.fromiter((F2_point(x, q2) for x in x_arr), count=x_arr.size, dtype=np.float64) + return out + + # slicing + + xK_cent = 0.5 * (xK_edges[:-1] + xK_edges[1:]) + Q2_cent = 0.5 * (Q2_edges[:-1] + Q2_edges[1:]) + + fig, ax = plt.subplots(figsize=(5, 5)) + + for islc, (q2min, q2max) in enumerate(q2_slices): + + selQ = (Q2_cent >= q2min) & (Q2_cent < q2max) + idxQ = np.where(selQ)[0] + if idxQ.size == 0: + continue + + F2_x_q = F2_on_grid(xK_cent, Q2_cent[idxQ]) + W = Nexp_xK_Q2[:, idxQ].astype(float) + good = np.isfinite(F2_x_q) & np.isfinite(W) & (W > 0) + W_eff = np.where(good, W, 0.0) + F2_eff = np.where(good, F2_x_q, 0.0) + Ntot_slice = np.sum(W_eff, axis=1) + F2_slice = np.full_like(xK_cent, np.nan, dtype=float) + rel_slice = np.full_like(xK_cent, np.nan, dtype=float) + validx = Ntot_slice > 0 + if np.any(validx): + F2_slice[validx] = np.sum(W_eff[validx] * F2_eff[validx], axis=1) / Ntot_slice[validx] + rel_slice[validx] = 1.0 / np.sqrt(Ntot_slice[validx]) + yerr = F2_slice * rel_slice + + # sanity check + med_rel = np.nanmedian(100 * rel_slice) if np.any(np.isfinite(rel_slice)) else np.nan + med_ye = np.nanmedian(yerr) if np.any(np.isfinite(yerr)) else np.nan + print(f"[DEBUG slice {q2min:g}-{q2max:g}] valid xK frac = {validx.mean():.3f} " + f"median rel% = {med_rel} median yerr = {med_ye}") + + lab = rf"${q2min:g} 0) + if np.any(okp): + ax.errorbar( + xK_cent[okp], + F2_slice[okp], + yerr=yerr[okp], + fmt="s-", + linewidth=1.5, + markersize=4, + capsize=2.5, + label=lab, + ) + + ax.set_xlabel(r"$x_K$") + ax.set_ylabel(r"$F_2^K(x_K)$ with projected uncertainties") + ax.set_xlim(xK_range[0], xK_range[1]) + ax.set_ylim(0,1) + + if title is None: + title = r"$F_2^K(x_K)$ with attached uncertainties from propagated $(x_K,Q^2)$ rel. errors" + ax.set_title(title) + + if logx: + ax.set_xscale("log") + if logy: + ax.set_yscale("log") + + if ymin is not None or ymax is not None: + ax.set_ylim(bottom=ymin, top=ymax) + + ax.legend(frameon=False) + fig.tight_layout() + + outpath = outdir / outname + savefig(fig, outpath, dpi=200) + return outpath + +# high level plotting : chi2 between models + +def plot_discriminant_kaon_sf_xK_Q2( + *, + beam: str, + nfiles: int, + root_base_dir: Path, + suffix: str, + gen_base_dir: Path, + outdir: Path, + tag: str, + xB_bins: int = 20, + xB_range: tuple[float, float] = (0.0, 1.0), + xK_bins: int = 20, + xK_range: tuple[float, float] = (0.0, 2.0), + Q2_bins: int = 20, + Q2_range: tuple[float, float] = (1.0, 500.0), + log_Q2: bool = True, + tmax: float | None = None, + lumin_fb: float = 1.0, + vmax: float | None = None, + modelA: str = "toy_baseline", + modelB: str = "toy_hard_valence", + toyA: ToyParams = ToyParams(), + toyB: ToyParams | None = None, + kparamsA: KaonModelParams = KaonModelParams(), + kparamsB: KaonModelParams | None = None, + sigma_ref: str = "A", # "A" | "B" | "mean" +): + """ + Build and plot a per-bin discrimination metric in (xK, Q2): + chi2_ij = (F2_A - F2_B)^2 / sigma_ij^2 + + where sigma_ij is built from the projected relative uncertainty relerr_xK_Q2 [%] + returned by plot_relerr_kaon_sf_xK_Q2: + sigma_ij = (relerr_ij/100) * F2_ref_ij + with F2_ref = F2_A (default), or F2_B, or (F2_A+F2_B)/2. + + Returns: chi2_map, relerr_map, (F2_A, F2_B) + """ + apply_mpl_style() + outdir = ensure_outdir(outdir) + + relerr_xB_Q2, relerr_xK_Q2, extra = plot_relerr_kaon_sf_xK_Q2( + beam=beam, + nfiles=nfiles, + root_base_dir=root_base_dir, + suffix=suffix, + gen_base_dir=gen_base_dir, + outdir=outdir, + tag=tag, + xB_bins=xB_bins, + xB_range=xB_range, + xK_bins=xK_bins, + xK_range=xK_range, + Q2_bins=Q2_bins, + Q2_range=Q2_range, + log_Q2=log_Q2, + tmax=tmax, + lumin_fb=lumin_fb, + ) + + # binning + + xK_edges = np.linspace(xK_range[0], xK_range[1], xK_bins + 1) + if log_Q2: + Q2_edges = np.logspace(np.log10(Q2_range[0]), np.log10(Q2_range[1]), Q2_bins + 1) + else: + Q2_edges = np.linspace(Q2_range[0], Q2_range[1], Q2_bins + 1) + + xK_cent = 0.5 * (xK_edges[:-1] + xK_edges[1:]) + if log_Q2: + Q2_cent = np.sqrt(Q2_edges[:-1] * Q2_edges[1:]) + else: + Q2_cent = 0.5 * (Q2_edges[:-1] + Q2_edges[1:]) + + # model pdf + + from config import setup_lhapdf_env, get_pion_pdfset, get_pion_pdfmember + setup_lhapdf_env() + pion_set = get_pion_pdfset() + pion_member = get_pion_pdfmember() + + toyB = toyB or toyA + kparamsB = kparamsB or kparamsA + + pdfA = make_kaon_pdf_func( + model=modelA, toy=toyA, kparams=kparamsA, + pion_set=pion_set, pion_member=pion_member, require_lhapdf=True + ) + pdfB = make_kaon_pdf_func( + model=modelB, toy=toyB, kparams=kparamsB, + pion_set=pion_set, pion_member=pion_member, require_lhapdf=True + ) + + # pdf mapping + + F2A = np.full((xK_bins, Q2_bins), np.nan, dtype=np.float64) + F2B = np.full((xK_bins, Q2_bins), np.nan, dtype=np.float64) + + for ix, x in enumerate(xK_cent): + for iQ, Q2 in enumerate(Q2_cent): + F2A[ix, iQ] = F2_from_pdf_func(pdfA, float(x), float(Q2)) + F2B[ix, iQ] = F2_from_pdf_func(pdfB, float(x), float(Q2)) + + # relative error mapping + + rel = relerr_xK_Q2 / 100.0 + + if sigma_ref == "A": + Fref = F2A + elif sigma_ref == "B": + Fref = F2B + elif sigma_ref == "mean": + Fref = 0.5 * (F2A + F2B) + else: + raise ValueError("sigma_ref must be 'A', 'B', or 'mean'") + + sigma = rel * Fref + + # sanity check + + chi2 = np.full_like(F2A, np.nan, dtype=np.float64) + ok = np.isfinite(F2A) & np.isfinite(F2B) & np.isfinite(sigma) & (sigma > 0) & np.isfinite(rel) & (rel > 0) + chi2[ok] = (F2A[ok] - F2B[ok])**2 / (sigma[ok]**2) + + # plot chi2 maps + + q2_tag = "logQ2" if log_Q2 else "linQ2" + title = rf"{beam}, $\mathcal{{L}}$={lumin_fb:g} fb$^{{-1}}$: {modelA} vs {modelB}" + + plot_relerr_map( + np.sqrt(chi2), + xK_edges, + Q2_edges, + xlabel=r"$x_K$", + ylabel=r"$Q^2\ (\mathrm{GeV}^2)$", + cbarlabel=r"Discrimination power $\sqrt{\chi^2}$ with projected uncertainties", + title=title, + outpath=outdir / f"chi2_xK_Q2_{beam}_{tag}_{modelA}_vs_{modelB}_L{lumin_fb:g}fb_{q2_tag}.png", + log_y=log_Q2, + vmax_percent=vmax, + ) + + return chi2, relerr_xK_Q2, (F2A, F2B) + +# models ratio computation + +def plot_F2K_ratio_xK_slices_to_reference( + relerr_xK_Q2: np.ndarray, + Nexp_xK_Q2: np.ndarray, + xK_edges: np.ndarray, + Q2_edges: np.ndarray, + ref_model: str, + models_to_compare: list[str], + q2_slices: list[tuple[float, float]] | None = None, + eval_cfg=None, + outdir: str | Path = "./outputs", + outname: str = "Discri_ratio_xK_slices.png", + title: str | None = None, + logx: bool = False, + ymin: float | None = None, + ymax: float | None = None, + xK_range=(0, 1), + Q2_range=(1, 500), +): + """ + Plot F2K(xK) ratios in Q2 slices: + ratio_model(xK) = F2_model(xK) / F2_ref(xK) + + The reference model appears as a horizontal line at 1. + The projected experimental relative uncertainty is shown as a band around 1: + 1 ± relerr_slice + + For each Q2 slice, both the ratio and the uncertainty band are built using + the same yield-weighted averaging as in plot_F2K_xK_slices_with_attached_errors. + """ + + apply_mpl_style() + outdir = ensure_outdir(outdir) + + if q2_slices is None: + q2_slices = [(0.0, 100.0), (100.0, 200.0), (200.0, 500.0)] + + if eval_cfg is None: + eval_cfg = EvalConfig() + + print("[DEBUG] Nexp_xK_Q2: finite frac =", np.isfinite(Nexp_xK_Q2).mean(), + "min =", np.nanmin(Nexp_xK_Q2), "max =", np.nanmax(Nexp_xK_Q2)) + print("[DEBUG] relerr_xK_Q2: finite frac =", np.isfinite(relerr_xK_Q2).mean(), + "min% =", np.nanmin(relerr_xK_Q2), "max% =", np.nanmax(relerr_xK_Q2)) + + setup_lhapdf_env() + if not is_lhapdf_available(): + print("WARNING: LHAPDF not available -> cannot compute F2K ratio slice plot.") + return None + + pion_set = get_pion_pdfset() + pion_member = get_pion_pdfmember() + + # helpers + + def _compute_model_map(model_name: str, x_arr: np.ndarray, q2_arr: np.ndarray) -> np.ndarray: + """ + Return F2(x_i, Q2_j) as array shape (len(x_arr), len(q2_arr)). + Works for both PDF-based models and direct JAM F2K replicas. + """ + x_arr = np.asarray(x_arr, dtype=np.float64) + q2_arr = np.asarray(q2_arr, dtype=np.float64) + out = np.full((x_arr.size, q2_arr.size), np.nan, dtype=np.float64) + + jam_replica_map = { + "jam25_rep165": 165, + "jam25_rep390": 390, + "jam25_rep400": 400, + } + + if model_name == "jam25_mean": + f2k = make_f2k_jam_mean() + + for j, q2 in enumerate(q2_arr): + for i, x in enumerate(x_arr): + out[i, j] = float(f2k(float(x), float(q2))) + + return out + + if model_name in jam_replica_map: + f2k = make_f2k_jam_replica(jam_replica_map[model_name]) + + for j, q2 in enumerate(q2_arr): + for i, x in enumerate(x_arr): + out[i, j] = float(f2k(float(x), float(q2))) + + return out + + pdf = make_kaon_pdf_func( + model=model_name, + toy=eval_cfg.toy, + kparams=KaonModelParams(), + pion_set=pion_set, + pion_member=pion_member, + require_lhapdf=True, + ) + + for j, q2 in enumerate(q2_arr): + for i, x in enumerate(x_arr): + out[i, j] = float(F2_from_pdf_func(pdf, float(x), float(q2))) + + return out + + def _weighted_q2_slice(A2d: np.ndarray, W2d: np.ndarray) -> np.ndarray: + good = np.isfinite(A2d) & np.isfinite(W2d) & (W2d > 0) + Aeff = np.where(good, A2d, 0.0) + Weff = np.where(good, W2d, 0.0) + denom = np.sum(Weff, axis=1) + out = np.full(A2d.shape[0], np.nan, dtype=float) + ok = denom > 0 + if np.any(ok): + out[ok] = np.sum(Weff[ok] * Aeff[ok], axis=1) / denom[ok] + return out + + all_models = [ref_model] + [m for m in models_to_compare if m != ref_model] + + # binning + + xK_cent = 0.5 * (xK_edges[:-1] + xK_edges[1:]) + Q2_cent = 0.5 * (Q2_edges[:-1] + Q2_edges[1:]) + + # Q2 slice pannel + + nslices = len(q2_slices) + fig, axes = plt.subplots( + nslices, 1, + figsize=(5.0, 5.0 * nslices), + sharex=True, + squeeze=False, + ) + axes = axes[:, 0] + + FONT = 12 + + model_colors = { + "jam25_rep165": "#5E3C99", + "jam25_rep400": "#B24C7C", + "jam25_rep390": "#3B5BA9", + "toy_baseline": "#1F3A5F", + "toy_soft_valence": "#3B5BA9", + "toy_hard_valence": "#B24C7C", + "toy_sea_enhanced": "#5E3C99", + "toy_su3_breaking": "#2E8B57", + "cosmao22": "#2E8B57", + } + + model_markers = { + "jam25_rep390": "o", + "jam25_rep400": "s", + "cosmao22": "^", + } + + for ax, (q2min, q2max) in zip(axes, q2_slices): + selQ = (Q2_cent >= q2min) & (Q2_cent < q2max) + idxQ = np.where(selQ)[0] + if idxQ.size == 0: + ax.set_visible(False) + continue + + W = Nexp_xK_Q2[:, idxQ].astype(float) + good_rw = np.isfinite(W) & (W > 0) + W_eff_rw = np.where(good_rw, W, 0.0) + Ntot_slice = np.sum(W_eff_rw, axis=1) + rel_slice = np.full_like(xK_cent, np.nan, dtype=float) + validx_rw = Ntot_slice > 0 + if np.any(validx_rw): + rel_slice[validx_rw] = 1.0 / np.sqrt(Ntot_slice[validx_rw]) + + # reference model slice + F2_ref_x_q = _compute_model_map(ref_model, xK_cent, Q2_cent[idxQ]) + good_ref = np.isfinite(F2_ref_x_q) & np.isfinite(W) & (W > 0) + W_eff_ref = np.where(good_ref, W, 0.0) + F2_ref_eff = np.where(good_ref, F2_ref_x_q, 0.0) + sumW_ref = np.sum(W_eff_ref, axis=1) + + F2_ref_slice = np.full_like(xK_cent, np.nan, dtype=float) + validx_ref = sumW_ref > 0 + + if np.any(validx_ref): + F2_ref_slice[validx_ref] = ( + np.sum(W_eff_ref[validx_ref] * F2_ref_eff[validx_ref], axis=1) + / sumW_ref[validx_ref] + ) + + if ref_model=="jam25_mean": + replica_ids = get_all_jam_replica_ids("JAM25kaon_nlonll_F2K") + + Frep = evaluate_f2k_for_jam_replicas( + replica_ids=replica_ids, + x_arr=xK_cent, + q2_arr=Q2_cent[idxQ], + setname="JAM25kaon_nlonll_F2K", + ) + + mean_rep = np.nanmean(Frep, axis=0) + std_rep = np.nanstd(Frep, axis=0, ddof=1) + + mean_slice = _weighted_q2_slice(mean_rep, W) + std_slice = _weighted_q2_slice(std_rep, W) + + spread_ratio = np.full_like(xK_cent, np.nan, dtype=float) + ok_spread = ( + np.isfinite(mean_slice) + & np.isfinite(std_slice) + & np.isfinite(F2_ref_slice) + & (F2_ref_slice > 0) + ) + + spread_ratio[ok_spread] = std_slice[ok_spread] / F2_ref_slice[ok_spread] + + ax.fill_between( + xK_cent[ok_spread], + 1.0 - spread_ratio[ok_spread], + 1.0 + spread_ratio[ok_spread], + color="#868E96", + alpha=0.28, + linewidth=0.0, + #label=r"JAM replicas $1\sigma$ spread" if ax is axes[0] else None, + zorder=0, + ) + + # horizontal reference line + ax.axhline( + 1.0, + linestyle="--", + linewidth=1.5, + label=f"{ref_model}" if ax is axes[0] else None, + color='black' + ) + + # compare requested models to reference + for model_name in models_to_compare: + + F2_mod_x_q = _compute_model_map(model_name, xK_cent, Q2_cent[idxQ]) + good_mod = np.isfinite(F2_mod_x_q) & np.isfinite(W) & (W > 0) + W_eff_mod = np.where(good_mod, W, 0.0) + F2_mod_eff = np.where(good_mod, F2_mod_x_q, 0.0) + sumW_mod = np.sum(W_eff_mod, axis=1) + + F2_mod_slice = np.full_like(xK_cent, np.nan, dtype=float) + validx_mod = sumW_mod > 0 + if np.any(validx_mod): + F2_mod_slice[validx_mod] = ( + np.sum(W_eff_mod[validx_mod] * F2_mod_eff[validx_mod], axis=1) + / sumW_mod[validx_mod] + ) + + ratio = np.full_like(xK_cent, np.nan, dtype=float) + ok_ratio = ( + np.isfinite(F2_mod_slice) + & np.isfinite(F2_ref_slice) + & (F2_ref_slice > 0) + ) + ratio[ok_ratio] = F2_mod_slice[ok_ratio] / F2_ref_slice[ok_ratio] + + yerr = ratio * rel_slice + okp = np.isfinite(ratio) & np.isfinite(yerr) & (yerr >= 0) + if np.any(okp): + ax.errorbar( + xK_cent[okp], + ratio[okp], + yerr=yerr[okp], + fmt="-" + model_markers.get(model_name, "o"), + color=model_colors.get(model_name, "#222222"), + ecolor=model_colors.get(model_name, "#222222"), + linewidth=2.0 if model_name == "cosmao22" else 1.8, + elinewidth=2.0 if model_name == "cosmao22" else 1.8, + markersize=6 if model_name == "cosmao22" else 5, + capsize=2.5, + capthick=1.1, + label=( + "CoSMAO22 (DSE)" if model_name == "cosmao22" + else model_name + ), + ) + + ax.set_ylabel("Ratio", fontsize=FONT) + ax.set_xlim(xK_range[0], xK_range[1]) + + if ymin is not None or ymax is not None: + ax.set_ylim(bottom=ymin, top=ymax) + + ax.legend( + title=rf"${q2min:g} np.ndarray: + good = np.isfinite(A2d) & np.isfinite(W2d) & (W2d > 0) + Aeff = np.where(good, A2d, 0.0) + Weff = np.where(good, W2d, 0.0) + denom = np.sum(Weff, axis=1) + + out = np.full(A2d.shape[0], np.nan, dtype=float) + ok = denom > 0 + if np.any(ok): + out[ok] = np.sum(Weff[ok] * Aeff[ok], axis=1) / denom[ok] + return out + + post_results = {} + + for L in lumin_list: + relerr_xB_Q2, relerr_xK_Q2, extra = plot_relerr_kaon_sf_xK_Q2( + beam=beam, + nfiles=nfiles, + root_base_dir=root_base_dir, + suffix=suffix, + gen_base_dir=gen_base_dir, + outdir=outdir, + tag=tag, + xB_bins=xB_bins, + xB_range=xB_range, + xK_bins=xK_bins, + xK_range=xK_range, + Q2_bins=Q2_bins, + Q2_range=Q2_range, + log_x=logx, + log_Q2=log_Q2, + tmax=tmax, + lumin_fb=L, + vmax_percent=100.0, + kinmethod=kinmethod, + ) + + Nexp_xK_Q2 = extra["Nexp_xK_Q2_reco"] + + sigma = (relerr_xK_Q2 / 100.0) * F_data + valid = np.isfinite(F_data) & np.isfinite(sigma) & (sigma > 0) + + chi2 = np.full(Frep_total.shape[0], np.nan, dtype=float) + for k in range(Frep_total.shape[0]): + ok = valid & np.isfinite(Frep_total[k]) + if np.any(ok): + chi2[k] = np.sum(((Frep_total[k][ok] - F_data[ok]) / sigma[ok]) ** 2) + + chi2_min = np.nanmin(chi2) + w = np.exp(-0.5 * (chi2 - chi2_min)) + w[~np.isfinite(w)] = 0.0 + + if np.sum(w) <= 0: + raise RuntimeError(f"All JAM replica weights vanished for L = {L:g} fb^-1.") + + w /= np.sum(w) + + mean_post = np.nansum(w[:, None, None] * Frep_plot, axis=0) + var_post = np.nansum( + w[:, None, None] * (Frep_plot - mean_post[None, :, :])**2, + axis=0, + ) + std_post = np.sqrt(np.maximum(var_post, 0.0)) + + post_results[L] = { + "relerr_xK_Q2": relerr_xK_Q2, + "Nexp_xK_Q2": Nexp_xK_Q2, + "weights": w, + "chi2": chi2, + "mean_post": mean_post, + "std_post": std_post, + } + + prior_color = "#495057" + base_colors = [ + "#A5D8FF", + "#4DABF7", + "#1C7ED6", + "#1864AB", + ] + + post_colors = { + L: base_colors[i % len(base_colors)] + for i, L in enumerate(lumin_list) + } + + nslices = len(q2_slices) + fig, axes = plt.subplots( + nslices, + 1, + figsize=(5, 5 * nslices), + sharex=True, + squeeze=False, + ) + axes = axes[:, 0] + + for ax, (q2min, q2max) in zip(axes, q2_slices): + idxQ = np.where((Q2_cent >= q2min) & (Q2_cent < q2max))[0] + + if idxQ.size == 0: + ax.set_visible(False) + continue + + W_prior = post_results[lumin_list[-1]]["Nexp_xK_Q2"][:, idxQ].astype(float) + + prior_mean_slice = weighted_q2_slice(mean_prior[:, idxQ], W_prior) + prior_std_slice = weighted_q2_slice(std_prior[:, idxQ], W_prior) + + ok_prior = np.isfinite(prior_mean_slice) & np.isfinite(prior_std_slice) + + if np.any(ok_prior): + ax.fill_between( + xK_cent[ok_prior], + prior_mean_slice[ok_prior] - prior_std_slice[ok_prior], + prior_mean_slice[ok_prior] + prior_std_slice[ok_prior], + color=prior_color, + alpha=0.35, + linewidth=0.0, + label="Before EIC", + zorder=1, + ) + + for iL, L in enumerate(lumin_list): + res = post_results[L] + W = res["Nexp_xK_Q2"][:, idxQ].astype(float) + + post_mean_slice = weighted_q2_slice(res["mean_post"][:, idxQ], W) + post_std_slice = weighted_q2_slice(res["std_post"][:, idxQ], W) + + ok_post = np.isfinite(post_mean_slice) & np.isfinite(post_std_slice) + + if np.any(ok_post): + ax.fill_between( + xK_cent[ok_post], + post_mean_slice[ok_post] - post_std_slice[ok_post], + post_mean_slice[ok_post] + post_std_slice[ok_post], + color=post_colors[L], + alpha=0.6, + linewidth=0.0, + label=rf"After EIC, ${L:g}\,\mathrm{{fb}}^{{-1}}$", + zorder=2 + iL, + ) + + ax.set_ylabel(ylabel, fontsize=13) + + ax.set_xlim(xK_range) + ax.set_ylim(0.0, 1.0) + ax.set_yticks([0.2, 0.4, 0.6, 0.8]) + + if logx: + ax.set_xscale("log") + ax.xaxis.set_major_locator(FixedLocator([1e-3, 1e-2, 1e-1, 1.0])) + ax.xaxis.set_major_formatter( + FixedFormatter([r"$10^{-3}$", r"$10^{-2}$", r"$10^{-1}$", r"$1$"]) + ) + ax.xaxis.set_minor_locator(LogLocator(base=10.0, subs=np.arange(2, 10) * 0.1)) + ax.xaxis.set_minor_formatter(NullFormatter()) + + ax.tick_params( + axis="x", + which="major", + direction="in", + bottom=True, + top=False, + length=7, + width=1.2, + labelsize=13, + ) + ax.tick_params( + axis="x", + which="minor", + direction="in", + bottom=True, + top=False, + length=4, + width=1.0, + ) + ax.tick_params( + axis="y", + which="major", + direction="in", + left=True, + right=True, + length=7, + width=1.2, + labelsize=13, + ) + ax.tick_params( + axis="y", + which="minor", + direction="in", + left=True, + right=True, + length=4, + width=1.0, + ) + + for spine in ax.spines.values(): + spine.set_linewidth(1.1) + + ax.legend( + title=rf"${q2min:g} < Q^2 < {q2max:g}\,\mathrm{{GeV}}^2$", + loc="upper right", + frameon=False, + fontsize=11, + title_fontsize=11, + ) + + axes[-1].set_xlabel(r"$x_K$", fontsize=13) + + if title is not None: + fig.suptitle( + title, + fontsize=13, + x=0.545, + y=0.985, + ha="center", + ) + + fig.subplots_adjust( + left=0.14, + right=0.95, + top=0.86, + bottom=0.12, + ) + + outpath = outdir / outname + savefig(fig, outpath, dpi=250) + + return { + "outpath": outpath, + "replica_ids": replica_ids, + "lumin_list": lumin_list, + "component": component, + "post_results": post_results, + "mean_prior": mean_prior, + "std_prior": std_prior, + "F_data": F_data, + } + +# high level plotting : reduction JAM replicas + +def plot_neff_reduction_vs_lumi_by_beam( + results_by_beam: dict, + *, + outdir: Path, + outname: str, + title: str | None = None, +): + apply_mpl_style() + outdir = ensure_outdir(outdir) + + beam_colors = { + "5x41": "#2B8A3E", + "10x100": "#1C7ED6", + "18x275": "#C92A2A", + } + + beam_markers = { + "5x41": "o", + "10x100": "s", + "18x275": "^", + } + + fig, ax = plt.subplots(figsize=(5.2,5.2)) + + all_reductions = [] + + for beam, res in results_by_beam.items(): + replica_ids = res["replica_ids"] + lumin_list = np.asarray(res["lumin_list"], dtype=float) + post_results = res["post_results"] + + nrep = len(replica_ids) + + reductions = [] + + for L in lumin_list: + w = np.asarray(post_results[L]["weights"], dtype=float) + w = w / np.sum(w) + + neff = 1.0 / np.sum(w**2) + reduction = 100.0 * (1.0 - neff / nrep) + + reductions.append(reduction) + + reductions = np.asarray(reductions) + all_reductions.extend(reductions[np.isfinite(reductions)]) + + ax.plot( + lumin_list, + reductions, + marker=beam_markers.get(beam, "o"), + markersize=6.5, + linewidth=2.2, + color=beam_colors.get(beam, None), + label=beam, + ) + + ax.set_xlabel(r"Luminosity (fb$^{-1}$)", fontsize=13) + ax.set_ylabel(r"Replica reduction (%)", fontsize=13) + + ax.set_xscale("log") + ax.set_xticks([5, 10, 50]) + ax.set_xticklabels([r"$5$", r"$10$", r"$50$"]) + + ymin = -2.0 + ymax = max(all_reductions) if len(all_reductions) else 5.0 + ax.set_ylim(ymin, 1.25 * ymax + 1.0) + + # Style ticks + ax.tick_params( + axis="both", + which="major", + direction="in", + top=True, + right=True, + length=6, + width=1.1, + labelsize=12, + ) + + ax.tick_params( + axis="both", + which="minor", + direction="in", + top=True, + right=True, + length=3, + width=0.9, + ) + + for spine in ax.spines.values(): + spine.set_linewidth(1.1) + + ax.legend( + frameon=False, + fontsize=12, + loc="upper left", + ) + + if title is not None: + fig.suptitle(title, fontsize=13, y=0.945) + + fig.subplots_adjust(left=0.16, right=0.96, top=0.84, bottom=0.16) + + outpath = outdir / outname + savefig(fig, outpath, dpi=250) + + return outpath \ No newline at end of file diff --git a/analysis/multicalo-lambda/studies_lambda.py b/analysis/multicalo-lambda/lambda_studies.py similarity index 81% rename from analysis/multicalo-lambda/studies_lambda.py rename to analysis/multicalo-lambda/lambda_studies.py index c124bcc..2e12c0c 100644 --- a/analysis/multicalo-lambda/studies_lambda.py +++ b/analysis/multicalo-lambda/lambda_studies.py @@ -5,11 +5,11 @@ import awkward as ak import matplotlib.pyplot as plt import matplotlib as mpl - -from config import PhysicsConstants, Paths, DEFAULT_CONST, DEFAULT_PATHS +import uproot +from config import PhysicsConstants, Paths, CONST, PATHS from utils import read_lambda_geant4, read_lambda_eicrecon, read_lambda_afterburner, iter_reco_files from physics import direct_energy_window, E_to_L, L_to_E, proton_kinematics_for_beam, xk_from_xb_xl -from plotting import apply_mpl_style, ensure_outdir, savefig +from plotting import apply_mpl_style, ensure_outdir, savefig, get_color, get_style def plot_lambda_spectra( @@ -21,7 +21,7 @@ def plot_lambda_spectra( suffix: str, outdir: Path, afterburner_pattern: str, - c: PhysicsConstants = DEFAULT_CONST, + c: PhysicsConstants = CONST, logy: bool = True, ) -> Path: apply_mpl_style() @@ -41,23 +41,31 @@ def plot_lambda_spectra( fig, ax = plt.subplots(figsize=(5, 5)) - if E_reco.size: - ax.hist( - E_reco, bins=bins, range=(0, xmax), - histtype="step", linewidth=2.0, - label=f"Reco ({E_reco.size:.1e})", - ) if E_geant4.size: ax.hist( E_geant4, bins=bins, range=(0, xmax), histtype="step", + color=get_color("geant4"), label=rf"Geant4 ({n_geant4_label:.1e})", + **get_style("geant4"), ) + if E_after.size: ax.hist( E_after, bins=bins, range=(0, xmax), - histtype="step", linestyle="--", + histtype="step", + color=get_color("afterburner"), label="Afterburner", + **get_style("afterburner"), + ) + + if E_reco.size: + ax.hist( + E_reco, bins=bins, range=(0, xmax), + histtype="step", + color=get_color("reco"), + label=f"Reco ({E_reco.size:.1e})", + **get_style("reco"), ) ax.set_xlabel("Energy (GeV)") @@ -80,7 +88,7 @@ def plot_angle_distrib( outdir: Path, afterburner_pattern: str, energies: tuple[str, ...] = ("5x41", "10x100", "18x275"), - c: PhysicsConstants = DEFAULT_CONST, + c: PhysicsConstants = CONST, ) -> Path: apply_mpl_style() outdir = ensure_outdir(outdir) @@ -93,12 +101,17 @@ def plot_angle_distrib( _, cos_after, sinphi_after = read_lambda_afterburner(beam=beam, nfiles=nfiles, pattern=afterburner_pattern, pid=c.pid_lambda) ax = axes[row, 0] + if cos_true.size: - ax.hist(cos_true, bins=bins, range=(-1, 1), histtype="step", label="Geant4") + ax.hist(cos_true, bins=bins, range=(-1, 1), histtype="step", + color=get_color("geant4"), label="Geant4", **get_style("geant4")) if cos_reco.size: - ax.hist(cos_reco, bins=bins, range=(-1, 1), histtype="step", linewidth=2.5, label="Reco") + ax.hist(cos_reco, bins=bins, range=(-1, 1), histtype="step", + color=get_color("reco"), label="Reco", **get_style("reco")) if cos_after.size: - ax.hist(cos_after, bins=bins, range=(-1, 1), histtype="step", linestyle="--", label="Afterburner") + ax.hist(cos_after, bins=bins, range=(-1, 1), histtype="step", + color=get_color("afterburner"), label="Afterburner", **get_style("afterburner")) + ax.set_xlabel(r"$\cos\theta_\Lambda$") ax.set_ylabel("Counts") ax.set_yscale("log") @@ -106,17 +119,21 @@ def plot_angle_distrib( ax.legend(frameon=False) ax = axes[row, 1] + if sinphi_true.size: - ax.hist(sinphi_true, bins=bins, range=(-1, 1), histtype="step", label="Geant4") + ax.hist(sinphi_true, bins=bins, range=(-1, 1), histtype="step", + color=get_color("geant4"), label="Geant4", **get_style("geant4")) if sinphi_reco.size: - ax.hist(sinphi_reco, bins=bins, range=(-1, 1), histtype="step", linewidth=2.5, label="Reco") + ax.hist(sinphi_reco, bins=bins, range=(-1, 1), histtype="step", + color=get_color("reco"), label="Reco", **get_style("reco")) if sinphi_after.size: - ax.hist(sinphi_after, bins=bins, range=(-1, 1), histtype="step", linestyle="--", label="Afterburner") + ax.hist(sinphi_after, bins=bins, range=(-1, 1), histtype="step", + color=get_color("afterburner"), label="Afterburner", **get_style("afterburner")) ax.set_xlabel(r"$\sin\phi_\Lambda$") ax.set_ylabel("Counts") ax.set_yscale("log") ax.set_title(beam) - ax.legend(frameon=False) + #ax.legend(frameon=False) fig.tight_layout() out = outdir / f"angledistrib_lambda_cos_sinphi_{suffix}.png" @@ -129,8 +146,9 @@ def efficiency_vs_energy( bins: int, root_base_dir: Path, suffix: str, - c: PhysicsConstants = DEFAULT_CONST, + c: PhysicsConstants = CONST, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + Emin, Emax = direct_energy_window(beam) E_truth, _, _ = read_lambda_geant4(beam, nfiles=nfiles, root_base_dir=root_base_dir, suffix=suffix, c=c) E_reco, _, _ = read_lambda_eicrecon(beam, nfiles=nfiles, root_base_dir=root_base_dir, suffix=suffix) @@ -167,8 +185,9 @@ def plot_efficiencies( outdir: Path, with_beta: bool, tag: str, - c: PhysicsConstants = DEFAULT_CONST, + c: PhysicsConstants = CONST, ) -> Path: + apply_mpl_style() outdir = ensure_outdir(outdir) @@ -184,25 +203,55 @@ def plot_efficiencies( scale = (100.0 / c.beta) if with_beta else 100.0 - ax.errorbar(c5, scale * e5, yerr=scale * s5, fmt="s", markersize=4, linewidth=1, label="5x41 (new)") - ax.errorbar(c5o, scale * e5o, yerr=scale * s5o, fmt="x", markersize=4, linewidth=1, label="5x41 (ZDC only)") - - ax.errorbar(c10, scale * e10, yerr=scale * s10, fmt="s", markersize=4, linewidth=1, label="10x100 (new)") - ax.errorbar(c10o, scale * e10o, yerr=scale * s10o, fmt="x", markersize=4, linewidth=1, label="10x100 (ZDC only)") - - ax.errorbar(c18, scale * e18, yerr=scale * s18, fmt="s", markersize=4, linewidth=1, label="18x275 (new)") - ax.errorbar(c18o, scale * e18o, yerr=scale * s18o, fmt="x", markersize=4, linewidth=1, label="18x275 (ZDC only)") + ax.errorbar( + c5, scale * e5, yerr=scale * s5, + fmt="s", markersize=4, linewidth=1, + color="green", ecolor="green", + label="5x41 (new)" + ) + + ax.errorbar( + c5o, scale * e5o, yerr=scale * s5o, + fmt="x", markersize=4, linewidth=1, + color="green", ecolor="green", + label="5x41 (ZDC only)" + ) + + ax.errorbar( + c10, scale * e10, yerr=scale * s10, + fmt="s", markersize=4, linewidth=1, + color="blue", ecolor="blue", + label="10x100 (new)" + ) + + ax.errorbar( + c10o, scale * e10o, yerr=scale * s10o, + fmt="x", markersize=4, linewidth=1, + color="blue", ecolor="blue", + label="10x100 (ZDC only)" + ) + + ax.errorbar( + c18, scale * e18, yerr=scale * s18, + fmt="s", markersize=4, linewidth=1, + color="red", ecolor="red", + label="18x275 (new)" + ) + + ax.errorbar( + c18o, scale * e18o, yerr=scale * s18o, + fmt="x", markersize=4, linewidth=1, + color="red", ecolor="red", + label="18x275 (ZDC only)" + ) ax.set_xlabel(r"$\Lambda^0$ energy: $E_\Lambda$ (GeV)") ax.set_ylabel(r"Reconstruction efficiency $\epsilon(E_\Lambda)$ (%)") - secax = ax.secondary_xaxis("top", functions=(lambda x: E_to_L(x, c), lambda x: L_to_E(x, c))) secax.set_xticks([5, 10, 15, 20]) secax.set_xlabel(r"Average decay vertex: $\langle z \rangle_\Lambda$ (m)") - - ax.legend(frameon=False) + ax.legend(frameon=False, loc="upper left") fig.tight_layout() - outname = f"epsilon_vs_energy_{tag}" + ("_beta" if with_beta else "") + ".png" return savefig(fig, outdir / outname, dpi=300) @@ -225,13 +274,17 @@ def plot_nlambda_vs_kinematics_all( t_bins: int = 50, t_range: tuple[float, float] | None = None, tmax: float | None = None, - c: PhysicsConstants = DEFAULT_CONST, + c: PhysicsConstants = CONST, + kinmethod: str = 'Truth' ) -> None: apply_mpl_style() outdir = ensure_outdir(outdir) - xB_branch = "InclusiveKinematicsTruth.x" - Q2_branch = "InclusiveKinematicsTruth.Q2" + if kinmethod not in ('Truth', 'Electron', 'JB', 'Sigma', 'DA', 'ML'): + raise ValueError(f"Invalid kinmethod: {kinmethod}. Choose Truth, Electron, JB, Sigma, DA, or ML.") + + xB_branch = f"InclusiveKinematics{kinmethod}.x" + Q2_branch = f"InclusiveKinematics{kinmethod}.Q2" lamE_branch = "ReconstructedFarForwardLambdas.energy" lampx_branch = "ReconstructedFarForwardLambdas.momentum.x" lampy_branch = "ReconstructedFarForwardLambdas.momentum.y" @@ -355,7 +408,7 @@ def q2_transform(Q2arr: np.ndarray): ax.set_ylabel(y_label) ax.set_title(rf"$\Lambda$ yield vs $(x_B,Q^2)$ at {beam}") fig.tight_layout() - savefig(fig, outdir / f"kinematics_xB_Q2_{beam}_{tag}.png", dpi=300) + savefig(fig, outdir / f"kinematics{kinmethod}_xB_Q2_{beam}_{tag}.png", dpi=300) if has_cand: Q2k_plot, mk = q2_transform(Q2_cand) @@ -376,4 +429,4 @@ def q2_transform(Q2arr: np.ndarray): ax.set_ylabel(y_label) ax.set_title(rf"$\Lambda$ candidates vs $(x_K,Q^2)$ at {beam}" + extra) fig.tight_layout() - savefig(fig, outdir / f"kinematics_xK_Q2_{beam}_{tag}.png", dpi=300) \ No newline at end of file + savefig(fig, outdir / f"kinematics{kinmethod}_xK_Q2_{beam}_{tag}.png", dpi=300) \ No newline at end of file diff --git a/analysis/multicalo-lambda/physics.py b/analysis/multicalo-lambda/physics.py index 6931880..928f985 100644 --- a/analysis/multicalo-lambda/physics.py +++ b/analysis/multicalo-lambda/physics.py @@ -1,9 +1,23 @@ from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Callable + import numpy as np -from config import PhysicsConstants, DEFAULT_CONST + +from config import PhysicsConstants, CONST +# ============================================================================= +# Beam helpers +# ============================================================================= + def parse_proton_energy(beam_str: str) -> float: + """ + Extract proton energy (GeV) from a beam string like '5x41', '10x100', '18x275'. + If the string doesn't contain 'x', it is interpreted as a float. + """ s = str(beam_str) if "x" in s: return float(s.split("x")[-1]) @@ -11,6 +25,9 @@ def parse_proton_energy(beam_str: str) -> float: def direct_energy_window(beam: str) -> tuple[float, float]: + """ + Energy window (GeV) used for direct-energy selections in afterburner studies. + """ if beam == "5x41": return 20.0, 41.0 if beam == "10x100": @@ -20,34 +37,55 @@ def direct_energy_window(beam: str) -> tuple[float, float]: raise ValueError(f"Unknown beam config: {beam}") -def E_to_L(E_GeV: np.ndarray | float, c: PhysicsConstants = DEFAULT_CONST) -> np.ndarray: +# ============================================================================= +# Lambda/proton kinematics utilities +# ============================================================================= + +def E_to_L(E_GeV: np.ndarray | float, c: PhysicsConstants = CONST) -> np.ndarray: + """ + Convert Lambda energy E (GeV) to average decay length in meters: + = (p/m) * c*tau = ctau * sqrt(E^2 - m^2) / m + """ E = np.asarray(E_GeV, dtype=np.float64) - m = c.m_lambda_gev - # = (sqrt(E^2 - m^2)/m) * c*tau - return c.ctau_lambda_m * np.sqrt(np.maximum(E * E - m * m, 0.0)) / m + m = float(c.m_lambda_gev) + return float(c.ctau_lambda_m) * np.sqrt(np.maximum(E * E - m * m, 0.0)) / m -def L_to_E(L_m: np.ndarray | float, c: PhysicsConstants = DEFAULT_CONST) -> np.ndarray: +def L_to_E(L_m: np.ndarray | float, c: PhysicsConstants = CONST) -> np.ndarray: + """ + Inverse of E_to_L: + E = sqrt((m*L/ctau)^2 + m^2) + """ L = np.asarray(L_m, dtype=np.float64) - m = c.m_lambda_gev - # E = sqrt((m*L/ctau)^2 + m^2) - return np.sqrt((m * L / c.ctau_lambda_m) ** 2 + m * m) + m = float(c.m_lambda_gev) + return np.sqrt((m * L / float(c.ctau_lambda_m)) ** 2 + m * m) -def proton_kinematics_for_beam(beam: str, c: PhysicsConstants = DEFAULT_CONST) -> dict[str, float | np.ndarray]: +def proton_kinematics_for_beam(beam: str, c: PhysicsConstants = CONST) -> dict[str, float | np.ndarray]: """ - Returns: - - Ep, p_mag, nhat, pplus_p - - proton 4-vector components (ppx,ppy,ppz,ppE) for t computation + Compute basic proton kinematics given a beam string (e.g. '18x275'). + + Returns a dict with: + Ep : proton energy (GeV) + p_mag : proton momentum magnitude (GeV) + nhat : unit direction vector (3,) + pplus_p : E + |p| (GeV) + ppx,ppy,ppz,ppE : proton 4-vector components (GeV) """ Ep = parse_proton_energy(beam) - p_mag = np.sqrt(max(Ep * Ep - c.mp_gev * c.mp_gev, 0.0)) - th = c.theta_proton_rad - nhat = np.array([np.sin(th), 0.0, np.cos(th)], dtype=np.float64) + mp = float(c.mp_gev) + p_mag = math.sqrt(max(Ep * Ep - mp * mp, 0.0)) + + th = float(c.theta_proton_rad) + sin_th = math.sin(th) + cos_th = math.cos(th) + + nhat = np.array([sin_th, 0.0, cos_th], dtype=np.float64) pplus_p = Ep + p_mag - ppx = p_mag * np.sin(th) + + ppx = p_mag * sin_th ppy = 0.0 - ppz = p_mag * np.cos(th) + ppz = p_mag * cos_th ppE = Ep return dict( @@ -63,8 +101,85 @@ def proton_kinematics_for_beam(beam: str, c: PhysicsConstants = DEFAULT_CONST) - def xk_from_xb_xl(xb: np.ndarray, xl: np.ndarray) -> np.ndarray: + """ + Compute x_K from x_B and x_L: + x_K = x_B / (1 - x_L) + + Returns NaN where undefined/invalid. + """ + xb = np.asarray(xb, dtype=np.float64) + xl = np.asarray(xl, dtype=np.float64) + denom = 1.0 - xl xk = np.full_like(xb, np.nan, dtype=np.float64) - ok = np.isfinite(xb) & np.isfinite(denom) & (denom > 0) + + ok = np.isfinite(xb) & np.isfinite(denom) & (denom > 0.0) xk[ok] = xb[ok] / denom[ok] - return xk \ No newline at end of file + return xk + +# ============================================================================= +# Kaon SF model parameters +# ============================================================================= + +@dataclass(frozen=True) +class ToyParams: + a: float = 0.5 + b: float = 1.5 + Q0: float = 0.5 + +@dataclass(frozen=True) +class KaonModelParams: + db_hard_valence: float = -0.5 + db_soft_valence: float = +0.5 + da_sbar_su3_break: float = +0.2 + db_sbar_su3_break: float = -0.4 + sea_amp: float = 0.05 + sea_a: float = 0.2 + sea_b: float = 4.0 + +# ============================================================================= +# Kaon structure function (LO) +# ============================================================================= + +# e_f^2 +_E2 = { + 1: 1.0 / 9.0, # d + 2: 4.0 / 9.0, # u + 3: 1.0 / 9.0, # s + 4: 4.0 / 9.0, # c + 5: 1.0 / 9.0, # b + 6: 4.0 / 9.0, # t +} + + +def beta_func(p: float, q: float) -> float: + """Euler beta function: B(p,q) = Γ(p) Γ(q) / Γ(p+q).""" + return math.gamma(p) * math.gamma(q) / math.gamma(p + q) + + +def toy_norm(a: float, b: float) -> float: + """Normalization so that ∫_0^1 norm * x^a (1-x)^b dx = 1.""" + return 1.0 / beta_func(a + 1.0, b + 1.0) + + +def toy_shape_x(x: float, a: float, b: float, norm: float) -> float: + """Normalized x-shape: norm * x^a (1-x)^b.""" + if x <= 0.0 or x >= 1.0: + return 0.0 + return norm * (x**a) * ((1.0 - x) ** b) + + +def F2_from_pdf_func(pdf_func: Callable[[int, float, float], float], x: float, Q2: float) -> float: + """ + Generic LO F2: + F2(x,Q2) = x * Σ_f e_f^2 [ q_f(x,Q2) + qbar_f(x,Q2) ] + using flavors 1..6. + + pdf_func(pid, x, Q2) must return f(x,Q2) (not x*f). + """ + s = 0.0 + for flav in range(1, 7): + q = pdf_func(+flav, x, Q2) + qb = pdf_func(-flav, x, Q2) + s += _E2[flav] * (q + qb) + return x * s \ No newline at end of file diff --git a/analysis/multicalo-lambda/plotting.py b/analysis/multicalo-lambda/plotting.py index 1e54669..55b1bfd 100644 --- a/analysis/multicalo-lambda/plotting.py +++ b/analysis/multicalo-lambda/plotting.py @@ -1,8 +1,52 @@ from __future__ import annotations + from pathlib import Path +from typing import Any, Final + import matplotlib as mpl import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import LogNorm + +from config import PATHS +from utils import Grid + +from typing import Literal + +# ============================================================================= +# Plot style registry +# ============================================================================= + +Scale = Literal["lin", "log"] + +COLORS: Final[dict[str, str]] = { + "geant4": "orange", + "reco": "tab:blue", + "afterburner": "black", +} + +STYLES: Final[dict[str, dict[str, Any]]] = { + "geant4": {"linestyle": "-", "linewidth": 2.0}, + "reco": {"linestyle": "-", "linewidth": 2.5}, + "afterburner": {"linestyle": "--", "linewidth": 2.0}, +} + + +def get_color(which: str) -> str: + try: + return COLORS[which] + except KeyError as e: + raise KeyError(f"Unknown color key '{which}'. Available: {sorted(COLORS)}") from e + + +def get_style(which: str) -> dict[str, Any]: + # styles are optional; unknown keys just return empty style dict + return dict(STYLES.get(which, {})) + +# ============================================================================= +# Global Matplotlib style +# ============================================================================= def apply_mpl_style() -> None: mpl.rcParams.update( @@ -15,15 +59,186 @@ def apply_mpl_style() -> None: ) +# ============================================================================= +# File helpers +# ============================================================================= + def ensure_outdir(outdir: str | Path) -> Path: p = Path(outdir) p.mkdir(parents=True, exist_ok=True) return p -def savefig(fig, path: str | Path, dpi: int = 300) -> Path: +def savefig(fig: mpl.figure.Figure, path: str | Path, dpi: int = 300) -> Path: p = Path(path) p.parent.mkdir(parents=True, exist_ok=True) fig.savefig(p, dpi=dpi) plt.close(fig) - return p \ No newline at end of file + return p + + +# ============================================================================= +# Plots +# ============================================================================= + +def plot_relerr_map( + Z: np.ndarray, + x_edges: np.ndarray, + y_edges: np.ndarray, + xlabel: str, + ylabel: str, + cbarlabel: str, + title: str, + outpath: str | Path, + *, + log_y: bool = False, + log_x: bool = False, + log_z: bool = False, + vmax_percent: float | None = None +) -> Path: + """ + Plot a relative statistical uncertainty map (in percent). + NOTE: expects Z to be shaped like (nx, ny) and will plot Z.T to match edges. + """ + fig, ax = plt.subplots(figsize=(6, 5)) + + Zm = np.ma.masked_invalid(Z) + Zm = np.ma.masked_where(Zm < 0, Zm) + + if log_z: + norm = LogNorm(vmin=np.nanmin(Zm), vmax=np.nanmax(Zm)) + pcm = ax.pcolormesh(x_edges, y_edges, Zm.T, shading="auto", cmap="coolwarm", norm=norm) + if vmax_percent is not None: + pcm.set_clim(vmin=1.0, vmax=float(vmax_percent)) + + else: + pcm = ax.pcolormesh(x_edges, y_edges, Zm.T, shading="auto", cmap="coolwarm") + if vmax_percent is not None: + pcm.set_clim(vmin=0.0, vmax=float(vmax_percent)) + + cb = fig.colorbar(pcm, ax=ax) + cb.set_label(cbarlabel) + + ax.set_xlabel(xlabel) + ax.set_ylabel(ylabel) + + if log_y: + ax.set_yscale("log") + + if log_x: + ax.set_xscale("log") + + ax.set_title(title) + fig.tight_layout() + return savefig(fig, outpath, dpi=300) + + +def plot_F2_map( + prefix: str, + grid: Grid, + F2: np.ndarray, + title: str, + cbar_label: str, + *, + xscale: Scale = "lin", + q2scale: Scale = "lin", + dpi: int = 220, +) -> Path: + """ + Plot F2 on (x, Q2) edges stored in Grid and save into PATHS.outputs. + """ + fig, ax = plt.subplots(figsize=(6.6, 5.6)) + + pcm = ax.pcolormesh(grid.x_edges, grid.q2_edges, F2, shading="auto", cmap="viridis") + cbar = fig.colorbar(pcm, ax=ax) + cbar.set_label(cbar_label) + + ax.set_xlabel(r"$x_K$") + ax.set_ylabel(r"$Q^2$ (GeV$^2$)") + ax.set_title(title) + + if xscale == "log": + ax.set_xscale("log") + ax.set_xlim(0.1,1) + + if q2scale == "log": + ax.set_yscale("log") + + fig.tight_layout() + + out_png = PATHS.outputs / f"{prefix}.png" + return savefig(fig, out_png, dpi=dpi) + +def plot_fischer_matrix_maps( + *, + info_aa: np.ndarray, + info_ab: np.ndarray, + info_bb: np.ndarray, + xK_range: tuple[float, float] = (0.0, 1.0), + xK_bins: int = 10, + Q2_range: tuple[float, float] = (1.0, 500.0), + Q2_bins: int = 10, + log_Q2: bool = True, + outdir: str | Path = "./outputs", + outname: str = "fischer_matrix_maps.png", + title: str | None = None, +): + """ + Plot Fisher information density maps arranged as a 2x2 matrix: + + [ f_aa f_ab ] + [ f_ba f_bb ] + + where each element is a map in (xK, Q2). + """ + + apply_mpl_style() + outdir = ensure_outdir(outdir) + + info_ba = info_ab + + xK_edges = np.linspace(xK_range[0], xK_range[1], xK_bins + 1) + + if log_Q2: + Q2_edges = np.logspace(np.log10(Q2_range[0]), np.log10(Q2_range[1]), Q2_bins + 1) + else: + Q2_edges = np.linspace(Q2_range[0], Q2_range[1], Q2_bins + 1) + + fig, axes = plt.subplots(2, 2, figsize=(9, 8), constrained_layout=True) + + panels = [ + (axes[0,0], info_aa, r"$f_{aa}(x_K,Q^2)$", "viridis"), + (axes[0,1], info_ab, r"$f_{ab}(x_K,Q^2)$", "viridis"), + (axes[1,0], info_ba, r"$f_{ba}(x_K,Q^2)$", "viridis"), + (axes[1,1], info_bb, r"$f_{bb}(x_K,Q^2)$", "viridis"), + ] + + for ax, z, lab, cmap in panels: + + m = ax.pcolormesh( + xK_edges, + Q2_edges, + z.T, + shading="auto", + cmap=cmap, + vmin=np.nanmin(z), + vmax=np.nanmax(z), + ) + + cb = fig.colorbar(m, ax=ax) + cb.set_label("Local Fisher information") + + ax.set_xlabel(r"$x_K$") + ax.set_ylabel(r"$Q^2$ (GeV$^2$)") + ax.set_title(lab) + + if log_Q2: + ax.set_yscale("log") + + if title is not None: + fig.suptitle(title) + + outpath = Path(outdir) / outname + savefig(fig, outpath, dpi=200) + + return outpath \ No newline at end of file diff --git a/analysis/multicalo-lambda/run.py b/analysis/multicalo-lambda/run.py index 3aa5e32..a5bcf61 100644 --- a/analysis/multicalo-lambda/run.py +++ b/analysis/multicalo-lambda/run.py @@ -2,36 +2,84 @@ import argparse from pathlib import Path +import numpy as np -from config import DEFAULT_CONST, DEFAULT_PATHS -from studies_lambda import ( +from config import CONST, PATHS + +from lambda_studies import ( plot_angle_distrib, plot_lambda_spectra, plot_efficiencies, plot_nlambda_vs_kinematics_all, ) -from studies_kaon_sf import plot_relerr_kaon_sf_xK_Q2 +from kaon_models import ( + evaluate_and_plot_F2K, + EvalConfig +) +from kaon_studies import ( + plot_relerr_kaon_sf_xK_Q2, + plot_F2K_xK_slices_with_attached_errors, + plot_discriminant_kaon_sf_xK_Q2, + plot_ratio_kaon_sf_xK_Q2, + plot_jam_reweighting_constraining_power_xK_slices, + plot_neff_reduction_vs_lumi_by_beam +) +from uq import ( + compute_local_fischer_maps, + compute_local_fischer, +) +from plotting import ( + apply_mpl_style, + get_color, get_style, + plot_fischer_matrix_maps +) + +KAON_MODELS = [ + # toy models + "toy_baseline", + "toy_soft_valence", + "toy_hard_valence", + "toy_sea_enhanced", + "toy_su3_breaking", + # JAM models + "jam25_rep165", + "jam25_rep390", + "jam25_rep400", + "jam25_mean", + # theoretical models + "cosmao22", + ] def build_parser() -> argparse.ArgumentParser: p = argparse.ArgumentParser(description="Multi-calo Lambda / Kaon-SF studies") - p.add_argument("--inputs", type=Path, default=DEFAULT_PATHS.inputs) - p.add_argument("--outputs", type=Path, default=DEFAULT_PATHS.outputs) + p.add_argument("--inputs", type=Path, default=PATHS.inputs) + p.add_argument("--outputs", type=Path, default=PATHS.outputs) p.add_argument("--suffix", type=str, default="testall_calocalib_verification") p.add_argument("--nfiles", type=int, default=1) - p.add_argument("--bins", type=int, default=50) - p.add_argument("--beam", type=str, default="5x41", choices=list(DEFAULT_CONST.energies)) + p.add_argument("--bins", type=int, default=10) + p.add_argument("--beam", type=str, default="5x41", choices=list(CONST.energies)) p.add_argument("--task", type=str, default="all", - choices=["all", "angles", "spectra", "eff", "kin", "relerr"]) + choices=["all", + "angles", + "spectra", + "eff", + "kin", + "models", + "relerr", + "discri", + "toymodels", + "replicas"]) p.add_argument("--with-beta", action="store_true") p.add_argument("--tmax", type=float, default=None) p.add_argument("--lumin-fb", type=float, default=1.0) p.add_argument("--logQ2", action="store_true") - p.add_argument("--gen-dir", type=Path, default=DEFAULT_PATHS.gen_base_dir) + p.add_argument("--gen-dir", type=Path, default=PATHS.gen_base_dir) return p def main() -> None: + args = build_parser().parse_args() inputs = args.inputs @@ -41,10 +89,12 @@ def main() -> None: nbins = args.bins beam = args.beam - after_pattern = DEFAULT_PATHS.afterburner_pattern + after_pattern = PATHS.afterburner_pattern # 1) angles + if args.task in ("angles", "all"): + out = plot_angle_distrib( nfiles=nfiles, bins=30, @@ -52,13 +102,16 @@ def main() -> None: suffix=suffix, outdir=outputs, afterburner_pattern=after_pattern, - energies=tuple(DEFAULT_CONST.energies), + energies=tuple(CONST.energies), ) - print("Saved:", out) + print("Saved angular distributions") # 2) spectra + if args.task in ("spectra", "all"): - for b in DEFAULT_CONST.energies: + + for b in CONST.energies: + xmax = float(b.split("x")[-1]) * 2.0 out = plot_lambda_spectra( beam=b, @@ -71,10 +124,12 @@ def main() -> None: afterburner_pattern=after_pattern, logy=True, ) - print("Saved:", out) + print(f"Saved Lambda spectra | {b}") # 3) efficiencies + if args.task in ("eff", "all"): + out = plot_efficiencies( nfiles_5x41=nfiles, nfiles_10x100=nfiles, @@ -89,53 +144,316 @@ def main() -> None: with_beta=args.with_beta, tag=suffix, ) - print("Saved:", out) + print(f"Saved efficiencies maps") # 4) kinematics + if args.task in ("kin", "all"): - for b in DEFAULT_CONST.energies: - plot_nlambda_vs_kinematics_all( - beam=b, - nfiles=nfiles, - root_base_dir=inputs, - suffix=suffix, - outdir=outputs, - tag=suffix, - xB_bins=nbins, - xB_range=(0, 2), - xK_bins=nbins, - xK_range=(0, 2), - Q2_bins=nbins, - Q2_range=(1, 3e3), - log_Q2=True, - logz=True, - tmax=args.tmax, - ) - print("Saved kinematics for", b) - # 5) relerr kaon SF + for b in CONST.energies: + + for kinme in ['Truth', 'Electron', 'JB']: + + plot_nlambda_vs_kinematics_all( + beam=b, + nfiles=nfiles, + root_base_dir=inputs, + suffix=suffix, + outdir=outputs, + tag=suffix, + xB_bins=nbins, + xB_range=(0, 1), + xK_bins=nbins, + xK_range=(0, 1), + Q2_bins=nbins, + Q2_range=(1, 500), + log_Q2=False, + logz=True, + tmax=args.tmax, + kinmethod=kinme + ) + print(f"Saved kinematics maps | {b}") + + # 5) kaon SF evaluation + + if args.task in ("models", "all"): + + for m in KAON_MODELS: + + out_eval = evaluate_and_plot_F2K(EvalConfig(model=m)) + + if out_eval: + print("Saved:", out_eval) + + # 6) relerr kaon SF + if args.task in ("relerr", "all"): - for b in DEFAULT_CONST.energies: - plot_relerr_kaon_sf_xK_Q2( - beam=b, - nfiles=nfiles, - root_base_dir=inputs, - suffix=suffix, - gen_base_dir=args.gen_dir, - outdir=outputs, - tag=suffix, - xK_bins=nbins, - xK_range=(0, 3.0), - Q2_bins=nbins, - Q2_range=(1.0, 550.0), - xB_bins=nbins, - xB_range=(0.0, 1.0), - lumin_fb=args.lumin_fb, - log_Q2=args.logQ2, - tmax=args.tmax, - ) - print("Saved relerr for", b) + for b in CONST.energies: + + for kime in ['Truth', 'Electron', 'JB']: + + # relative error heatmap + + relerr_xB_Q2, relerr_xK_Q2, meta = plot_relerr_kaon_sf_xK_Q2( + beam=b, + nfiles=nfiles, + root_base_dir=inputs, + suffix=suffix, + gen_base_dir=args.gen_dir, + outdir=outputs, + tag=suffix, + xK_bins=nbins, + xK_range=(1e-3, 1.0), + Q2_bins=nbins, + Q2_range=(1, 100.0), + xB_bins=nbins, + xB_range=(1e-3, 1.0), + lumin_fb=args.lumin_fb, + log_Q2=True,#args.logQ2, + log_x=True, + tmax=args.tmax, + vmax_percent=100.0, + kinmethod=kime, + ) + print(f"Saved relative error maps | {b}") + + # fischer matrix and covariance + + # res_fisher = compute_local_fischer_maps( + # a0=0.5, + # b0=1.5, + # relerr_xK_Q2=relerr_xK_Q2, + # xK_range=(0, 1), + # xK_bins=nbins, + # Q2_range=(1, 500), + # Q2_bins=nbins, + # model="toy_baseline", + # da=1e-2, + # db=1e-2, + # log_Q2_centers=False, #args.logQ2, + # ) + + # plot_fischer_matrix_maps( + # info_aa=res_fisher["info_aa"], + # info_ab=res_fisher["info_ab"], + # info_bb=res_fisher["info_bb"], + # xK_range=(0, 1), + # xK_bins=nbins, + # Q2_range=(1, 500), + # Q2_bins=nbins, + # log_Q2=False, #args.logQ2, + # outdir=outputs, + # outname=f"fisher_matrix_maps_{b}_L{args.lumin_fb}fb.png", + # title=r"Local Fisher information for toy-model around $\theta=(a,b)=(0.5,1.5)$", + # ) + + # 7) models discrimination heatmaps + + if args.task in ("discri", "all"): + + for b in CONST.energies: + + for modelb in KAON_MODELS: + + chi2, relerr_xK_Q2, (F2A, F2B) = plot_discriminant_kaon_sf_xK_Q2( + beam=b, + nfiles=nfiles, + root_base_dir=inputs, + suffix=suffix, + gen_base_dir=args.gen_dir, + outdir=outputs, + tag=suffix, + xK_bins=nbins, + xK_range=(0.0, 1.0), + Q2_bins=nbins, + Q2_range=(1.0, 500.0), + xB_bins=nbins, + xB_range=(0.0, 1.0), + lumin_fb=args.lumin_fb, + log_Q2=args.logQ2, + tmax=args.tmax, + modelA="toy_baseline", + modelB=modelb, + sigma_ref="mean", + vmax=3, + ) + print(f"Saved chi2 map | toy_baseline vs. {modelb} | {b}") + + # 8) toy models study + + if args.task in ("toymodels", "all"): + + for b in CONST.energies: + + for slicing in [(10,100)]: + + for kime in ['Truth', 'Electron', 'JB']: + + # toy models + plot_ratio_kaon_sf_xK_Q2( + beam=b, + nfiles=nfiles, + root_base_dir=inputs, + suffix=suffix, + gen_base_dir=args.gen_dir, + outdir=outputs, + tag=suffix, + ref_model="toy_baseline", + models_to_compare=[ + "toy_soft_valence", + "toy_hard_valence", + "toy_sea_enhanced", + "toy_su3_breaking", + ], + q2_slices=[slicing], + xK_bins=nbins, + xK_range=(1e-3, 1.0), + Q2_bins=nbins, + Q2_range=(1e-3, 500.0), + xB_bins=nbins, + xB_range=(1e-3, 1.0), + log_Q2=args.logQ2, + logx=True, + tmax=args.tmax, + lumin_fb=args.lumin_fb, + title = f"{b} GeV | {args.lumin_fb} fb$^{{-1}}$ | {kime} kinematics", + outname=f"Ratio_toymodels_{b}_{args.lumin_fb}fb_Q2_{slicing}_{kime}.png", + ymin=0., + ymax=2.5, + kinmethod=kime + ) + print(f"Saved ratio toy models | Q2 {slicing} @ {b} | {kime} kinematicss") + + # 9) literature models (JAM, DSE) + + if args.task in ("replicas", "all"): + + components = ["total", "valence", "sea"] + + component_title = { + "total": "Total F2K", + "valence": "Valence F2K", + "sea": "Quark-sea", + "gluon": "Gluon-sea", + } + + for slicing in [(10, 100)]: + + for kime in ["Electron", "JB"]: + + reweighting_results_by_beam = {} + + for b in CONST.energies: + + # discrimination JAM replicas + + plot_ratio_kaon_sf_xK_Q2( + beam=b, + nfiles=nfiles, + root_base_dir=inputs, + suffix=suffix, + gen_base_dir=args.gen_dir, + outdir=outputs, + tag=suffix, + ref_model="jam25_mean", + models_to_compare=["jam25_rep400"], + q2_slices=[slicing], + xK_bins=nbins, + xK_range=(1e-3, 1.0), + Q2_bins=nbins, + Q2_range=(1e-3, 500.0), + xB_bins=nbins, + xB_range=(1e-3, 1.0), + log_Q2=args.logQ2, + logx=True, + tmax=args.tmax, + lumin_fb=args.lumin_fb, + title=f"{b} GeV | {args.lumin_fb} fb$^{{-1}}$ | {kime} kinematics", + outname=f"Ratio_JAM-JAM_{b}_{args.lumin_fb}fb_Q2_{slicing}_{kime}.png", + ymin=0.0, + ymax=2.5, + kinmethod=kime, + ) + + print(f"Saved ratio replicas | Q2 {slicing} @ {b} | {kime} kinematics") + + # discrimination JAM/DSE + + plot_ratio_kaon_sf_xK_Q2( + beam=b, + nfiles=nfiles, + root_base_dir=inputs, + suffix=suffix, + gen_base_dir=args.gen_dir, + outdir=outputs, + tag=suffix, + ref_model="jam25_mean", + models_to_compare=["cosmao22"], + q2_slices=[slicing], + xK_bins=nbins, + xK_range=(1e-3, 1.0), + Q2_bins=nbins, + Q2_range=(1e-3, 500.0), + xB_bins=nbins, + xB_range=(1e-3, 1.0), + log_Q2=args.logQ2, + logx=True, + tmax=args.tmax, + lumin_fb=args.lumin_fb, + title=f"{b} GeV | {args.lumin_fb} fb$^{{-1}}$ | {kime} kinematics", + outname=f"Ratio_JAM-DSE_{b}_{args.lumin_fb}fb_Q2_{slicing}_{kime}.png", + ymin=0.0, + ymax=2.5, + kinmethod=kime, + ) + + print(f"Saved ratio DSE | Q2 {slicing} @ {b} | {kime} kinematics") + + # constraining power JAM replicas + + # for comp in components: + + # reweighting_res = plot_jam_reweighting_constraining_power_xK_slices( + # beam=b, + # nfiles=nfiles, + # root_base_dir=inputs, + # suffix=suffix, + # gen_base_dir=args.gen_dir, + # outdir=outputs, + # tag=suffix, + # q2_slices=[slicing], + # xK_bins=nbins, + # xK_range=(1e-3, 1.0), + # Q2_bins=nbins, + # Q2_range=(1e-3, 500.0), + # xB_bins=nbins, + # xB_range=(1e-3, 1.0), + # log_Q2=args.logQ2, + # logx=True, + # tmax=args.tmax, + # lumin_fb=[5, 10, 50], + # kinmethod=kime, + # component=comp, + # title=f"{component_title[comp]} | {b} GeV | {kime} kinematics", + # outname=f"Reweighting_replicas_{comp}_{b}_5-10-50fb_Q2_{slicing}_{kime}.png", + # ) + + # print(f"Saved JAM reweighting | component {comp} | Q2 {slicing} @ {b} | {kime} kinematics") + + # if comp == "total": + # reweighting_results_by_beam[b] = reweighting_res + + # plot_neff_reduction_vs_lumi_by_beam( + # reweighting_results_by_beam, + # outdir=outputs, + # title=f"$Q^2$ = {slicing[0]}–{slicing[1]} GeV$^2$ | {kime} kinematics", + # outname=f"Neff_reduction_by_beam_5-10-50fb_Q2_{slicing}_{kime}.png", + # ) + + # print(f"Saved replicas reduction | Q2 {slicing} | {kime} kinematics") + +# main if __name__ == "__main__": main() \ No newline at end of file diff --git a/analysis/multicalo-lambda/studies_kaon_sf.py b/analysis/multicalo-lambda/studies_kaon_sf.py deleted file mode 100644 index 333aaf0..0000000 --- a/analysis/multicalo-lambda/studies_kaon_sf.py +++ /dev/null @@ -1,302 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -import numpy as np -import awkward as ak -import uproot -import matplotlib.pyplot as plt -from matplotlib.colors import LogNorm - -from plotting import apply_mpl_style, ensure_outdir, savefig -from physics import proton_kinematics_for_beam, xk_from_xb_xl -from config import PhysicsConstants, DEFAULT_CONST - - -def build_sigma_map_xb_q2( - beam: str, - base_dir: str | Path, - xb_range=(0.0, 1.0), - xb_bins=10, - q2_range=(1.0, 500.0), - q2_bins=10, - chunk=2_000_000, - savefig: bool = True, - fig_dir: str | Path = "./outputs", - plot_density: bool = True, -): - Ee_str, Ep_str = beam.lower().split("x") - Ee, Ep = float(Ee_str), float(Ep_str) - - fname = f"k_lambda_crossing_0.000-{Ee:.1f}on{Ep:.1f}_x0.0001-1.0000_q1.0-500.0.root" - fpath = Path(base_dir) / fname - if not fpath.exists(): - raise FileNotFoundError(fpath) - - xb_edges = np.linspace(xb_range[0], xb_range[1], xb_bins + 1) - q2_edges = np.logspace(np.log10(q2_range[0]), np.log10(q2_range[1]), q2_bins + 1) - - sumw_map = np.zeros((xb_bins, q2_bins), dtype=np.float64) - Ngen_map = np.zeros((xb_bins, q2_bins), dtype=np.float64) - - with uproot.open(fpath) as f: - meta = f["Meta"] - evnts = f["Evnts"] - n_use = min(meta.num_entries, evnts.num_entries) - - mc0 = meta["MC"].array(entry_start=0, entry_stop=1, library="np") - Ngen_declared = int(mc0["nEvts"][0]) - - for start in range(0, n_use, chunk): - stop = min(start + chunk, n_use) - - jac = meta["Jacob"].array(entry_start=start, entry_stop=stop, library="np") - mc = meta["MC"].array(entry_start=start, entry_stop=stop, library="np") - inv = evnts["invts"].array(entry_start=start, entry_stop=stop, library="np") - - w = mc["PhSpFct"] * jac # nb - xb = inv["xBj"] - q2 = inv["Q2"] - - ok = np.isfinite(w) & np.isfinite(xb) & np.isfinite(q2) & (q2 > 0) - if not np.any(ok): - continue - - w = w[ok] - xb = xb[ok] - q2 = q2[ok] - - sumw_map += np.histogram2d(xb, q2, bins=(xb_edges, q2_edges), weights=w)[0] - Ngen_map += np.histogram2d(xb, q2, bins=(xb_edges, q2_edges))[0] - - sigma_map_nb = sumw_map / float(Ngen_declared) - sigma_tot_nb = float(sumw_map.sum()) / float(Ngen_declared) - - if savefig: - apply_mpl_style() - outdir = ensure_outdir(fig_dir) - outpath = outdir / f"sigma_xb_q2_{beam}.png" - - if plot_density: - dx = np.diff(xb_edges)[:, None] - dQ2 = np.diff(q2_edges)[None, :] - z = sigma_map_nb / (dx * dQ2) - cbar_label = r"$d\sigma/(dx_B\,dQ^2)$ [nb/GeV$^2$]" - title = rf"Sullivan $d\sigma/(dx_B dQ^2)$ — {beam}" - else: - z = sigma_map_nb - cbar_label = r"$\sigma$ [nb]" - title = rf"Sullivan $\sigma(x_B,Q^2)$ — {beam}" - - pos = z[z > 0] - norm = LogNorm(vmin=float(pos.min()), vmax=float(pos.max())) if pos.size else None - - fig, ax = plt.subplots(figsize=(6, 5)) - m = ax.pcolormesh(xb_edges, q2_edges, z.T, shading="auto", norm=norm, cmap="coolwarm") - cb = fig.colorbar(m, ax=ax) - cb.set_label(cbar_label) - ax.set_xlabel(r"$x_B$") - ax.set_ylabel(r"$Q^2$ [GeV$^2$]") - ax.set_yscale("log") - ax.set_title(title) - fig.tight_layout() - savefig(fig, outpath, dpi=200) - - return sigma_map_nb, Ngen_map, sigma_tot_nb - - -def plot_relerr_kaon_sf_xK_Q2( - beam: str, - nfiles: int, - root_base_dir: Path, - suffix: str, - gen_base_dir: Path, - outdir: Path, - tag: str, - xB_bins: int = 20, - xB_range: tuple[float, float] = (0.0, 1.0), - xK_bins: int = 20, - xK_range: tuple[float, float] = (0.0, 2.0), - Q2_bins: int = 20, - Q2_range: tuple[float, float] = (1.0, 500.0), - log_Q2: bool = True, - tmax: float | None = None, - lumin_fb: float = 1.0, - vmax_percent: float | None = None, - c: PhysicsConstants = DEFAULT_CONST, -): - """ - - sigma_gen(xB,Q2) from generator (nb/bin) - - eps(xB,Q2) from reco sample: Nreco_lambda / Ngen_evt - - project to (xK,Q2) with migration P(xK|xB,Q2) - - relerr = 100/sqrt(Nexp) [%], with Nexp = sigma*eps*L - """ - apply_mpl_style() - outdir = ensure_outdir(outdir) - - xB_edges = np.linspace(xB_range[0], xB_range[1], xB_bins + 1) - xK_edges = np.linspace(xK_range[0], xK_range[1], xK_bins + 1) - if log_Q2: - Q2_edges = np.logspace(np.log10(Q2_range[0]), np.log10(Q2_range[1]), Q2_bins + 1) - else: - Q2_edges = np.linspace(Q2_range[0], Q2_range[1], Q2_bins + 1) - - sigma_gen_nb, _, sigma_tot_nb = build_sigma_map_xb_q2( - beam=beam, - base_dir=gen_base_dir, - xb_range=xB_range, - xb_bins=xB_bins, - q2_range=Q2_range, - q2_bins=Q2_bins, - chunk=2_000_000, - savefig=False, - ) - - xB_branch = "InclusiveKinematicsTruth.x" - Q2_branch = "InclusiveKinematicsTruth.Q2" - lamE_branch = "ReconstructedFarForwardLambdas.energy" - lampx_branch = "ReconstructedFarForwardLambdas.momentum.x" - lampy_branch = "ReconstructedFarForwardLambdas.momentum.y" - lampz_branch = "ReconstructedFarForwardLambdas.momentum.z" - - pk = proton_kinematics_for_beam(beam, c=c) - nhat = pk["nhat"] - pplus_p = float(pk["pplus_p"]) - ppx, ppy, ppz, ppE = float(pk["ppx"]), float(pk["ppy"]), float(pk["ppz"]), float(pk["ppE"]) - - Ngen_evt = np.zeros((xB_bins, Q2_bins), dtype=np.float64) - Nreco_lam = np.zeros((xB_bins, Q2_bins), dtype=np.float64) - M = np.zeros((xK_bins, xB_bins, Q2_bins), dtype=np.float64) - - for i in range(1, nfiles + 1): - fpath = root_base_dir / f"k_lambda_{beam}_5000evt_{i:03d}_{suffix}.root" - if not fpath.exists(): - continue - - try: - with uproot.open(fpath) as f: - tree = f["events"] - - xB_evt = ak.to_numpy(ak.flatten(tree[xB_branch].array(library="ak"))) - Q2_evt = ak.to_numpy(ak.flatten(tree[Q2_branch].array(library="ak"))) - n_evt = min(len(xB_evt), len(Q2_evt)) - if n_evt == 0: - continue - xB_evt = xB_evt[:n_evt] - Q2_evt = Q2_evt[:n_evt] - - Ngen_evt += np.histogram2d(xB_evt, Q2_evt, bins=(xB_edges, Q2_edges))[0] - - lamE_j = tree[lamE_branch].array(library="ak")[:n_evt] - px_j = tree[lampx_branch].array(library="ak")[:n_evt] - py_j = tree[lampy_branch].array(library="ak")[:n_evt] - pz_j = tree[lampz_branch].array(library="ak")[:n_evt] - - xB_bc, _ = ak.broadcast_arrays(ak.Array(xB_evt), lamE_j) - Q2_bc, _ = ak.broadcast_arrays(ak.Array(Q2_evt), lamE_j) - - xB_c = ak.to_numpy(ak.flatten(xB_bc)) - Q2_c = ak.to_numpy(ak.flatten(Q2_bc)) - E_c = ak.to_numpy(ak.flatten(lamE_j)) - px_c = ak.to_numpy(ak.flatten(px_j)) - py_c = ak.to_numpy(ak.flatten(py_j)) - pz_c = ak.to_numpy(ak.flatten(pz_j)) - if xB_c.size == 0: - continue - - pL_par = px_c * nhat[0] + py_c * nhat[1] + pz_c * nhat[2] - xL = (E_c + pL_par) / pplus_p - xK_c = xk_from_xb_xl(xB_c, xL) - - ok = np.isfinite(xB_c) & np.isfinite(Q2_c) & np.isfinite(xK_c) & (Q2_c > 0) - - if tmax is not None: - dE = ppE - E_c - dpx = ppx - px_c - dpy = ppy - py_c - dpz = ppz - pz_c - t = dE * dE - (dpx * dpx + dpy * dpy + dpz * dpz) - tneg = -t - ok &= np.isfinite(tneg) & (tneg < tmax) - - if not np.any(ok): - continue - - xB_sel = xB_c[ok] - Q2_sel = Q2_c[ok] - xK_sel = xK_c[ok] - - Nreco_lam += np.histogram2d(xB_sel, Q2_sel, bins=(xB_edges, Q2_edges))[0] - - ixB = np.searchsorted(xB_edges, xB_sel, side="right") - 1 - iQ = np.searchsorted(Q2_edges, Q2_sel, side="right") - 1 - ixK = np.searchsorted(xK_edges, xK_sel, side="right") - 1 - good = (ixB >= 0) & (ixB < xB_bins) & (iQ >= 0) & (iQ < Q2_bins) & (ixK >= 0) & (ixK < xK_bins) - np.add.at(M, (ixK[good], ixB[good], iQ[good]), 1.0) - - except Exception as e: - print(f"[plot_relerr_kaon_sf_xK_Q2] Warning {fpath.name}: {e}") - - eps = np.zeros_like(Ngen_evt, dtype=np.float64) - m = Ngen_evt > 0 - eps[m] = Nreco_lam[m] / Ngen_evt[m] - eps = np.clip(eps, 0.0, 1.0) - - L_nb_inv = float(lumin_fb) * 1e6 # fb^-1 -> nb^-1 - Nexp_xB_Q2 = sigma_gen_nb * eps * L_nb_inv - - relerr_xB_Q2 = np.full_like(Nexp_xB_Q2, np.nan, dtype=np.float64) - okN = Nexp_xB_Q2 > 10 - relerr_xB_Q2[okN] = 100.0 / np.sqrt(Nexp_xB_Q2[okN]) - - denom = np.sum(M, axis=0) - P = np.zeros_like(M, dtype=np.float64) - ok = denom > 0 - P[:, ok] = M[:, ok] / denom[ok] - - Nexp_xK_Q2 = np.zeros((xK_bins, Q2_bins), dtype=np.float64) - for iQ in range(Q2_bins): - Nexp_xK_Q2[:, iQ] = P[:, :, iQ] @ Nexp_xB_Q2[:, iQ] - - relerr_xK_Q2 = np.full_like(Nexp_xK_Q2, np.nan, dtype=np.float64) - okNk = Nexp_xK_Q2 > 10 - relerr_xK_Q2[okNk] = 100.0 / np.sqrt(Nexp_xK_Q2[okNk]) - - print(f"σ_tot(gen) = {sigma_tot_nb:.6g} nb") - print("Total expected Lambdas (xB,Q2):", float(np.nansum(Nexp_xB_Q2))) - print("Total expected Lambdas (xK,Q2):", float(np.nansum(Nexp_xK_Q2))) - - extra = "" if tmax is None else rf" ($-t<{tmax}$)" - q2_tag = "logQ2" if log_Q2 else "linQ2" - - def plot_map(Z, x_edges, y_edges, xlabel, outname): - fig, ax = plt.subplots(figsize=(6, 5)) - Zm = np.ma.masked_invalid(Z) - pcm = ax.pcolormesh(x_edges, y_edges, Zm.T, shading="auto", cmap="viridis") - cb = fig.colorbar(pcm, ax=ax) - cb.set_label(r"Relative stat. uncertainty $\delta N/N$ [%]" + extra) - if vmax_percent is not None and np.isfinite(vmax_percent): - pcm.set_clim(vmin=0.0, vmax=float(vmax_percent)) - ax.set_xlabel(xlabel) - ax.set_ylabel(r"$Q^2\ (\mathrm{GeV}^2)$") - if log_Q2: - ax.set_yscale("log") - ax.set_title(rf"{beam}, $\mathcal{{L}}$={lumin_fb:g} fb$^{{-1}}$") - fig.tight_layout() - savefig(fig, outdir / outname, dpi=300) - - plot_map( - relerr_xB_Q2, xB_edges, Q2_edges, r"$x_B$", - f"relerr_xB_Q2_{beam}_{tag}_L{lumin_fb:g}fb_{q2_tag}.png", - ) - plot_map( - relerr_xK_Q2, xK_edges, Q2_edges, r"$x_K \simeq x_B/(1-x_\Lambda)$", - f"relerr_xK_Q2_{beam}_{tag}_L{lumin_fb:g}fb_{q2_tag}.png", - ) - - return relerr_xB_Q2, relerr_xK_Q2, dict( - eps=eps, - Nexp_xB_Q2=Nexp_xB_Q2, - Nexp_xK_Q2=Nexp_xK_Q2, - sigma_gen_nb=sigma_gen_nb, - sigma_tot_nb=sigma_tot_nb, - ) \ No newline at end of file diff --git a/analysis/multicalo-lambda/test.py b/analysis/multicalo-lambda/test.py new file mode 100644 index 0000000..d3274a2 --- /dev/null +++ b/analysis/multicalo-lambda/test.py @@ -0,0 +1,154 @@ +# from pathlib import Path +# import os +# import lhapdf + + +# name = "JAM19FF_kaon_nlo" +# pdf = lhapdf.mkPDF(name, 0) + +# print("Set:", name) +# print("x range:", pdf.xMin(), pdf.xMax()) +# print("Q2 range:", pdf.q2Min(), pdf.q2Max()) +# print("flavors:", pdf.flavors()) + +# # Try to locate .info file +# datapath = os.environ.get("LHAPDF_DATA_PATH", "") +# candidates = [Path(p) for p in datapath.split(":") if p] +# for base in candidates: +# info = base / name / f"{name}.info" +# if info.exists(): +# print("\nINFO FILE:", info) +# print(info.read_text()[:1200]) +# break +# else: +# print("\nCould not locate .info file via LHAPDF_DATA_PATH.") + + +from __future__ import annotations + +from pathlib import Path +import numpy as np +import matplotlib.pyplot as plt +from kaon_studies import make_kaon_pdf_func, F2_from_pdf_func +from physics import ToyParams, KaonModelParams +from config import setup_lhapdf_env, get_pion_pdfset, get_pion_pdfmember +setup_lhapdf_env() +pion_set = get_pion_pdfset() +pion_member = get_pion_pdfmember() + +# suppose que tu as déjà : +# - build_kaon_model(model_kind, ...) +# - F2_from_pdf_func(pdf_func, x, Q2) +# - apply_mpl_style() + +# colorss = ['black', 'darkorange', 'cadetblue', 'darkred'] +colorss = ['black', 'blue', 'gold', 'darkgreen'] + +TOY_MODEL_LABELS = { + # "toy_baseline": "Baseline", + # "toy_soft_valence": "Soft valence", + # "toy_hard_valence": "Hard valence", + # "toy_sea_enhanced": "Sea enhanced", + # "toy_su3_breaking": "SU(3) breaking", + + "smrs": "SMRS (ansatz)", + "grv_grs": "GRV-GRS (ansatz)", + "jam_grs": "JAM-GRS (ansatz)", + "dse": "DSE (ansatz)", +} + +def plot_toy_models_vs_x( + Q2: float, + outpath: str | Path | None = None, + xmin: float = 1e-3, + xmax: float = 0.95, + nx: int = 400, + logx: bool = True, +): + """ + Plot F2^K(x,Q2) vs x for the toy models on a single figure. + """ + + # apply_mpl_style() # décommente si tu veux ton style global + + model_kinds = [ + # "toy_baseline", + # "toy_hard_valence", + # "toy_sea_enhanced", + # "toy_su3_breaking", + "smrs", + "grv_grs", + "jam_grs", + "dse" + ] + + if logx: + xgrid = np.logspace(np.log10(xmin), np.log10(xmax), nx) + else: + xgrid = np.linspace(xmin, xmax, nx) + + fig, ax = plt.subplots(figsize=(7,5)) + + i = 0 + + for model_kind in model_kinds: + + # adapte cette ligne à ton builder exact + pdf_func = make_kaon_pdf_func( + model=model_kind, toy=ToyParams, kparams=KaonModelParams, + pion_set=pion_set, pion_member=pion_member, require_lhapdf=True) + + y = np.array([F2_from_pdf_func(pdf_func, x, Q2) for x in xgrid]) + # y_u = np.array([x * (4.0/9.0) * pdf_func(2, x, Q2) for x in xgrid]) + # y_sbar = np.array([x * (1.0/9.0) * pdf_func(-3, x, Q2) for x in xgrid]) + # y_tot = y_u + y_sbar + + ax.plot( + xgrid, + y, + lw=2.2, + label=TOY_MODEL_LABELS.get(model_kind, model_kind), + color=colorss[i] + ) + + # ax.plot( + # xgrid, + # y_sbar, + # lw=2.2, + # # label=TOY_MODEL_LABELS.get(model_kind, model_kind), + # color='gray' + # ) + + # ax.plot( + # xgrid, + # y_u, + # lw=2.2, + # # label=TOY_MODEL_LABELS.get(model_kind, model_kind), + # color='black' + # ) + + i+=1 + + ax.set_xlabel(r"$x$") + ax.set_ylabel(r"$F$") + ax.set_xticks([0,0.5,1]) + # ax.set_title(rf"Toy kaon structure-function models at $Q^2 = {Q2:g}\ \mathrm{{GeV}}^2$") + ax.legend(frameon=False) + + # if logx: + # ax.set_xscale("log") + + # ax.grid(True, alpha=0.3) + fig.tight_layout() + + if outpath is not None: + outpath = Path(outpath) + outpath.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(outpath, dpi=300, bbox_inches="tight") + + return fig, ax + +plot_toy_models_vs_x( + Q2=10.0, + outpath="outputs/vizu_realmodels.png", +) \ No newline at end of file diff --git a/analysis/multicalo-lambda/uq.py b/analysis/multicalo-lambda/uq.py new file mode 100644 index 0000000..af4e74a --- /dev/null +++ b/analysis/multicalo-lambda/uq.py @@ -0,0 +1,446 @@ +from __future__ import annotations + +import numpy as np + +from physics import ( + KaonModelParams, + ToyParams, + F2_from_pdf_func +) + +from config import ( + setup_lhapdf_env, + get_pion_pdfset, + get_pion_pdfmember, +) + +from kaon_models import ( + make_kaon_pdf_func, +) + +def efficiency_from_counts( + n_reco: np.ndarray, + n_gen: np.ndarray, + *, + clip: bool = True, +) -> np.ndarray: + """ + eps = n_reco / n_gen with safe handling (eps=0 when n_gen=0). + """ + eps = np.zeros_like(n_gen, dtype=np.float64) + m = n_gen > 0 + eps[m] = n_reco[m] / n_gen[m] + if clip: + eps = np.clip(eps, 0.0, 1.0) + return eps + + +def expected_yields_nb( + sigma_nb: np.ndarray, + eps: np.ndarray, + lumin_fb: float, +) -> np.ndarray: + """ + Nexp = sigma(nb/bin) * eps * L(nb^-1), with L = lumin_fb * 1e6 (fb^-1 -> nb^-1). + """ + L_nb_inv = float(lumin_fb) * 1e6 + return sigma_nb * eps * L_nb_inv + + +def relerr_combined( + n_exp: np.ndarray, + eps: np.ndarray, + n_gen: np.ndarray, +) -> np.ndarray: + """ + Relative uncertainty [%] combining: + - Poisson counting on expected data: 1/sqrt(Nexp) + - Binomial MC stat on eps: sqrt((1-eps)/(eps*Ngen)) + - ... + + Returns an array with NaN where undefined. + """ + out = np.full_like(n_exp, np.nan, dtype=np.float64) + + okN = n_exp > 0 + ok_eps = (eps > 0) & (n_gen > 0) + mask = okN & ok_eps + if not np.any(mask): + return out + + stat_term = 1.0 / np.sqrt(n_exp[mask]) + relerr_eps = np.sqrt((1.0 - eps[mask]) / (eps[mask] * n_gen[mask])) + + out[mask] = 100.0 * np.sqrt(stat_term**2) #100.0 * np.sqrt(stat_term**2 + relerr_eps**2) + + return out + + +def migration_probability_xk_given_xb_q2(M: np.ndarray) -> np.ndarray: + """ + Build P(xK | xB, Q2) from counts M with shape (xK, xB, Q2). + Normalizes over xK for each (xB,Q2). Where denom=0, leaves P=0. + """ + denom = np.sum(M, axis=0) # (xB,Q2) + P = np.zeros_like(M, dtype=np.float64) + ok = denom > 0 + P[:, ok] = M[:, ok] / denom[ok] + return P + + +def project_xbq2_to_xkq2( + P_xk_xb_q2: np.ndarray, + A_xb_q2: np.ndarray, +) -> np.ndarray: + """ + For each Q2 bin: A_xk(Q2) = P(Q2) @ A_xb(Q2) + P has shape (xK, xB, Q2), A has shape (xB, Q2). + Returns shape (xK, Q2). + """ + xk_bins, xb_bins, q2_bins = P_xk_xb_q2.shape + out = np.zeros((xk_bins, q2_bins), dtype=np.float64) + for iQ in range(q2_bins): + out[:, iQ] = P_xk_xb_q2[:, :, iQ] @ A_xb_q2[:, iQ] + return out + + +def projected_efficiency_from_migration( + M: np.ndarray, + P_xk_xb_q2: np.ndarray, + n_gen_xb_q2: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Build effective projected quantities: + - Ngen_xK_Q2 = P @ Ngen_xB_Q2 + - Nreco_xK_Q2 = sum_xB M (counts already in xK,Q2) + - eps_xK_Q2 = Nreco_xK_Q2 / Ngen_xK_Q2 + """ + n_gen_xK_Q2 = project_xbq2_to_xkq2(P_xk_xb_q2, n_gen_xb_q2) + n_reco_xK_Q2 = np.sum(M, axis=1) # (xK,Q2) + eps_xK_Q2 = efficiency_from_counts(n_reco_xK_Q2, n_gen_xK_Q2, clip=True) + return eps_xK_Q2, n_gen_xK_Q2, n_reco_xK_Q2 + + +def _compute_toy_F2_map_xK_Q2( + *, + model: str, + toy: ToyParams, + kparams: KaonModelParams, + xK_cent: np.ndarray, + Q2_cent: np.ndarray, +) -> np.ndarray: + """ + Compute F2^K(xK, Q2) on the bin-center grid for a given toy model. + + Returns + ------- + F2 : np.ndarray + Array of shape (len(xK_cent), len(Q2_cent)). + """ + setup_lhapdf_env() + pion_set = get_pion_pdfset() + pion_member = get_pion_pdfmember() + + pdf = make_kaon_pdf_func( + model=model, + toy=toy, + kparams=kparams, + pion_set=pion_set, + pion_member=pion_member, + require_lhapdf=True, + ) + + F2 = np.empty((len(xK_cent), len(Q2_cent)), dtype=np.float64) + + for ix, x in enumerate(xK_cent): + for iQ, Q2 in enumerate(Q2_cent): + F2[ix, iQ] = float(F2_from_pdf_func(pdf, float(x), float(Q2))) + + return F2 + + +def compute_local_fischer( + *, + a0: float, + b0: float, + relerr_xK_Q2: np.ndarray, + xK_range: tuple[float, float] = (0, 1), + xK_bins: int = 10, + Q2_range: tuple[float, float] = (1, 500), + Q2_bins: int = 10, + model: str = "toy_baseline", + kparams: KaonModelParams = KaonModelParams(), + da: float = 1e-2, + db: float = 1e-2, + log_Q2_centers: bool = True, + eps: float = 1e-15, +) -> dict: + """ + Compute the local Fisher matrix for the toy-model parameters (a, b), + using F2^K(xK, Q2) as observable and relerr_xK_Q2 as projected relative + uncertainty map. + + Parameters + ---------- + a0, b0 + Fiducial toy-model parameters where the Fisher matrix is evaluated. + + relerr_xK_Q2 + Relative uncertainty map in percent, shape (xK_bins, Q2_bins). + + xK_edges, Q2_edges + Bin edges used to define the observable grid. + + model + Kaon model name passed to make_kaon_pdf_func. For your first study, + this should typically stay "toy_baseline". + + da, db + Finite-difference steps for numerical derivatives. + + log_Q2_centers + If True, use geometric bin centers for Q2, which is the natural choice + when Q2 bins are logarithmic. + + eps + Numerical floor to avoid divisions by zero. + + Returns + ------- + dict with keys: + "F" : Fisher matrix, shape (2,2) + "Cov" : inverse Fisher matrix if invertible, else None + "sigma_map" : absolute sigma_ij used in the calculation + "F2_0" : central F2 map + "dF_da" : numerical derivative wrt a + "dF_db" : numerical derivative wrt b + "xK_cent" : xK bin centers + "Q2_cent" : Q2 bin centers + "mask" : boolean mask of bins effectively used + """ + + xK_edges = np.linspace(xK_range[0], xK_range[1], xK_bins + 1) + + if log_Q2_centers: + Q2_edges = np.logspace(np.log10(Q2_range[0]), np.log10(Q2_range[1]), Q2_bins + 1) + else: + Q2_edges = np.linspace(Q2_range[0], Q2_range[1], Q2_bins + 1) + + relerr_xK_Q2 = np.asarray(relerr_xK_Q2, dtype=np.float64) + xK_edges = np.asarray(xK_edges, dtype=np.float64) + Q2_edges = np.asarray(Q2_edges, dtype=np.float64) + + xK_cent = 0.5 * (xK_edges[:-1] + xK_edges[1:]) + + if log_Q2_centers: + Q2_cent = np.sqrt(Q2_edges[:-1] * Q2_edges[1:]) + else: + Q2_cent = 0.5 * (Q2_edges[:-1] + Q2_edges[1:]) + + if relerr_xK_Q2.shape != (xK_cent.size, Q2_cent.size): + raise ValueError( + f"relerr_xK_Q2 has shape {relerr_xK_Q2.shape}, expected {(xK_cent.size, Q2_cent.size)}" + ) + + # --- Central model + toy_0 = ToyParams(a=a0, b=b0) + F2_0 = _compute_toy_F2_map_xK_Q2( + model=model, + toy=toy_0, + kparams=kparams, + xK_cent=xK_cent, + Q2_cent=Q2_cent, + ) + + # --- Shifted models for finite differences + toy_ap = ToyParams(a=a0 + da, b=b0) + toy_am = ToyParams(a=a0 - da, b=b0) + toy_bp = ToyParams(a=a0, b=b0 + db) + toy_bm = ToyParams(a=a0, b=b0 - db) + + F2_ap = _compute_toy_F2_map_xK_Q2( + model=model, + toy=toy_ap, + kparams=kparams, + xK_cent=xK_cent, + Q2_cent=Q2_cent, + ) + F2_am = _compute_toy_F2_map_xK_Q2( + model=model, + toy=toy_am, + kparams=kparams, + xK_cent=xK_cent, + Q2_cent=Q2_cent, + ) + F2_bp = _compute_toy_F2_map_xK_Q2( + model=model, + toy=toy_bp, + kparams=kparams, + xK_cent=xK_cent, + Q2_cent=Q2_cent, + ) + F2_bm = _compute_toy_F2_map_xK_Q2( + model=model, + toy=toy_bm, + kparams=kparams, + xK_cent=xK_cent, + Q2_cent=Q2_cent, + ) + + # --- Numerical derivatives + dF_da = (F2_ap - F2_am) / (2.0 * da) + dF_db = (F2_bp - F2_bm) / (2.0 * db) + + # --- Absolute sigma map from relative errors + rel = relerr_xK_Q2 / 100.0 + Fref = F2_0 + sigma_map = rel * Fref + + # --- Valid bins mask + mask = ( + np.isfinite(F2_0) + & np.isfinite(dF_da) + & np.isfinite(dF_db) + & np.isfinite(sigma_map) + & np.isfinite(rel) + & (sigma_map > 0.0) + & (rel > 0.0) + ) + + if not np.any(mask): + raise RuntimeError("No valid bins available to compute the local Fisher matrix.") + + w = np.zeros_like(F2_0, dtype=np.float64) + w[mask] = 1.0 / np.maximum(sigma_map[mask] ** 2, eps) + + F_aa = np.sum(w[mask] * dF_da[mask] * dF_da[mask]) + F_ab = np.sum(w[mask] * dF_da[mask] * dF_db[mask]) + F_bb = np.sum(w[mask] * dF_db[mask] * dF_db[mask]) + + F = np.array( + [ + [F_aa, F_ab], + [F_ab, F_bb], + ], + dtype=np.float64, + ) + + try: + Cov = np.linalg.inv(F) + except np.linalg.LinAlgError: + Cov = None + + return { + "F": F, + "Cov": Cov, + "sigma_map": sigma_map, + "F2_0": F2_0, + "dF_da": dF_da, + "dF_db": dF_db, + "xK_cent": xK_cent, + "Q2_cent": Q2_cent, + "mask": mask, + } + +def compute_local_fischer_maps( + *, + a0: float, + b0: float, + relerr_xK_Q2: np.ndarray, + xK_range: tuple[float, float] = (0, 1), + xK_bins: int = 10, + Q2_range: tuple[float, float] = (1, 500), + Q2_bins: int = 10, + model: str = "toy_baseline", + kparams: KaonModelParams = KaonModelParams(), + da: float = 1e-2, + db: float = 1e-2, + log_Q2_centers: bool = True, + eps: float = 1e-15, +) -> dict: + """ + Compute the bin-by-bin Fisher information density maps associated with the + local Fisher matrix in parameter space (a, b). + + For each bin (xK, Q2), define + f_aa = (1/sigma^2) * (dF/da)^2 + f_ab = (1/sigma^2) * (dF/da) * (dF/db) + f_bb = (1/sigma^2) * (dF/db)^2 + + so that the total Fisher matrix is obtained by summing over bins: + F_aa = sum f_aa + F_ab = sum f_ab + F_bb = sum f_bb + + Returns + ------- + dict with keys: + "F" : total Fisher matrix, shape (2, 2) + "Cov" : inverse Fisher matrix if invertible, else None + "maps" : array of shape (2, 2, xK_bins, Q2_bins) + "info_aa" : Fisher density map for aa, shape (xK_bins, Q2_bins) + "info_ab" : Fisher density map for ab, shape (xK_bins, Q2_bins) + "info_ba" : same as info_ab + "info_bb" : Fisher density map for bb, shape (xK_bins, Q2_bins) + "sigma_map" : absolute uncertainty map + "F2_0" : central F2 map + "dF_da" : derivative wrt a + "dF_db" : derivative wrt b + "xK_cent" : xK bin centers + "Q2_cent" : Q2 bin centers + "mask" : boolean validity mask + """ + res = compute_local_fischer( + a0=a0, + b0=b0, + relerr_xK_Q2=relerr_xK_Q2, + xK_range=xK_range, + xK_bins=xK_bins, + Q2_range=Q2_range, + Q2_bins=Q2_bins, + model=model, + kparams=kparams, + da=da, + db=db, + log_Q2_centers=log_Q2_centers, + eps=eps, + ) + + sigma_map = res["sigma_map"] + dF_da = res["dF_da"] + dF_db = res["dF_db"] + mask = res["mask"] + + w = np.zeros_like(sigma_map, dtype=np.float64) + w[mask] = 1.0 / np.maximum(sigma_map[mask] ** 2, eps) + + info_aa = np.full_like(sigma_map, np.nan) + info_ab = np.full_like(sigma_map, np.nan) + info_bb = np.full_like(sigma_map, np.nan) + + info_aa[mask] = w[mask] * dF_da[mask] * dF_da[mask] + info_ab[mask] = w[mask] * dF_da[mask] * dF_db[mask] + info_bb[mask] = w[mask] * dF_db[mask] * dF_db[mask] + + maps = np.zeros((2, 2) + sigma_map.shape, dtype=np.float64) + maps[0, 0] = info_aa + maps[0, 1] = info_ab + maps[1, 0] = info_ab + maps[1, 1] = info_bb + + return { + "F": res["F"], + "Cov": res["Cov"], + "maps": maps, + "info_aa": info_aa, + "info_ab": info_ab, + "info_ba": info_ab, + "info_bb": info_bb, + "sigma_map": sigma_map, + "F2_0": res["F2_0"], + "dF_da": dF_da, + "dF_db": dF_db, + "xK_cent": res["xK_cent"], + "Q2_cent": res["Q2_cent"], + "mask": mask, + } \ No newline at end of file diff --git a/analysis/multicalo-lambda/utils.py b/analysis/multicalo-lambda/utils.py index e1e9ceb..117f8ae 100644 --- a/analysis/multicalo-lambda/utils.py +++ b/analysis/multicalo-lambda/utils.py @@ -1,13 +1,30 @@ from __future__ import annotations + +import os +from dataclasses import dataclass from pathlib import Path -from typing import Iterable +from typing import Callable, Iterable, TYPE_CHECKING + import numpy as np import awkward as ak import uproot -from config import PhysicsConstants, Paths, DEFAULT_CONST, DEFAULT_PATHS +import math + +from config import CONST +from physics import ToyParams, toy_norm, toy_shape_x +if TYPE_CHECKING: + from config import PhysicsConstants + + +# ============================================================================= +# Files helpers +# ============================================================================= def iter_reco_files(beam: str, nfiles: int, root_base_dir: Path, suffix: str) -> Iterable[Path]: + """ + Yield existing non-empty ROOT files following your naming convention. + """ for i in range(1, nfiles + 1): fpath = root_base_dir / f"k_lambda_{beam}_5000evt_{i:03d}_{suffix}.root" if not fpath.exists() or fpath.stat().st_size == 0: @@ -15,6 +32,10 @@ def iter_reco_files(beam: str, nfiles: int, root_base_dir: Path, suffix: str) -> yield fpath +# ============================================================================= +# ROOT readers +# ============================================================================= + def read_lambda_eicrecon( beam: str, nfiles: int, @@ -22,6 +43,7 @@ def read_lambda_eicrecon( suffix: str, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ + Read reconstructed Lambda from ReconstructedFarForwardLambdas. Returns: E [GeV], cos(theta), sin(phi) """ e_branch = "ReconstructedFarForwardLambdas.energy" @@ -80,7 +102,7 @@ def read_lambda_geant4( nfiles: int, root_base_dir: Path, suffix: str, - c: PhysicsConstants = DEFAULT_CONST, + c: "PhysicsConstants" = CONST, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ Read truth Lambda from MCParticles. @@ -123,7 +145,7 @@ def read_lambda_geant4( cos_theta = pz / p phi = np.arctan2(py, px) sin_phi = np.sin(phi) - E = np.sqrt(p * p + (c.m_lambda_gev**2)) + E = np.sqrt(p * p + (c.m_lambda_gev ** 2)) E_list.append(ak.to_numpy(ak.flatten(E))) cos_list.append(ak.to_numpy(ak.flatten(cos_theta))) @@ -138,6 +160,10 @@ def read_lambda_geant4( return E_out, cos_out, sinphi_out +# ============================================================================= +# Afterburner reader (HepMC3 ASCII) +# ============================================================================= + def read_lambda_afterburner( beam: str, nfiles: int, @@ -145,6 +171,7 @@ def read_lambda_afterburner( pid: int, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ + Read Lambda from HepMC v3 afterburner file. Returns: E [GeV], cos(theta), sin(phi) """ E_list: list[np.ndarray] = [] @@ -196,4 +223,140 @@ def read_lambda_afterburner( E_all = np.concatenate(E_list) if E_list else np.array([], dtype=np.float64) cos_all = np.concatenate(cos_list) if cos_list else np.array([], dtype=np.float64) sinphi_all = np.concatenate(sinphi_list) if sinphi_list else np.array([], dtype=np.float64) - return E_all, cos_all, sinphi_all \ No newline at end of file + return E_all, cos_all, sinphi_all + + +# ============================================================================= +# LHAPDF helpers (kaon toy × pion evolution) +# ============================================================================= + +try: + import lhapdf # type: ignore +except Exception: + lhapdf = None + + +def is_lhapdf_available() -> bool: + return lhapdf is not None + + +def ensure_lhapdf_set_installed(setname: str) -> None: + if lhapdf is None: + raise RuntimeError("python lhapdf is not available in this environment.") + installed = set(lhapdf.availablePDFSets()) + if setname not in installed: + raise RuntimeError( + f"Pion PDF set '{setname}' is not installed in LHAPDF_DATA_PATH={os.environ.get('LHAPDF_DATA_PATH')}\n" + f"Install it e.g.:\n" + f" lhapdf install --pdfdir {os.environ.get('LHAPDF_DATA_PATH')} {setname}\n" + ) + + +def pion_light_quark_sum(pdf, x: float, Q: float) -> float: + """ + Using (u+ubar + d+dbar) as a proxy for meson evolution. + LHAPDF returns x*f, so divide by x to get f. + """ + if x <= 0.0: + return 0.0 + u = pdf.xfxQ(2, x, Q) / x + ub = pdf.xfxQ(-2, x, Q) / x + d = pdf.xfxQ(1, x, Q) / x + db = pdf.xfxQ(-1, x, Q) / x + return (u + ub + d + db) + + +def make_pdf_kaon_toy_times_pion_evolution( + p: ToyParams, + pion_set: str, + member: int = 0, +) -> Callable[[int, float, float], float]: + """ + q_K(pid,x,Q2) = q_K^toy(pid,x,Q0) * R_pi(x,Q2) + where R_pi is derived from pion PDFs. + K+ content: u and sbar only (valence-only toy). + """ + if lhapdf is None: + raise RuntimeError("python lhapdf is not available in this environment.") + ensure_lhapdf_set_installed(pion_set) + + pion_pdf = lhapdf.mkPDF(pion_set, member) + norm = toy_norm(p.a, p.b) + Q0 = float(np.sqrt(p.Q0)) + + def pdf_kaon(pid: int, x: float, Q2: float) -> float: + base = toy_shape_x(x, p.a, p.b, norm) + + Q = float(np.sqrt(Q2)) + num = pion_light_quark_sum(pion_pdf, x, Q) + den = pion_light_quark_sum(pion_pdf, x, Q0) + + R = 1.0 if den <= 0.0 else (num / den) + val = base * R + + if pid == 2: # u + return val + if pid == -3: # sbar + return val + return 0.0 + + return pdf_kaon + + +# ============================================================================= +# Grids +# ============================================================================= + +@dataclass(frozen=True) +class Grid: + x_edges: np.ndarray + q2_edges: np.ndarray + x_cent: np.ndarray + q2_cent: np.ndarray + Xc: np.ndarray + Q2c: np.ndarray + + +def make_log_grid( + x_min: float, + x_max: float, + q2_min: float, + q2_max: float, + nx: int, + nq2: int, +) -> Grid: + x_edges = np.logspace(np.log10(x_min), np.log10(x_max), nx + 1) + q2_edges = np.logspace(np.log10(q2_min), np.log10(q2_max), nq2 + 1) + x_cent = 0.5 * (x_edges[:-1] + x_edges[1:]) + q2_cent = 0.5 * (q2_edges[:-1] + q2_edges[1:]) + Xc, Q2c = np.meshgrid(x_cent, q2_cent, indexing="xy") + return Grid(x_edges, q2_edges, x_cent, q2_cent, Xc, Q2c) + + +def make_lin_grid( + x_min: float, + x_max: float, + q2_min: float, + q2_max: float, + nx: int, + nq2: int, +) -> Grid: + x_edges = np.linspace(x_min, x_max, nx + 1) + q2_edges = np.linspace(q2_min, q2_max, nq2 + 1) + x_cent = 0.5 * (x_edges[:-1] + x_edges[1:]) + q2_cent = 0.5 * (q2_edges[:-1] + q2_edges[1:]) + Xc, Q2c = np.meshgrid(x_cent, q2_cent, indexing="xy") + return Grid(x_edges, q2_edges, x_cent, q2_cent, Xc, Q2c) + + +def compute_F2_grid(grid: Grid, f2_point: Callable[[float, float], float]) -> np.ndarray: + """ + f2_point(x, Q2) -> F2 + """ + F2 = np.empty_like(grid.Xc, dtype=np.float64) + for i in range(grid.Q2c.shape[0]): + for j in range(grid.Xc.shape[1]): + x = float(grid.Xc[i, j]) + Q2 = float(grid.Q2c[i, j]) + F2[i, j] = f2_point(x, Q2) + return F2 \ No newline at end of file