{ "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 }