{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "76fec76e-f945-4851-85a0-fceac3752075",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "# Generating precompute tables"
   ]
  },
  {
   "cell_type": "raw",
   "id": "e3be28ef-8b49-40bf-9b4a-80f16c61054e",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Here we precompute the Bragg pulses needed to simulate an SCI.\n",
    "\n",
    "First we define functions to call :py:meth:`mwave.integrate.gbragg` in a parallel loop with the targeted parameters. Also define a test function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a9c984db-1aab-4805-93d2-40224cceb15b",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from mwave.integrate import make_kvec, make_phi, gbragg, pops_vs_time\n",
    "from mwave.precompute import write_bragg_precompute, read_bragg_precompute\n",
    "import numpy as np\n",
    "from scipy.interpolate import RegularGridInterpolator as RGI\n",
    "from matplotlib import pyplot as plt\n",
    "from tqdm import tqdm\n",
    "from joblib import Parallel, delayed\n",
    "import h5py\n",
    "\n",
    "# Define output file name\n",
    "fname = 'sig0.247.h5'\n",
    "\n",
    "# Define simulation parameters\n",
    "n_init = 0\n",
    "n_bragg = 5\n",
    "N_bloch = 10\n",
    "sigma = 0.247\n",
    "\n",
    "# Define function to perform precomputation\n",
    "def precompute_gbragg(n0, nf, n_bragg, N_bloch=None):\n",
    "\n",
    "    # Make kvec\n",
    "    kvec, n0_idx, nf_idx = make_kvec(n0,nf)\n",
    "    \n",
    "    # Compute over grid\n",
    "    omegas = np.linspace(0, 43, 800)\n",
    "    deltas = np.linspace(-2, 2, 400) + 4*n_bragg\n",
    "    \n",
    "    # Create array to store output\n",
    "    phi = np.full((len(omegas), len(deltas), len(kvec)), np.nan, dtype=np.complex128)\n",
    "    \n",
    "    # Define the grid shape\n",
    "    grid_shape = [len(omegas), len(deltas)]\n",
    "    \n",
    "    # Define function to compute single Bragg pulse\n",
    "    def do_gbragg(i):\n",
    "        idx1, idx2 = np.unravel_index(i, grid_shape)\n",
    "        if N_bloch is not None:\n",
    "            sol = gbragg(kvec, make_phi(kvec, n0), 6*sigma, deltas[idx2], omegas[idx1], sigma, omega_mod=8*(N_bloch+n_bragg))\n",
    "        else:\n",
    "            sol = gbragg(kvec, make_phi(kvec, n0), 6*sigma, deltas[idx2], omegas[idx1], sigma)\n",
    "        return sol.y[:,-1]\n",
    "    \n",
    "    # Compute function in parallel\n",
    "    out = Parallel(n_jobs=-1)(delayed(do_gbragg)(i) for i in tqdm(range(np.prod(grid_shape))))\n",
    "    \n",
    "    # Put all of the output into the phi array\n",
    "    for i in range(np.prod(grid_shape)):\n",
    "        idx1, idx2 = np.unravel_index(i, grid_shape)\n",
    "        phi[idx1, idx2, :] = out[i]\n",
    "\n",
    "    # Write output to HDF5 file\n",
    "    write_bragg_precompute(fname, phi, kvec, ((omegas, 'omegas'), (deltas, 'deltas')), n0, nf, n_bragg, N_bloch=N_bloch)\n",
    "\n",
    "def test_precomp_grid(n0, nf, multifreq=False):\n",
    "    # Load grid\n",
    "    phi, kvec, grid = read_bragg_precompute(fname, n0, nf, n_bragg, N_bloch if multifreq else None)\n",
    "    omegas, deltas = grid[0][0], grid[1][0]\n",
    "    \n",
    "    # Interpolate\n",
    "    rgi = RGI([omegas, deltas], phi, method='cubic')\n",
    "    \n",
    "    # Create vectorized gbragg function\n",
    "    def vgbragg(omegas, deltas):\n",
    "        if len(omegas) != len(deltas):\n",
    "            raise ValueError('omegas and deltas must have the same length')\n",
    "        phi = np.full((len(omegas), len(kvec)), np.nan, dtype=np.complex128)\n",
    "        for i in tqdm(range(len(omegas))):\n",
    "            if multifreq:\n",
    "                sol = gbragg(kvec, make_phi(kvec, n0), 6*sigma, deltas[i], omegas[i], sigma, omega_mod=8*(N_bloch+n_bragg))\n",
    "            else:\n",
    "                sol = gbragg(kvec, make_phi(kvec, n0), 6*sigma, deltas[i], omegas[i], sigma)\n",
    "            phi[i, :] = sol.y[:,-1]\n",
    "        \n",
    "        return phi\n",
    "    \n",
    "    # Drop random points on grid\n",
    "    npoints = 4000\n",
    "    rng = np.random.default_rng()\n",
    "    omega_pnts = rng.uniform(omegas[0], omegas[-1], npoints)\n",
    "    delta_pnts = rng.uniform(deltas[0], deltas[-1], npoints)\n",
    "    \n",
    "    # Compute real value on points\n",
    "    phi_actual = vgbragg(omega_pnts, delta_pnts)\n",
    "    \n",
    "    # Compute interpolated value on points\n",
    "    phi_interp = rgi((omega_pnts, delta_pnts))\n",
    "    \n",
    "    # Take the difference\n",
    "    phi_diff = phi_actual - phi_interp\n",
    "    \n",
    "    # Sum the error along each wavefunction\n",
    "    err = np.sum(phi_diff, axis=-1)\n",
    "    \n",
    "    # Plot the real and imaginary parts of the error\n",
    "    plt.figure()\n",
    "    plt.scatter(omega_pnts, delta_pnts, c=np.real(err))\n",
    "    cbar = plt.colorbar()\n",
    "    cbar.ax.set_ylabel('re abs err')\n",
    "    plt.show()\n",
    "    \n",
    "    plt.figure()\n",
    "    plt.scatter(omega_pnts, delta_pnts, c=np.imag(err))\n",
    "    cbar = plt.colorbar()\n",
    "    cbar.ax.set_ylabel('im abs err')\n",
    "    plt.show()\n",
    "\n",
    "    # Plot histograms of the error\n",
    "    plt.figure()\n",
    "    plt.hist(np.real(err), range=(-2e-6,2e-6), bins=200)\n",
    "    plt.xlabel('re abs err')\n",
    "    plt.show()\n",
    "    \n",
    "    plt.figure()\n",
    "    plt.hist(np.imag(err), range=(-2e-6,2e-6), bins=200)\n",
    "    plt.xlabel('im abs err')\n",
    "    plt.show()\n",
    "\n",
    "    # Print the error mean and std\n",
    "    print(f'Re[error] mean={np.mean(np.real(err))}, error std={np.std(np.real(err))}')\n",
    "    print(f'Im[error] mean={np.mean(np.imag(err))}, error std={np.std(np.imag(err))}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db742e21-d0df-47ca-a4a0-3b6d3c050866",
   "metadata": {},
   "source": [
    "Precompute the single frequency Bragg pulses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3a0cb1ba-76c4-425b-8bc0-2ce9923ef0a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "n0, nf = n_init, n_init + n_bragg\n",
    "# Uncomment the below code when you want to perform the computation and/or test the precompute accuracy\n",
    "# precompute_gbragg(n0, nf, n_bragg)\n",
    "# test_precomp_grid(n0,nf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d814d013-1f40-4cd7-b501-d44fd7a23a2a",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "n0, nf = n_init + n_bragg, n_init\n",
    "# Uncomment the below code when you want to perform the computation and/or test the precompute accuracy\n",
    "# precompute_gbragg(n0, nf, n_bragg)\n",
    "# test_precomp_grid(n0,nf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bb66bda-ec49-46b5-afe4-b5f3b4a17a86",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Precompute the multifrequency Bragg pulses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c408cb1b-6ca3-40cd-99eb-4d700aba7b6b",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "n0, nf = -N_bloch, -N_bloch-n_bragg\n",
    "# Uncomment the below code when you want to perform the computation and/or test the precompute accuracy\n",
    "# precompute_gbragg(n0, nf, n_bragg, N_bloch=N_bloch)\n",
    "# test_precomp_grid(n0,nf,multifreq=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "53aa0736-2ed9-40f0-9263-32ecd8373cde",
   "metadata": {},
   "outputs": [],
   "source": [
    "n0, nf = n_bragg+N_bloch, 2*n_bragg+N_bloch\n",
    "# Uncomment the below code when you want to perform the computation and/or test the precompute accuracy\n",
    "# precompute_gbragg(n0, nf, n_bragg, N_bloch=N_bloch)\n",
    "# test_precomp_grid(n0,nf,multifreq=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "72b4d4f4-9658-47e3-9ef8-67588b2adc6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "n0, nf = 2*n_bragg+N_bloch, n_bragg+N_bloch\n",
    "# Uncomment the below code when you want to perform the computation and/or test the precompute accuracy\n",
    "# precompute_gbragg(n0, nf, n_bragg, N_bloch=N_bloch)\n",
    "# test_precomp_grid(n0,nf,multifreq=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "db451518-6dfe-4726-854c-c072bc61ff0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "n0, nf = -n_bragg-N_bloch, -N_bloch\n",
    "# Uncomment the below code when you want to perform the computation and/or test the precompute accuracy\n",
    "# precompute_gbragg(n0, nf, n_bragg, N_bloch=N_bloch)\n",
    "# test_precomp_grid(n0,nf,multifreq=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}