{ "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", "from scipy.special import erf\n", "import matplotlib.pyplot as plt\n", "from matplotlib.gridspec import GridSpec\n", "from aare import FallingScurve\n" ] }, { "cell_type": "code", "execution_count": null, "id": "2536939d-8f34-438e-850a-6f3bdf48f658", "metadata": {}, "outputs": [], "source": [ "import aare\n", "aare.__file__" ] }, { "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 # fraction of step height (p4)\n", "SEED = 42\n", "N_THREADS = 4\n", "\n", "N_REPEATS = 10\n", "N_WARMUP = 3 # untimed iterations (icache + branch predictor warmup)\n", "COOLDOWN = 2.0 # seconds between (method, thread_count) pairs" ] }, { "cell_type": "markdown", "id": "1d35abe9-8d28-4417-ad49-ff1e18e298c1", "metadata": {}, "source": [ "## Synthetic data genrator" ] }, { "cell_type": "code", "execution_count": null, "id": "06c8fddb-56d9-4f84-8b21-4ffd8b7bd26f", "metadata": {}, "outputs": [], "source": [ "# Models in use\n", "def scurve(x, p0, p1, p2, p3, p4, p5): # rising Scurve\n", " z = (x - p2) / (np.sqrt(2) * p3)\n", " return (p0 + p1 * x) + 0.5 * (1 + erf(z)) * (p4 + p5 * (x - p2))\n", "\n", "def scurve2(x, p0, p1, p2, p3, p4, p5): # Falling Scurve\n", " z = (x - p2) / (np.sqrt(2) * p3)\n", " return (p0 + p1 * x) + 0.5 * (1 - erf(z)) * (p4 + p5 * (x - p2))\n", " \n", "def generate_3d_scurve_data(fn, rows, cols, n_scan, noise_frac, seed):\n", " \"\"\"\n", " Synthetic detector image stack with per-pixel rising S-curves.\n", " \n", " Returns\n", " -------\n", " x : (n_scan,)\n", " y : (rows, cols, n_scan)\n", " y_err : (rows, cols, n_scan)\n", " truths : dict {name: (rows, cols) array}\n", " \"\"\"\n", " rng = np.random.default_rng(seed)\n", " \n", " # Per-pixel ground truth - shapes (rows, cols)\n", " p0_true = rng.uniform(-5, 5, (rows, cols)) # baseline offset\n", " p1_true = rng.uniform(-0.02, 0.02, (rows, cols)) # baseline slope\n", " p2_true = rng.uniform(30, 70, (rows, cols)) # threshold\n", " p3_true = rng.uniform(2, 8, (rows, cols)) # width\n", " p4_true = rng.uniform(200, 800, (rows, cols)) # step height\n", " p5_true = rng.uniform(-0.5, 0.5, (rows, cols)) # post-step slope\n", " \n", " x = np.linspace(0, 100, n_scan)\n", " \n", " # broadcast: x -> (1, 1, n_scan), params -> (rows, cols, 1)\n", " y_clean = fn(\n", " x[None, None, :],\n", " p0_true[:, :, None], p1_true[:, :, None],\n", " p2_true[:, :, None], p3_true[:, :, None],\n", " p4_true[:, :, None], p5_true[:, :, None],\n", " )\n", " \n", " noise_sigma = noise_frac * p4_true[:, :, None] * np.ones_like(y_clean)\n", " y = y_clean + rng.normal(0, noise_sigma)\n", " y_err = noise_sigma.copy()\n", " \n", " truths = dict(p0=p0_true, p1=p1_true, p2=p2_true,\n", " p3=p3_true, p4=p4_true, p5=p5_true)\n", " return x, y, y_err, truths " ] }, { "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", " for _ in range(n_warmup):\n", " res = fn()\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": "24ec20a3-8117-4718-9236-19d37897e04e", "metadata": {}, "source": [ "## Data generation" ] }, { "cell_type": "code", "execution_count": null, "id": "13bcd5a6-8f95-4300-a002-04023d0f9bab", "metadata": {}, "outputs": [], "source": [ "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, truths = generate_3d_scurve_data(\n", " scurve2, ROWS, COLS, N_SCAN, NOISE_FRAC, SEED\n", ")" ] }, { "cell_type": "markdown", "id": "e16e2559-059b-4058-88c9-a9632988fb74", "metadata": {}, "source": [ "## Define the tested/benched methods" ] }, { "cell_type": "code", "execution_count": null, "id": "1f6bc651-80c1-41dd-8a05-7f15aba006aa", "metadata": {}, "outputs": [], "source": [ "model = FallingScurve()\n", "METHOD_DEFS = [\n", " (\"Minuit2 (analytic)\",\n", " lambda nt: lambda: model.fit(x, y, n_threads=nt),\n", " \"#4CAF50\", {\"linewidth\": 4.0, \"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": "markdown", "id": "bc2f50a4-ea9d-4b0d-86f6-39b2f71e9592", "metadata": {}, "source": [ "## Single-call benchmard" ] }, { "cell_type": "code", "execution_count": null, "id": "5a417145-ce42-4c3a-a7bf-05ff97ba0450", "metadata": {}, "outputs": [], "source": [ "PARAM_NAMES = [\"p0\", \"p1\", \"p2\", \"p3\", \"p4\", \"p5\"]\n", "ndf = N_SCAN - len(PARAM_NAMES)\n", "def extract_result(label, res):\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", "# ---- summary table ----\n", "header = f\"{'Method':24s} {'time (ms)':>10s}\"\n", "for pn in PARAM_NAMES:\n", " header += f\" {'med|d' + pn + '|':>10s}\"\n", "print(header)\n", "print(\"-\" * (26 + 12 + 12 * len(PARAM_NAMES)))\n", " \n", "for name, m in methods.items():\n", " par = m[\"par\"]\n", " med_t = np.median(m[\"times\"]) * 1e3\n", " \n", " deltas = \" \".join(\n", " f\"{np.median(np.abs(par[:, :, i] - truths[pn])):10.4f}\"\n", " for i, pn in enumerate(PARAM_NAMES)\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 {deltas}{chi2_str}\")" ] }, { "cell_type": "markdown", "id": "7b6a3419-52a7-4b1b-92b6-b160d4227f64", "metadata": {}, "source": [ "## Thread scaling" ] }, { "cell_type": "code", "execution_count": null, "id": "55ecc77c-0823-408d-a811-8e4f4900f332", "metadata": { "scrolled": true }, "outputs": [], "source": [ "thread_counts = [1, 2, 4, 8]\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", " 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} \u00b5s/pixel)\")\n", " \n", " print(\"\\n\")" ] }, { "cell_type": "markdown", "id": "7aedaf85-0c70-449f-94c1-4acd6ff991eb", "metadata": {}, "source": [ "## Visualization: Residuals & Performance" ] }, { "cell_type": "code", "execution_count": null, "id": "d5f3152d-be84-420e-8045-9c42ac5c24cb", "metadata": {}, "outputs": [], "source": [ "# FIGURE 1: Residual histograms (6 panels)\n", "truth_arrays = [truths[pn] for pn in PARAM_NAMES]\n", " \n", "fig1, axes1 = plt.subplots(2, 3, figsize=(16, 9))\n", "fig1.suptitle(\n", " f\"Parameter Residuals \u2014 {ROWS}\u00d7{COLS} pixels, {N_SCAN} scan points\",\n", " fontsize=14, fontweight=\"bold\")\n", " \n", "for idx, (pname, truth) in enumerate(zip(PARAM_NAMES, truth_arrays)):\n", " ax = axes1.flat[idx]\n", " \n", " res_by_method = {}\n", " all_res = []\n", " for mname, m in methods.items():\n", " residual = (m[\"par\"][:, :, idx] - 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", "\n", "# FIGURE 2: bar chart + thread scaling\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", "fig2 = plt.figure(figsize=(14, 5))\n", "gs = GridSpec(1, 2, figure=fig2, width_ratios=[1, 1.3])\n", " \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", "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", " 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", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "id": "cda3dc45-eec8-466a-a0e1-6d62fa37d5da", "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 }