{ "cells": [ { "cell_type": "markdown", "id": "0", "metadata": {}, "source": [ "# Ray Tune generic black-box optimiser\n", "\n", "[Ray Tune](https://docs.ray.io/en/latest/tune/index.html) is a hyperparameter tuning library primarily developed for machine learning.\n", "However, it is suitable for generic black-box optimisation as well.\n", "For our purpose, it provides an interface for running simulations inside a given *search space* and optimising for a given *loss function* $\\mathcal{L}$ using a given *algorithm*.\n", "It automatically manages checkpointing, logging and, importantly, parallel (or even distributed) computing.\n", "\n", "You can see installation instructions [here](https://docs.ray.io/en/latest/ray-overview/installation.html), but the regular pip install should work for most. Notably, ARM-based macOS support is experimental.\n", "\n", "```console\n", "pip install \"ray[tune,air]\" hyperopt\n", "```\n", "\n", "You can optimise a `mmi1x2` component for a transmission of $|S_{21}|^2 = 0.5$ (50% power) for a given wavelength using MEEP." ] }, { "cell_type": "code", "execution_count": null, "id": "1", "metadata": {}, "outputs": [], "source": [ "from functools import partial\n", "\n", "import gdsfactory as gf\n", "import numpy as np\n", "import ray\n", "import ray.air\n", "import ray.air.session\n", "from gdsfactory.config import PATH\n", "from gdsfactory.generic_tech import get_generic_pdk\n", "from ray import tune\n", "\n", "import gplugins.gmeep as gm\n", "\n", "gf.config.rich_output()\n", "PDK = get_generic_pdk()\n", "PDK.activate()\n", "\n", "tmp = PATH.optimiser\n", "tmp.mkdir(exist_ok=True)" ] }, { "cell_type": "markdown", "id": "2", "metadata": { "lines_to_next_cell": 2 }, "source": [ "## Loss function $\\mathcal{L}$\n", "\n", "The loss function is very important and should be designed to be meaningful for your need.\n", "\n", "The easiest method to optimise for a specific value is to use $L_1$ or $L_2$ (MSE) loss. Different optimisation algorithms might prefer more or less aggressive behaviour close to target, so choose depending on that.\n", "$$\n", "\\begin{align*}\n", " L_1(x) &= |x_\\text{target} - x|, \\\\\n", " L_2(x) &= \\left(x_\\text{target} - x\\right)^2\n", " .\n", "\\end{align*}\n", "$$\n" ] }, { "cell_type": "code", "execution_count": null, "id": "3", "metadata": {}, "outputs": [], "source": [ "def loss_S21_L1(x, target):\n", " r\"\"\"Loss function. Returns :math:`$\\sum_i L_1(x_i)$` and :math:`$x$` as a tuple\"\"\"\n", " return np.abs(target - x), x" ] }, { "cell_type": "markdown", "id": "4", "metadata": {}, "source": [ "Let's select a target of $0.7$ for $S_{21}$" ] }, { "cell_type": "code", "execution_count": null, "id": "5", "metadata": {}, "outputs": [], "source": [ "loss = partial(loss_S21_L1, target=0.5)" ] }, { "cell_type": "markdown", "id": "6", "metadata": {}, "source": [ "## Optimisation settings\n", "\n", "Here we specify the search space, the optimiser and its settings.\n", "\n", "
\n", " Note Choosing a new optimiser often requires you to install a separate package, see Ray Tune → Search Algorithms for details. Here one needs to install Hyperopt.\n", "
" ] }, { "cell_type": "code", "execution_count": null, "id": "7", "metadata": {}, "outputs": [], "source": [ "search_config = {\n", " \"length_mmi\": tune.uniform(0.05, 2),\n", " \"width_mmi\": tune.uniform(0.05, 2),\n", "}\n", "\n", "# pylint: disable=wrong-import-position,ungrouped-imports\n", "from ray.tune.search.hyperopt import HyperOptSearch\n", "\n", "tune_config = tune.TuneConfig(\n", " metric=\"loss\",\n", " mode=\"min\",\n", " search_alg=HyperOptSearch(),\n", " max_concurrent_trials=2, # simulations to run in parallel\n", " num_samples=-1, # max iterations, can be -1 for infinite\n", " time_budget_s=60\n", " * 20, # time after which optimisation is stopped. May be useful along with ``num_samples=-1``.\n", ")" ] }, { "cell_type": "markdown", "id": "8", "metadata": {}, "source": [ "## Implement a *trainable* function\n", "\n", "You need to implement a function which can be *trained* to be improved w.r.t. our $\\mathcal{L}$.\n", "In other words, we create a function for a single training step, which generates, runs, and returns output $\\mathcal{L}(\\vec{x})$ from simulations for given parameters $\\vec{x}$. This may require a bit more effort and some shell scripting to get right depending on your simulations.\n", "\n", "Here we demonstrate a trainable for S-parameter simulations. The `write_sparameters_meep` returns $\\mathbf{S}$ as a function of $\\lambda$ given in $\\text{µm}$. From this, we select $S_{21}(\\lambda)$ and try to optimise for $\\min_\\text{geometry} \\sum_\\lambda (S_{21}(\\lambda) - \\text{target})$. In other words, that the transmission from 1 to 2 would be as close to target as possible for the given wavelength (or range of wavelengths).\n" ] }, { "cell_type": "code", "execution_count": null, "id": "9", "metadata": {}, "outputs": [], "source": [ "use_mpi = False # change this to true if you have MPI support" ] }, { "cell_type": "code", "execution_count": null, "id": "10", "metadata": {}, "outputs": [], "source": [ "def trainable_simulations(config):\n", " \"\"\"Training step, or `trainable`, function for Ray Tune to run simulations and return results.\"\"\"\n", "\n", " # Component to optimise\n", " component = gf.components.mmi1x2(**config)\n", "\n", " # Simulate and get output\n", " dirpath = tmp / ray.air.session.get_trial_id()\n", "\n", " meep_params = dict(\n", " component=component,\n", " run=True,\n", " dirpath=dirpath,\n", " wavelength_start=1.5,\n", " # wavelength_stop=1.6,\n", " wavelength_points=1,\n", " )\n", "\n", " if use_mpi: # change this to false if no MPI support\n", " s_params = gm.write_sparameters_meep_mpi(\n", " cores=2,\n", " **meep_params, # set this to be same as in `tune.Tuner`\n", " )\n", " s_params = np.load(s_params) # parallel version returns filepath to npz instead\n", " else:\n", " s_params = gm.write_sparameters_meep(**meep_params)\n", "\n", " s_params_abs = np.abs(s_params[\"o2@0,o1@0\"]) ** 2\n", "\n", " loss_x, x = loss(s_params_abs)\n", " if not np.isscalar(x): # for many wavelengths, consider sum and mean\n", " loss_x, x = loss_x.sum(), x.mean()\n", "\n", " return {\"loss\": loss_x, \"value\": x}\n", "\n", " # ALTERNATIVE\n", " # For a shell-based solution to more software, subprocess.run is recommended roughly as below\n", " # interpreter = shutil.which('bash')\n", " # subprocess.run(\n", " # [interpreter, '-c', './generated_simulation.sh'],\n", " # cwd=dirpath,\n", " # check=True,\n", " # )" ] }, { "cell_type": "markdown", "id": "11", "metadata": {}, "source": [ "## Run optimiser\n", "In the cell below, we gather all the code above to a [`tune.Tuner`](https://docs.ray.io/en/latest/tune/api_docs/execution.html#tuner) instance and start the optimisation by calling `tuner.fit()`." ] }, { "cell_type": "code", "execution_count": null, "id": "12", "metadata": {}, "outputs": [], "source": [ "tuner = tune.Tuner(\n", " tune.with_resources(\n", " trainable_simulations, {\"cpu\": 2}\n", " ), # maximum resources given to a worker, it also supports 'gpu'\n", " param_space=search_config,\n", " tune_config=tune_config,\n", " run_config=ray.air.RunConfig(\n", " local_dir=tmp / \"ray_results\",\n", " checkpoint_config=ray.air.CheckpointConfig(checkpoint_frequency=1),\n", " log_to_file=True,\n", " verbose=2, # Intermediate results in Jupyter\n", " ),\n", ")\n", "\n", "# Previous simulations can be restored with, see https://docs.ray.io/en/latest/tune/tutorials/tune-stopping.html\n", "# tuner = Tuner.restore(path=tmp / \"ray_results/my_experiment\")\n", "\n", "results = tuner.fit()" ] }, { "cell_type": "markdown", "id": "13", "metadata": {}, "source": [ "The results can be seen and manipulated in DataFrame" ] }, { "cell_type": "code", "execution_count": null, "id": "14", "metadata": {}, "outputs": [], "source": [ "df = results.get_dataframe()\n", "df" ] }, { "cell_type": "markdown", "id": "15", "metadata": {}, "source": [ "There are clearly many possible solutions, so making a [Pareto front](https://en.wikipedia.org/wiki/Pareto_front) plot w.r.t. some other parameter like overall size would make sense here." ] }, { "cell_type": "code", "execution_count": null, "id": "16", "metadata": {}, "outputs": [], "source": [ "best_params = results.get_best_result(metric=\"loss\", mode=\"min\").metrics\n", "best_params[\"loss\"], best_params[\"config\"]" ] } ], "metadata": { "jupytext": { "cell_metadata_filter": "-all", "main_language": "python", "notebook_metadata_filter": "-all" } }, "nbformat": 4, "nbformat_minor": 5 }