examples/noisy_optimization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🦾🤖 Noise-aware optimization scheme with ✨ XLuminA ✨: \n",
    "\n",
    "This notebook is a step-by-step guide for building a robust (noise-aware) optimization scheme in XLuminA.\n",
    "\n",
    "We will set-up an optimization scheme for *the sharp focus for a radially polarized light beam* - (we use **robust_discovery** from **optical_elements.py**)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "# Setting the path for XLuminA modules:\n",
    "current_path = os.path.abspath(os.path.join('..'))\n",
    "module_path = os.path.join(current_path)\n",
    "\n",
    "if module_path not in sys.path:\n",
    "    sys.path.append(module_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from xlumina.__init__ import um, nm, mm\n",
    "from xlumina.vectorized_optics import *\n",
    "from xlumina.optical_elements import robust_discovery\n",
    "from xlumina.toolbox import space, softmin\n",
    "from xlumina.loss_functions import vectorized_loss_hybrid\n",
    "import jax.numpy as jnp\n",
    "import jax\n",
    "from jax import random\n",
    "import optax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## System specs, define light sources, output dimensions and static parameters during optimization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. System specs:\n",
    "sensor_lateral_size = 512  # Resolution\n",
    "wavelength1 = 650*nm\n",
    "x_total = 2500*um\n",
    "x, y = space(x_total, sensor_lateral_size)\n",
    "shape = jnp.shape(x)[0]\n",
    "\n",
    "# 2. Define the optical functions: two orthogonally polarized beams:\n",
    "w0 = (1200*um, 1200*um)  \n",
    "ls1 = PolarizedLightSource(x, y, wavelength1)\n",
    "ls1.gaussian_beam(w0=w0, jones_vector=(1, 1))\n",
    "\n",
    "# 3. Define the output (High Resolution) detection:\n",
    "x_out, y_out = jnp.array(space(10*um, 400)) # Pixel size detector: 20 um / 400 pix \n",
    "\n",
    "# 4. High NA objective lens specs:\n",
    "NA = 0.9 \n",
    "radius_lens = 3.6*mm/2 \n",
    "f_lens = radius_lens / NA\n",
    "\n",
    "# 5. Static parameters - don't change during optimization:\n",
    "fixed_params = [radius_lens, f_lens, x_out, y_out]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the optical setup:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Vectorized version of the optical setup [`optical_elements.py` > `robust_discovery`] over a new axis (defined by the noise).\n",
    "\n",
    "    Here the args of the function are: `light source` (ls1 - ls6, common to all tables), `parameters` (common to all tables), `fixed_params` (common to all tables), `noise` (DIFFERENT for each table). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def batch_robust_discovery(ls1, ls2, ls3, ls4, ls5, ls6, parameters, fixed_params, noise_distances, noise_slms, noise_wps, noise_amps, distance_offset):\n",
    "    \"\"\"\n",
    "    Vectorized (efficient) version of robust_discovery() for batch optimization. \n",
    "    \n",
    "    Parameters: \n",
    "        ls1, ls2, ls3, ls4, ls5, ls6 (PolarizedLightSource)\n",
    "        parameters (jnp.array): parameters to pass to the optimizer\n",
    "            BB 1: [phase1_1, phase1_2, eta1, theta1, z1_1, z1_2]\n",
    "            BB 2: [phase2_1, phase2_2, eta2, theta2, z2_1, z2_2] \n",
    "            BB 3: [phase3_1, phase3_2, eta3, theta3, z3_1, z3_2]\n",
    "            BS ratios: [bs1, bs2, bs3, bs4, bs5, bs6, bs7, bs8, bs9] <- automatically contains \n",
    "            Extra distances: [z4, z5]\n",
    "        fixed_params (jnp.array): parameters to maintain fixed during optimization [r, f, xout and yout]; that is radius and focal length of the objective lens.\n",
    "         \n",
    "        noise_distances (jnp.array): Misalignment (in microns): [noise_z1_1, noise_z1_2, ...]\n",
    "        noise_slms (jnp.array): Noise (in radians): [noise_phase1_1, noise_phase1_2, ...]\n",
    "        noise_wps (jnp.array): Noise (in radians): [noise_eta1, noise_theta1, ...]\n",
    "        noise_amps (jnp.array): Noise (in AU): [noise_A1, noise_A2, ...]\n",
    "        \n",
    "    Returns vectorized version of detected light (intensity tensor): (# tables, (6, resolution, resolution))\n",
    "    \"\"\"\n",
    "    # Noise shapes are: \n",
    "    # distance: (#tables, 1, 8); \n",
    "    # slms (amp and phase): (#tables, 6, (resolution,resolution)); \n",
    "    # wp: (#tables, 1, 6)\n",
    "    # vmap in axes 0 -> across optical tables\n",
    "    detected_intensities_z = vmap(robust_discovery, in_axes=(None, None, None, None, None, None, None, None, \n",
    "                                                             0, 0, 0, 0, None))(ls1, ls2, ls3, ls4, ls5, ls6, parameters, fixed_params, \n",
    "                                                                             noise_distances, noise_slms, noise_wps, noise_amps, distance_offset)\n",
    "    return detected_intensities_z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. Vectorize the loss function: `vmap` the computation of the loss function across different optical tables.(imported from `loss_functions.py` > `vectorized_loss_hybrid`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mean_batch_discover(detected_light):\n",
    "    \"\"\" vmap loss over optical tables: \"\"\"\n",
    "    # detected_light with shape (#optical tables, (6, N, N))\n",
    "    # vmap loss in axis #optical tables \n",
    "    return vmap(vectorized_loss_hybrid, in_axes=(0,))(detected_light)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. Define the loss: first computes light from paralellized optical tables (`batch_robust_discovery`), later compute the loss function for each detector in each optical table. Finally, compute the mean loss value across the optical tables and get the minimum value using `softmin`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit \n",
    "def loss_batch_discovery(parameters, noise_d, noise_slm, noise_wp, noise_amp):\n",
    "    \"\"\"\n",
    "    Loss function. It computes L= Area/I_{epsilon} across detectors.\n",
    "    \n",
    "    Parameters:\n",
    "        parameters (list): Optimized parameters.\n",
    "        noise_distances (jnp.array): Misalignment (in microns): [noise_d1, noise_d2, ...]\n",
    "        noise_slms (jnp.array): Noise (in radians): [noise_slm_1, noise_slm_2, ...]\n",
    "        noise_wps (jnp.array): Noise (in radians): [noise_eta, noise_theta, ...]\n",
    "        noise_amp (jnp.array): Noise (in AU): [noise_A1, noise_A2, ...]\n",
    "\n",
    "    Returns the mean value of the loss computed for all the inputs. \n",
    "    \"\"\"\n",
    "    # Output from batch_robust_discovery is (#optical tables, (6, N, N)): for 6 detectors each\n",
    "    detected_z_intensities = batch_robust_discovery(ls1, ls1, ls1, ls1, ls1, ls1, parameters, fixed_params, noise_d, noise_slm, noise_wp, noise_amp, distance_offset = 15) \n",
    "\n",
    "    # Get the minimum value within loss value array\n",
    "    # output from mean_batch_discover is (#optical tables, (6, 1)). \n",
    "    # Compute the mean across #optical tables and get the minimum value using softmin.\n",
    "    loss_val = softmin(jnp.mean(mean_batch_discover(detected_z_intensities), axis=0, keepdims=True))\n",
    "    return loss_val # shape (#optical tables, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimizer settings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Global variable\n",
    "shape = jnp.array([sensor_lateral_size, sensor_lateral_size])\n",
    "# Define the loss function:\n",
    "loss_function = loss_batch_discovery"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optimization settings\n",
    "OS = {'n_best': 500,\n",
    "      'best_loss': 3*1e2,\n",
    "      'num_iterations': 50000,\n",
    "      'num_samples': 1,\n",
    "      'WEIGHT_DECAY': 1e-3,\n",
    "      'BASE_lr': 0.05,\n",
    "      'END_lr': 0.001,\n",
    "      'DECAY_STEPS': 4000\n",
    "     }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [!!]  define noise settings dictionary:\n",
    "\n",
    "**NS (dict)**: \n",
    "\n",
    "    NS = {'n_tables': __, \n",
    "          'number of distances': __,\n",
    "          'number of sSLM': __, \n",
    "          'number of wps': __,\n",
    "          'noise_level': __, \n",
    "          'misalignment': (minval, maxval), \n",
    "          'phase_noise': (minval, maxval),\n",
    "          'discretize': __}\n",
    "\n",
    "where, \n",
    "\n",
    "**n_tables (int)**: number of optical tables to compute in parallel\n",
    "\n",
    "**number of distances, number of sSLM, number of wps (str)**: number of distances, sSLM and wave plates in the optical setup.\n",
    "\n",
    "**level (str)**:\n",
    "\n",
    "    1. low: noise in SLMs and WPs $\\pm$(0.01 to 0.05) rads and misalignment of $\\pm$(0.01 to 0.05) mm \n",
    "    2. mild: noise in SLMs and WPs $\\pm$(0.05 to 0.5) rads and misalignment of $\\pm$(0.05 to 0.5) mm \n",
    "    3. high: noise in SLMs and WPs $\\pm$(0.5 to 1) rads and misalignment of $\\pm$(0.5 to 1) mm \n",
    "    4. all: noise in SLMs and WPs $\\pm$(0.01 to 1) rads and misalignment of $\\pm$(0.01 to 1) mm \n",
    "    5. tunable: tunable noise via NS dictionary\n",
    "\n",
    "**discretize (bool)**: if true, discretize SLM noise to 8-bit.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Noise settings:\n",
    "NS = {'n_tables': 3, \n",
    "      'number of distances': 8,\n",
    "      'number of sSLM': 3, \n",
    "      'number of wps': 3,\n",
    "      'noise_level': 'tunable',\n",
    "      'misalignment': (10, 100), \n",
    "      'phase_noise': (0.01, 0.1),\n",
    "      'discretize': False}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Define keychain for the number of optical tables specified in NS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Keychain for optical tables\n",
    "def keychain_optical_tables(seed, number_of_tables):\n",
    "    \"\"\" \n",
    "    Generates keychain for # optical tables especified \n",
    "    \"\"\"\n",
    "    keychain = []\n",
    "    for num in range(number_of_tables):\n",
    "        key_table = random.PRNGKey(seed + num)\n",
    "        keychain.append(key_table)\n",
    "            \n",
    "    return jnp.array(keychain)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. Define `shake_setup()` and `batch_shake_setup()` as functions to include noise in the setup per iteration. \n",
    "\n",
    "    Two types of shaking functions are provided in `optical_elements.py`. \n",
    "\n",
    "    1. `shake_setup` takes noise settings NS:dict as argument. Thus, `batch_shake_setup` can't be used with @jit or @partial(jit). \n",
    "\n",
    "    2. However, if you want to @jit `batch_shake_setup`, copy-paste `shake_setup_jit` in your optimizer file, as it doesn't have NS:dict as an argument. \n",
    "\n",
    "Here we will copy-paste `shake_setup_jit` and use NS as global variable. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def shake_setup_jit(key, resolution):\n",
    "    \"\"\"\n",
    "    [THIS FUNCTION IS INTENDED TO BE PASTED IN THE OPTIMIZER FILE TO ENABLE @jit COMPILATION FOR `batch_shake_setup`]\n",
    "    \n",
    "    Creates noise for all the different optical variables on an optical table.\n",
    "    \n",
    "    Parameters:\n",
    "        key (PRNGKey): JAX random key for reproducibility\n",
    "        resolution (int): number of pixels for space\n",
    "        \n",
    "        global variable NS (dict): noise settings as\n",
    "\n",
    "            NS = {'n_tables': __, 'number of distances': __,\n",
    "                  'number of sSLM': __, 'number of wps': __,\n",
    "                  'noise_level': __, \n",
    "                  'misalignment': (minval, maxval), \n",
    "                  'phase_noise': (minval, maxval),\n",
    "                  'discretize': __}\n",
    "    \n",
    "    Returns:\n",
    "        random_noise_distances, random_noise_slms, random_noise_wps, random_noise_amps, key0 (new key to split in the next iteration), key (old key)\n",
    "    \"\"\"\n",
    "    num_physical_variables = 4 # of physical variables (e.g., distance, slm phase, ,...) to optimize.\n",
    "    # split as many times as variables + 1 to renew the key0 each step\n",
    "    key0, key1, key2, key3, key4 = random.split(key, num_physical_variables+1)\n",
    "    d_type = 'int8'\n",
    "\n",
    "    # NS is not an input to ensure vmap during optimization.\n",
    "    level = NS['noise_level']\n",
    "    discretize = NS['discretize']\n",
    "    \n",
    "    # level can be: 'low' == 0, 'mild' == 1, 'high' == 2, 'all' == 3 and 'tunable' == 4\n",
    "\n",
    "    if level == 'low':\n",
    "        # Misalignment (um)\n",
    "        minval_d = 10*um  # 0.01 mm\n",
    "        maxval_d = 50*um # 0.05 mm\n",
    "        # SLM / WP phase and amplitude (rads and AU, respectively)\n",
    "        minval_phase = 0.01 \n",
    "        maxval_phase = 0.05\n",
    "\n",
    "    if level == 'mild':\n",
    "        # Misalignment (um)\n",
    "        minval_d = 50*um  # 0.05 mm\n",
    "        maxval_d = 500*um # 0.5 mm\n",
    "        # SLM / WP phase and amplitude (rads and AU, respectively)\n",
    "        minval_phase = 0.05 \n",
    "        maxval_phase = 0.5\n",
    "\n",
    "    if level == 'high':\n",
    "        # Misalignment (um)\n",
    "        minval_d = 500*um # 0.5 mm\n",
    "        maxval_d = 1000*um # 1 mm\n",
    "        # SLM / WP phase and amplitude (rads and AU, respectively)\n",
    "        minval_phase = 0.5\n",
    "        maxval_phase = 1\n",
    "\n",
    "    if level == 'all':\n",
    "        # Misalignment (um)\n",
    "        minval_d = 10*um  # 0.01 mm\n",
    "        maxval_d = 1000*um # 0.15 mm\n",
    "        # SLM / WP phase and amplitude (rads and AU, respectively)\n",
    "        minval_phase = 0.01 \n",
    "        maxval_phase = 1\n",
    "\n",
    "    if level == 'tunable':\n",
    "        # Misalignment (um)\n",
    "        minval_d, maxval_d  = NS['misalignment']  # in um\n",
    "        # SLM / WP phase and amplitude (rads and AU, respectively)\n",
    "        minval_phase, maxval_phase = NS['phase_noise']\n",
    "\n",
    "    if discretize: \n",
    "        d_type = 'uint8'\n",
    "\n",
    "    # noise for distances (d1 and d2): shape = (1, NS['number of distances'])\n",
    "    random_noise_distances = jnp.squeeze(random.uniform(key1, shape=(1, NS['number of distances']), minval=minval_d, maxval=maxval_d), axis=0) \n",
    "    # noise for SLMs phases and amplitude (slm1 and slm2): shape = (2, (resolution, resolution))\n",
    "    random_noise_amps = random.choice(key4, jnp.array([-1,1]), shape=(2*NS['number of sSLM'], resolution, resolution)).astype(d_type) * random.uniform(key2, shape=(2*NS['number of sSLM'], resolution, resolution), minval=minval_phase, maxval=maxval_phase) \n",
    "    random_noise_slms = random.choice(key2, jnp.array([-1,1]), shape=(2*NS['number of sSLM'], resolution, resolution)).astype(d_type) * random.uniform(key2, shape=(2*NS['number of sSLM'], resolution, resolution), minval=minval_phase, maxval=maxval_phase) \n",
    "    # noise for WP angles (eta and theta): shape = (1, 2)\n",
    "    random_noise_wps = jnp.squeeze(random.choice(key3, jnp.array([-1,1]), shape=(1, 2*NS['number of wps'])).astype('int8'), axis=0) * jnp.squeeze(random.uniform(key3, shape=(1, 2*NS['number of wps']), minval=minval_phase, maxval=maxval_phase), axis=0) \n",
    "    \n",
    "    return random_noise_distances, random_noise_slms, random_noise_wps, random_noise_amps, key0, key\n",
    "\n",
    "\n",
    "@jit\n",
    "def batch_shake_setup(key_array, array_for_shape):\n",
    "    \"\"\"\n",
    "    Creates noise for all the different optimizable variables on multiple optical tables given by size(key_array).\n",
    "    \n",
    "    Parameters: \n",
    "        key_array (PRNGKey): Array with different keys -- will change for each step in the optimization. \n",
    "        The dimension of this array is decided by # of optical tables to compute in parallel.\n",
    "        array_for_shape (jnp.array): array of shape [resolution, resolution] to make it jit. \n",
    "        \n",
    "    Returns:\n",
    "        random_noise_distances [with shape = (size(key_array), 1, NS['number of distances'])], \n",
    "        random_noise_amps [with shape = (size(key_array), 2*NS['number of sSLM'], resolution, resolution)], \n",
    "        random_noise_slms [with shape = (size(key_array), 2*NS['number of sSLM'], resolution, resolution)], \n",
    "        random_noise_wps [with shape = (size(key_array), 1, 2*NS['number of wps'])]\n",
    "        key0 (PRNGKey): array with key0 to split in the next iteration step\n",
    "    \"\"\"\n",
    "    return vmap(shake_setup_jit, in_axes = (0, None))(key_array, jnp.shape(array_for_shape)[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define optimizer (adamw with schedule)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def adamw_schedule(base_lr, end_lr, decay_steps, weight_decay) -> optax.GradientTransformation:\n",
    "    \"\"\"\n",
    "    Custom optimizer - adamw: applies several transformations in sequence\n",
    "    1) Apply ADAMW\n",
    "    2) Apply lr schedule\n",
    "    \"\"\"\n",
    "    lr_schedule = base_lr\n",
    "    #lr_schedule = optax.linear_schedule(init_value= base_lr, end_value = end_lr, transition_steps = decay_steps, transition_begin = 500)                                           \n",
    "    return optax.adamw(learning_rate=lr_schedule, weight_decay=weight_decay)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimization loop: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit(params: optax.Params, optimizer: optax.GradientTransformation, num_iterations, keys, x) -> optax.Params:\n",
    "    \n",
    "    # Init the optimizer with initial parameters\n",
    "    opt_state = optimizer.init(params)\n",
    "\n",
    "    @jit\n",
    "    def update(parameters, opt_state, noise_d, noise_slm, noise_wp, noise_amp):\n",
    "        # Define single update step - contains noise_array: \n",
    "        loss_value, grads = jax.value_and_grad(loss_function)(parameters, noise_d, noise_slm, noise_wp, noise_amp)\n",
    "        # Update the state of the optimizer\n",
    "        updates, state = optimizer.update(grads, opt_state, parameters)\n",
    "        # Update the parameters\n",
    "        new_params = optax.apply_updates(parameters, updates)\n",
    "        \n",
    "        return new_params, parameters, state, loss_value, updates\n",
    "\n",
    "    # Initialize some parameters    \n",
    "    n_best = OS['n_best']\n",
    "    best_loss = OS['best_loss']\n",
    "    best_params = None\n",
    "    best_keys = None\n",
    "    best_step = 0\n",
    "    \n",
    "    print('Starting Optimization', flush=True)\n",
    "    \n",
    "    for step in range(num_iterations):\n",
    "        \n",
    "        # Add noise: update noise and keys each iteration\n",
    "        noise_d, noise_slm, noise_wp, noise_amp, keys, old_keys = batch_shake_setup(keys, x) # 'x' is the space variable from optical table\n",
    "        \n",
    "        # Apply update step\n",
    "        params, old_params, opt_state, loss_value, grads = update(params, opt_state, noise_d, noise_slm, noise_wp, noise_amp)\n",
    "        \n",
    "        print(f\"Step {step}\")\n",
    "        print(f\"Loss {loss_value}\")\n",
    "        \n",
    "        # Update the `best_loss` value:\n",
    "        if loss_value < best_loss:\n",
    "            # Best loss value\n",
    "            best_loss = loss_value\n",
    "            # Best optimized parameters\n",
    "            best_params = old_params\n",
    "            # Keys for best params\n",
    "            best_keys = old_keys\n",
    "            best_step = step\n",
    "            print('Best loss value is updated')\n",
    "\n",
    "        if step % 100 == 0:\n",
    "            # Stopping criteria: if best_loss has not changed every 500 steps, stop.\n",
    "            if step - best_step > n_best:\n",
    "                print(f'Stopping criterion: no improvement in loss value for {n_best} steps')\n",
    "                break\n",
    "    \n",
    "    print(f'Best loss: {best_loss} at step {best_step}')\n",
    "    print(f'Best parameters: {best_params}')  \n",
    "    return best_params, best_loss, best_keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optimizer settings\n",
    "num_iterations = OS['num_iterations']\n",
    "num_samples = OS['num_samples']\n",
    "\n",
    "for i in range(num_samples):\n",
    "    tic = time.perf_counter()\n",
    "    \n",
    "    # seed1 to ensure randomness among samples\n",
    "    seed1 = np.random.randint(9999)\n",
    "    \n",
    "    # Init keychain for noise -- as many init keys as optical tables to parallelize\n",
    "    keys = keychain_optical_tables(seed1, NS['n_tables'])\n",
    "        \n",
    "    # Optimizer settings\n",
    "    WEIGHT_DECAY = OS['WEIGHT_DECAY']\n",
    "    BASE_lr = OS['BASE_lr']\n",
    "    END_lr = OS['END_lr']\n",
    "    DECAY_STEPS = OS['DECAY_STEPS']\n",
    "    \n",
    "    # Random init parameters:\n",
    "    phase1_1 = jnp.array([np.random.uniform(0, 1, shape)], dtype=jnp.float64)[0]\n",
    "    phase1_2 = jnp.array([np.random.uniform(0, 1, shape)], dtype=jnp.float64)[0]\n",
    "    a1_1 = jnp.array([np.random.uniform(0, 1, shape)], dtype=jnp.float64)[0]\n",
    "    a1_2 = jnp.array([np.random.uniform(0, 1, shape)], dtype=jnp.float64)[0]\n",
    "    \n",
    "    phase2_1 = jnp.array([np.random.uniform(0, 1, shape)], dtype=jnp.float64)[0]\n",
    "    phase2_2 = jnp.array([np.random.uniform(0, 1, shape)], dtype=jnp.float64)[0]\n",
    "    a2_1 = jnp.array([np.random.uniform(0, 1, shape)], dtype=jnp.float64)[0]\n",
    "    a2_2 = jnp.array([np.random.uniform(0, 1, shape)], dtype=jnp.float64)[0]\n",
    "    \n",
    "    phase3_1 = jnp.array([np.random.uniform(0, 1, shape)], dtype=jnp.float64)[0]\n",
    "    phase3_2 = jnp.array([np.random.uniform(0, 1, shape)], dtype=jnp.float64)[0]\n",
    "    a3_1 = jnp.array([np.random.uniform(0, 1, shape)], dtype=jnp.float64)[0]\n",
    "    a3_2 = jnp.array([np.random.uniform(0, 1, shape)], dtype=jnp.float64)[0]\n",
    "    \n",
    "    eta1 = jnp.array([np.random.uniform(0, 1, 1)], dtype=jnp.float64)[0]\n",
    "    theta1 = jnp.array([np.random.uniform(0, 1, 1)], dtype=jnp.float64)[0]\n",
    "    eta2 = jnp.array([np.random.uniform(0, 1, 1)], dtype=jnp.float64)[0]\n",
    "    theta2 = jnp.array([np.random.uniform(0, 1, 1)], dtype=jnp.float64)[0]\n",
    "    eta3 = jnp.array([np.random.uniform(0, 1, 1)], dtype=jnp.float64)[0]\n",
    "    theta3 = jnp.array([np.random.uniform(0, 1, 1)], dtype=jnp.float64)[0]\n",
    "    eta4 = jnp.array([np.random.uniform(0, 1, 1)], dtype=jnp.float64)[0]\n",
    "    theta4 = jnp.array([np.random.uniform(0, 1, 1)], dtype=jnp.float64)[0]\n",
    "    \n",
    "    z1_1 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    z1_2 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    z2_1 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    z2_2 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    z3_1 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    z3_2 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    z4 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    z5 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    \n",
    "    bs1 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    bs2 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    bs3 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    bs4 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    bs5 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    bs6 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    bs7 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    bs8 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    bs9 = jnp.array([np.random.uniform(0, 1)], dtype=jnp.float64)\n",
    "    \n",
    "    # Init params for 3x3 robust discovery\n",
    "    init_params = [phase1_1, phase1_2, a1_1, a1_2, eta1, theta1, z1_1, z1_2, \n",
    "                   phase2_1, phase2_2, a2_1, a2_2, eta2, theta2, z2_1, z2_2, \n",
    "                   phase3_1, phase3_2, a3_1, a3_2, eta3, theta3, z3_1, z3_2, \n",
    "                   bs1, bs2, bs3, bs4, bs5, bs6, bs7, bs8, bs9, \n",
    "                   z4, z5]\n",
    "                   \n",
    "    # Init optimizer:\n",
    "    optimizer = adamw_schedule(BASE_lr, END_lr, DECAY_STEPS, WEIGHT_DECAY)\n",
    "\n",
    "    # Apply fit function:\n",
    "    best_params, best_loss, iteration_steps, loss_list, keys_noise = fit(init_params, optimizer, num_iterations, keys, x)\n",
    "    \n",
    "    print(f\"Time taken to optimize one sample - in seconds {(time.perf_counter() - tic):.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dummy_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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}