{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parallelization\n",
    "\n",
    "A `DataFit` evaluates its objectives on a pool of worker processes. This example\n",
    "shows how to choose the size of that pool, how the size you ask for relates to\n",
    "the size you get, and when a pool is worth having at all.\n",
    "\n",
    "Parallelism is a property of *how a fit runs*, not of the fit itself, so it is\n",
    "configured with an `ExecutionConfig` passed to `run` rather than with a keyword\n",
    "on `DataFit`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "import ionworkspipeline as iwp\n",
    "from ionworkspipeline.execution import (\n",
    "    ExecutionConfig,\n",
    "    PrintSink,\n",
    "    default_local_workers,\n",
    ")\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pybamm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up a fit worth parallelizing\n",
    "\n",
    "We fit two active material volume fractions of a DFN against three discharge\n",
    "curves. Three objectives and an expensive model is a shape where a worker pool\n",
    "can actually help — we will measure how much, and see the conditions under which\n",
    "it does not.\n",
    "\n",
    "First, synthetic data at three C-rates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = pybamm.lithium_ion.DFN()\n",
    "parameter_values = iwp.ParameterValues(\"Chen2020\")\n",
    "parameter_values.update(\n",
    "    {\n",
    "        \"Negative electrode active material volume fraction\": 0.75,\n",
    "        \"Positive electrode active material volume fraction\": 0.665,\n",
    "    }\n",
    ")\n",
    "\n",
    "C_RATES = (0.2, 0.5, 1.0)\n",
    "\n",
    "\n",
    "def generate_data(c_rate):\n",
    "    values = parameter_values.copy()\n",
    "    values[\"Current function [A]\"] = c_rate * 5.0\n",
    "    duration = 3300 / c_rate\n",
    "    t = np.arange(0, duration, duration / 300)\n",
    "    sim = iwp.Simulation(\n",
    "        model, parameter_values=values, t_eval=[0, duration], t_interp=t\n",
    "    )\n",
    "    sol = sim.solve()\n",
    "    return pd.DataFrame(\n",
    "        {x: sol[x](sol.t) for x in [\"Time [s]\", \"Current [A]\", \"Voltage [V]\"]}\n",
    "    )\n",
    "\n",
    "\n",
    "objectives = {\n",
    "    f\"{c_rate}C\": iwp.objectives.CurrentDriven(\n",
    "        generate_data(c_rate), options={\"model\": model}\n",
    "    )\n",
    "    for c_rate in C_RATES\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We fit with CMA-ES, a population optimizer: each iteration proposes a whole\n",
    "generation of candidate points, and every point in that generation can be\n",
    "evaluated at the same time. That is what makes a fit worth parallelizing — a\n",
    "optimizer that produces one point at a time can only ever occupy one worker per\n",
    "start, however large the pool.\n",
    "\n",
    "Bounds are kept tight around the true values so every sampled point is a real\n",
    "solve rather than a cheap failure, which keeps the timings below comparable.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = {\n",
    "    \"Negative electrode active material volume fraction\": iwp.Parameter(\n",
    "        \"Negative electrode active material volume fraction\", bounds=(0.70, 0.80)\n",
    "    ),\n",
    "    \"Positive electrode active material volume fraction\": iwp.Parameter(\n",
    "        \"Positive electrode active material volume fraction\", bounds=(0.62, 0.72)\n",
    "    ),\n",
    "}\n",
    "known_parameters = {k: v for k, v in parameter_values.items() if k not in parameters}\n",
    "\n",
    "\n",
    "def build_data_fit(multistarts=10):\n",
    "    return iwp.DataFit(\n",
    "        objectives,\n",
    "        parameters=parameters,\n",
    "        multistarts=multistarts,\n",
    "        options={\"seed\": 0},\n",
    "        optimizer=iwp.optimizers.CMAES(population_size=8, max_iterations=15),\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How many workers can this fit use?\n",
    "\n",
    "Asking for more workers than the fit can occupy does not make it faster — the\n",
    "extra processes sit idle. `max_useful_workers()` reports that ceiling, and\n",
    "`tasks_per_point()` reports how a single parameter vector fans out.\n",
    "\n",
    "The pool is sized to `min(max_useful_workers(), num_workers)`, so these two\n",
    "numbers tell you the largest `num_workers` worth requesting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_fit = build_data_fit()\n",
    "data_fit.setup(known_parameters)\n",
    "\n",
    "print(f\"multistarts:          {data_fit.num_jobs}\")\n",
    "print(f\"tasks per point:      {data_fit.tasks_per_point()}\")\n",
    "print(f\"max useful workers:   {data_fit.max_useful_workers()}\")\n",
    "print(f\"default pool on this machine: {default_local_workers()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`tasks_per_point` is 3 because each parameter vector is scored against three\n",
    "objectives, and those solves are independent, so they are dispatched as separate\n",
    "tasks. `max_useful_workers` multiplies that fan-out by the width of one round of\n",
    "the search: `multistarts x population_size x tasks_per_point`, or\n",
    "`10 x 8 x 3 = 240` here. Beyond that the fit has no work to hand out.\n",
    "\n",
    "The default pool is a fraction of the machine's cores rather than all of them,\n",
    "so a fit started from an interactive session leaves the machine usable.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What one evaluation costs\n",
    "\n",
    "Before timing the pool, it is worth knowing what it is dividing. The *first*\n",
    "evaluation in a process pays the model build and solver compile; every\n",
    "subsequent one is just a solve. Timing both on the freshly set-up fit shows the\n",
    "gap.\n",
    "\n",
    "This has to be measured now, before any fit has run — the compiled model is\n",
    "cached for the lifetime of the process, so a \"cold\" call made later is not cold."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x0 = np.array(data_fit.x_initial_values_list)\n",
    "\n",
    "start = time.perf_counter()\n",
    "data_fit.objective_function(x0)\n",
    "cold = time.perf_counter() - start\n",
    "\n",
    "# A nearby, distinct point — the objective caches on exact inputs, so re-running\n",
    "# x0 would report a cache hit rather than a solve.\n",
    "start = time.perf_counter()\n",
    "data_fit.objective_function(x0 * 1.001)\n",
    "warm = time.perf_counter() - start\n",
    "\n",
    "print(f\"cold evaluation (build + solve): {cold:.3f}s\")\n",
    "print(f\"warm evaluation (solve only):    {warm:.3f}s\")\n",
    "print(f\"ratio: {cold / warm:.0f}x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each worker pays that build cost once. A pool of `N` workers therefore only wins\n",
    "when the warm work it removes exceeds `N` builds plus process startup — which is\n",
    "why a cheap objective (an analytic model, or a short ODE solve) usually runs\n",
    "*slower* on a local pool than in-process. The expensive model here is what makes\n",
    "the pool worth starting."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Choosing the worker count\n",
    "\n",
    "Pass `num_workers` on an `ExecutionConfig`. This is the only knob needed for\n",
    "local execution: the pool, the in-flight task limit and the choice of executor\n",
    "are all derived from it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def time_fit(num_workers):\n",
    "    fit = build_data_fit()\n",
    "    start = time.perf_counter()\n",
    "    result = fit.run(\n",
    "        known_parameters, execution=ExecutionConfig(num_workers=num_workers)\n",
    "    )\n",
    "    return time.perf_counter() - start, float(np.min(np.atleast_1d(result.costs)))\n",
    "\n",
    "\n",
    "worker_counts = [1, 4, 8]\n",
    "results = {n: time_fit(n) for n in worker_counts}\n",
    "timings = {n: seconds for n, (seconds, _) in results.items()}\n",
    "\n",
    "baseline = timings[worker_counts[0]]\n",
    "for n, (seconds, cost) in results.items():\n",
    "    print(\n",
    "        f\"num_workers={n:>2}: {seconds:6.1f}s  ({baseline / seconds:.1f}x)  cost={cost:.4e}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(list(timings), list(timings.values()), \"o-\")\n",
    "ax.set_xlabel(\"num_workers\")\n",
    "ax.set_ylabel(\"Fit wall time [s]\")\n",
    "ax.set_ylim(bottom=0)\n",
    "ax.grid(alpha=0.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The speedup is real but sub-linear, and it flattens well before\n",
    "`max_useful_workers`. Two costs explain the gap, and both are worth knowing\n",
    "before choosing a number:\n",
    "\n",
    "- **Each worker builds the model once.** The pool ships the objectives to every\n",
    "  process, and the first solve in each pays the model build and solver compile.\n",
    "- **The pool is bounded by physical cores.** Requesting more workers than the\n",
    "  machine has oversubscribes it, and the curve turns back up.\n",
    "\n",
    "So the useful range is roughly \"up to the core count, and no more than\n",
    "`max_useful_workers()`\". Note that the fitted cost is identical at every worker\n",
    "count: the pool changes how fast the search runs, never what it finds.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Watching a parallel fit make progress\n",
    "\n",
    "A `print` inside an objective is useless once a pool is running: it executes in\n",
    "a worker process, whose standard output is discarded. Progress has to come from\n",
    "the driver instead, and that is what a **progress sink** is for.\n",
    "\n",
    "`PrintSink` is the built-in one — a line per generation, and a line when the fit\n",
    "finishes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit = build_data_fit(multistarts=2)\n",
    "fit.run(\n",
    "    known_parameters,\n",
    "    execution=ExecutionConfig(num_workers=4, progress_sink=PrintSink()),\n",
    ")\n",
    "fig, axes = fit.plot_trace()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each line is tagged with the start that produced it, so the two starts above\n",
    "interleave as they progress. The sink runs on the driver, so this output is\n",
    "identical whether the fit runs on one worker or fifty — which is exactly what a\n",
    "`print` in an objective cannot give you.\n",
    "\n",
    "The sink only decides where progress is shown. The fit records its own trace\n",
    "whichever sink you pass, which is why `plot_trace()` works above.\n",
    "\n",
    "For an `ArrayDataFit`, each line also names its element — `[25.0:0]` is start\n",
    "`0` of the element at `25.0` — since every element restarts its start ids from\n",
    "`\"0\"`. To show your own names instead, pass `labels`; one is used per element,\n",
    "in run order, including an element that is canceled before any start finishes:\n",
    "\n",
    "```python\n",
    "PrintSink(labels=[f\"{T} K\" for T in temperatures])\n",
    "```\n",
    "\n",
    "A sink is also the place to send progress somewhere other than the terminal —\n",
    "implement `emit(event)` and pass it as `progress_sink` to record generations to\n",
    "a file, a database, or a live plot. Its last event for each fit is an\n",
    "`ExecutionFinished`, delivered however the fit ended.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Objective-level parallelism\n",
    "\n",
    "`tasks_per_point` was 3 because the three objectives were dispatched separately.\n",
    "That is the default for a multi-objective fit, and it multiplies the width of\n",
    "every generation. It is set on the `DataFit` rather than the `ExecutionConfig`,\n",
    "because whether objectives can be split depends on the cost function, not on the\n",
    "machine.\n",
    "\n",
    "Turning it off batches all of a point's objectives into one task: fewer, larger\n",
    "tasks and a narrower fit. That is the right trade when the objectives are cheap\n",
    "enough that per-task overhead dominates.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for mode in (\"auto\", \"off\"):\n",
    "    fit = iwp.DataFit(\n",
    "        objectives,\n",
    "        parameters=parameters,\n",
    "        multistarts=4,\n",
    "        options={\"seed\": 0},\n",
    "        optimizer=iwp.optimizers.CMAES(population_size=8, max_iterations=15),\n",
    "        objective_parallelism=mode,\n",
    "    )\n",
    "    fit.setup(known_parameters)\n",
    "    print(\n",
    "        f\"objective_parallelism={mode!r:>7}: \"\n",
    "        f\"tasks_per_point={fit.tasks_per_point()}, \"\n",
    "        f\"max_useful_workers={fit.max_useful_workers()}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "- `ExecutionConfig(num_workers=N)` is the one setting for local pool size.\n",
    "- The pool is `min(max_useful_workers(), N)`; check `max_useful_workers()`\n",
    "  before asking for more. For a population optimizer that ceiling is\n",
    "  `multistarts x population_size x tasks_per_point`.\n",
    "- A pool pays off only when the warm evaluation cost is large relative to the\n",
    "  per-worker model build. Measure both before scaling up.\n",
    "- Watch progress with `progress_sink=PrintSink()`; a `print` inside an objective\n",
    "  runs in a worker and is discarded.\n",
    "- `num_workers=1` skips the pool entirely and runs in-process, which is the\n",
    "  easiest way to debug an objective that raises.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".env",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}