{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1471b1cc",
   "metadata": {},
   "source": [
    "### Helicity minimization\n",
    "\n",
    "In this tutorial, demonstrate how to use the auto-differentiation features"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0c02017",
   "metadata": {},
   "source": [
    "The setup is similar to the tokamak equilibrium tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "5fda129b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax.numpy import cos, pi, sin\n",
    "from mrx.derham_sequence import DeRhamSequence\n",
    "from mrx.mappings import toroid_map\n",
    "\n",
    "jax.config.update(\"jax_enable_x64\", True)\n",
    "\n",
    "Phi = toroid_map(epsilon=1/3)\n",
    "Seq = DeRhamSequence(\n",
    "    (5, 5, 5),                         # nb. of splines in (r, θ, ζ)\n",
    "    (3, 3, 3),                           # degree of splines in (r, θ, ζ)\n",
    "    5,                                   # nb. of quadrature points per spline\n",
    "    (\"clamped\", \"periodic\", \"periodic\"), # spline type in (r, θ, ζ)\n",
    "    Phi,                                 # mapping from (r, θ, ζ) to (x, y, z)\n",
    "    polar=True,                          # domain has a polar singularity\n",
    "    dirichlet=True                       # impose Dirichlet BCs on r=1 boundary\n",
    ")\n",
    "\n",
    "Seq.evaluate_1d()\n",
    "Seq.assemble_all()\n",
    "\n",
    "def B_0(p):\n",
    "    x, y, z = Phi(p)\n",
    "    R, phi = (x**2 + y**2)**0.5, jnp.arctan2(y, x)\n",
    "    BR, Bphi, Bz = -z/R, 1/R, (R-1)/R\n",
    "    Bx = BR * cos(phi) - Bphi * sin(phi)\n",
    "    By = BR * sin(phi) + Bphi * cos(phi)\n",
    "    return jnp.array([Bx, By, Bz])\n",
    "    \n",
    "B = jnp.linalg.solve(Seq.M2, Seq.P2(B_0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1858e851",
   "metadata": {},
   "source": [
    "Next, we compute the (generalized) helicity of the magnetic field, defined as\n",
    "$$\n",
    "    \\mathcal H = \\int_\\Omega A \\cdot (B + B_\\mathfrak{H}) \\, \\mathrm dx,\n",
    "$$\n",
    "where $\\mathrm{curl} \\, A = B - B_\\mathfrak{H}$ and $B_\\mathfrak{H}$ is the harmonic part of the magnetic field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8712a8e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_helicity(B, Seq):\n",
    "    A = jnp.linalg.solve(Seq.dd1, Seq.weak_curl @ B)\n",
    "    return A @ Seq.M12 @ B"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7536fe75",
   "metadata": {},
   "source": [
    "We now want to minimize the absolute value of the helicity of the magnetic field while keeping the magnetic energy fixed. To do so with a simple gradient descent method plus re-normalization. To compute the gradient of the helicity with respect to the magnetic field, we can use JAX's auto-differentiation capabilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "851d105f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial helicity: -0.062102\n"
     ]
    }
   ],
   "source": [
    "H = compute_helicity(B, Seq)\n",
    "dHdB = jax.grad(compute_helicity)(B, Seq)\n",
    "print(f\"Initial helicity: {H:.6f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "29f77b1e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.06333230830297515"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "2 * pi**2/3 * (2 - (2+1/9)*(1-1/9)**0.5)  # analytical value for comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e631f2c5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(0.22377362, dtype=float64)"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "helicity_value_and_grad = jax.jit(jax.value_and_grad(compute_helicity, argnums=0), static_argnums=1)\n",
    "\n",
    "i = 0\n",
    "helicity = 1e10\n",
    "while (helicity > 1e-6):\n",
    "    helicity, grad_helicity = helicity_value_and_grad(B, Seq)\n",
    "    B -= 0.01 * grad_helicity\n",
    "    B /= (B @ Seq.M2 @ B)**0.5\n",
    "    i += 1\n",
    "    if i % 100 == 0:\n",
    "        print(f\"helicity and energy after {i} iterations: {helicity:.2e}, {0.5 * B @ Seq.M2 @ B:.2e}\")\n",
    "print(f\"helicity and energy after {i} iterations: {helicity:.2e}, {0.5 * B @ Seq.M2 @ B:.2e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66dc64ec",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}