what_we_actually_want.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Aaaahhhhhhhh\n",
    "Alle guten Dinge sind zwei, also mal schauen ob ichs jetzt richtig verstanden hab.\n",
    "\n",
    "- Fix geometry\n",
    "- Vary the B.n created by plasma (B_external_normal)\n",
    "- Analyze the impact"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from datetime import datetime\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "from scipy.optimize import minimize\n",
    "from simsopt.field import BiotSavart, Current, coils_via_symmetries\n",
    "import simsopt.geo\n",
    "from simsopt.mhd import VirtualCasing, Vmec\n",
    "from simsopt.objectives import QuadraticPenalty, SquaredFlux\n",
    "from simsopt.util import in_github_actions\n",
    "\n",
    "\n",
    "# Major radius for the initial circular coils:\n",
    "R0 = 5.5\n",
    "\n",
    "# Minor radius for the initial circular coils:\n",
    "R1 = 1.25\n",
    "R1_reactor = R0/5\n",
    "\n",
    "# Number of Fourier modes describing each Cartesian component of each coil:\n",
    "order = 9\n",
    "\n",
    "# Number of iterations to perform:\n",
    "MAXITER = 100\n",
    "\n",
    "# File for the desired boundary magnetic surface:\n",
    "TEST_DIR = \"tests/test_files\"\n",
    "filename = 'wout_W7-X_without_coil_ripple_beta0p05_d23p4_tm_reference.nc'\n",
    "vmec_file = TEST_DIR +\"/\"+ filename\n",
    "\n",
    "# Resolution on the plasma boundary surface:\n",
    "# nphi is the number of grid points in 1/2 a field period.\n",
    "nphi = 32\n",
    "ntheta = 64\n",
    "\n",
    "# Resolution for the virtual casing calculation:\n",
    "vc_src_nphi = 80\n",
    "# (For the virtual casing src_ resolution, only nphi needs to be\n",
    "# specified; the theta resolution is computed automatically to\n",
    "# minimize anisotropy of the grid.)\n",
    "\n",
    "#######################################################\n",
    "# End of input parameters.\n",
    "#######################################################\n",
    "\n",
    "# Directory for output\n",
    "out_dir = Path(\"output\")\n",
    "out_dir.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# Once the virtual casing calculation has been run once, the results\n",
    "# can be used for many coil optimizations. Therefore here we check to\n",
    "# see if the virtual casing output file alreadys exists. If so, load\n",
    "# the results, otherwise run the virtual casing calculation and save\n",
    "# the results.\n",
    "head, tail = os.path.split(vmec_file)\n",
    "vc_filename = os.path.join(head, tail.replace('wout', 'vcasing'))\n",
    "print('virtual casing data file:', vc_filename)\n",
    "if os.path.isfile(vc_filename):\n",
    "    print('Loading saved virtual casing result')\n",
    "    vc = VirtualCasing.load(vc_filename)\n",
    "else:\n",
    "    # Virtual casing must not have been run yet.\n",
    "    print('Running the virtual casing calculation')\n",
    "    vc = VirtualCasing.from_vmec(vmec_file, src_nphi=vc_src_nphi, trgt_nphi=nphi, trgt_ntheta=ntheta)\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "stacked_img = np.vstack([vc.B_external_normal, np.flip(vc.B_external_normal, axis=0)])\n",
    "extended_img = np.hstack([stacked_img, np.flip(stacked_img, axis=1)])\n",
    "plt.imshow(extended_img)\n",
    "plt.colorbar()\n",
    "plt.show()\n",
    "plt.imshow(np.real(np.fft.ifftshift(np.fft.fft2(extended_img))[20:-20,20:-20]))\n",
    "plt.colorbar()\n",
    "plt.title(\"Spectrum of the loaded B.n\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nfp = 3\n",
    "s = simsopt.geo.SurfaceRZFourier.from_nphi_ntheta(32, 32, \"field period\", nfp )\n",
    "s.set_rc(0,0,R0)\n",
    "s.set_rc(1,0,R1_reactor)\n",
    "s.set_zs(1,0,R1_reactor)\n",
    "s.plot()\n",
    "print(\"Aspect ratio\",R0/R1_reactor)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_modes(image_mn, **kwargs):\n",
    "  extent = (-image_mn.shape[1]/2, image_mn.shape[1]/2, image_mn.shape[0]-0.5,-0.5)\n",
    "\n",
    "  plt.imshow(image_mn, extent=extent, **kwargs)\n",
    "  plt.colorbar()\n",
    "  plt.ylabel(\"m\")\n",
    "  plt.xlabel(\"n\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Only contains the even terms, for odd terms the imaginary equivalent has to be used! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.axes_grid1 import ImageGrid\n",
    "\n",
    "Mpol_max = 5\n",
    "Ntor_max = 4\n",
    "\n",
    "images = []\n",
    "ms = []\n",
    "ns = []\n",
    "for m in range(Mpol_max):\n",
    "  for sig_n in [-1, 1]:\n",
    "    if m==0 and sig_n==-1:\n",
    "      images.extend([None]*(Ntor_max-1))\n",
    "      continue\n",
    "    for n in range(Ntor_max) if sig_n>0 else reversed(range(1,Ntor_max)):\n",
    "      if m==0 and n==0:\n",
    "         images.append(None)\n",
    "         continue\n",
    "      Bn_fspace = np.zeros((nphi, ntheta))\n",
    "      Bn_fspace[m][n*sig_n] = nphi*10\n",
    "      B_external_normal = np.fft.ifft2(Bn_fspace)\n",
    "\n",
    "      images.append(np.real(B_external_normal))\n",
    "      ms.append(m)\n",
    "      ns.append(n*sig_n)\n",
    "\n",
    "fig = plt.figure(figsize=(Mpol_max*2, Ntor_max*3))\n",
    "grid = ImageGrid(fig, 111,  # similar to subplot(111)\n",
    "                 nrows_ncols=(Mpol_max, Ntor_max*2-1),  # creates 2x2 grid of axes\n",
    "                 axes_pad=0.1,  # pad between axes in inch.\n",
    "                 )\n",
    "\n",
    "for ax, im in zip(grid, images):\n",
    "    # Iterating over the grid returns the Axes.\n",
    "    if im is not None:\n",
    "      ax.imshow(im)\n",
    "plt.show()\n",
    "\n",
    "images = list(filter(lambda x: x is not None, images))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets\n",
    "\n",
    "def plot_coil_complexities(campaign_folder):\n",
    "    resfilepaths = os.listdir(campaign_folder)\n",
    "\n",
    "    results = [simsopt.load(os.path.join(campaign_folder, resfile)) for resfile in resfilepaths]\n",
    "    mmax = 0\n",
    "    nmax = 0\n",
    "    for res in results:\n",
    "        mmax = max(mmax, res[\"m\"])\n",
    "        nmax = max(nmax, res[\"n\"])\n",
    "    heatmap_convergence_mean = np.zeros((mmax+1, nmax*2+1))\n",
    "    heatmap_convergence_max = np.zeros((mmax+1, nmax*2+1))\n",
    "    coil_complexity_map = np.zeros_like(heatmap_convergence_mean)\n",
    "    J_map = np.zeros_like(heatmap_convergence_mean)\n",
    "    \n",
    "    for res in results:\n",
    "        heatmap_convergence_mean[res[\"m\"]][nmax+res[\"n\"]] = np.mean(np.abs(res[\"BdotN\"]))\n",
    "        heatmap_convergence_max[res[\"m\"]][nmax+res[\"n\"]] = np.max(res[\"BdotN\"])\n",
    "        J_map[res[\"m\"]][nmax+res[\"n\"]] = res[\"J\"]\n",
    "        \n",
    "        if res[\"J\"] < B_DOT_N_THRESHOLD:\n",
    "            coil_complexity_map[res[\"m\"]][nmax+res[\"n\"]] = res[\"complexity\"]\n",
    "        else:\n",
    "            coil_complexity_map[res[\"m\"]][nmax+res[\"n\"]] = np.nan \n",
    "    plt.subplot(221)\n",
    "    plot_modes(heatmap_convergence_mean)\n",
    "    plt.title(\"Mean B.n error\")\n",
    "    plt.subplot(222)\n",
    "    plot_modes(heatmap_convergence_max)\n",
    "    plt.title(\"Max B.n error\")\n",
    "    plt.subplot(223)\n",
    "    plot_modes(coil_complexity_map)\n",
    "    plt.title(\"coil_complexity_map\")\n",
    "    plt.subplot(224)\n",
    "    plot_modes(J_map)\n",
    "    plt.title(\"J\")\n",
    "    plt.savefig(\"complexity.png\")\n",
    "\n",
    "    # for res in results:\n",
    "    #     simsopt.geo.plot(res[\"coils\"])\n",
    "\n",
    "directories = [os.path.join(out_dir, name) for name in os.listdir(out_dir) if os.path.isdir(os.path.join(out_dir, name))]\n",
    "ipywidgets.interact(plot_coil_complexities, campaign_folder=ipywidgets.Dropdown(options=reversed(directories)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sensitivity analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spectrum generation functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def image_from_magn_and_phase(magnitudes, phases, image_size=(64, 64)):\n",
    "    # Construct frequency domain signal\n",
    "    # magnitudes[0] = 0\n",
    "    freq_signal = magnitudes * np.exp(1j * phases)\n",
    "    \n",
    "    # Perform 2D inverse Fourier transform\n",
    "    image = np.fft.ifft2(freq_signal, s=image_size)\n",
    "\n",
    "    # plt.subplot(131)\n",
    "    # plt.imshow(np.fft.fftshift(np.imag(image)))\n",
    "    # plt.subplot(132)\n",
    "    # plt.imshow(np.fft.fftshift(np.real(image)))\n",
    "\n",
    "    # Take real part of the image\n",
    "    return np.imag(image)\n",
    "\n",
    "# Generate random magnitudes and phases\n",
    "magnitudes = np.random.rand(Mpol_max, Ntor_max)\n",
    "phases = np.random.rand(Mpol_max, Ntor_max)  * 2 * np.pi\n",
    "image_from_magn_and_phase(magnitudes, phases)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from skimage.filters import window\n",
    "\n",
    "def image_from_spectral_radius(spectral_radius:float, magnitude = 1, image_size=(64, 64), normalization=None):\n",
    "    if normalization is None:\n",
    "        normalization = np.ones(image_size)\n",
    "    # Construct frequency domain signal\n",
    "    magnitudes = np.random.rand(*image_size)\n",
    "    magnitudes[0,0] = 0\n",
    "    phases = np.random.rand(*image_size) * 2 * np.pi\n",
    "    freq_signal = magnitudes * np.exp(1j * phases)\n",
    "\n",
    "    y, x = np.indices(image_size)\n",
    "    r = np.sqrt(x**2 + y**2)\n",
    "    # w = window(spectral_radius, image_size)\n",
    "    # plt.imshow(np.fft.ifftshift(w))\n",
    "    # plt.show()\n",
    "    freq_signal = np.where(r<=spectral_radius, freq_signal, np.zeros_like(freq_signal))\n",
    "    \n",
    "    # Perform 2D inverse Fourier transform\n",
    "    image = np.real(np.fft.ifft2(1j*freq_signal.imag, s=image_size)/normalization)\n",
    "\n",
    "    return image * magnitude * (image_size[0]* image_size[1]) / np.sum(np.abs(image))\n",
    "\n",
    "\n",
    "plt.imshow(image_from_spectral_radius(5))\n",
    "plt.colorbar()\n",
    "plt.show()\n",
    "np.sum(np.abs(image_from_spectral_radius(3)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define and run the optimization problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_dir_campaign = out_dir / datetime.now().strftime(\"%Y-%m-%d--%H-%M\") #For the current experimental campaign\n",
    "out_dir_campaign.mkdir(parents=True, exist_ok=True)\n",
    "print(out_dir_campaign)\n",
    "\n",
    "def problem(spectral_radius, magnitude, complex_surface=False, stellsym=True, batch=False):\n",
    "    nfp = 4\n",
    "    if complex_surface:\n",
    "        nfp = 1\n",
    "    \n",
    "\n",
    "    if stellsym:\n",
    "        nphi_asym = nphi\n",
    "        ncoils = 3\n",
    "    else:\n",
    "        nphi_asym = nphi*2\n",
    "        ncoils = 6\n",
    "\n",
    "    # Generate Geometry from X2 parameter\n",
    "    if complex_surface:\n",
    "        filename = \"tests/test_files/input.LandremanPaul2021_QA\"\n",
    "        s = simsopt.geo.SurfaceRZFourier.from_vmec_input(filename, nphi=nphi_asym, ntheta=ntheta)\n",
    "        s.extend_via_projected_normal(0.1)\n",
    "    else:\n",
    "        s = simsopt.geo.SurfaceRZFourier.from_nphi_ntheta(nphi_asym, ntheta, range=\"half period\" if stellsym else \"field period\", nfp=nfp)\n",
    "        # s = simsopt.geo.SurfaceRZFourier.from_nphi_ntheta(nphi_asym, ntheta, range=\"field period\", nfp=nfp)\n",
    "        \n",
    "        s.change_resolution(5, 5)\n",
    "        \n",
    "        s.set_rc(0,0,R0)\n",
    "        s.set_rc(1,0,R1_reactor)\n",
    "        s.set_zs(1,0,R1_reactor)\n",
    "        # mask = s.rc < 0.0\n",
    "        # mask[0,:s.ntor+1] = True\n",
    "        # s.rc[~mask] += X2\n",
    "    \n",
    "    normalization_weight = np.sqrt(s.area()/(2 * np.pi**2 *np.linalg.norm(s.normal(), axis=-1)))\n",
    "    # , normalization = normalization_weight\n",
    "    B_external_normal = image_from_spectral_radius(spectral_radius, magnitude, normalization_weight.shape)\n",
    "\n",
    "    if(not batch):\n",
    "        plt.figure()\n",
    "        plt.subplot(221)\n",
    "        plt.imshow(B_external_normal)\n",
    "        plt.colorbar()\n",
    "        plt.title(\"Target field\")\n",
    "\n",
    "    CS_THRESHOLD = 0.05\n",
    "    # Create the initial coils:\n",
    "    base_curves = simsopt.geo.create_equally_spaced_curves(ncoils, s.nfp, stellsym=stellsym, R0=s.major_radius(), R1=s.minor_radius()*(1+5*CS_THRESHOLD), order=order)\n",
    "    # base_curves = simsopt.geo.create_equally_spaced_curves(ncoils, s.nfp, stellsym=stellsym, R0=1, R1=0.5, order=order)\n",
    "    base_currents = [Current(5e6 / ncoils * 1e-5) * 1e5 for _ in range(ncoils)]\n",
    "\n",
    "    # Above, the factors of 1e-5 and 1e5 are included so the current\n",
    "    # degrees of freedom are O(1) rather than ~ MA.  The optimization\n",
    "    # algorithm may not perform well if the dofs are scaled badly.\n",
    "\n",
    "    coils = coils_via_symmetries(base_curves, base_currents, s.nfp, stellsym=stellsym)\n",
    "    bs = BiotSavart(coils)\n",
    "    bs.set_points(s.gamma().reshape((-1, 3)))\n",
    "    curves = [c.curve for c in coils]\n",
    "\n",
    "    # Form the total objective function.\n",
    "    B_DOT_N_THRESHOLD = 0.05  # 10% error allowed\n",
    "    LENGTH_WEIGHT = 0.05\n",
    "    TARGET_LENGTH = 1.0*s.minor_radius() * 2 * np.pi\n",
    "    CC_WEIGHT = 1.0\n",
    "    CC_THRESHOLD = 0.02\n",
    "    CURVATURE_THRESHOLD = 0.02\n",
    "    CURVATURE_WEIGHT = 1e-3\n",
    "\n",
    "    MSC_THRESHOLD = 0.2\n",
    "    MSC_WEIGHT = 1e-3\n",
    "        \n",
    "    Jf_initial = SquaredFlux(s, BiotSavart([]), target=B_external_normal).J() # Scalar evaluation to create the percentage error threshold\n",
    "\n",
    "    Jf = 2*SquaredFlux(s, bs, target=B_external_normal)\n",
    "    Jf_threshold =  QuadraticPenalty(Jf, B_DOT_N_THRESHOLD*Jf_initial, \"max\")\n",
    "    Jls = LENGTH_WEIGHT*sum([QuadraticPenalty(simsopt.geo.CurveLength(c), TARGET_LENGTH) for c in base_curves])\n",
    "    Jccdist = CC_WEIGHT*simsopt.geo.CurveCurveDistance(curves, CC_THRESHOLD, num_basecurves=ncoils)\n",
    "    Jcsdist = simsopt.geo.CurveSurfaceDistance(curves, s, CS_THRESHOLD)\n",
    "    Jcs = CURVATURE_WEIGHT * sum([simsopt.geo.LpCurveCurvature(c, 2, CURVATURE_THRESHOLD) for c in base_curves])\n",
    "    Jmscs = MSC_WEIGHT * sum( [QuadraticPenalty(simsopt.geo.MeanSquaredCurvature(c), MSC_THRESHOLD, \"max\") for c in base_curves])\n",
    "    Jcs_nothresh = CURVATURE_WEIGHT * sum([simsopt.geo.LpCurveCurvature(c, 2) for c in base_curves])\n",
    "    Jmscs_nothresh = MSC_WEIGHT * sum([simsopt.geo.MeanSquaredCurvature(c) for c in base_curves])\n",
    "    \n",
    "    JF = Jf  + Jls + Jccdist  + Jcsdist  + Jcs + Jmscs\n",
    "       \n",
    "\n",
    "    def fun(dofs):\n",
    "        JF.x = dofs\n",
    "        J = JF.J()\n",
    "        grad = JF.dJ()\n",
    "        return J, grad\n",
    "\n",
    "    res = minimize(fun, JF.x, jac=True, method='L-BFGS-B', options={'maxiter': MAXITER, 'maxcor': 300, 'ftol': 1e-20, 'gtol': 1e-20}, tol=1e-20)\n",
    "\n",
    "    # Writing to file\n",
    "    Bbs = bs.B().reshape((nphi_asym, ntheta, 3))\n",
    "    BdotN = B_external_normal - np.sum(Bbs * s.unitnormal(), axis=2)\n",
    "    pointData = {\"B_N\": BdotN[:, :, None]}\n",
    "    s.to_vtk(out_dir / \"surf_opt\", extra_data=pointData)\n",
    "    simsopt.geo.curves_to_vtk(curves, out_dir / \"curves_opt\")\n",
    "    \n",
    "    def get_unique_filename(filename):\n",
    "        if not os.path.exists(filename):\n",
    "            return filename\n",
    "\n",
    "        # Split the filename and its extension\n",
    "        name, ext = os.path.splitext(filename)\n",
    "\n",
    "        # Initialize counter\n",
    "        counter = 1\n",
    "\n",
    "        # Increment the filename until it is unique\n",
    "        while True:\n",
    "            new_filename = f\"{name}_{counter}{ext}\"\n",
    "            if not os.path.exists(new_filename):\n",
    "                return new_filename\n",
    "            counter += 1\n",
    "\n",
    "    complexity = (Jls.J()+Jccdist.J()+Jcsdist.J()+ Jcs_nothresh.J() + Jmscs_nothresh.J()) # only the coil complexity terms, without the magnetic field residual\n",
    "    simsopt.save({  \"spectral_radius\":spectral_radius,\n",
    "                    \"magnitude\":magnitude,\n",
    "                    \"J\": JF.J(),\n",
    "                    \"complexity\": complexity,\n",
    "                    \"coils\":coils,\n",
    "                    \"surf\":s,\n",
    "                    \"B_external_normal\":B_external_normal,\n",
    "                    \"BdotN\":BdotN}, \n",
    "                    get_unique_filename(out_dir_campaign/\"sens_result.json\"))\n",
    "    \n",
    "    if(not batch):\n",
    "        plt.subplot(223)\n",
    "        plt.imshow(np.real(np.sum(Bbs * s.unitnormal(), axis=2)))\n",
    "        plt.colorbar()\n",
    "        plt.title(\"B.n realized\")\n",
    "\n",
    "        plt.subplot(122)\n",
    "        plt.imshow(np.real(BdotN))\n",
    "        plt.colorbar()\n",
    "        plt.title(\"B.n Diff\")\n",
    "\n",
    "        plt.show()\n",
    "\n",
    "        simsopt.geo.plot([s]+curves)\n",
    "\n",
    "        print(\"Jf\",Jf.J())\n",
    "        print(\"Jf_threshold\",Jf_threshold.J())\n",
    "        print(\"Jls\",Jls.J())\n",
    "        print(\"Jccdist\",Jccdist.J())\n",
    "        print(\"Jcsdist\",Jcsdist.J())\n",
    "        print(\"Jcs\",Jcs.J())\n",
    "        print(\"Jmscs\",Jmscs.J())\n",
    "\n",
    "    return complexity, np.mean(np.abs(BdotN))\n",
    "\n",
    "counter = 0\n",
    "def debug_func(X,X2, **kwargs):\n",
    "    global counter\n",
    "\n",
    "    plt.subplot(4,4,(counter%16+1))\n",
    "    plt.imshow(image_from_magn_and_phase(X, X2))\n",
    "    counter+=1\n",
    "    if(counter%16 == 0):\n",
    "        plt.colorbar()\n",
    "        plt.show()\n",
    "        \n",
    "    return np.mean(X), np.max(X)\n",
    "\n",
    "def wrapped_problem(X: np.ndarray, func=problem) -> np.ndarray:\n",
    "    N, D = X.shape\n",
    "    results = []\n",
    "    for i in range(N):\n",
    "        print(f\"{i+1}/{N}\")\n",
    "        results.append(func(X[i,0], 0.25, batch=True))\n",
    "\n",
    "    return np.array(results)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "problem(1, 0.25)\n",
    "problem(2, 0.25)\n",
    "problem(3, 0.25)\n",
    "problem(4, 0.25)\n",
    "problem(5, 0.25)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from SALib import ProblemSpec\n",
    "\n",
    "sp = ProblemSpec({\n",
    "        \"names\":  [\"spectral_radius\"],\n",
    "        \"bounds\": [[1, 12]],\n",
    "        \"outputs\": [\"coil_complexity\", \"BdotN_mean\"],\n",
    "    })\n",
    "\n",
    "sp.sample_sobol(1024, calc_second_order=True).evaluate(wrapped_problem).analyze_sobol()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget\n",
    "import pandas as pd\n",
    "# https://gsa-module.readthedocs.io/en/stable/implementation/sobol_indices.html\n",
    "sp.plot()\n",
    "sp.heatmap()\n",
    "complexity_Si, BdotN_Si = sp.to_df()\n",
    "complexity_Si[0].to_csv(out_dir_campaign/\"total_Si_complex_landreman.csv\")\n",
    "BdotN_Si[0].to_csv(out_dir_campaign/\"total_Si_BdotN_landreman.csv\")\n",
    "complexity_Si[1].to_csv(out_dir_campaign/\"first_Si_complex_landreman.csv\")\n",
    "BdotN_Si[1].to_csv(out_dir_campaign/\"first_Si_BdotN_landreman.csv\")\n",
    "complexity_Si[2].to_csv(out_dir_campaign/\"second_Si_complex_landreman.csv\")\n",
    "BdotN_Si[2].to_csv(out_dir_campaign/\"second_Si_BdotN_landreman.csv\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sp\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CSV analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import plotly.express as px\n",
    "import ipywidgets\n",
    "\n",
    "def implot_df(df:pd.DataFrame, **kwargs):\n",
    "  mmesh, nmesh = np.meshgrid(np.arange(Mpol_max), np.arange(Ntor_max))\n",
    "  map_to_mn = {f\"M{i}\": (m, n) for m, n, i in zip(mmesh.flatten(), nmesh.flatten(), range(Ntor_max*Mpol_max))}\n",
    "\n",
    "  df[\"mn\"] = df.index.to_series().map(map_to_mn)\n",
    "  df[\"m\"] = df[\"mn\"].apply(lambda x: x[0])\n",
    "  df[\"n\"] = df[\"mn\"].apply(lambda x: x[1])\n",
    "\n",
    "  return px.bar(df, x=\"m\", y=df.columns[0], color=\"n\",facet_col=\"n\", error_y=df.columns[1], barmode=\"group\", **kwargs)\n",
    "  return px.scatter(df, \"n\",\"m\", color=df.columns[0], size=df.columns[1], title=\"Impact of geometry coefficients\")\n",
    "\n",
    "analysis_folder = \"\"\n",
    "def load_dataframes(campaign_folder):\n",
    "  global analysis_folder, first_Si_complexity, total_Si_complexity, total_Si_BdotN, first_Si_BdotN\n",
    "  try:\n",
    "    first_Si_complexity = pd.read_csv(os.path.join(campaign_folder,\"first_Si_complex_landreman.csv\"), index_col=0)\n",
    "    total_Si_complexity = pd.read_csv(os.path.join(campaign_folder,\"total_Si_complex_landreman.csv\"), index_col=0)\n",
    "    total_Si_BdotN = pd.read_csv(os.path.join(campaign_folder,\"total_Si_BdotN_landreman.csv\"), index_col=0)\n",
    "    first_Si_BdotN = pd.read_csv(os.path.join(campaign_folder,\"first_Si_BdotN_landreman.csv\"), index_col=0)\n",
    "\n",
    "    analysis_folder = campaign_folder\n",
    "\n",
    "    implot_df(total_Si_complexity, title=\"Coil Complexity\").show()\n",
    "  except:\n",
    "    print(\"could not find files at\", campaign_folder)\n",
    "out_dir = \"output\"\n",
    "directories = [os.path.join(out_dir, name) for name in os.listdir(out_dir) if os.path.isdir(os.path.join(out_dir, name))]\n",
    "ipywidgets.interact(load_dataframes, campaign_folder=ipywidgets.Dropdown(options=reversed(directories)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "implot_df(total_Si_complexity, title=\"Coil Complexity\").show()\n",
    "implot_df(total_Si_BdotN, title=\"B.n Residual\").show()\n",
    "implot_df(first_Si_complexity, title=\"Coil Complexity\").show()\n",
    "implot_df(first_Si_BdotN, title=\"B.n Residual\").show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget\n",
    "simsopt.geo.plot([results[100][\"surf\"]]+results[100][\"coils\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "import simsopt\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "plt.close(\"all\")\n",
    "\n",
    "\n",
    "results = []\n",
    "i = 0\n",
    "for path in os.listdir(analysis_folder):\n",
    "  if(path.endswith(\".json\")):\n",
    "    optimization_res = simsopt.load(analysis_folder + \"/\" +path)\n",
    "    results.append(optimization_res)\n",
    "\n",
    "    # simsopt.geo.plot([optimization_res[\"surf\"]]+optimization_res[\"coils\"])\n",
    "    if i<=320:\n",
    "      plt.subplot(8,8,(i%64+1))\n",
    "      plt.imshow(optimization_res[\"B_external_normal\"])\n",
    "      i+=1\n",
    "      if(i%64 == 0):\n",
    "        plt.show()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import plotly.express as px\n",
    "\n",
    "df = pd.DataFrame(results)\n",
    "df[\"J\"] = df[\"J\"].astype(float)\n",
    "df[\"B target max\"] = df[\"B_external_normal\"].apply(np.max)\n",
    "df[\"B.n residual max\"] = df[\"BdotN\"].apply(lambda x: min(5,np.max(x)))\n",
    "df[\"spectral_radius\"] = df[\"spectral_radius\"].astype(float)\n",
    "df[\"complexity\"] = df[\"complexity\"].astype(float)\n",
    "df[\"magnitude\"] = df[\"magnitude\"].astype(float)\n",
    "df[\"run_id\"] = df.index\n",
    "\n",
    "df[\"success\"] = df[\"B.n residual max\"]<df[\"B target max\"]/2\n",
    "fig = px.scatter(df.select_dtypes(exclude=[\"object\"]), \"spectral_radius\", \"complexity\", color=\"B.n residual max\", symbol=\"success\",\n",
    "                 hover_data={\"run_id\":True}, template=\"plotly_dark\")\n",
    "fig.update_layout(height=600, clickmode='event+select')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "second_Si_complexity = pd.read_csv(\"second_Si_complex_landreman.csv\", index_col=0)\n",
    "second_Si_BdotN = pd.read_csv(\"second_Si_BdotN_landreman.csv\", index_col=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def implot_second_order_df(df:pd.DataFrame):\n",
    "  df[\"mn\"] = df.index.to_series().apply(lambda x: x.split(\"', '\")[0][2:]).map(map_to_mn)\n",
    "  df[\"mn2\"] = df.index.to_series().apply(lambda x: x.split(\"', '\")[1][:-2]).map(map_to_mn)\n",
    "  df[\"m\"] = df[\"mn\"].apply(lambda x: x[0])\n",
    "  df[\"n\"] = df[\"mn\"].apply(lambda x: x[1])\n",
    "  df[\"m2\"] = df[\"mn2\"].apply(lambda x: x[0])\n",
    "  df[\"n2\"] = df[\"mn2\"].apply(lambda x: x[1])\n",
    "\n",
    "  return px.bar(df, x=\"m\", y=df.columns[0], color=\"m2\",facet_col=\"n\", facet_row=\"n2\", error_y=df.columns[1], barmode=\"group\")\n",
    "\n",
    "implot_second_order_df(second_Si_complexity)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Can we find a pattern?\n",
    "Summarize m and n terms of a certain \"radius\" into bins. Check for correlations, i.e. if ANY high order frequencies exist."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis_df = pd.DataFrame(results)\n",
    "analysis_df[\"J\"] = analysis_df[\"J\"].astype(float)\n",
    "analysis_df[\"complexity\"] = analysis_df[\"complexity\"].astype(float)\n",
    "# analysis_df[\"BdotN_mean\"] = analysis_df[\"BdotN\"].map(np.max)\n",
    "analysis_df.plot()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def radial_profile(data, center):\n",
    "    y, x = np.indices((data.shape))\n",
    "    r = np.sqrt((x - center[0])**2 + (y - center[1])**2)\n",
    "    r = r.astype(int)\n",
    "\n",
    "    tbin = np.bincount(r.ravel(), data.ravel())\n",
    "    nr = np.bincount(r.ravel())\n",
    "    radialprofile = tbin / nr\n",
    "    return radialprofile \n",
    "\n",
    "def radial_profile_center(data):\n",
    "    return radial_profile(data, np.array(data.shape)/2)\n",
    "\n",
    "num_bins = len(analysis_df[\"Bmagnitude\"].apply(radial_profile_center)[0])\n",
    "res = analysis_df[\"Bmagnitude\"].apply(radial_profile_center)\n",
    "analysis_df = analysis_df.join(pd.DataFrame(res.tolist()).add_prefix(\"bin_\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget\n",
    "plt.close()\n",
    "plt.imshow(analysis_df.select_dtypes(exclude=[\"object\"]).corr()-np.eye(num_bins+2)) \n",
    "plt.xlabel([*analysis_df.select_dtypes(exclude=[\"object\"]).columns])\n",
    "# plt.ylabel([*analysis_df.select_dtypes(exclude=[\"object\"]).columns])\n",
    "plt.colorbar()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Efficient Fields (Landreman basis functions)\n",
    "$f_j(\\theta, \\phi) = \\sqrt{\\frac{A}{2 \\pi^2 \\|\\vec{N}(\\theta, \\phi)\\|}} sin(m_j \\theta - n_j \\phi)$\n",
    "\n",
    "$f_j(\\theta, \\phi) = \\sqrt{\\frac{A}{2 \\pi^2 \\|\\vec{N}(\\theta, \\phi)\\|}} cos(m_j \\theta - n_j \\phi)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import simsopt\n",
    "import simsopt.geo\n",
    "import numpy as np\n",
    "\n",
    "surf = simsopt.geo.SurfaceRZFourier.from_nphi_ntheta(nfp=3, mpol=5, ntor=4)\n",
    "normalization_weight = np.sqrt(surf.area()/(2 * np.pi**2 *np.linalg.norm(surf.normal(), axis=-1)))\n",
    "\n",
    "print(surf.dmean_cross_sectional_area_by_dcoeff().shape)\n",
    "print(surf.gamma().shape)\n",
    "print(surf.mean_cross_sectional_area())\n",
    "# px.scatter_3d(x=surf.gamma()[:,:,0].flatten(), y=surf.gamma()[:,:,1].flatten(), z=surf.gamma()[:,:,2].flatten())\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}