mirror of
https://github.com/slsdetectorgroup/aare.git
synced 2026-09-21 07:22:09 +02:00
Erik Fröjdh
4f6b69fba6
Removed the lmfit based fitting and the old API for calling functions closes #296
433 lines
19 KiB
Plaintext
433 lines
19 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "efef8e20-6571-4561-8f0b-6048b57907de",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import time\n",
|
|
"import random\n",
|
|
"import numpy as np\n",
|
|
"import matplotlib.pyplot as plt\n",
|
|
"from matplotlib.gridspec import GridSpec\n",
|
|
"from aare import Gaussian\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "35772e2c-37c6-4986-a9c0-7dca4a034827",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"ROWS = 100\n",
|
|
"COLS = 100\n",
|
|
"N_SCAN = 100\n",
|
|
"NOISE_FRAC = 0.05\n",
|
|
"SEED = 42\n",
|
|
"N_THREADS = 4\n",
|
|
"\n",
|
|
"N_REPEATS = 7\n",
|
|
"N_WARMUP = 3 # untimed iterations (icache + branch predictor warmup)\n",
|
|
"COOLDOWN = 2.0 # seconds between (method, thread_count) pairs"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "be455445-7df8-47cc-9f39-d2a38856e0bd",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Data generator"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "06c8fddb-56d9-4f84-8b21-4ffd8b7bd26f",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def generate_3d_data(rows, cols, n_scan, noise_frac, seed):\n",
|
|
" \"\"\"\n",
|
|
" Generate a synthetic detector image stack where each pixel has a\n",
|
|
" Gaussian response curve with per-pixel variation in A, mu, sigma.\n",
|
|
"\n",
|
|
" Returns x (n_scan,), y (rows, cols, n_scan), y_err (rows, cols, n_scan),\n",
|
|
" and the ground-truth parameter arrays.\n",
|
|
" \"\"\"\n",
|
|
" rng = np.random.default_rng(seed)\n",
|
|
"\n",
|
|
" # Per-pixel true params each of shape: [rows, cols, 1]\n",
|
|
" A_true = rng.uniform(200, 1000, size=(rows, cols))\n",
|
|
" mu_true = rng.uniform(20, 80, size=(rows, cols))\n",
|
|
" sig_true = rng.uniform(3, 12, size=(rows, cols))\n",
|
|
" \n",
|
|
" # One common binned energy array\n",
|
|
" x = np.linspace(0, 100, n_scan) # shape [1, 1, nscan]\n",
|
|
"\n",
|
|
" # Build ground truth signals per-pixel\n",
|
|
" exponent = -0.5 * ((x[None, None, :] - mu_true[:, :, None]) / sig_true[:,:, None])**2 # shape [rows, cols, nscan]\n",
|
|
" y_clean = A_true[:, :, None] * np.exp(exponent)\n",
|
|
"\n",
|
|
" # Perturb with noise\n",
|
|
" noise_sigma = noise_frac * A_true[:, :, None] * np.ones_like(y_clean) # shape [rows, cols, nscan]\n",
|
|
" noise = rng.normal(0, noise_sigma)\n",
|
|
" y = y_clean + noise\n",
|
|
"\n",
|
|
" y_err = noise_sigma.copy()\n",
|
|
"\n",
|
|
" return x, y, y_err, A_true, mu_true, sig_true "
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "30ed6bb8-4798-497f-9990-4c518f885855",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Profiling function"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "dc0b52b6-9fc3-453d-b6eb-e325a2b8342e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def bench(fn, n_warmup=N_WARMUP, n_repeats=N_REPEATS):\n",
|
|
" \"\"\"\n",
|
|
" Warmup then time `fn` over `n_repeats` calls.\n",
|
|
" Returns (last_result, list_of_walltimes_in_seconds).\n",
|
|
" \"\"\"\n",
|
|
" # warmup: primes icache, branch predictor, and lets CPU ramp to boost clock\n",
|
|
" for _ in range(n_warmup):\n",
|
|
" res = fn()\n",
|
|
"\n",
|
|
" times = []\n",
|
|
" for _ in range(n_repeats):\n",
|
|
" t0 = time.perf_counter()\n",
|
|
" res = fn()\n",
|
|
" t1 = time.perf_counter()\n",
|
|
" times.append(t1 - t0)\n",
|
|
" return res, times"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "377fc820-95b2-48aa-b104-a272c50e4103",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Quick check on small (2x2) frame"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "9cc999b9-e0d8-4deb-ae29-e8bd0141534b",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Generate 2 x 2 dataset of Gaussian-like profiles for each pixel\n",
|
|
"x2, y2, yerr2, true_A2, true_mu2, true_sig2 = generate_3d_data(\n",
|
|
" 2, 2, N_SCAN, NOISE_FRAC, SEED\n",
|
|
")\n",
|
|
"model_g = Gaussian()\n",
|
|
"model_g.compute_errors = True\n",
|
|
"result = model_g.fit(x2, y2, yerr2)\n",
|
|
"\n",
|
|
"from pprint import pprint\n",
|
|
"print(\"== True Gaussian params == \")\n",
|
|
"print(\"A_true = \\n\", true_A2)\n",
|
|
"print(\"mu_true = \\n\", true_mu2)\n",
|
|
"print(\"sig_true = \\n\",true_sig2)\n",
|
|
"print(\"\\n\")\n",
|
|
"\n",
|
|
"print(\"== Fit results ==\")\n",
|
|
"par = result['par']\n",
|
|
"# print(par)\n",
|
|
"A_fit = par[:, :, 0]\n",
|
|
"mu_fit = par[:, :, 1]\n",
|
|
"sig_fit = par[:, :, 2]\n",
|
|
"print(\"A_fit = \\n\", A_fit)\n",
|
|
"print(\"mu_fit = \\n\", mu_fit)\n",
|
|
"print(\"sig_fit = \\n\", sig_fit)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "a98a41f3-fd2e-4bfc-9ec6-0d23dd38e896",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"fig, ax = plt.subplots(2, 2, figsize=(12,8))\n",
|
|
"\n",
|
|
"# Gaussians in 2x2 frame: True vs Fit\n",
|
|
"for row in range(2):\n",
|
|
" for col in range(2):\n",
|
|
" ax[row, col].plot(x2, y2[row, col,:], label=\"data\")\n",
|
|
" ax[row, col].plot(x2, model_g(x2, result['par'][row, col,:]), linewidth=1, color=\"green\", label=\"minuit\")\n",
|
|
" ax[row, col].set_title(f\"Gaussian Fit to data in pixel [{row}, {col}]\")\n",
|
|
" ax[row, col].legend()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "fcf52481-0278-4f95-8676-829a6d61eff8",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Fit data with different backends"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "1f6bc651-80c1-41dd-8a05-7f15aba006aa",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# ===============\n",
|
|
"# DATA GENERATION\n",
|
|
"# ===============\n",
|
|
"print(f\"Generating synthetic data: {ROWS}x{COLS} pixels, \"\n",
|
|
" f\"{N_SCAN} scan points, noise_frac={NOISE_FRAC}\\n\")\n",
|
|
"\n",
|
|
"x, y, yerr, true_A, true_mu, true_sig = generate_3d_data(\n",
|
|
" ROWS, COLS, N_SCAN, NOISE_FRAC, SEED\n",
|
|
")\n",
|
|
"\n",
|
|
"model = Gaussian()\n",
|
|
"print(f\"model.max_calls = {model.max_calls}\")\n",
|
|
"print(f\"model.tolerance = {model.tolerance}\")\n",
|
|
"print(\"model.compute_errors =\", model.compute_errors)\n",
|
|
"METHOD_DEFS = [\n",
|
|
" (\"Minuit2 (obj API)\",\n",
|
|
" lambda nt: lambda: model.fit(x, y, n_threads=nt),\n",
|
|
" \"#FF9800\", {\"linewidth\": 2.5, \"linestyle\": \":\"}),\n",
|
|
"]\n",
|
|
"\n",
|
|
"colors = {label: c for label, _, c, _ in METHOD_DEFS}\n",
|
|
"styles = {label: s for label, _, _, s in METHOD_DEFS}"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "5a417145-ce42-4c3a-a7bf-05ff97ba0450",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# ====================================\n",
|
|
"# SINGLE-CALL BENCHMARK (at N_THREADS)\n",
|
|
"# ====================================\n",
|
|
"def extract_result(label, res):\n",
|
|
" \"\"\"Normalize return values across fitters into a common dict.\"\"\"\n",
|
|
" if isinstance(res, dict):\n",
|
|
" out = {\"par\": res[\"par\"]}\n",
|
|
" if \"par_err\" in res:\n",
|
|
" out[\"par_err\"] = res[\"par_err\"]\n",
|
|
" if \"chi2\" in res:\n",
|
|
" out[\"chi2\"] = res[\"chi2\"]\n",
|
|
" return out\n",
|
|
" \n",
|
|
"methods = {}\n",
|
|
"for label, factory, _, _ in METHOD_DEFS:\n",
|
|
" time.sleep(COOLDOWN)\n",
|
|
" res, times = bench(factory(N_THREADS))\n",
|
|
" entry = extract_result(label, res)\n",
|
|
" entry[\"times\"] = times\n",
|
|
" methods[label] = entry\n",
|
|
"\n",
|
|
"# ---- Print summary ----\n",
|
|
"ndf = N_SCAN - 3\n",
|
|
"print(f\"{'Method':24s} {'time (ms)':>10s} {'med|dA|':>10s} {'med|dMu|':>10s} {'med|dSig|':>10s}\")\n",
|
|
"print(\"-\" * 80)\n",
|
|
"for name, m in methods.items():\n",
|
|
" par = m[\"par\"]\n",
|
|
" med_t = np.median(m[\"times\"]) * 1e3\n",
|
|
" dA = np.median(np.abs(par[:,:,0] - true_A))\n",
|
|
" dMu = np.median(np.abs(par[:,:,1] - true_mu))\n",
|
|
" dSig = np.median(np.abs(par[:,:,2] - true_sig))\n",
|
|
"\n",
|
|
" chi2_str = \"\"\n",
|
|
" if \"chi2\" in m:\n",
|
|
" chi2_str = f\" chi2/ndf={np.median(m['chi2'] / ndf):.4f}\"\n",
|
|
"\n",
|
|
" print(f\"[{name:22s}] {med_t:8.2f} ms \"\n",
|
|
" f\"{dA:10.3f} {dMu:10.4f} {dSig:10.4f}{chi2_str}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "55ecc77c-0823-408d-a811-8e4f4900f332",
|
|
"metadata": {
|
|
"scrolled": true
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# ===============\n",
|
|
"# THREAD SCALING\n",
|
|
"# ===============\n",
|
|
"thread_counts = [1, 2, 4, 8, 16]\n",
|
|
"\n",
|
|
"thread_times = {label: [] for label, _, _, _ in METHOD_DEFS}\n",
|
|
"ttimes_stddev = {label: [] for label, _, _, _ in METHOD_DEFS}\n",
|
|
"\n",
|
|
"for nt in thread_counts:\n",
|
|
" # shuffle method order per thread count to decorrelate thermal bias\n",
|
|
" run_order = list(METHOD_DEFS)\n",
|
|
" random.shuffle(run_order)\n",
|
|
"\n",
|
|
" for label, factory, _, _ in run_order:\n",
|
|
" time.sleep(COOLDOWN)\n",
|
|
" _, times = bench(factory(nt))\n",
|
|
"\n",
|
|
" med = np.median(times) * 1e3\n",
|
|
" std = np.std(times) * 1e3\n",
|
|
" thread_times[label].append(med)\n",
|
|
" ttimes_stddev[label].append(std)\n",
|
|
"\n",
|
|
" per_px = med / (ROWS * COLS) * 1e3\n",
|
|
" per_px_std = std / (ROWS * COLS) * 1e3\n",
|
|
" print(f\" {label:22s} n_threads={nt:2d} \"\n",
|
|
" f\"{med:8.2f} \u00b1 {std:6.2f} ms \"\n",
|
|
" f\"({per_px:.4f} \u00b1 {per_px_std:.4f} \u03bcs/pixel)\")\n",
|
|
" print(\"\\n\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "d5f3152d-be84-420e-8045-9c42ac5c24cb",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# =============================\n",
|
|
"# FIGURE 1: Residual histograms\n",
|
|
"# =============================\n",
|
|
"param_names = [\"A\", \"\u03bc\", \"\u03c3\"]\n",
|
|
"param_truths = [true_A, true_mu, true_sig]\n",
|
|
"\n",
|
|
"fig1, axes1 = plt.subplots(1, 3, figsize=(15, 5))\n",
|
|
"fig1.suptitle(f\"Parameter Residuals \u2014 {ROWS}\u00d7{COLS} pixels, {N_SCAN} scan points\",\n",
|
|
" fontsize=14, fontweight=\"bold\")\n",
|
|
"\n",
|
|
"for col, (pname, truth) in enumerate(zip(param_names, param_truths)):\n",
|
|
" ax = axes1[col]\n",
|
|
"\n",
|
|
" # collect residuals across all methods for shared bin edges\n",
|
|
" res_by_method = {}\n",
|
|
" all_res = []\n",
|
|
" for mname, m in methods.items():\n",
|
|
" residual = (m[\"par\"][:, :, col] - truth).ravel()\n",
|
|
" res_by_method[mname] = residual\n",
|
|
" all_res.append(residual)\n",
|
|
" all_res = np.concatenate(all_res)\n",
|
|
"\n",
|
|
" lo, hi = np.percentile(all_res, [0.5, 99.5])\n",
|
|
" edges = np.linspace(lo, hi, 101)\n",
|
|
"\n",
|
|
" for mname, r in res_by_method.items():\n",
|
|
" ax.hist(r, bins=edges, histtype=\"step\", label=mname,\n",
|
|
" color=colors[mname],\n",
|
|
" linewidth=styles[mname][\"linewidth\"],\n",
|
|
" linestyle=styles[mname][\"linestyle\"])\n",
|
|
"\n",
|
|
" ax.axvline(0, color=\"k\", linestyle=\"--\", linewidth=1, alpha=0.7)\n",
|
|
" ax.set_xlabel(f\"Fitted {pname} \u2212 True {pname}\")\n",
|
|
" ax.set_ylabel(\"Pixel count\")\n",
|
|
" ax.set_title(f\"\u0394{pname}\")\n",
|
|
" ax.legend(fontsize=8)\n",
|
|
" ax.grid(alpha=0.3)\n",
|
|
"\n",
|
|
"fig1.tight_layout()\n",
|
|
"# fig1.savefig(\"fig1_residual_histograms.png\", dpi=150, bbox_inches=\"tight\")\n",
|
|
"# print(\"\\nSaved fig1_residual_histograms.png\")\n",
|
|
"\n",
|
|
"# ====================================================\n",
|
|
"# FIGURE 2: Performance \u2014 bar chart + thread scaling\n",
|
|
"# ====================================================\n",
|
|
"fig2 = plt.figure(figsize=(14, 5))\n",
|
|
"gs = GridSpec(1, 2, figure=fig2, width_ratios=[1, 1.3])\n",
|
|
"\n",
|
|
"# -- Left: bar chart at N_THREADS --\n",
|
|
"ax2a = fig2.add_subplot(gs[0])\n",
|
|
"names = list(methods.keys())\n",
|
|
"medians = [np.median(methods[n][\"times\"]) * 1e3 for n in names]\n",
|
|
"bars = ax2a.barh(names, medians,\n",
|
|
" color=[colors[n] for n in names],\n",
|
|
" edgecolor=\"white\", height=0.5)\n",
|
|
"ax2a.set_xlabel(\"Median wall time (ms)\")\n",
|
|
"ax2a.set_title(f\"Single call \u2014 {ROWS}\u00d7{COLS} px, {N_THREADS} threads\")\n",
|
|
"for bar, val in zip(bars, medians):\n",
|
|
" ax2a.text(bar.get_width() + max(medians) * 0.02,\n",
|
|
" bar.get_y() + bar.get_height() / 2,\n",
|
|
" f\"{val:.1f} ms\", va=\"center\", fontsize=10)\n",
|
|
"ax2a.grid(axis=\"x\", alpha=0.3)\n",
|
|
"ax2a.set_xlim(0, max(medians) * 1.25)\n",
|
|
"\n",
|
|
"# -- Right: thread scaling with error bars --\n",
|
|
"ax2b = fig2.add_subplot(gs[1])\n",
|
|
"for label, _, _, _ in METHOD_DEFS:\n",
|
|
" tt = thread_times[label]\n",
|
|
" sd = ttimes_stddev[label]\n",
|
|
" speedup = [tt[0] / t for t in tt]\n",
|
|
" # propagate uncertainty: S = t0/t \u2192 \u03b4S/S = sqrt((\u03b4t0/t0)\u00b2 + (\u03b4t/t)\u00b2)\n",
|
|
" speedup_err = [\n",
|
|
" s * np.sqrt((sd[0] / tt[0])**2 + (sd[i] / tt[i])**2)\n",
|
|
" for i, s in enumerate(speedup)\n",
|
|
" ]\n",
|
|
" ax2b.errorbar(thread_counts, speedup, yerr=speedup_err,\n",
|
|
" fmt=\"o-\", label=label, color=colors[label],\n",
|
|
" linewidth=2, markersize=7, capsize=4)\n",
|
|
"\n",
|
|
"ax2b.plot(thread_counts, thread_counts, \"k--\", alpha=0.4, label=\"Ideal linear\")\n",
|
|
"ax2b.set_xlabel(\"Number of threads\")\n",
|
|
"ax2b.set_ylabel(\"Speedup vs 1 thread\")\n",
|
|
"ax2b.set_title(\"Thread scaling\")\n",
|
|
"ax2b.set_xticks(thread_counts)\n",
|
|
"ax2b.legend(fontsize=9)\n",
|
|
"ax2b.grid(alpha=0.3)\n",
|
|
"\n",
|
|
"fig2.tight_layout()\n",
|
|
"# fig2.savefig(\"fig2_performance.png\", dpi=150, bbox_inches=\"tight\")\n",
|
|
"# print(\"Saved fig2_performance.png\")\n",
|
|
"\n",
|
|
"plt.show()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "92a8c724-e139-45d4-a354-1b8408637ead",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": []
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "Python 3 (ipykernel)",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.11.15"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|