examples/test_diffractio_vs_xlumina.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2dae72b8-43d4-40e7-8dce-6f3458e42818",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Testing XLuminA: \n",
    "\n",
    "(1) *Diffractio vs XLuminA propagation functions.*\n",
    "\n",
    "(2) *JAX (autodiff-based ADAM) vs SciPy's (BFGS) optimization using XLuminA.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "e3b8dee2-f6c2-4383-8444-e193942adfd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "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": "markdown",
   "id": "cca4db6a-5b8c-4933-830b-5b3ed45c54a6",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## (1) Speed test on propagation functions:\n",
    "\n",
    "    1. Scalar Rayleigh-Sommerfeld (RS)\n",
    "    2. Scalar Chirped z-transform (CZT)\n",
    "    3. Vectorial Rayleigh-Sommerfeld (VRS)\n",
    "    4. Vectorial Chirped z-transform (VCZT)\n",
    "    \n",
    "*Test for a computational window size of 2048x2048. You can change the sampling via \"resolution\".* "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bb59d384-9481-403c-b86e-2459f7b80ee8",
   "metadata": {},
   "outputs": [],
   "source": [
    "resolution = 2048"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "628755b6-5572-4593-8300-92dbdbed3ef2",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Diffractio: *runs on CPU device*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "65196a3c-41d5-4a0d-98c4-732d3893a8f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffractio import np, degrees, um, mm\n",
    "from diffractio.scalar_sources_XY import Scalar_source_XY\n",
    "from diffractio.vector_sources_XY import Vector_source_XY\n",
    "from diffractio.scalar_fields_XY import Scalar_field_XY\n",
    "from diffractio.vector_fields_XY import Vector_field_XY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "79843f19-9128-4417-a150-bd2566e43a11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# System specs:\n",
    "wavelength = .6328 * um\n",
    "w0 = (1200*um , 1200*um)\n",
    "x = np.linspace(-15 * mm, 15 * mm, resolution)\n",
    "y = np.linspace(-15 * mm, 15 * mm, resolution)\n",
    "x_out = np.linspace(-15 * mm, 15 * mm, resolution)\n",
    "y_out = np.linspace(-15 * mm, 15 * mm, resolution)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc6204ac-b6ad-4a2d-a6eb-bab8ea31a60d",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Scalar optics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b51a4438-1dfb-4083-99b3-1962838389ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time taken for RS propagation - in seconds 4.051304008811712\n",
      "Time taken for CZT propagation - in seconds 1.9945963714271784\n"
     ]
    }
   ],
   "source": [
    "ls = Scalar_source_XY(x, y, wavelength, info='Light source')\n",
    "ls.gauss_beam(r0=(0 * um, 0 * um), w0=w0, z0=(0,0), A=1, theta=0 * degrees, phi=0 * degrees)\n",
    "\n",
    "# Rayleigh-Sommerfeld:\n",
    "tic = time.perf_counter()\n",
    "ls_propagated = ls.RS(z=5*mm)\n",
    "print(\"Time taken for RS propagation - in seconds\", time.perf_counter() - tic)\n",
    "\n",
    "# Chirped z-transform:\n",
    "tic = time.perf_counter()\n",
    "ls_propagated = ls.CZT(z=5*mm, xout=x_out, yout=y_out, verbose=False)\n",
    "print(\"Time taken for CZT propagation - in seconds\", time.perf_counter() - tic)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfc48484-3781-4800-bbdd-a3e5fd583c0f",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Vectorial optics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b3ea108d-9513-49d1-8a07-9fcbfbccac08",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time taken for VRS propagation - in seconds 12.461026553064585\n",
      "Time taken for VCZT propagation - in seconds 6.241402745246887\n"
     ]
    }
   ],
   "source": [
    "vls = Vector_source_XY(x, y, wavelength=wavelength, info='Light source polarization')\n",
    "vls.constant_polarization(u=ls, v=(1, 0), has_normalization=False, radius=(15*mm, 15*mm)) \n",
    "\n",
    "# Rayleigh-Sommerfeld:\n",
    "tic = time.perf_counter()\n",
    "vls.VRS(z=5*mm, n=1, new_field=False, verbose=False, amplification=(1, 1))\n",
    "print(\"Time taken for VRS propagation - in seconds\", time.perf_counter() - tic)\n",
    "\n",
    "# Chirped z-transform:\n",
    "tic = time.perf_counter()\n",
    "vls.CZT(z=5*mm, xout=x_out, yout=y_out, verbose=False)\n",
    "print(\"Time taken for VCZT propagation - in seconds\", time.perf_counter() - tic)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdb9398d-0b54-4dd5-839b-3047599a6a64",
   "metadata": {
    "tags": []
   },
   "source": [
    "### XLuminA: *runs on CPU and GPU*\n",
    "\n",
    "*If you are running this notebook in a* **CPU** *node, the software will automatically fall back to CPU. You don't need to change anything :)*\n",
    "\n",
    "*If you are running this notebook in a* **GPU** *node, you can enable the* **CPU** *via adding these lines of code at the beginning of the cell:*\n",
    "\n",
    "    import os\n",
    "    os.environ['CUDA_VISIBLE_DEVICES'] = '[]'\n",
    "    \n",
    "*For XLuminA's test we run the functions* **twice** *: the 1st time a function that contains the* **just-in-time decorator** *(@jit) is called, JAX takes extra time to store the function in cache. The second (and subsequent) runs just call the pre-compiled function so the computation is faster.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "39aad6ae-8181-4218-b46a-473b4ad3f8e0",
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'xlumina'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[3], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mxlumina\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m__init__\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m mm\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mxlumina\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mwave_optics\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;241m*\u001b[39m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mxlumina\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mvectorized_optics\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;241m*\u001b[39m\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'xlumina'"
     ]
    }
   ],
   "source": [
    "from xlumina.__init__ import mm\n",
    "from xlumina.wave_optics import *\n",
    "from xlumina.vectorized_optics import *\n",
    "from xlumina.optical_elements import *\n",
    "from xlumina.loss_functions import *\n",
    "from xlumina.toolbox import space\n",
    "import jax.numpy as jnp\n",
    "\n",
    "wavelength = .6328 * um\n",
    "w0 = (1200*um , 1200*um)\n",
    "x, y = space(15*mm, resolution)\n",
    "x_out, y_out = jnp.array(space(15*mm, resolution))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a65c47b3-2ca3-49a6-8b77-9bbacc8ebd68",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Scalar optics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5f6041ef-9ba2-4e98-842f-933211cdefd5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time taken to perform one RS propagation (in seconds): 1.6592820473015308\n",
      "Time taken for 1st RS propagation - in seconds 1.6601537950336933\n",
      "Time taken to perform one RS propagation (in seconds): 0.009246926754713058\n",
      "Time taken for 2nd RS propagation - in seconds 0.00949065387248993\n",
      "Time taken to perform one CZT propagation (in seconds): 2.8923826217651367\n",
      "Time taken for 1st CZT propagation - in seconds 2.8929686844348907\n",
      "Time taken to perform one CZT propagation (in seconds): 0.030335675925016403\n",
      "Time taken for 2nd CZT propagation - in seconds 0.030557945370674133\n"
     ]
    }
   ],
   "source": [
    "gb = LightSource(x, y, wavelength)\n",
    "gb.gaussian_beam(w0=w0, E0=1)\n",
    "\n",
    "# Rayleigh-Sommerfeld:\n",
    "tic = time.perf_counter() \n",
    "ls_propagated = gb.RS_propagation(z=5*mm)\n",
    "print(\"Time taken for 1st RS propagation - in seconds\", time.perf_counter() - tic)\n",
    "\n",
    "tic = time.perf_counter() \n",
    "ls_propagated = gb.RS_propagation(z=5*mm)\n",
    "print(\"Time taken for 2nd RS propagation - in seconds\", time.perf_counter() - tic)\n",
    "\n",
    "# Chirped z-transform:\n",
    "tic = time.perf_counter() \n",
    "ls_propagated = gb.CZT(z=5*mm, xout=x_out, yout=y_out)\n",
    "print(\"Time taken for 1st CZT propagation - in seconds\", time.perf_counter() - tic)\n",
    "\n",
    "tic = time.perf_counter() \n",
    "ls_propagated = gb.CZT(z=5*mm, xout=x_out, yout=y_out)\n",
    "print(\"Time taken for 2nd CZT propagation - in seconds\", time.perf_counter() - tic)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76b503cc-2d29-4845-be11-7ed6ba03aeb0",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Vectorial optics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bd6ca97f-3d3c-4990-bcb4-26ca3ce30fae",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time taken to perform one VRS propagation (in seconds): 1.3885535262525082\n",
      "Time taken for 1st VRS propagation - in seconds 1.3888123966753483\n",
      "Time taken to perform one VRS propagation (in seconds): 0.14746328070759773\n",
      "Time taken for 2nd VRS propagation - in seconds 0.14762601256370544\n",
      "Time taken to perform one VCZT propagation (in seconds): 1.1658967472612858\n",
      "Time taken for 1st VCZT propagation - in seconds 1.1661826558411121\n",
      "Time taken to perform one VCZT propagation (in seconds): 0.07554645836353302\n",
      "Time taken for 2nd VCZT propagation - in seconds 0.07574645429849625\n"
     ]
    }
   ],
   "source": [
    "gb_lp = PolarizedLightSource(x, y, wavelength)\n",
    "gb_lp.gaussian_beam(w0=w0, jones_vector=(1, 0))\n",
    "\n",
    "# Rayleigh-Sommerfeld:\n",
    "tic = time.perf_counter() \n",
    "gb_propagated = gb_lp.VRS_propagation(z=5*mm)\n",
    "print(\"Time taken for 1st VRS propagation - in seconds\", time.perf_counter() - tic)\n",
    "\n",
    "tic = time.perf_counter() \n",
    "gb_propagated = gb_lp.VRS_propagation(z=5*mm)\n",
    "print(\"Time taken for 2nd VRS propagation - in seconds\", time.perf_counter() - tic)\n",
    "\n",
    "# Chirped z-transform:\n",
    "tic = time.perf_counter() \n",
    "gb_propagated = gb_lp.VCZT(z=5*mm, xout=x_out, yout=y_out)\n",
    "print(\"Time taken for 1st VCZT propagation - in seconds\", time.perf_counter() - tic)\n",
    "\n",
    "tic = time.perf_counter() \n",
    "gb_propagated = gb_lp.VCZT(z=5*mm, xout=x_out, yout=y_out)\n",
    "print(\"Time taken for 2nd VCZT propagation - in seconds\", time.perf_counter() - tic)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee25250c-d579-43b4-85da-1b0ad73290df",
   "metadata": {
    "tags": []
   },
   "source": [
    "## (2) Comparison of XLuminA's autodiff optimization with numerical methods.\n",
    "\n",
    "*For this test we code an optical setup consisting of a gaussian beam interacting with a spatial light modulator (SLM) after some propagation. The ground truth is a gaussian beam with a spiral phase. The loss function is the mean squared error (MSE) of the phase between the outcome of the setup and the ground truth. Via optimization, we aim to rediscover the phase mask in the SLM device.*\n",
    "\n",
    "*We optimize using: (1) Diffractio's optical simulator with SciPy's BFGS optimizer (computing the gradients* **numerically** *) and (2) XLuminA with the ADAM optimizer from JAX, which due to autodiff, computes the gradients* **analytically**. *We run the optimization for different* **resolution**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "c08c8556-0d31-45c7-b519-687ac156abb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "resolution = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "909fb323-dfbe-48db-a076-f9ae48c11903",
   "metadata": {
    "tags": []
   },
   "source": [
    "### XLuminA with autodiff optimization (ADAM):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "cdbebd1e-9374-4d0f-8fd6-479b1a7cd361",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xlumina.__init__ import um, nm, mm, degrees, radians\n",
    "from xlumina.wave_optics import *\n",
    "from xlumina.optical_elements import SLM\n",
    "from xlumina.toolbox import space\n",
    "import jax.numpy as jnp\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "beddd43f-c1e7-4708-8dc1-2a64ff0a0f0b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x1477247e21f0>"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAZgAAAGdCAYAAAAv9mXmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAWH0lEQVR4nO3dbYyUhf3v4d+ysLNgl7ViIBIR8aQJCBqVNY2Cto2GxKfUpMdWo9Zom9QjKkiOUaptoy1u7IMh0YpZ0xhbg/KiNdKktiU2glaNiPiQtpHTmshWaqgeu4uiCzsz58X5u/3TEdwBftz3rNeVzAsmM843s8t8uHfWudvq9Xo9AOAgG1f0AADGJoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFOMP9QPWarXYtm1bdHV1RVtb26F+eAAOQL1ejx07dsT06dNj3Lh9H6Mc8sBs27YtZsyYcagfFoCDqL+/P44++uh93uaQB6arqysiIr5w3P+K8eMqh/rh96reMaHoCQ3qEw/5l+cTVUu4abizfJt2H1a+nz5XO8v3E4PhEm6qdpRvU61E3+LVXR/Glp/dPvJavi+HfPZHPxYbP64S49tLFJj2EgamhJvaxpfoO/0jE8q3qT6hfIGJEr5w1ku4qYzPU1v5XgpG9RZHCf8WADAWCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkEBgAUggMACkEBoAU+xWYe++9N2bNmhWdnZ0xf/78eOqppw72LgBaXNOBWbNmTSxdujRuueWW2Lx5c5xxxhlxzjnnxNatWzP2AdCimg7MXXfdFd/4xjfim9/8ZsyZMydWrlwZM2bMiFWrVmXsA6BFNRWYXbt2xaZNm2LRokV7XL9o0aJ45plnPvY+Q0NDMTg4uMcFgLGvqcC8/fbbUa1WY9q0aXtcP23atHjrrbc+9j69vb3R3d09cnE2S4BPh/16k/8/TzRTr9f3evKZ5cuXx8DAwMilv79/fx4SgBbT1KkAjzzyyGhvb284Wtm+fXvDUc1HKpVKVCrlOXMlAIdGU0cwHR0dMX/+/Fi3bt0e169bty5OP/30gzoMgNbW9MnMly1bFpdffnn09PTEaaedFn19fbF169a4+uqrM/YB0KKaDszXvva1eOedd+L222+Pf/zjHzFv3rz4zW9+EzNnzszYB0CLajowERHXXHNNXHPNNQd7CwBjiM8iAyCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEixX59FdjC07d4dbVV926f2jz+JW5HaJpTwa9ZZL3pBo/J96aJWwu+nWmGvQHtXm1D0gka1jqIX/Futib9uJXy1AGAsEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFOOLeuD6B0NRH1cv6uEbtFVrRU9oCePGl+/fJG2TCvs23qv6uLaiJzSodRS9oFG1UsLnqVL0gka1En2L15p4qSzfqwUAY4LAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKZoKTG9vb5x66qnR1dUVU6dOjQsvvDBee+21rG0AtLCmArN+/fpYvHhxPPfcc7Fu3boYHh6ORYsWxfvvv5+1D4AW1dRpbH7729/u8ecHHnggpk6dGps2bYozzzzzoA4DoLUd0HnSBgYGIiLiiCOO2OtthoaGYmhoaOTPg4ODB/KQALSI/X6Tv16vx7Jly2LhwoUxb968vd6ut7c3uru7Ry4zZszY34cEoIXsd2CuvfbaeOWVV+Lhhx/e5+2WL18eAwMDI5f+/v79fUgAWsh+/Yjsuuuui7Vr18aGDRvi6KOP3udtK5VKVCqV/RoHQOtqKjD1ej2uu+66ePTRR+PJJ5+MWbNmZe0CoMU1FZjFixfH6tWr47HHHouurq546623IiKiu7s7Jk6cmDIQgNbU1Hswq1atioGBgfjiF78YRx111MhlzZo1WfsAaFFN/4gMAEbDZ5EBkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApDigUyYfiPr770e9bXdRD9+oo0Rb/kvb8HDRExqU8V8k7Z/pKHpCg2r5JsXuSW1FT2gwPKnoBY1qE4pe0Kg+vjyfA1mN0W8p4+sFAGOAwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkGF/UA9fe3xm1tt1FPXyjnTuLXtCgrb296AkN2gZ3FD2hwYSiB3yMD06aVvSEBkNT6kVPaFAr4RevXsZ/drcVPeDfak3ctoxPJQBjgMAAkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApDigwvb290dbWFkuXLj1IcwAYK/Y7MBs3boy+vr448cQTD+YeAMaI/QrMe++9F5deemncf//98dnPfvZgbwJgDNivwCxevDjOO++8OPvssz/xtkNDQzE4OLjHBYCxr+lTJj/yyCPx4osvxsaNG0d1+97e3rjtttuaHgZAa2vqCKa/vz+WLFkSDz30UHR2do7qPsuXL4+BgYGRS39//34NBaC1NHUEs2nTpti+fXvMnz9/5LpqtRobNmyIe+65J4aGhqK9vX2P+1QqlahUKgdnLQAto6nAnHXWWfHqq6/ucd2VV14Zs2fPjptuuqkhLgB8ejUVmK6urpg3b94e1x122GExZcqUhusB+HTzf/IDkKLp3yL7T08++eRBmAHAWOMIBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASDFAX8W2f56dMurMbmrPH2r1mtFT2gwHNWiJzTYUdtV9IQGLwwdUfSEBv/7lf9Z9IQGQ+9MKnpCo+G2ohc0aKuWb1PUix7wb7UY/WtleV7hARhTBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEgxfiiHvid6vuxq1qevu2s14ue0OBftcK+PHu1bfiIoic0eGHnrKInNDj68H8VPaHB34se8DE+2NlR9IQGtd3leV0aMVyeTfV6bdS3Lc9qAMYUgQEghcAAkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFI0HZg333wzLrvsspgyZUpMmjQpTjrppNi0aVPGNgBaWFMnHHn33XdjwYIF8aUvfSkef/zxmDp1avztb3+Lww8/PGkeAK2qqcDceeedMWPGjHjggQdGrjv22GMP9iYAxoCmfkS2du3a6OnpiYsuuiimTp0aJ598ctx///37vM/Q0FAMDg7ucQFg7GsqMK+//nqsWrUqPve5z8Xvfve7uPrqq+P666+Pn//853u9T29vb3R3d49cZsyYccCjASi/tnp99Cej7+joiJ6ennjmmWdGrrv++utj48aN8eyzz37sfYaGhmJoaGjkz4ODgzFjxoz461+mRVdXeX6Jbefon4ZD5l+1pn6CeUhsG+4uekKDF3bOKnpCg6ff/h9FT2jw938dXvSEBh/s7Ch6QoPa7vK8Lo0YLs+m2gcfxt+v+14MDAzE5MmT93nbplYfddRRcfzxx+9x3Zw5c2Lr1q17vU+lUonJkyfvcQFg7GsqMAsWLIjXXnttj+u2bNkSM2fOPKijAGh9TQXmhhtuiOeeey7uuOOO+Otf/xqrV6+Ovr6+WLx4cdY+AFpUU4E59dRT49FHH42HH3445s2bF9///vdj5cqVcemll2btA6BFNf0u8vnnnx/nn39+xhYAxpDy/GoCAGOKwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEhR2Bmt3q62xYfVtqIevsGOevlOfPR/q58pekKDN3d/tugJDXZUO4ue0OAzE4Y++UaHWNfED4ue0BJ2DZXvRH/VEp1wLGrDo75piVYDMJYIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0CK8UU98Lu1zthVK0/fdtQmFj2hwTvVzxQ9ocFAdVLRExoMl+j76CPjx9WKntBg0oTdRU9oMFxtL3pCS9jVVp7nqW1CddS3Ld/fTADGBIEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSNBWY4eHhuPXWW2PWrFkxceLEOO644+L222+PWq18H00OQLGaOh/MnXfeGffdd188+OCDMXfu3HjhhRfiyiuvjO7u7liyZEnWRgBaUFOBefbZZ+PLX/5ynHfeeRERceyxx8bDDz8cL7zwQso4AFpXUz8iW7hwYTzxxBOxZcuWiIh4+eWX4+mnn45zzz13r/cZGhqKwcHBPS4AjH1NHcHcdNNNMTAwELNnz4729vaoVquxYsWKuOSSS/Z6n97e3rjtttsOeCgAraWpI5g1a9bEQw89FKtXr44XX3wxHnzwwfjxj38cDz744F7vs3z58hgYGBi59Pf3H/BoAMqvqSOYG2+8MW6++ea4+OKLIyLihBNOiDfeeCN6e3vjiiuu+Nj7VCqVqFQqB74UgJbS1BHMzp07Y9y4Pe/S3t7u15QBaNDUEcwFF1wQK1asiGOOOSbmzp0bmzdvjrvuuiuuuuqqrH0AtKimAnP33XfHd77znbjmmmti+/btMX369PjWt74V3/3ud7P2AdCimgpMV1dXrFy5MlauXJk0B4CxwmeRAZBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKRo6rPIDqb3ap1Rq7UX9fANBmudRU9osKNavk07ax1FT2hQLeG/k8ZFvegJDSaMqxY9oUHH+OGiJzSo1tqKnlBq1eHRn56lfH8zARgTBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0CK8Yf6Aev1ekRE7Hyveqgfep921sq1JyLig+pw0RMafDi8u+gJDXbtLt+m3UO7ip7QYHiovegJDYbL9zRFdXet6AkNqsPleX2q7hyKiH+/lu/LIQ/Mjh07IiLiioX/51A/NAAHyY4dO6K7u3uft2mrjyZDB1GtVott27ZFV1dXtLW17fd/Z3BwMGbMmBH9/f0xefLkg7hwbPE8jY7naXQ8T6Mzlp+ner0eO3bsiOnTp8e4cft+l+WQH8GMGzcujj766IP235s8efKY+wJm8DyNjudpdDxPozNWn6dPOnL5iDf5AUghMACkaNnAVCqV+N73vheVSqXoKaXmeRodz9PoeJ5Gx/P0/x3yN/kB+HRo2SMYAMpNYABIITAApBAYAFK0bGDuvffemDVrVnR2dsb8+fPjqaeeKnpSqfT29sapp54aXV1dMXXq1LjwwgvjtddeK3pWqfX29kZbW1ssXbq06Cml8+abb8Zll10WU6ZMiUmTJsVJJ50UmzZtKnpWqQwPD8ett94as2bNiokTJ8Zxxx0Xt99+e9Rq5ftss0OlJQOzZs2aWLp0adxyyy2xefPmOOOMM+Kcc86JrVu3Fj2tNNavXx+LFy+O5557LtatWxfDw8OxaNGieP/994ueVkobN26Mvr6+OPHEE4ueUjrvvvtuLFiwICZMmBCPP/54/PnPf46f/OQncfjhhxc9rVTuvPPOuO++++Kee+6Jv/zlL/HDH/4wfvSjH8Xdd99d9LTCtOSvKX/+85+PU045JVatWjVy3Zw5c+LCCy+M3t7eApeV1z//+c+YOnVqrF+/Ps4888yi55TKe++9F6ecckrce++98YMf/CBOOumkWLlyZdGzSuPmm2+OP/7xj35K8AnOP//8mDZtWvzsZz8bue4rX/lKTJo0KX7xi18UuKw4LXcEs2vXrti0aVMsWrRoj+sXLVoUzzzzTEGrym9gYCAiIo444oiCl5TP4sWL47zzzouzzz676CmltHbt2ujp6YmLLroopk6dGieffHLcf//9Rc8qnYULF8YTTzwRW7ZsiYiIl19+OZ5++uk499xzC15WnEP+YZcH6u23345qtRrTpk3b4/pp06bFW2+9VdCqcqvX67Fs2bJYuHBhzJs3r+g5pfLII4/Eiy++GBs3bix6Smm9/vrrsWrVqli2bFl8+9vfjueffz6uv/76qFQq8fWvf73oeaVx0003xcDAQMyePTva29ujWq3GihUr4pJLLil6WmFaLjAf+c+P+q/X6wf08f9j2bXXXhuvvPJKPP3000VPKZX+/v5YsmRJ/P73v4/Ozs6i55RWrVaLnp6euOOOOyIi4uSTT44//elPsWrVKoH5b9asWRMPPfRQrF69OubOnRsvvfRSLF26NKZPnx5XXHFF0fMK0XKBOfLII6O9vb3haGX79u0NRzVEXHfddbF27drYsGHDQT1NwliwadOm2L59e8yfP3/kumq1Ghs2bIh77rknhoaGor29fGeBPNSOOuqoOP744/e4bs6cOfHLX/6yoEXldOONN8bNN98cF198cUREnHDCCfHGG29Eb2/vpzYwLfceTEdHR8yfPz/WrVu3x/Xr1q2L008/vaBV5VOv1+Paa6+NX/3qV/GHP/whZs2aVfSk0jnrrLPi1VdfjZdeemnk0tPTE5deemm89NJL4vJfFixY0PAr7lu2bImZM2cWtKicdu7c2XACrvb29k/1rym33BFMRMSyZcvi8ssvj56enjjttNOir68vtm7dGldffXXR00pj8eLFsXr16njssceiq6tr5Iivu7s7Jk6cWPC6cujq6mp4T+qwww6LKVOmeK/qv7nhhhvi9NNPjzvuuCO++tWvxvPPPx99fX3R19dX9LRSueCCC2LFihVxzDHHxNy5c2Pz5s1x1113xVVXXVX0tOLUW9RPf/rT+syZM+sdHR31U045pb5+/fqiJ5VKRHzs5YEHHih6Wql94QtfqC9ZsqToGaXz61//uj5v3rx6pVKpz549u97X11f0pNIZHBysL1mypH7MMcfUOzs768cdd1z9lltuqQ8NDRU9rTAt+f/BAFB+LfceDACtQWAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkEBgAUvw/kZ70PmbdWnkAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wavelength = 650*nm\n",
    "x_total = 1000*um \n",
    "x, y = space(x_total, resolution)\n",
    "shape = jnp.shape(x)[0]\n",
    "\n",
    "# Light source specs\n",
    "w0 = (1200*um , 1200*um)\n",
    "gb = LightSource(x, y, wavelength)\n",
    "gb.gaussian_beam(w0=w0, E0=1)\n",
    "gb_gt = ScalarLight(x, y, wavelength)\n",
    "# Set spiral phase for the ground truth\n",
    "phase_mask = jnp.arctan2(gb.Y,gb.X)\n",
    "gb_gt.field = gb.field * jnp.exp(1j * phase_mask)\n",
    "\n",
    "plt.imshow(phase_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "39398ddc-36ce-42b7-bf40-245383f37bc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the setup\n",
    "def setup(gb, parameters):\n",
    "    gb_propagated = gb.RS_propagation(z=25000*mm)\n",
    "    gb_modulated, _ = SLM(gb, parameters, gb.x.shape[0])\n",
    "    return gb_modulated\n",
    "\n",
    "def mse_phase(input_light, target_light):\n",
    "    return jnp.sum((jnp.angle(input_light.field) - jnp.angle(target_light.field)) ** 2) / resolution**2\n",
    "def loss(parameters):\n",
    "    out = setup(gb, parameters)\n",
    "    loss_val = mse_phase(out, gb_gt)\n",
    "    return loss_val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "8d363ec1-e4e6-4023-8e2b-4382884291d3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[gpu(id=0), gpu(id=1)]\n",
      "Starting Optimization\n",
      "Stopping criterion: no improvement in loss value for 50 steps\n",
      "Best loss: 9.932328871996077e-30 at step 5741; actual step 5800\n",
      "Time taken to optimize one sample - in seconds 3.0919757653027773\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x1477246df220>"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAZgAAAGdCAYAAAAv9mXmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAWH0lEQVR4nO3dbYyUhf3v4d+ysLNgl7ViIBIR8aQJCBqVNY2Cto2GxKfUpMdWo9Zom9QjKkiOUaptoy1u7IMh0YpZ0xhbg/KiNdKktiU2glaNiPiQtpHTmshWaqgeu4uiCzsz58X5u/3TEdwBftz3rNeVzAsmM843s8t8uHfWudvq9Xo9AOAgG1f0AADGJoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFOMP9QPWarXYtm1bdHV1RVtb26F+eAAOQL1ejx07dsT06dNj3Lh9H6Mc8sBs27YtZsyYcagfFoCDqL+/P44++uh93uaQB6arqysiIr5w3P+K8eMqh/rh96reMaHoCQ3qEw/5l+cTVUu4abizfJt2H1a+nz5XO8v3E4PhEm6qdpRvU61E3+LVXR/Glp/dPvJavi+HfPZHPxYbP64S49tLFJj2EgamhJvaxpfoO/0jE8q3qT6hfIGJEr5w1ku4qYzPU1v5XgpG9RZHCf8WADAWCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkEBgAUggMACkEBoAU+xWYe++9N2bNmhWdnZ0xf/78eOqppw72LgBaXNOBWbNmTSxdujRuueWW2Lx5c5xxxhlxzjnnxNatWzP2AdCimg7MXXfdFd/4xjfim9/8ZsyZMydWrlwZM2bMiFWrVmXsA6BFNRWYXbt2xaZNm2LRokV7XL9o0aJ45plnPvY+Q0NDMTg4uMcFgLGvqcC8/fbbUa1WY9q0aXtcP23atHjrrbc+9j69vb3R3d09cnE2S4BPh/16k/8/TzRTr9f3evKZ5cuXx8DAwMilv79/fx4SgBbT1KkAjzzyyGhvb284Wtm+fXvDUc1HKpVKVCrlOXMlAIdGU0cwHR0dMX/+/Fi3bt0e169bty5OP/30gzoMgNbW9MnMly1bFpdffnn09PTEaaedFn19fbF169a4+uqrM/YB0KKaDszXvva1eOedd+L222+Pf/zjHzFv3rz4zW9+EzNnzszYB0CLajowERHXXHNNXHPNNQd7CwBjiM8iAyCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEixX59FdjC07d4dbVV926f2jz+JW5HaJpTwa9ZZL3pBo/J96aJWwu+nWmGvQHtXm1D0gka1jqIX/Futib9uJXy1AGAsEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFOOLeuD6B0NRH1cv6uEbtFVrRU9oCePGl+/fJG2TCvs23qv6uLaiJzSodRS9oFG1UsLnqVL0gka1En2L15p4qSzfqwUAY4LAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKZoKTG9vb5x66qnR1dUVU6dOjQsvvDBee+21rG0AtLCmArN+/fpYvHhxPPfcc7Fu3boYHh6ORYsWxfvvv5+1D4AW1dRpbH7729/u8ecHHnggpk6dGps2bYozzzzzoA4DoLUd0HnSBgYGIiLiiCOO2OtthoaGYmhoaOTPg4ODB/KQALSI/X6Tv16vx7Jly2LhwoUxb968vd6ut7c3uru7Ry4zZszY34cEoIXsd2CuvfbaeOWVV+Lhhx/e5+2WL18eAwMDI5f+/v79fUgAWsh+/Yjsuuuui7Vr18aGDRvi6KOP3udtK5VKVCqV/RoHQOtqKjD1ej2uu+66ePTRR+PJJ5+MWbNmZe0CoMU1FZjFixfH6tWr47HHHouurq546623IiKiu7s7Jk6cmDIQgNbU1Hswq1atioGBgfjiF78YRx111MhlzZo1WfsAaFFN/4gMAEbDZ5EBkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApDigUyYfiPr770e9bXdRD9+oo0Rb/kvb8HDRExqU8V8k7Z/pKHpCg2r5JsXuSW1FT2gwPKnoBY1qE4pe0Kg+vjyfA1mN0W8p4+sFAGOAwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkGF/UA9fe3xm1tt1FPXyjnTuLXtCgrb296AkN2gZ3FD2hwYSiB3yMD06aVvSEBkNT6kVPaFAr4RevXsZ/drcVPeDfak3ctoxPJQBjgMAAkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApDigwvb290dbWFkuXLj1IcwAYK/Y7MBs3boy+vr448cQTD+YeAMaI/QrMe++9F5deemncf//98dnPfvZgbwJgDNivwCxevDjOO++8OPvssz/xtkNDQzE4OLjHBYCxr+lTJj/yyCPx4osvxsaNG0d1+97e3rjtttuaHgZAa2vqCKa/vz+WLFkSDz30UHR2do7qPsuXL4+BgYGRS39//34NBaC1NHUEs2nTpti+fXvMnz9/5LpqtRobNmyIe+65J4aGhqK9vX2P+1QqlahUKgdnLQAto6nAnHXWWfHqq6/ucd2VV14Zs2fPjptuuqkhLgB8ejUVmK6urpg3b94e1x122GExZcqUhusB+HTzf/IDkKLp3yL7T08++eRBmAHAWOMIBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASDFAX8W2f56dMurMbmrPH2r1mtFT2gwHNWiJzTYUdtV9IQGLwwdUfSEBv/7lf9Z9IQGQ+9MKnpCo+G2ohc0aKuWb1PUix7wb7UY/WtleV7hARhTBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEgxfiiHvid6vuxq1qevu2s14ue0OBftcK+PHu1bfiIoic0eGHnrKInNDj68H8VPaHB34se8DE+2NlR9IQGtd3leV0aMVyeTfV6bdS3Lc9qAMYUgQEghcAAkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFI0HZg333wzLrvsspgyZUpMmjQpTjrppNi0aVPGNgBaWFMnHHn33XdjwYIF8aUvfSkef/zxmDp1avztb3+Lww8/PGkeAK2qqcDceeedMWPGjHjggQdGrjv22GMP9iYAxoCmfkS2du3a6OnpiYsuuiimTp0aJ598ctx///37vM/Q0FAMDg7ucQFg7GsqMK+//nqsWrUqPve5z8Xvfve7uPrqq+P666+Pn//853u9T29vb3R3d49cZsyYccCjASi/tnp99Cej7+joiJ6ennjmmWdGrrv++utj48aN8eyzz37sfYaGhmJoaGjkz4ODgzFjxoz461+mRVdXeX6Jbefon4ZD5l+1pn6CeUhsG+4uekKDF3bOKnpCg6ff/h9FT2jw938dXvSEBh/s7Ch6QoPa7vK8Lo0YLs+m2gcfxt+v+14MDAzE5MmT93nbplYfddRRcfzxx+9x3Zw5c2Lr1q17vU+lUonJkyfvcQFg7GsqMAsWLIjXXnttj+u2bNkSM2fOPKijAGh9TQXmhhtuiOeeey7uuOOO+Otf/xqrV6+Ovr6+WLx4cdY+AFpUU4E59dRT49FHH42HH3445s2bF9///vdj5cqVcemll2btA6BFNf0u8vnnnx/nn39+xhYAxpDy/GoCAGOKwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEhR2Bmt3q62xYfVtqIevsGOevlOfPR/q58pekKDN3d/tugJDXZUO4ue0OAzE4Y++UaHWNfED4ue0BJ2DZXvRH/VEp1wLGrDo75piVYDMJYIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0CK8UU98Lu1zthVK0/fdtQmFj2hwTvVzxQ9ocFAdVLRExoMl+j76CPjx9WKntBg0oTdRU9oMFxtL3pCS9jVVp7nqW1CddS3Ld/fTADGBIEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSNBWY4eHhuPXWW2PWrFkxceLEOO644+L222+PWq18H00OQLGaOh/MnXfeGffdd188+OCDMXfu3HjhhRfiyiuvjO7u7liyZEnWRgBaUFOBefbZZ+PLX/5ynHfeeRERceyxx8bDDz8cL7zwQso4AFpXUz8iW7hwYTzxxBOxZcuWiIh4+eWX4+mnn45zzz13r/cZGhqKwcHBPS4AjH1NHcHcdNNNMTAwELNnz4729vaoVquxYsWKuOSSS/Z6n97e3rjtttsOeCgAraWpI5g1a9bEQw89FKtXr44XX3wxHnzwwfjxj38cDz744F7vs3z58hgYGBi59Pf3H/BoAMqvqSOYG2+8MW6++ea4+OKLIyLihBNOiDfeeCN6e3vjiiuu+Nj7VCqVqFQqB74UgJbS1BHMzp07Y9y4Pe/S3t7u15QBaNDUEcwFF1wQK1asiGOOOSbmzp0bmzdvjrvuuiuuuuqqrH0AtKimAnP33XfHd77znbjmmmti+/btMX369PjWt74V3/3ud7P2AdCimgpMV1dXrFy5MlauXJk0B4CxwmeRAZBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKRo6rPIDqb3ap1Rq7UX9fANBmudRU9osKNavk07ax1FT2hQLeG/k8ZFvegJDSaMqxY9oUHH+OGiJzSo1tqKnlBq1eHRn56lfH8zARgTBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkEBgAUggMACkEBoAUAgNACoEBIIXAAJBCYABIITAApBAYAFIIDAApBAaAFAIDQAqBASCFwACQQmAASCEwAKQQGABSCAwAKQQGgBQCA0CK8Yf6Aev1ekRE7Hyveqgfep921sq1JyLig+pw0RMafDi8u+gJDXbtLt+m3UO7ip7QYHiovegJDYbL9zRFdXet6AkNqsPleX2q7hyKiH+/lu/LIQ/Mjh07IiLiioX/51A/NAAHyY4dO6K7u3uft2mrjyZDB1GtVott27ZFV1dXtLW17fd/Z3BwMGbMmBH9/f0xefLkg7hwbPE8jY7naXQ8T6Mzlp+ner0eO3bsiOnTp8e4cft+l+WQH8GMGzcujj766IP235s8efKY+wJm8DyNjudpdDxPozNWn6dPOnL5iDf5AUghMACkaNnAVCqV+N73vheVSqXoKaXmeRodz9PoeJ5Gx/P0/x3yN/kB+HRo2SMYAMpNYABIITAApBAYAFK0bGDuvffemDVrVnR2dsb8+fPjqaeeKnpSqfT29sapp54aXV1dMXXq1LjwwgvjtddeK3pWqfX29kZbW1ssXbq06Cml8+abb8Zll10WU6ZMiUmTJsVJJ50UmzZtKnpWqQwPD8ett94as2bNiokTJ8Zxxx0Xt99+e9Rq5ftss0OlJQOzZs2aWLp0adxyyy2xefPmOOOMM+Kcc86JrVu3Fj2tNNavXx+LFy+O5557LtatWxfDw8OxaNGieP/994ueVkobN26Mvr6+OPHEE4ueUjrvvvtuLFiwICZMmBCPP/54/PnPf46f/OQncfjhhxc9rVTuvPPOuO++++Kee+6Jv/zlL/HDH/4wfvSjH8Xdd99d9LTCtOSvKX/+85+PU045JVatWjVy3Zw5c+LCCy+M3t7eApeV1z//+c+YOnVqrF+/Ps4888yi55TKe++9F6ecckrce++98YMf/CBOOumkWLlyZdGzSuPmm2+OP/7xj35K8AnOP//8mDZtWvzsZz8bue4rX/lKTJo0KX7xi18UuKw4LXcEs2vXrti0aVMsWrRoj+sXLVoUzzzzTEGrym9gYCAiIo444oiCl5TP4sWL47zzzouzzz676CmltHbt2ujp6YmLLroopk6dGieffHLcf//9Rc8qnYULF8YTTzwRW7ZsiYiIl19+OZ5++uk499xzC15WnEP+YZcH6u23345qtRrTpk3b4/pp06bFW2+9VdCqcqvX67Fs2bJYuHBhzJs3r+g5pfLII4/Eiy++GBs3bix6Smm9/vrrsWrVqli2bFl8+9vfjueffz6uv/76qFQq8fWvf73oeaVx0003xcDAQMyePTva29ujWq3GihUr4pJLLil6WmFaLjAf+c+P+q/X6wf08f9j2bXXXhuvvPJKPP3000VPKZX+/v5YsmRJ/P73v4/Ozs6i55RWrVaLnp6euOOOOyIi4uSTT44//elPsWrVKoH5b9asWRMPPfRQrF69OubOnRsvvfRSLF26NKZPnx5XXHFF0fMK0XKBOfLII6O9vb3haGX79u0NRzVEXHfddbF27drYsGHDQT1NwliwadOm2L59e8yfP3/kumq1Ghs2bIh77rknhoaGor29fGeBPNSOOuqoOP744/e4bs6cOfHLX/6yoEXldOONN8bNN98cF198cUREnHDCCfHGG29Eb2/vpzYwLfceTEdHR8yfPz/WrVu3x/Xr1q2L008/vaBV5VOv1+Paa6+NX/3qV/GHP/whZs2aVfSk0jnrrLPi1VdfjZdeemnk0tPTE5deemm89NJL4vJfFixY0PAr7lu2bImZM2cWtKicdu7c2XACrvb29k/1rym33BFMRMSyZcvi8ssvj56enjjttNOir68vtm7dGldffXXR00pj8eLFsXr16njssceiq6tr5Iivu7s7Jk6cWPC6cujq6mp4T+qwww6LKVOmeK/qv7nhhhvi9NNPjzvuuCO++tWvxvPPPx99fX3R19dX9LRSueCCC2LFihVxzDHHxNy5c2Pz5s1x1113xVVXXVX0tOLUW9RPf/rT+syZM+sdHR31U045pb5+/fqiJ5VKRHzs5YEHHih6Wql94QtfqC9ZsqToGaXz61//uj5v3rx6pVKpz549u97X11f0pNIZHBysL1mypH7MMcfUOzs768cdd1z9lltuqQ8NDRU9rTAt+f/BAFB+LfceDACtQWAASCEwAKQQGABSCAwAKQQGgBQCA0AKgQEghcAAkEJgAEghMACkEBgAUvw/kZ70PmbdWnkAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import jax\n",
    "from jax import grad, jit\n",
    "from jax.example_libraries import optimizers\n",
    "\n",
    "# Print device info (GPU or CPU)\n",
    "print(jax.devices(), flush=True)\n",
    "\n",
    "# Define the update:\n",
    "@jit\n",
    "def update(step_index, optimizer_state):\n",
    "    # define single update step\n",
    "    parameters = get_params(optimizer_state)\n",
    "    # Call the loss function and compute the gradients\n",
    "    computed_loss = loss_value(parameters)\n",
    "    computed_gradients = grad(loss_value, allow_int=True)(parameters)\n",
    "\n",
    "    return opt_update(step_index, computed_gradients, optimizer_state), computed_loss, computed_gradients\n",
    "\n",
    "# Define the loss function and compute its gradients:\n",
    "loss_value = jit(loss)\n",
    "\n",
    "# Optimizer settings\n",
    "STEP_SIZE = 0.1\n",
    "num_iterations = 50000\n",
    "n_best = 50\n",
    "best_loss = 1e10\n",
    "best_params = None\n",
    "best_step = 0\n",
    "\n",
    "# Init parameters:\n",
    "init_params = jnp.array([np.random.uniform(-jnp.pi, jnp.pi, (resolution, resolution))], dtype=jnp.float64)\n",
    "   \n",
    "# Define the optimizer and initialize it\n",
    "opt_init, opt_update, get_params = optimizers.adam(STEP_SIZE)\n",
    "opt_state = opt_init(init_params)\n",
    "\n",
    "print('Starting Optimization', flush=True)\n",
    "tic = time.perf_counter()\n",
    "        \n",
    "# Optimize in a loop\n",
    "for step in range(num_iterations):\n",
    "        \n",
    "    opt_state, loss_value, gradients = update(step, opt_state)\n",
    "\n",
    "    if loss_value < best_loss:\n",
    "        best_loss = loss_value\n",
    "        best_params = get_params(opt_state)\n",
    "        best_step = step\n",
    "\n",
    "    # Stopping criteria:\n",
    "    if step % 100 == 0:\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}; actual step {step}')\n",
    "print(\"Time taken to optimize one sample - in seconds\", time.perf_counter() - tic)\n",
    "\n",
    "plt.imshow(best_params[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf99a66b-434f-40e4-b362-448301924e9e",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Diffractio with SciPy's optimizer (BFGS):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "636e23cd-5c27-4b31-8f22-923e0e7988f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffractio import np, degrees, um, mm, nm\n",
    "from diffractio.scalar_sources_XY import Scalar_source_XY\n",
    "from diffractio.scalar_fields_XY import Scalar_field_XY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "09ece06c-5aad-46f8-b6b1-bc4acc3d24c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Light source settings\n",
    "wavelength = .6328 * um\n",
    "w0 = (1200*um , 1200*um)\n",
    "\n",
    "wavelength = 650*nm\n",
    "x_total = 1000*um \n",
    "x = np.linspace(-x_total , x_total , resolution)\n",
    "y = np.linspace(-x_total , x_total , resolution)\n",
    "gb = Scalar_source_XY(x, y, wavelength, info='Light source')\n",
    "gb.gauss_beam(r0=(0 * um, 0 * um), w0=w0, z0=(0,0), A=1, theta=0 * degrees, phi=0 * degrees)\n",
    "\n",
    "# Spiral phase for ground truth\n",
    "gb_gt = Scalar_field_XY(x, y, wavelength)\n",
    "phase_mask = np.arctan2(gb.Y,gb.X)\n",
    "gb_gt.u = gb.u * np.exp(1j * phase_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "18134d1b-c4dc-4afa-ba98-ab010cac38d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optical setup\n",
    "def setup(gb, parameters):\n",
    "    gb_propagated = gb.RS(25000*mm)\n",
    "    gb_modulated, _ = npSLM(gb_propagated.u, parameters, gb.x.shape[0])\n",
    "    return gb_modulated\n",
    "\n",
    "def mse_phase(input_light, target_light):\n",
    "    return np.sum((np.angle(input_light) - np.angle(target_light.u)) ** 2) / resolution**2\n",
    "\n",
    "def phase(phase):\n",
    "    return np.exp(1j * phase)\n",
    "\n",
    "def npSLM(input_field, phase_array, shape):\n",
    "    slm = np.fromfunction(lambda i, j: phase(phase_array[i, j]),\n",
    "                           (shape, shape), dtype=int)\n",
    "    light_out = input_field * slm  # Multiplies element-wise\n",
    "    return light_out, slm\n",
    "\n",
    "def loss(parameters_flat):\n",
    "    parameters = parameters_flat.reshape(resolution, resolution)\n",
    "    out = setup(gb, parameters)\n",
    "    loss_val = mse_phase(out, gb_gt)\n",
    "    return loss_val\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "2ddae8c7-368e-4770-97d4-317601f838d0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Optimization terminated successfully.\n",
      "         Current function value: 0.000000\n",
      "         Iterations: 4\n",
      "         Function evaluations: 606\n",
      "         Gradient evaluations: 6\n",
      "Time taken to optimize one sample - in seconds 0.3504616217687726\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x14771fb8a670>"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAeEAAAGdCAYAAAAlqsu0AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAhVklEQVR4nO3df3BU1f3/8ddmIRvQbBTyTZRhwVD7LZBohYTpl18WRpsxKiMzHaqOUvw5UgKCmXEkYv1BC1ur5ZOOlGgcB2kdlOlYKp1RasYWkFq+QgSlv2DwV3aEFLU2gSAbsvd+/0DydV0Ss3t395xkn4+ZM0Pu7Nn7zgJ55/0+597rc13XFQAAyLo80wEAAJCrSMIAABhCEgYAwBCSMAAAhpCEAQAwhCQMAIAhJGEAAAwhCQMAYMiQbJ/QcRwdPnxYhYWF8vl82T49AMAD13V17NgxjRo1Snl5mavjTp48qa6uLs/vk5+fr4KCgjRElBlZT8KHDx9WKBTK9mkBAGkUiUQ0evTojLz3yZMnVTb2XLUdjXl+rwsuuEDvv/++tYk460m4sLBQkvSN2gflD9jzofhPmo4g0ZCT9t1RdMjnpiNINORzx3QICYacsDCmk95/oKVbXtTCmLrt+7uTY8/Pgu5YVK+/8z89P8szoaurS21HY3q/ZayChalX2x3HHJVVfqiuri6S8BlnWtD+QIFdSdief+M9/Bb9xzvDb9/PTA2x8IfmkKEWxtRt319eXszCmFz7/u7ksy+mbCwnBgvzPCXhZDU2NqqxsVEffPCBJKm8vFwPPvigampqMnbOrCdhAAD6I+Y6inmoRWJJ/kI1evRo/exnP9PFF18sSdqwYYOuu+467d27V+Xl5akH0geSMADASo5cOUo9Cyc7d86cOXFfr1q1So2Njdq1axdJGACQWxw58tKIPzO7o6Mj7nggEFAgEOhzbiwW029/+1t1dnZq6tSpHqLoG9cJAwAGtVAopKKiop4RDod7fe3+/ft17rnnKhAIaOHChdq8ebMmTpyYsdiohAEAVoq5rmJu6u3oM3MjkYiCwWDP8b6q4G9961vat2+f/vvf/+rFF1/UggULtH379owlYpIwAMBK6VoTDgaDcUm4L/n5+T0bs6qqqrR792798pe/1FNPPZVyHH2hHQ0AQC9c11U0Gs3Y+1MJAwCs5MhVLIu7o++//37V1NQoFArp2LFjeuGFF7Rt2zZt3bo15Ri+DkkYAGClbF+i9O9//1vz58/XkSNHVFRUpEsvvVRbt27V9773vZRj+DokYQAAJD3zzDNZPydJGABgpXTtjrZZShuz1q1bp7KyMhUUFKiyslKvv/56uuMCAOQ4Jw3Ddkkn4U2bNmnZsmVasWKF9u7dq5kzZ6qmpkatra2ZiA8AgEEr6SS8Zs0a3X777brjjjs0YcIENTQ0KBQKqbGxMRPxAQByVOyL3dFehu2SWhPu6upSS0uLli9fHne8urpab7zxxlnnRKPRuGusvnoPTwAAzibmyuNTlNIXS6YkVQl/8sknisViKi0tjTteWlqqtra2s84Jh8Nx9+wMhUKpRwsAyBmsCffiqw9zdl231wc819fXq729vWdEIpFUTgkAwKCTVDu6uLhYfr8/oeo9evRoQnV8Rn8eGQUAwFc58immsxd4/Z1vu6Qq4fz8fFVWVqq5uTnueHNzs6ZNm5bWwAAAuc1xvQ/bJX2zjrq6Os2fP19VVVWaOnWqmpqa1NraqoULF2YiPgAABq2kk/D111+vTz/9VCtXrtSRI0dUUVGhl19+WWPHjs1EfACAHBXz2I72MjdbUrpt5aJFi7Ro0aJ0xwIAQI9cSMI8TxgAAEN4gAMAwEqO65Pjetgd7WFutpCEAQBWoh0NAAAyhkoYAGClmPIU81ArxtIYS6aQhAEAVnI9rgm7rAkDAJAa1oQBAEDGUAkDAKwUc/MUcz2sCQ/Ge0cDAJANjnxyPDRsHdmfhWlHAwBgiLFK2PWdHrbw0PHIGBtjspHPMR1Borwu+4LKO2nfBRv+z0+ZDiGBr6vbdAiJYvb8e8qLRbN2rlzYmEU7GgBgJe9rwrSjAQBAL6iEAQBWOr0xy8MDHGhHAwCQGsfjbSvZHQ0AAHpFJQwAsFIubMwiCQMArOQob9DfrIMkDACwUsz1KebhhhJe5mYLa8IAABhCJQwAsFLM4+7oGO1oAABS47h5cjxszHIGwMYs2tEAABhCJQwAsBLtaAAADHHkbYezPc+e6h3taAAADKESBgBYyfvNOuyvM0nCAAAreb9tpf1J2P4IAQAYpKiEAQBW4nnCAAAYkgvtaJIwAMBK3q8Ttj8J2x8hAACDFJUwAMBKjuuT4+VmHQPgUYYkYQCAlRyP7eiBcJ2w/RECADBIUQkDAKzk/VGG9teZ9kcIAMhJMfk8j2SEw2FNmTJFhYWFKikp0dy5c3XgwIEMfXenkYQBAJC0fft21dbWateuXWpublZ3d7eqq6vV2dmZsXPSjgYAWCnb7eitW7fGfb1+/XqVlJSopaVFl19+ecpx9IUkDACwUkxKuqX81fmS1NHREXc8EAgoEAh87fz29nZJ0ogRI1KO4evQjgYADGqhUEhFRUU9IxwOf+0c13VVV1enGTNmqKKiImOxUQkDAKyUrnZ0JBJRMBjsOd6fKnjx4sV65513tHPnzpTP3x8kYQCAldL1AIdgMBiXhL/OkiVLtGXLFu3YsUOjR49O+fz9QRIGAFjJ9fgoQzfJua7rasmSJdq8ebO2bdumsrKylM/dXyRhAAAk1dbWauPGjXrppZdUWFiotrY2SVJRUZGGDRuWkXOShAEAVsr284QbGxslSbNmzYo7vn79et1yyy0px9EXY0nYGSr5hpo6eyKfYzqCRM4p+54A4gxxTYcwIOR12/c5+U90mQ4hge9E1HQICXxdp0yHkChmzw8on5O9v7NsP0XJdbP//5ZLlAAAMIR2NADASjGPjzL0MjdbSMIAACtlux1tgv2/JgAAMEhRCQMArOQoT46HWtHL3GwhCQMArBRzfYp5aCl7mZst9v+aAADAIEUlDACwUi5szCIJAwCs5Hp8ipLrYW62kIQBAFaKyaeYhwc4eJmbLfb/mgAAwCBFJQwAsJLjelvXdey7hXsCkjAAwEqOxzVhL3Ozxf4IAQAYpJJKwuFwWFOmTFFhYaFKSko0d+5cHThwIFOxAQBymCOf52G7pJLw9u3bVVtbq127dqm5uVnd3d2qrq5WZ2dnpuIDAOSoM3fM8jJsl9Sa8NatW+O+Xr9+vUpKStTS0qLLL788rYEBADDYedqY1d7eLkkaMWJEr6+JRqOKRqM9X3d0dHg5JQAgR7Axqw+u66qurk4zZsxQRUVFr68Lh8MqKirqGaFQKNVTAgByiCNfz60rUxqDbU34yxYvXqx33nlHzz//fJ+vq6+vV3t7e8+IRCKpnhIAgEElpXb0kiVLtGXLFu3YsUOjR4/u87WBQECBQCCl4AAAucv1uMPZHQCVcFJJ2HVdLVmyRJs3b9a2bdtUVlaWqbgAADmOpyh9RW1trTZu3KiXXnpJhYWFamtrkyQVFRVp2LBhGQkQAJCb2Jj1FY2NjWpvb9esWbN04YUX9oxNmzZlKj4AAAatpNvRAABkA+1oAAAM8XrryUF9iRIAAPCGShgAYCXa0QAAGJILSZh2NAAAhlAJAwCslAuVMEkYAGClXEjCtKMBADCEShgAYCVX3q71HQi3lyIJAwCslAvtaJIwAMBKJOEMcoe6cvLtaRb4LPzLitnz8fTwxez7nJwh9sXkP9ltOoQEee2dpkNI4HZ+bjqEBO6pLtMhWM11+XzSiUoYAGAlKmEAAAzJhSTMJUoAABhCJQwAsJLr+uR6qGa9zM0WkjAAwEo8TxgAAGQMlTAAwEq5sDGLJAwAsFIurAnTjgYAwBAqYQCAlWhHAwBgSC60o0nCAAAruR4r4YGQhFkTBgDgCzt27NCcOXM0atQo+Xw+/f73v8/o+UjCAAAruZJc18NI4ZydnZ369re/rbVr16b72zkr2tEAACs58smX5Ttm1dTUqKamJuVzJoskDAAY1Do6OuK+DgQCCgQChqKJRzsaAGClM7ujvQxJCoVCKioq6hnhcNjwd/b/UQkDAKzkuD750nCdcCQSUTAY7DluSxUskYQBAINcMBiMS8I2IQkDAKx0Zpezl/m2IwkDAKxk4o5Zx48f16FDh3q+fv/997Vv3z6NGDFCY8aMSTmW3pCEAQD4wp49ezR79uyer+vq6iRJCxYs0LPPPpv285GEAQBWMlEJz5o1S24W+9gkYQCAldK1O9pmJGEAgJVyYWMWN+sAAMAQKmEAgJVOV8Je1oTTGEyGkIQBAFYysTEr22hHAwBgCJUwAMBKrlJ7JvCX59uOJAwAsBLtaAAAkDFUwgAAO+VAP5okDACwk8d2tAZAO5okDACwEnfMAgAAGWOsEo4FJDdg6uyJXL99vzLF8k1HkCiWb197x8sN3jPlvP/7mekQEsQ+OmI6hARud7fpEAYGnz3/xmPuqaydKxd2R9OOBgDYyfV5W9cdAEmYdjQAAIZQCQMArJQLG7NIwgAAO+XAdcK0owEAMIRKGABgJXZHAwBg0gBoKXtBOxoAAEOohAEAVqIdDQCAKTmwO5okDACwlO+L4WW+3VgTBgDAECphAICdaEcDAGBIDiRhT+3ocDgsn8+nZcuWpSkcAAByR8qV8O7du9XU1KRLL700nfEAAHAajzI8u+PHj+umm27S008/rfPPPz/dMQEA0PMUJS/Ddikl4draWl1zzTW68sorv/a10WhUHR0dcQMAAKTQjn7hhRf01ltvaffu3f16fTgc1iOPPJJ0YACAHMfGrHiRSERLly7Vc889p4KCgn7Nqa+vV3t7e8+IRCIpBQoAyDFn1oS9DMslVQm3tLTo6NGjqqys7DkWi8W0Y8cOrV27VtFoVH6/P25OIBBQIBBIT7QAAAwiSSXhK664Qvv37487duutt2r8+PG67777EhIwAACp8rmnh5f5tksqCRcWFqqioiLu2DnnnKORI0cmHAcAwJMcWBPmjlkAADvlwHXCnpPwtm3b0hAGAAC5h0oYAGAn2tEAABiSA0mY5wkDAGAIlTAAwE45UAmThAEAdsqB3dG0owEAMIRKGABgJe6YBQCAKTmwJkw7GgAAQ0jCAAAYQjsaAGAlnzyuCactkswxloSdAkca5pg6fQLHnlCs1jXEvkWWUxO7TIeQYNnSv5oOIcH/Gfah6RASnGdhL+5c31DTISQYnpdvOoQeHcccnf+/s3QyQ5corVu3To899piOHDmi8vJyNTQ0aObMmanH0QcL/wsAAGDGpk2btGzZMq1YsUJ79+7VzJkzVVNTo9bW1oycjyQMALCTm4aRpDVr1uj222/XHXfcoQkTJqihoUGhUEiNjY3ev5+zIAkDAOyUpiTc0dERN6LR6FlP19XVpZaWFlVXV8cdr66u1htvvJHu704SSRgAMMiFQiEVFRX1jHA4fNbXffLJJ4rFYiotLY07Xlpaqra2tozExu5oAICV0nXHrEgkomAw2HM8EAj0Pc8Xv6HLdd2EY+lCEgYA2ClNd8wKBoNxSbg3xcXF8vv9CVXv0aNHE6rjdKEdDQCApPz8fFVWVqq5uTnueHNzs6ZNm5aRc1IJAwDsZODe0XV1dZo/f76qqqo0depUNTU1qbW1VQsXLvQQSO9IwgAAK5l4itL111+vTz/9VCtXrtSRI0dUUVGhl19+WWPHjk09kD6QhAEA+JJFixZp0aJFWTkXSRgAYCdDt63MJpIwAMBOOfA8YZIwAMBKJtaEs41LlAAAMIRKGABgJ9rRAAAY4rEdPRCSMO1oAAAMoRIGANiJdjQAAIbkQBKmHQ0AgCFUwgAAK3GdMAAAyBiSMAAAhtCOBgDYKQc2ZpGEAQBWyoU1YZIwAMBeAyCResGaMAAAhlAJAwDsxJowAABm5MKaMO1oAAAMoRIGANiJdjQAAGbQjgYAABlDJQwAsBPtaAAADMmBJEw7GgAAQ4xVwu6wmNxhMVOnHxB8fvt+jRtacMp0CAkuOP+Y6RAS7D8RMh1Cgnyfff/fLhr6iekQEvwv/+emQ0hQ5Nrz/+5YzMnauXJhYxbtaACAnXKgHU0SBgDYKQeSMGvCAAAYQiUMALASa8IAAJhCOxoAAGQKlTAAwEq0owEAMIV2NAAAyBQqYQCAnXKgEiYJAwCs5PtieJlvO9rRAAAYQiUMALAT7WgAAMzIhUuUkm5Hf/TRR7r55ps1cuRIDR8+XJdddplaWloyERsAIJe5aRiWS6oS/uyzzzR9+nTNnj1br7zyikpKSvTuu+/qvPPOy1B4AAAMXkkl4UcffVShUEjr16/vOXbRRRelOyYAAE4bANWsF0m1o7ds2aKqqirNmzdPJSUlmjRpkp5++uk+50SjUXV0dMQNAAC+zpk1YS/Ddkkl4ffee0+NjY365je/qT/+8Y9auHCh7r77bv3617/udU44HFZRUVHPCIVCnoMGAGAwSCoJO46jyZMna/Xq1Zo0aZLuuusu3XnnnWpsbOx1Tn19vdrb23tGJBLxHDQAIAewMSvehRdeqIkTJ8YdmzBhgl588cVe5wQCAQUCgdSiAwDkLC5R+orp06frwIEDcccOHjyosWPHpjUoAAByQVJJ+J577tGuXbu0evVqHTp0SBs3blRTU5Nqa2szFR8AIFdZ3o5etWqVpk2bpuHDh6d8qW5SSXjKlCnavHmznn/+eVVUVOgnP/mJGhoadNNNN6V0cgAAemP77uiuri7NmzdPP/rRj1J+j6RvW3nttdfq2muvTfmEAAAMBo888ogk6dlnn035Pbh3NADATml6gMNX709h04ZhHmUIALBTmtaEQ6FQ3P0qwuFwdr+PPpCEAQBWSteacCQSibtfRX19fa/nfPjhh+Xz+foce/bsSdv3SDsaADCoBYNBBYPBfr128eLFuuGGG/p8TTqfmUASBgDYKU1rwskoLi5WcXGxh5MmhyQMALCSz3Xlc1PPwl7m9kdra6v+85//qLW1VbFYTPv27ZMkXXzxxTr33HP79R4kYQAAUvDggw9qw4YNPV9PmjRJkvTnP/9Zs2bN6td7sDELAGAny++Y9eyzz8p13YTR3wQsUQkDACzFAxwAAEDGUAkDAOxkYHd0thlLwv6CbuUN6zZ1+gR5FvYt/EMc0yEkGF4QNR1CguFDu0yHkODzWL7pEBK0dReZDiFBQd4p0yEMEJ+bDqDHcSd7PytpRwMAgIyhHQ0AsBPtaAAAzMiFdjRJGABgpxyohFkTBgDAECphAIC1BkJL2QuSMADATq57eniZbzna0QAAGEIlDACwErujAQAwhd3RAAAgU6iEAQBW8jmnh5f5tiMJAwDsRDsaAABkCpUwAMBK7I4GAMCUHLhZB0kYAGClXKiEWRMGAMAQKmEAgJ1yYHc0SRgAYCXa0QAAIGOohAEAdmJ3NAAAZtCOBgAAGUMlDACwE7ujAQAwg3Y0AADIGCphAICdHPf08DLfciRhAICdWBMGAMAMnzyuCactksxhTRgAAEOohAEAduKOWQAAmMElSgAAIGOohAEAdmJ3NAAAZvhcVz4P67pe5maLsSQcKOiWv+CUqdMnyMtzTIeQIDAkZjqEBOfkd5kOIcHwIfbFZKOoM9R0CAk6nYDpEBJ05uWbDiHBcNeen5Wdrn0/KwcyKmEAgJ2cL4aX+ZYjCQMArJQL7Wh2RwMAYAiVMADATuyOBgDAEO6YBQCAGdwxCwAAZAyVMADATrSjAQAww+ecHl7m2452NAAASfrggw90++23q6ysTMOGDdM3vvENPfTQQ+rqSu4OflTCAAA7WdyO/te//iXHcfTUU0/p4osv1t/+9jfdeeed6uzs1OOPP97v9yEJAwDsZPF1wldddZWuuuqqnq/HjRunAwcOqLGxkSQMAMAZHR0dcV8HAgEFAul/eEh7e7tGjBiR1BzWhAEAVjpz72gvQ5JCoZCKiop6RjgcTnus7777rp544gktXLgwqXlJJeHu7m498MADPQvR48aN08qVK+U4A2ALGgBgYDmzJuxlSIpEImpvb+8Z9fX1vZ7y4Ycfls/n63Ps2bMnbs7hw4d11VVXad68ebrjjjuS+haTakc/+uijevLJJ7VhwwaVl5drz549uvXWW1VUVKSlS5cmdWIAALIhGAwqGAz267WLFy/WDTfc0OdrLrroop4/Hz58WLNnz9bUqVPV1NSUdGxJJeG//vWvuu6663TNNdf0BPL8888n/FYAAIBnrrw9EziFjVnFxcUqLi7u12s/+ugjzZ49W5WVlVq/fr3y8pJf4U1qxowZM/Taa6/p4MGDkqS3335bO3fu1NVXX93rnGg0qo6OjrgBAMDXSdeacCYcPnxYs2bNUigU0uOPP66PP/5YbW1tamtrS+p9kqqE77vvPrW3t2v8+PHy+/2KxWJatWqVbrzxxl7nhMNhPfLII0kFBQDA6UuUvFwnnLZIErz66qs6dOiQDh06pNGjR8efNomYk6qEN23apOeee04bN27UW2+9pQ0bNujxxx/Xhg0bep1TX18ftyAeiUSSOSUAANa55ZZb5LruWUcykqqE7733Xi1fvrxn0fqSSy7Rhx9+qHA4rAULFpx1TqauxwIADHIW3zErXZJKwidOnEhYePb7/VyiBABIP0eSz+N8yyWVhOfMmaNVq1ZpzJgxKi8v1969e7VmzRrddtttmYoPAIBBK6kk/MQTT+jHP/6xFi1apKNHj2rUqFG666679OCDD2YqPgBAjvK6wzmTu6PTJakkXFhYqIaGBjU0NGQoHAAAvpADa8LcOxoAAEN4ihIAwE45UAmThAEAdsqBJEw7GgAAQ6iEAQB24jphAADM4BIlAABMYU0YAABkCpUwAMBOjiv5PFSzjv2VMEkYAGAn2tEAACBTjFXCwwNd8ge87D1PL3+efXvZA/6Y6RASnJsfNR1CguFDTpkOIcGQPPv+7mx0yvWbDiHBKde+BuEp15566VRWi0uPlbDsr4Tt+9cGAIBEOxoAAGQOlTAAwE6OK08tZXZHAwCQItc5PbzMtxztaAAADKESBgDYKQc2ZpGEAQB2Yk0YAABDcqASZk0YAABDqIQBAHZy5bESTlskGUMSBgDYiXY0AADIFCphAICdHEeShxtuOPbfrIMkDACwE+1oAACQKVTCAAA75UAlTBIGANgpB+6YRTsaAABDqIQBAFZyXUeuh8cRepmbLSRhAICdXNdbS5k1YQAAUuR6XBMeAEmYNWEAAAyhEgYA2MlxJJ+HdV3WhAEASBHtaAAAkClUwgAAK7mOI9dDO5pLlAAASBXtaAAAkClUwgAAOzmu5BvclTBJGABgJ9eV5OUSJfuTMO1oAAAMoRIGAFjJdVy5HtrR7gCohEnCAAA7uY68taO5RAkAgJTkQiXMmjAAAIZkvRI+85tJ7EQ026fum5ebhGdItz9mOoQEp051mQ4hgY0xdZ06ZTqEBCfz7Yvp86HdpkNIcMLC/3fH/fb8fOo8fjqWbFSZ3W7UU0u5W/b9m/+qrCfhY8eOSZL2//BX2T41ACBNjh07pqKiooy8d35+vi644ALtbHvZ83tdcMEFys/PT0NUmeFzs9w0dxxHhw8fVmFhoXw+X8rv09HRoVAopEgkomAwmMYIBxc+p/7hc+ofPqf+Gcyfk+u6OnbsmEaNGqW8vMytaJ48eVJdXd67XPn5+SooKEhDRJmR9Uo4Ly9Po0ePTtv7BYPBQfePPBP4nPqHz6l/+Jz6Z7B+TpmqgL+soKDA6uSZLmzMAgDAEJIwAACGDNgkHAgE9NBDDykQCJgOxWp8Tv3D59Q/fE79w+eE/sr6xiwAAHDagK2EAQAY6EjCAAAYQhIGAMAQkjAAAIYM2CS8bt06lZWVqaCgQJWVlXr99ddNh2SVcDisKVOmqLCwUCUlJZo7d64OHDhgOiyrhcNh+Xw+LVu2zHQo1vnoo4908803a+TIkRo+fLguu+wytbS0mA7LKt3d3XrggQdUVlamYcOGady4cVq5cqUcx577PsM+AzIJb9q0ScuWLdOKFSu0d+9ezZw5UzU1NWptbTUdmjW2b9+u2tpa7dq1S83Nzeru7lZ1dbU6OztNh2al3bt3q6mpSZdeeqnpUKzz2Wefafr06Ro6dKheeeUV/eMf/9AvfvELnXfeeaZDs8qjjz6qJ598UmvXrtU///lP/fznP9djjz2mJ554wnRosNiAvETpO9/5jiZPnqzGxsaeYxMmTNDcuXMVDocNRmavjz/+WCUlJdq+fbsuv/xy0+FY5fjx45o8ebLWrVunn/70p7rsssvU0NBgOixrLF++XH/5y1/oNn2Na6+9VqWlpXrmmWd6jn3/+9/X8OHD9Zvf/MZgZLDZgKuEu7q61NLSourq6rjj1dXVeuONNwxFZb/29nZJ0ogRIwxHYp/a2lpdc801uvLKK02HYqUtW7aoqqpK8+bNU0lJiSZNmqSnn37adFjWmTFjhl577TUdPHhQkvT2229r586duvrqqw1HBptl/QEOXn3yySeKxWIqLS2NO15aWqq2tjZDUdnNdV3V1dVpxowZqqioMB2OVV544QW99dZb2r17t+lQrPXee++psbFRdXV1uv/++/Xmm2/q7rvvViAQ0A9/+EPT4VnjvvvuU3t7u8aPHy+/369YLKZVq1bpxhtvNB0aLDbgkvAZX30Mouu6nh6NOJgtXrxY77zzjnbu3Gk6FKtEIhEtXbpUr776ak48rSVVjuOoqqpKq1evliRNmjRJf//739XY2EgS/pJNmzbpueee08aNG1VeXq59+/Zp2bJlGjVqlBYsWGA6PFhqwCXh4uJi+f3+hKr36NGjCdUxpCVLlmjLli3asWNHWh8hORi0tLTo6NGjqqys7DkWi8W0Y8cOrV27VtFoVH6/32CEdrjwwgs1ceLEuGMTJkzQiy++aCgiO917771avny5brjhBknSJZdcog8//FDhcJgkjF4NuDXh/Px8VVZWqrm5Oe54c3Ozpk2bZigq+7iuq8WLF+t3v/ud/vSnP6msrMx0SNa54oortH//fu3bt69nVFVV6aabbtK+fftIwF+YPn16wuVtBw8e1NixYw1FZKcTJ04kPOTe7/dziRL6NOAqYUmqq6vT/PnzVVVVpalTp6qpqUmtra1auHCh6dCsUVtbq40bN+qll15SYWFhT+egqKhIw4YNMxydHQoLCxPWyM855xyNHDmStfMvueeeezRt2jStXr1aP/jBD/Tmm2+qqalJTU1NpkOzypw5c7Rq1SqNGTNG5eXl2rt3r9asWaPbbrvNdGiwmTtA/epXv3LHjh3r5ufnu5MnT3a3b99uOiSrSDrrWL9+venQrPbd737XXbp0qekwrPOHP/zBraiocAOBgDt+/Hi3qanJdEjW6ejocJcuXeqOGTPGLSgocMeNG+euWLHCjUajpkODxQbkdcIAAAwGA25NGACAwYIkDACAISRhAAAMIQkDAGAISRgAAENIwgAAGEISBgDAEJIwAACGkIQBADCEJAwAgCEkYQAADCEJAwBgyP8DV0RoQZQguvsAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from scipy.optimize import minimize \n",
    "\n",
    "# Init parameters\n",
    "parameters = np.random.uniform(-np.pi, np.pi, (resolution, resolution)).flatten()\n",
    "\n",
    "# Optimizer\n",
    "tic = time.perf_counter()\n",
    "res = minimize(loss, parameters, method='BFGS', options={'disp': True})\n",
    "print(\"Time taken to optimize one sample - in seconds\", time.perf_counter() - tic)\n",
    "\n",
    "result = (res.x).reshape(resolution, resolution)\n",
    "\n",
    "# Wrap phase (-pi, pi):\n",
    "plt.imshow(np.arctan2(np.cos(result), np.sin(result)))\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8401bb15-d445-418e-a458-9253a3887b2d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pretzel",
   "language": "python",
   "name": "pretzel"
  },
  "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}