{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Multi-angle sweeping\n",
    "\n",
    "This notebook shows how to do multi-angle sweeping, including fused strategies versus composed strategies."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "We first create the same circular domain as before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "jax.config.update(\"jax_enable_x64\", True)\n",
    "jax.config.update(\"jax_platforms\", \"cpu\")\n",
    "import jax.numpy as jnp\n",
    "from jax_sbgeom.flux_surfaces.miller_tokamak import MillerTokamak\n",
    "import jax_sbgeom as jsb\n",
    "from jax_sn.domain import Domain, CrossSectionsIndexed, CrossSectionData\n",
    "from jax_sn.solution_domain import BasixLagrangianSimplex\n",
    "from jax_sn.tools.plotly import plot_domain\n",
    "\n",
    "tokamak = MillerTokamak(rminor = 1.0, rmajor = 1.2, elongation = 0.5, triangularity = 0.2)\n",
    "\n",
    "mesh_tokamak = jsb.flux_surfaces.flux_surface_meshing.mesh_poloidal_plane(tokamak, jnp.linspace(0.0,1.0,4), 0.0, 7)\n",
    "# Convert it to a 2D mesh!\n",
    "mesh_tokamak           = (jnp.stack([mesh_tokamak[0][:,0], mesh_tokamak[0][:,2]], axis=-1), mesh_tokamak[1])\n",
    "\n",
    "cross_section          = CrossSectionData(\n",
    "                            total     = jnp.array([[5.0]]), \n",
    "                            scattering= jnp.array([[[[0.0]]]]))\n",
    "\n",
    "cross_sections_tokamak = CrossSectionsIndexed(indices = jnp.full((mesh_tokamak[1].shape[0],), 0), cross_section_data = cross_section)\n",
    "\n",
    "domain_tokamak         = Domain.from_mesh_and_cross_sections(mesh_tokamak, BasixLagrangianSimplex(1, 2).face_template, cross_sections_tokamak)\n",
    "plot_domain(domain_tokamak, width=800, height=800)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "Let's create a quadrature set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sn.quadrature_set import create_tn_quadrature_set\n",
    "\n",
    "quadrature_set = create_tn_quadrature_set(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "In general, composed operators have two possibilities (also for single group and multigroup operators): either we use a sub-operator and compose it using Python functions (e.g. vmap or scan), or we implement the entire operator immediately. We will now show how to construct multi-angle transport operator inversion using either approach, but always use Vmap forward strategy.\n",
    "\n",
    "\n",
    "First, we will just vmap the a BatchedSweep (which instead of first inverting all element matrices and then collapsing, does this batch-by-batch), and print what it requires of the sweep data (nothing: a plain topological sort will do):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sn.operators.transport_operator import TransportOperatorMultiAngle, VmapSweep, VmapAnglesInversion, VmapForward, VmapAnglesForward, BatchedSweep\n",
    "multi_angle_forward_strategy   = VmapAnglesForward(VmapForward())\n",
    "multi_angle_inversion_strategy = VmapAnglesInversion(BatchedSweep(5))\n",
    "print(\"Required sweep-data capabilities:\", multi_angle_inversion_strategy.required_capabilities())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "The second sweep type is a fused Warp kernel, which requires Levels data. This means data to parallelize over the levels of the elements as well (instead of \n",
    "just the topological sort)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sn.operators.transport_operator import WarpLevelAngles\n",
    "\n",
    "multi_angle_warp_strategy = WarpLevelAngles()\n",
    "print(\"Required sweep-data capabilities:\", multi_angle_warp_strategy.required_capabilities())\n",
    "\n",
    "plot_domain(domain_tokamak, show_levels=True, level_mode = \"levels\", graph_angle = jnp.array([1/ jnp.sqrt(2), 1/ jnp.sqrt(2)]), show_graph = True, width=800, height=800)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "Elements 0 and 1 can be computed in parallel; this Warp-based kernel parallelizes that too to achieve higher parallelism (mostly useful on GPU).\n",
    "\n",
    "Let's actually construct them (warp requires $\\leq 4$ basis functions, so limited to linear triangles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sn.solution_domain import SolutionDomain\n",
    "from jax_sn.operators.transport_operator import TransportData\n",
    "solution_domain = SolutionDomain.from_element_and_domain(BasixLagrangianSimplex(1, 2), domain_tokamak)\n",
    "\n",
    "transport_data_basic = TransportData(solution_domain.spatial_data, 0, quadrature_set.angles,\n",
    "                                      multi_angle_inversion_strategy.build_default_sweep_data(domain_tokamak, quadrature_set.angles))\n",
    "transport_data_level = TransportData(solution_domain.spatial_data, 0, quadrature_set.angles, \n",
    "                                     multi_angle_warp_strategy.build_default_sweep_data(domain_tokamak, quadrature_set.angles))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "The added sweep data is permissive: we can construct a multi-angle strategy requiring the basic data with the more specialized data (but obviously not vice-versa):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sn.operators.transport_operator import TransportOperatorMultiAngle\n",
    "\n",
    "TransportOperatorMultiAngle(transport_data_level, multi_angle_forward_strategy, multi_angle_inversion_strategy)\n",
    "try: \n",
    "    TransportOperatorMultiAngle(transport_data_basic, multi_angle_forward_strategy, multi_angle_warp_strategy)\n",
    "except Exception as e:\n",
    "    print(\"Expected error:\", e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "transport_operator_vmap = TransportOperatorMultiAngle(transport_data_basic, \n",
    "                                                      multi_angle_forward_strategy, \n",
    "                                                      multi_angle_inversion_strategy)\n",
    "transport_operator_warp = TransportOperatorMultiAngle(transport_data_level, \n",
    "                                                      multi_angle_forward_strategy, \n",
    "                                                      multi_angle_warp_strategy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "They have the same inversion call as single-angle transport operators:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "rhs = jnp.ones_like(transport_operator_vmap.out_structure())\n",
    "\n",
    "sol_1 = transport_operator_vmap.inv(rhs)\n",
    "sol_2 = transport_operator_warp.inv(rhs)\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "print(\"Shape [n_angles, n_elements, n_basis]:\", sol_1.shape)\n",
    "print(\"Difference between vmap and warp solutions:\", np.max(np.abs(sol_1 - sol_2)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "If the operator is composed or has a 'natural' single-angle version,  it supports indexing to get the single-angle operator for a given angle:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "transport_operator_vmap[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "transport_operator_warp[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "To recap, a composed strategy in this case is just either two nested strategies or one bigger:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sn.operators.transport_operator import MultiAngleInversionStrategy\n",
    "multi_angle_inversion_strategy = VmapAnglesInversion(BatchedSweep(5))\n",
    "multi_angle_warp_strategy      = WarpLevelAngles()\n",
    "\n",
    "print(isinstance(multi_angle_inversion_strategy, MultiAngleInversionStrategy))\n",
    "print(isinstance(multi_angle_warp_strategy, MultiAngleInversionStrategy))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {},
   "source": [
    "Naturally, in tests, we test that all such implementations are equivalent up to numerical precision, including VJP/JVPs and adjoint/transport/inversion idempotence for a set of representative dimension/domain combinations (and similarly for other operators!)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jaxsbgeom",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}