{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Geometry (mapped domains)\n",
    "\n",
    "In Struphy, a computational domain is defined by a **mapping** from logical coordinates to physical space. Concretely, we use a map\n",
    "$$\n",
    "F:(0,1)^3 \\rightarrow \\Omega,\\qquad \\boldsymbol{\\eta}=(\\eta_1,\\eta_2,\\eta_3) \\mapsto \\mathbf{x}=(x,y,z),\n",
    "$$\n",
    "where the logical unit cube is easy to mesh, and the mapping bends this mesh into the physical geometry of interest.\n",
    "\n",
    "This is why mapped domains are so useful: we can keep structured discretizations in logical space while representing curved geometries (cylinders, tori, tokamaks, stellarators) in physical space.\n",
    "\n",
    "A practical challenge is the **polar singularity** (the magnetic axis). In Struphy, two strategies are common:\n",
    "\n",
    "1. Use `polar splines` in the de Rham setup (more expensive).\n",
    "2. Cut out a small hole around the axis (cheap, but particles cannot live in that excluded region).\n",
    "\n",
    "In this tutorial we focus on the second strategy and build intuition from simple to more physics-driven mappings.\n",
    "\n",
    "## HollowCylinder\n",
    "\n",
    "We start with `HollowCylinder`, the simplest curvilinear example: a cylindrical shell obtained from the logical cube by an analytical mapping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from struphy import domains\n",
    "\n",
    "domain = domains.HollowCylinder()\n",
    "domain.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A 3D rendering helps to connect the mapping parameters to geometry: inner radius, outer radius, and length become visually obvious."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain.show_3d()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us inspect the default parameters of `HollowCylinder`. Read them as: `a1` = inner radius, `a2` = outer radius, `Lz` = axial length, and `poc` controls angular periodicity (\"piece of cake\")."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key, val in domain.params.items():\n",
    "    print(key, \"=\", val)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`poc` lets us model only a sector of the full cylinder and reconstruct periodicity through the mapping. This is useful when symmetry allows reducing the simulated angular extent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain_poc = domains.HollowCylinder(poc=3)\n",
    "domain_poc.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some key attributes to interpret are:\n",
    "\n",
    "- `kind_map`: internal mapping family identifier (`20-29` are cylinder/torus analytical mappings).\n",
    "- `pole`: `True` when the mapping includes a polar singularity on axis.\n",
    "- `periodic_eta3`: tells whether the logical direction $\\eta_3$ is periodic in this mapping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(domain.kind_map)\n",
    "print(domain.pole)\n",
    "print(domain.periodic_eta3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is also useful to inspect callable methods of a domain object: this shows what geometric operations (mapping, Jacobian, metric-related helpers, visualization, etc.) are available for downstream workflows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for attr in dir(domain):\n",
    "    if callable(getattr(domain, attr)) and \"__\" not in attr and attr[0] != \"_\":\n",
    "        print(attr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The domain object is itself callable: evaluating `domain(...)` applies the mapping $F(\\eta_1,\\eta_2,\\eta_3)=(x,y,z)$ from logical to physical coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(domain.__call__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we vary `a1` to control how much of the axis neighborhood is removed. Increasing `a1` enlarges the excluded core and avoids singular behavior at the center."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain = domains.HollowCylinder(a1=0.05)\n",
    "domain.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we set `a1=0`, the hole disappears and the axis is included. In that case, `domain.pole` becomes `True`, signaling the polar singularity discussed in the introduction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain = domains.HollowCylinder(a1=0.0)\n",
    "domain.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(domain.kind_map)\n",
    "print(domain.pole)\n",
    "print(domain.periodic_eta3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## HollowTorus\n",
    "\n",
    "`HollowTorus` extends the same idea to toroidal geometry: the logical cube is mapped to a torus with optional inner hole around the magnetic axis.\n",
    "\n",
    "Compared to `HollowCylinder`, important new parameters are:\n",
    "- `R0`: major radius (distance from torus center to tube center),\n",
    "- `sfl`: whether to use straight-field-line poloidal parametrization,\n",
    "- `tor_period`: built-in toroidal periodicity (e.g. `3` means one third of a full torus in the mapped geometry).\n",
    "\n",
    "Let us create [HollowTorus](https://struphy.pages.mpcdf.de/struphy/sections/domains.html#struphy.geometry.domains.HollowTorus) with default parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain = domains.HollowTorus()\n",
    "domain.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this mapping, `periodic_eta3` is `True` because the toroidal angle is encoded in $\\eta_3$ and treated periodically by construction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(domain.kind_map)\n",
    "print(domain.pole)\n",
    "print(domain.periodic_eta3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us inspect the defaults. In particular, compare `a1/a2` (minor radii), `R0` (major radius), and periodicity-related parameters to understand the torus shape and angular parametrization choices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key, val in domain.params.items():\n",
    "    print(key, \"=\", val)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we modify three knobs at once:\n",
    "\n",
    "- `a1` to change the hole around the axis,\n",
    "- `sfl=True` to switch poloidal coordinates to straight-field-line style,\n",
    "- `tor_period=1` to represent the full torus period in the map.\n",
    "\n",
    "This helps build intuition for how coordinate choices alter the geometry and periodicity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain = domains.HollowTorus(a1=0.05, sfl=True, tor_period=1)\n",
    "domain.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokamak\n",
    "\n",
    "For realistic tokamak geometry, Struphy uses mappings tied to an axisymmetric MHD equilibrium. Instead of prescribing a simple analytic torus, the domain is built from a poloidal flux function $\\psi$ and field-line tracing.\n",
    "\n",
    "So conceptually, the workflow is: equilibrium model $\\rightarrow$ flux surfaces $\\rightarrow$ mapped computational domain.\n",
    "\n",
    "[Tokamak](https://struphy.pages.mpcdf.de/struphy/sections/domains.html#struphy.geometry.domains.Tokamak) is the domain class implementing this idea.\n",
    "\n",
    "Let us create a `Tokamak` with default parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain = domains.Tokamak()\n",
    "domain.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain.show_3d()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us inspect the default tokamak mapping parameters. Notice how they now include options linked to equilibrium/flux-surface construction rather than only simple geometric radii."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key, val in domain.params.items():\n",
    "    if \"cx\" not in key and \"cy\" not in key:\n",
    "        print(key, \"=\", val)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A `Tokamak` domain is always coupled to an [AxisymmMHDequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/mhd_equils.html#struphy.fields_background.mhd_equil.base.AxisymmMHDequilibrium), which provides the flux function $\\psi$ used to construct the map.\n",
    "\n",
    "By default, this is [AdhocTorus](https://struphy.pages.mpcdf.de/struphy/sections/mhd_equils.html#struphy.fields_background.mhd_equil.equils.AdhocTorus), a convenient analytic equilibrium for testing. We can also switch to [EQDSKequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/mhd_equils.html#struphy.fields_background.mhd_equil.equils.EQDSKequilibrium), which represents a more realistic equilibrium source."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from struphy import equils\n",
    "\n",
    "mhd_eq = equils.EQDSKequilibrium()\n",
    "\n",
    "domain = domains.Tokamak(equilibrium=mhd_eq)\n",
    "domain.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we shrink the inner excluded region by adjusting `psi_shifts`, which changes how close the mapped domain gets to the magnetic axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain = domains.Tokamak(equilibrium=mhd_eq, psi_shifts=[0.2, 2])\n",
    "domain.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stellarator mappings\n",
    "\n",
    "Tokamak mappings above are axisymmetric. Stellarators, in contrast, are intrinsically 3D and typically require equilibrium data generated by dedicated external codes.\n",
    "\n",
    "Struphy provides interfaces to read equilibrium geometry from:\n",
    "\n",
    "- [GVEC equilibrium code](https://gitlab.mpcdf.mpg.de/gvec-group/gvec)\n",
    "- [DESC equilibrium code](https://desc-docs.readthedocs.io/en/latest/index.html)\n",
    "\n",
    "### GVEC interface\n",
    "\n",
    "The corresponding Struphy domain class is [GVECunit](https://struphy.pages.mpcdf.de/struphy/sections/domains.html?highlight=gvec#struphy.geometry.domains.GVECunit).\n",
    "Let us create an instance with default parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain = domains.GVECunit()\n",
    "domain.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inspecting defaults is useful here as well: they describe both geometry controls and how the external equilibrium information is interpreted by the mapping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key, val in domain.params.items():\n",
    "    if \"cx\" not in key and \"cy\" not in key and \"cz\" not in key:\n",
    "        print(key, \"=\", val)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we modify equilibrium-side parameters to open a hole around the magnetic axis (`rmin`) and use the whole stellarator (`use_nfp=False`). These settings are passed through [GVECequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/mhd_equils.html#struphy.fields_background.mhd_equil.equils.GVECequilibrium):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gvec_equil = equils.GVECequilibrium(rmin=0.1, use_nfp=False)\n",
    "domain = domains.GVECunit(gvec_equil)\n",
    "domain.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DESC interface\n",
    "\n",
    "The analogous interface for DESC is [DESCunit](https://struphy.pages.mpcdf.de/struphy/sections/domains.html?highlight=gvec#struphy.geometry.domains.DESCunit).\n",
    "As with GVEC, equilibrium choices are specified through [DESCequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/mhd_equils.html#struphy.fields_background.mhd_equil.equils.DESCequilibrium), then consumed by the domain mapping:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "desc_equil = equils.DESCequilibrium(use_nfp=False)\n",
    "domain = domains.DESCunit(desc_equil)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env (3.12.3)",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}