laddu 0.20.0

Amplitude analysis tools for Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "matplotlib==3.10.7",
#   "numpy==2.3.4",
# ]
# ///
"""Plot the binned K-short closure projection emitted by generate_and_fit_ksks."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import NotRequired, TypedDict

import matplotlib as mpl
import numpy as np

mpl.use('Agg')
import matplotlib.pyplot as plt


class StairStyle(TypedDict):
    color: NotRequired[str]
    linewidth: NotRequired[float]
    linestyle: NotRequired[str]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('input', type=Path, help='projection JSON from the laddu example')
    parser.add_argument(
        '-o',
        '--output',
        type=Path,
        help='output image (default: input path with a .png suffix)',
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    output = args.output or args.input.with_suffix('.png')
    with args.input.open(encoding='utf-8') as source:
        payload = json.load(source)
    if payload.get('schema_version') != 2:  # noqa: PLR2004
        msg = f'unsupported projection schema: {payload.get("schema_version")!r}'
        raise ValueError(msg)

    edges = np.asarray(payload['data']['histogram']['bin_edges'], dtype=float)
    centers = 0.5 * (edges[:-1] + edges[1:])
    widths = np.diff(edges)
    data = np.asarray(payload['data']['histogram']['counts'], dtype=float)
    if len(edges) != len(data) + 1:
        msg = 'bin_edges must contain one more entry than each histogram'
        raise ValueError(msg)

    fig, ax = plt.subplots(figsize=(8.0, 5.5), constrained_layout=True)
    ax.errorbar(
        centers,
        data,
        xerr=0.5 * widths,
        yerr=np.sqrt(np.clip(data, 0.0, None)),
        fmt='o',
        color='black',
        markersize=4,
        linewidth=1,
        capsize=0,
        label=payload['data']['label'],
        zorder=4,
    )

    styles: dict[str, StairStyle] = {
        'fit': {'color': '#d62728', 'linewidth': 2.2},
        'f0': {'color': '#1f77b4', 'linewidth': 1.8, 'linestyle': '--'},
        'f2': {'color': '#2ca02c', 'linewidth': 1.8, 'linestyle': ':'},
    }
    for series in payload['projections']:
        values = np.asarray(series['histogram']['counts'], dtype=float)
        if len(values) != len(data):
            msg = f'projection {series["id"]!r} has the wrong number of bins'
            raise ValueError(msg)
        ax.stairs(
            values,
            edges,
            label=series['label'],
            **styles.get(series['id'], {'linewidth': 1.8}),
        )

    ax.set(
        xlabel=rf'{payload["observable"]} [$\mathrm{{{payload["unit"]}}}$]',
        ylabel=rf'$\mathrm{{Events}}\,/\,({widths[0]:.3f}\,\mathrm{{{payload["unit"]}}})$',
        title=payload['title'],
        xlim=(edges[0], edges[-1]),
        ylim=(0.0, None),
    )
    ax.legend(frameon=False)
    ax.tick_params(direction='in', top=True, right=True)
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=180)
    print(output)


if __name__ == '__main__':
    main()