diff --git a/notebooks/searches/mle.ipynb b/notebooks/searches/mle.ipynb index 96482f8..3a6ddd9 100644 --- a/notebooks/searches/mle.ipynb +++ b/notebooks/searches/mle.ipynb @@ -12,11 +12,14 @@ " - `Drawer`: Draws a fixed number of samples uniformly from the priors (useful for sensitivity mapping and\n", " quantifying stochasticity of the likelihood function).\n", " - `LBFGS`: The scipy L-BFGS-B optimization algorithm.\n", + " - `MultiStartAdam`: A JAX / `optax` multi-start first-order gradient MAP optimizer (with `MultiStartADABelief`\n", + " and `MultiStartLion` as drop-in alternatives).\n", "\n", "Relevant links:\n", "\n", " - Drawer: https://github.com/PyAutoLabs/PyAutoFit/blob/main/autofit/non_linear/optimize/drawer/drawer.py\n", " - L-BFGS: https://docs.scipy.org/doc/scipy/reference/optimize.minimize-lbfgsb.html\n", + " - MultiStart: https://github.com/PyAutoLabs/PyAutoFit/blob/main/autofit/non_linear/search/mle/multi_start_gradient/search.py\n", "\n", "__Contents__\n", "\n", @@ -26,7 +29,8 @@ "- **Model + Analysis**: Setting up the model and analysis shared by every search below.\n", "- **Search: Drawer**: Configuring and running the Drawer search.\n", "- **Search: LBFGS**: Configuring and running the L-BFGS-B optimizer.\n", - "- **Search Internal**: Accessing the internal optimizer for advanced use (shown once for LBFGS)." + "- **Search Internal**: Accessing the internal optimizer for advanced use (shown once for LBFGS).\n", + "- **Search: MultiStartAdam**: Running the JAX multi-start gradient MAP optimizer." ] }, { @@ -337,7 +341,69 @@ "source": [ "search_internal = result.search_internal\n", "\n", - "print(search_internal)\n" + "print(search_internal)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search: MultiStartAdam__\n", + "\n", + "We now use `MultiStartAdam`, a JAX / `optax` multi-start first-order gradient MAP optimizer.\n", + "\n", + "Unlike a single-start optimizer such as `LBFGS`, it launches `n_starts` independent optimizations from broad,\n", + "randomly drawn starting points, all in parallel via `jax.vmap`, and returns the best one. Taking a fixed\n", + "self-normalised Adam step per start, this wide population of starts reliably escapes the local maxima that\n", + "trap single-start gradient and line-search methods, making it a robust maximum a posteriori (MAP) optimizer.\n", + "\n", + "The search is JAX-native, so unlike the `Drawer` and `LBFGS` examples above it requires a JAX-traceable\n", + "analysis (`use_jax=True`). No JAX pytree registration of the model is needed: the search builds each model\n", + "instance inside its own traced objective from a plain parameter vector, so the model never has to cross a\n", + "`jax.jit` boundary as a pytree.\n", + "\n", + "`MultiStartADABelief` and `MultiStartLion` are drop-in alternatives which simply swap the local `optax`\n", + "update rule; `Lion` is sign-based and therefore prefers a ~10x smaller `learning_rate`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_jax = af.ex.Analysis(data=data, noise_map=noise_map, use_jax=True)\n", + "\n", + "search = af.MultiStartAdam(\n", + " path_prefix=\"searches\",\n", + " name=\"MultiStartAdam\",\n", + " n_starts=16,\n", + " n_steps=500,\n", + " learning_rate=0.5,\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis_jax)\n", + "\n", + "model_data = result.max_log_likelihood_instance.model_data_from(\n", + " xvalues=np.arange(data.shape[0])\n", + ")\n", + "\n", + "plt.errorbar(\n", + " x=range(data.shape[0]),\n", + " y=data,\n", + " yerr=noise_map,\n", + " linestyle=\"\",\n", + " color=\"k\",\n", + " ecolor=\"k\",\n", + " elinewidth=1,\n", + " capsize=2,\n", + ")\n", + "plt.plot(range(data.shape[0]), model_data, color=\"r\")\n", + "plt.title(\"MultiStartAdam model fit to 1D Gaussian dataset.\")\n", + "plt.xlabel(\"x values of profile\")\n", + "plt.ylabel(\"Profile normalization\")\n", + "plt.show()\n", + "plt.close()\n" ], "outputs": [], "execution_count": null diff --git a/scripts/searches/mle.py b/scripts/searches/mle.py index 9e90f0d..f296ca4 100644 --- a/scripts/searches/mle.py +++ b/scripts/searches/mle.py @@ -7,11 +7,14 @@ - `Drawer`: Draws a fixed number of samples uniformly from the priors (useful for sensitivity mapping and quantifying stochasticity of the likelihood function). - `LBFGS`: The scipy L-BFGS-B optimization algorithm. + - `MultiStartAdam`: A JAX / `optax` multi-start first-order gradient MAP optimizer (with `MultiStartADABelief` + and `MultiStartLion` as drop-in alternatives). Relevant links: - Drawer: https://github.com/PyAutoLabs/PyAutoFit/blob/main/autofit/non_linear/optimize/drawer/drawer.py - L-BFGS: https://docs.scipy.org/doc/scipy/reference/optimize.minimize-lbfgsb.html + - MultiStart: https://github.com/PyAutoLabs/PyAutoFit/blob/main/autofit/non_linear/search/mle/multi_start_gradient/search.py __Contents__ @@ -22,6 +25,7 @@ - **Search: Drawer**: Configuring and running the Drawer search. - **Search: LBFGS**: Configuring and running the L-BFGS-B optimizer. - **Search Internal**: Accessing the internal optimizer for advanced use (shown once for LBFGS). +- **Search: MultiStartAdam**: Running the JAX multi-start gradient MAP optimizer. """ # from autoconf import setup_notebook; setup_notebook() @@ -214,3 +218,54 @@ search_internal = result.search_internal print(search_internal) + +""" +__Search: MultiStartAdam__ + +We now use `MultiStartAdam`, a JAX / `optax` multi-start first-order gradient MAP optimizer. + +Unlike a single-start optimizer such as `LBFGS`, it launches `n_starts` independent optimizations from broad, +randomly drawn starting points, all in parallel via `jax.vmap`, and returns the best one. Taking a fixed +self-normalised Adam step per start, this wide population of starts reliably escapes the local maxima that +trap single-start gradient and line-search methods, making it a robust maximum a posteriori (MAP) optimizer. + +The search is JAX-native, so unlike the `Drawer` and `LBFGS` examples above it requires a JAX-traceable +analysis (`use_jax=True`). No JAX pytree registration of the model is needed: the search builds each model +instance inside its own traced objective from a plain parameter vector, so the model never has to cross a +`jax.jit` boundary as a pytree. + +`MultiStartADABelief` and `MultiStartLion` are drop-in alternatives which simply swap the local `optax` +update rule; `Lion` is sign-based and therefore prefers a ~10x smaller `learning_rate`. +""" +analysis_jax = af.ex.Analysis(data=data, noise_map=noise_map, use_jax=True) + +search = af.MultiStartAdam( + path_prefix="searches", + name="MultiStartAdam", + n_starts=16, + n_steps=500, + learning_rate=0.5, +) + +result = search.fit(model=model, analysis=analysis_jax) + +model_data = result.max_log_likelihood_instance.model_data_from( + xvalues=np.arange(data.shape[0]) +) + +plt.errorbar( + x=range(data.shape[0]), + y=data, + yerr=noise_map, + linestyle="", + color="k", + ecolor="k", + elinewidth=1, + capsize=2, +) +plt.plot(range(data.shape[0]), model_data, color="r") +plt.title("MultiStartAdam model fit to 1D Gaussian dataset.") +plt.xlabel("x values of profile") +plt.ylabel("Profile normalization") +plt.show() +plt.close()