{ "cells": [ { "cell_type": "markdown", "id": "0", "metadata": {}, "source": [ "# Particle tracing\n", "\n", "In this tutorial you will build and compare particle-tracing simulations with the toy models [Vlasov](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.Vlasov) and [GuidingCenter](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.GuidingCenter).\n", "\n", "By the end, you will know how to:\n", "\n", "1. change the geometry,\n", "2. choose different particle loading strategies,\n", "3. add a static background magnetic field.\n", "\n", "## Part 1: Particles in a cylinder\n", "\n", "As in Tutorial 2, we recreate the setup directly in the notebook so each option is visible and easy to modify. If you prefer starting from default launch files in a terminal, run:\n", "\n", "```\n", "struphy params Vlasov\n", "struphy params GuidingCenter\n", "```\n", "\n", "Here we set the simulation domain $\\Omega$ to a cylinder and compare two sampling strategies for particle positions:\n", "\n", "- uniform in logical space $[0,1]^3 = F^{-1}(\\Omega)$,\n", "- uniform directly in the physical cylinder $\\Omega$.\n", "\n", "Start by importing the API and the model:" ] }, { "cell_type": "code", "execution_count": null, "id": "1", "metadata": {}, "outputs": [], "source": [ "from struphy import BaseUnits, EnvironmentOptions, Time\n", "from struphy import equils\n", "from struphy import domains\n", "from struphy import grids\n", "from struphy import DerhamOptions\n", "from struphy import maxwellians\n", "from struphy import (BoundaryParameters, \n", " LoadingParameters, \n", " WeightsParameters,\n", " SortingParameters,\n", " SavingParameters,\n", " )\n", "from struphy import Simulation\n", "\n", "# import model\n", "from struphy.models import Vlasov" ] }, { "cell_type": "markdown", "id": "2", "metadata": {}, "source": [ "### Step 1: define separate output folders\n", "\n", "We will run two comparable simulations and save them in different output folders. Set the folder names through environment variables:" ] }, { "cell_type": "code", "execution_count": null, "id": "3", "metadata": {}, "outputs": [], "source": [ "# light-weight model instances\n", "model = Vlasov()\n", "model_2 = Vlasov()\n", "\n", "# environment options\n", "env = EnvironmentOptions()\n", "env_2 = EnvironmentOptions(sim_folder=\"sim_2\")" ] }, { "cell_type": "markdown", "id": "4", "metadata": {}, "source": [ "### Step 2: set shared simulation options\n", "\n", "Keep the generic options identical in both runs so the loading strategy is the main difference. We perform only one time step, use a cylindrical geometry, and leave the equilibrium empty:" ] }, { "cell_type": "code", "execution_count": null, "id": "5", "metadata": {}, "outputs": [], "source": [ "# time stepping\n", "time_opts = Time(dt=0.2, Tend=0.2)\n", "\n", "# geometry\n", "a1 = 1e-6 # to avoid polar singularity, we cut a small hole around the axis.\n", "a2 = 5.0\n", "Lz = 20.0\n", "domain = domains.HollowCylinder(a1=a1, a2=a2, Lz=Lz)\n", "\n", "# fluid equilibrium (can be used as part of initial conditions)\n", "equil = None\n", "\n", "# grid\n", "grid = grids.TensorProductGrid()\n", "\n", "# derham options\n", "derham_opts = DerhamOptions()" ] }, { "cell_type": "markdown", "id": "6", "metadata": {}, "source": [ "### Step 3: instantiate the first simulation\n", "\n", "Create the first simulation object with the baseline particle loading:" ] }, { "cell_type": "code", "execution_count": null, "id": "7", "metadata": {}, "outputs": [], "source": [ "sim = Simulation(model,\n", " env=env,\n", " time_opts=time_opts,\n", " domain=domain,\n", " equil=equil,\n", " grid=grid,\n", " derham_opts=derham_opts,)" ] }, { "cell_type": "markdown", "id": "8", "metadata": {}, "source": [ "Struphy provides a convenient way to clone a simulation and change only selected parameters.\n", "\n", "Use that pattern here to create a second simulation with modified model/environment options:" ] }, { "cell_type": "code", "execution_count": null, "id": "9", "metadata": {}, "outputs": [], "source": [ "sim_2 = sim.spawn_sister(model= model_2, env=env_2)" ] }, { "cell_type": "markdown", "id": "10", "metadata": {}, "source": [ "### Step 4: inspect the domain\n", "\n", "Before changing the loading options, verify the simulation domain visually:" ] }, { "cell_type": "code", "execution_count": null, "id": "11", "metadata": {}, "outputs": [], "source": [ "domain.show()" ] }, { "cell_type": "markdown", "id": "12", "metadata": {}, "source": [ "For the second simulation, set `spatial=\"disc\"` in the loading parameters.\n", "\n", "This samples particles uniformly on the cylinder cross section, which makes it easier to compare against logical-space loading:" ] }, { "cell_type": "code", "execution_count": null, "id": "13", "metadata": {}, "outputs": [], "source": [ "# species parameters\n", "loading_params = LoadingParameters(Np=1000)\n", "loading_params_2 = LoadingParameters(Np=1000, spatial=\"disc\")\n", "\n", "weights_params = WeightsParameters()\n", "boundary_params = BoundaryParameters()\n", "saving_params = SavingParameters(n_markers=1.0)\n", "\n", "model.kinetic_ions.set_markers(\n", " loading_params=loading_params, \n", " weights_params=weights_params, \n", " boundary_params=boundary_params,\n", " saving_params=saving_params\n", ")\n", "model_2.kinetic_ions.set_markers(\n", " loading_params=loading_params_2, \n", " weights_params=weights_params, \n", " boundary_params=boundary_params,\n", " saving_params=saving_params\n", ")" ] }, { "cell_type": "markdown", "id": "14", "metadata": {}, "source": [ "Use identical propagator options and initial conditions in both runs so that only the loading choice changes the result:" ] }, { "cell_type": "code", "execution_count": null, "id": "15", "metadata": {}, "outputs": [], "source": [ "# propagator options\n", "model.propagators.push_vxb.options = model.propagators.push_vxb.Options()\n", "model.propagators.push_eta.options = model.propagators.push_eta.Options()\n", "\n", "model_2.propagators.push_vxb.options = model_2.propagators.push_vxb.Options()\n", "model_2.propagators.push_eta.options = model_2.propagators.push_eta.Options()" ] }, { "cell_type": "code", "execution_count": null, "id": "16", "metadata": {}, "outputs": [], "source": [ "# initial conditions (background + perturbation)\n", "perturbation = None\n", "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", "\n", "model.kinetic_ions.var.add_background(background)\n", "model_2.kinetic_ions.var.add_background(background)" ] }, { "cell_type": "markdown", "id": "17", "metadata": {}, "source": [ "### Step 5: run simulation A\n", "\n", "Execute the first simulation:" ] }, { "cell_type": "code", "execution_count": null, "id": "18", "metadata": {}, "outputs": [], "source": [ "sim.run()" ] }, { "cell_type": "markdown", "id": "19", "metadata": {}, "source": [ "### Step 6: run simulation B\n", "\n", "Now execute the second simulation with disc-based spatial loading:" ] }, { "cell_type": "code", "execution_count": null, "id": "20", "metadata": {}, "outputs": [], "source": [ "sim_2.run()" ] }, { "cell_type": "markdown", "id": "21", "metadata": {}, "source": [ "### Step 7: compare initial particle distributions\n", "\n", "Post-process both runs, load the generated data, and plot initial particle positions on a cylinder cross section.\n", "\n", "This plot is the key check for understanding how loading in logical space differs from loading in physical space:" ] }, { "cell_type": "code", "execution_count": null, "id": "22", "metadata": {}, "outputs": [], "source": [ "sim.pproc()" ] }, { "cell_type": "code", "execution_count": null, "id": "23", "metadata": {}, "outputs": [], "source": [ "sim_2.pproc()" ] }, { "cell_type": "code", "execution_count": null, "id": "24", "metadata": {}, "outputs": [], "source": [ "sim.load_plotting_data()" ] }, { "cell_type": "code", "execution_count": null, "id": "25", "metadata": {}, "outputs": [], "source": [ "sim_2.load_plotting_data()" ] }, { "cell_type": "code", "execution_count": null, "id": "26", "metadata": {}, "outputs": [], "source": [ "from matplotlib import pyplot as plt\n", "\n", "fig = plt.figure(figsize=(10, 6))\n", "\n", "orbits = sim.orbits.kinetic_ions\n", "orbits_uni = sim_2.orbits.kinetic_ions\n", "\n", "plt.subplot(1, 2, 1)\n", "plt.scatter(orbits[0, :, 0], orbits[0, :, 1], s=2.0)\n", "circle1 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n", "ax = plt.gca()\n", "ax.add_patch(circle1)\n", "ax.set_aspect(\"equal\")\n", "plt.xlabel(\"x\")\n", "plt.ylabel(\"y\")\n", "plt.title(\"sim_1: draw uniform in logical space\")\n", "\n", "plt.subplot(1, 2, 2)\n", "plt.scatter(orbits_uni[0, :, 0], orbits_uni[0, :, 1], s=2.0)\n", "circle2 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n", "ax = plt.gca()\n", "ax.add_patch(circle2)\n", "ax.set_aspect(\"equal\")\n", "plt.xlabel(\"x\")\n", "plt.ylabel(\"y\")\n", "plt.title(\"sim_2: draw uniform on disc\")" ] }, { "cell_type": "markdown", "id": "27", "metadata": {}, "source": [ "### Optional: quasi-uniform marker loading (Sobol)\n", "\n", "For a more homogeneous particle representation with limited particle count, use Sobol-based loading.\n", "\n", "Set `loading=\"sobol_standard\"` or `loading=\"sobol_antithetic\"` in `LoadingParameters`:" ] }, { "cell_type": "code", "execution_count": null, "id": "28", "metadata": {}, "outputs": [], "source": [ "env_3 = EnvironmentOptions(sim_folder=\"sim_3\")\n", "model_3 = Vlasov()\n", "sim_3 = Simulation.spawn_sister(sim, env=env_3, model=model_3)\n", "\n", "# species parameters\n", "loading_params_3 = LoadingParameters(Np=1000, spatial=\"disc\", loading=\"sobol_standard\")\n", "\n", "model_3.kinetic_ions.set_markers(\n", " loading_params=loading_params_3, \n", " weights_params=weights_params, \n", " boundary_params=boundary_params,\n", " saving_params=saving_params,\n", ")\n", "\n", "# propagator options\n", "model_3.propagators.push_vxb.options = model_3.propagators.push_vxb.Options()\n", "model_3.propagators.push_eta.options = model_3.propagators.push_eta.Options()\n", "\n", "# initial conditions (background + perturbation)\n", "model_3.kinetic_ions.var.add_background(background)\n", "\n", "sim_3.run()\n", "sim_3.pproc()\n", "sim_3.load_plotting_data()\n", "\n", "fig = plt.figure(figsize=(15, 6))\n", "\n", "orbits_standard = sim_3.orbits.kinetic_ions\n", "\n", "plt.subplot(1, 3, 1)\n", "plt.scatter(orbits[0, :, 0], orbits[0, :, 1], s=2.0)\n", "circle1 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n", "ax = plt.gca()\n", "ax.add_patch(circle1)\n", "ax.set_aspect(\"equal\")\n", "plt.xlabel(\"x\")\n", "plt.ylabel(\"y\")\n", "plt.title(\"sim_1: draw uniform in logical space\")\n", "\n", "plt.subplot(1, 3, 2)\n", "plt.scatter(orbits_uni[0, :, 0], orbits_uni[0, :, 1], s=2.0)\n", "circle2 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n", "ax = plt.gca()\n", "ax.add_patch(circle2)\n", "ax.set_aspect(\"equal\")\n", "plt.xlabel(\"x\")\n", "plt.ylabel(\"y\")\n", "plt.title(\"sim_2: draw uniform on disc\")\n", "\n", "plt.subplot(1, 3, 3)\n", "plt.scatter(orbits_standard[0, :, 0], orbits_standard[0, :, 1], s=2.0)\n", "circle3 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n", "ax = plt.gca()\n", "ax.add_patch(circle3)\n", "ax.set_aspect(\"equal\")\n", "plt.xlabel(\"x\")\n", "plt.ylabel(\"y\")\n", "plt.title(\"sim_3: draw using sobol_standard\")" ] }, { "cell_type": "markdown", "id": "29", "metadata": {}, "source": [ "### Optional: antithetic Sobol loading\n", "\n", "Repeat the same setup with `loading=\"sobol_antithetic\"` if you want a variance-reduced counterpart to standard Sobol loading:" ] }, { "cell_type": "code", "execution_count": null, "id": "30", "metadata": {}, "outputs": [], "source": [ "env_3 = EnvironmentOptions(sim_folder=\"sim_3\")\n", "model_3 = Vlasov()\n", "sim_3 = Simulation.spawn_sister(sim, env=env_3, model=model_3)\n", "\n", "# species parameters\n", "loading_params_3 = LoadingParameters(Np=1000, spatial=\"disc\", loading=\"sobol_standard\")\n", "\n", "model_3.kinetic_ions.set_markers(\n", " loading_params=loading_params_3, \n", " weights_params=weights_params, \n", " boundary_params=boundary_params,\n", " saving_params=saving_params,\n", ")\n", "\n", "# propagator options\n", "model_3.propagators.push_vxb.options = model_3.propagators.push_vxb.Options()\n", "model_3.propagators.push_eta.options = model_3.propagators.push_eta.Options()\n", "\n", "# initial conditions (background + perturbation)\n", "model_3.kinetic_ions.var.add_background(background)\n", "\n", "sim_3.run()\n", "sim_3.pproc()\n", "sim_3.load_plotting_data()\n", "\n", "fig = plt.figure(figsize=(15, 6))\n", "\n", "orbits_standard = sim_3.orbits.kinetic_ions\n", "\n", "plt.subplot(1, 3, 1)\n", "plt.scatter(orbits[0, :, 0], orbits[0, :, 1], s=2.0)\n", "circle1 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n", "ax = plt.gca()\n", "ax.add_patch(circle1)\n", "ax.set_aspect(\"equal\")\n", "plt.xlabel(\"x\")\n", "plt.ylabel(\"y\")\n", "plt.title(\"sim_1: draw uniform in logical space\")\n", "\n", "plt.subplot(1, 3, 2)\n", "plt.scatter(orbits_uni[0, :, 0], orbits_uni[0, :, 1], s=2.0)\n", "circle2 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n", "ax = plt.gca()\n", "ax.add_patch(circle2)\n", "ax.set_aspect(\"equal\")\n", "plt.xlabel(\"x\")\n", "plt.ylabel(\"y\")\n", "plt.title(\"sim_2: draw uniform on disc\")\n", "\n", "plt.subplot(1, 3, 3)\n", "plt.scatter(orbits_standard[0, :, 0], orbits_standard[0, :, 1], s=2.0)\n", "circle3 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n", "ax = plt.gca()\n", "ax.add_patch(circle3)\n", "ax.set_aspect(\"equal\")\n", "plt.xlabel(\"x\")\n", "plt.ylabel(\"y\")\n", "plt.title(\"sim_3: draw using sobol_standard\")" ] }, { "cell_type": "markdown", "id": "31", "metadata": {}, "source": [ "## Part 2: Reflecting boundary conditions\n", "\n", "Next, create a new simulation with 50 time steps and 15 particles in the cylinder.\n", "\n", "In addition, enable reflecting boundary conditions in the radial direction. In Struphy this corresponds to logical direction $\\eta_1$:" ] }, { "cell_type": "code", "execution_count": null, "id": "32", "metadata": {}, "outputs": [], "source": [ "time_opts = Time(dt=0.2, Tend=10.0)\n", "loading_params = LoadingParameters(Np=15, spatial=\"disc\")\n", "boundary_params = BoundaryParameters(bc=(\"reflect\", \"periodic\", \"periodic\"))" ] }, { "cell_type": "code", "execution_count": null, "id": "33", "metadata": {}, "outputs": [], "source": [ "# light-weight model instance\n", "model = Vlasov()\n", "\n", "sim = Simulation(model,\n", " env=env,\n", " time_opts=time_opts,\n", " domain=domain,\n", " equil=equil,\n", " grid=grid,\n", " derham_opts=derham_opts,)" ] }, { "cell_type": "code", "execution_count": null, "id": "34", "metadata": {}, "outputs": [], "source": [ "model.kinetic_ions.set_markers(\n", " loading_params=loading_params, \n", " weights_params=weights_params, \n", " boundary_params=boundary_params,\n", " saving_params=saving_params,\n", ")" ] }, { "cell_type": "markdown", "id": "35", "metadata": {}, "source": [ "Set propagator options and initial conditions as in the previous section:" ] }, { "cell_type": "code", "execution_count": null, "id": "36", "metadata": {}, "outputs": [], "source": [ "# propagator options\n", "model.propagators.push_vxb.options = model.propagators.push_vxb.Options()\n", "model.propagators.push_eta.options = model.propagators.push_eta.Options()\n", "\n", "# initial conditions (background + perturbation)\n", "perturbation = None\n", "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", "\n", "model.kinetic_ions.var.add_background(background)" ] }, { "cell_type": "markdown", "id": "37", "metadata": {}, "source": [ "Run the simulation, post-process the output, and plot the resulting trajectories to verify reflections at the radial boundary:" ] }, { "cell_type": "code", "execution_count": null, "id": "38", "metadata": {}, "outputs": [], "source": [ "sim.run()" ] }, { "cell_type": "code", "execution_count": null, "id": "39", "metadata": {}, "outputs": [], "source": [ "sim.pproc()" ] }, { "cell_type": "code", "execution_count": null, "id": "40", "metadata": {}, "outputs": [], "source": [ "sim.load_plotting_data()" ] }, { "cell_type": "markdown", "id": "41", "metadata": {}, "source": [ "Under `sim.orbits[]`, Struphy stores orbit data in a 3D NumPy array:\n", "\n", "- axis 0: time step,\n", "- axis 1: particle index,\n", "- axis 2: particle attributes.\n", "\n", "The first three attributes are particle positions, followed by velocities and then weights (initial and time-dependent)." ] }, { "cell_type": "code", "execution_count": null, "id": "42", "metadata": {}, "outputs": [], "source": [ "orbits = sim.orbits.kinetic_ions\n", "\n", "Nt = orbits.shape[0]\n", "Np = orbits.shape[1]\n", "Nattr = orbits.shape[2]" ] }, { "cell_type": "code", "execution_count": null, "id": "43", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "\n", "fig = plt.figure()\n", "ax = fig.gca()\n", "\n", "colors = [\"tab:blue\", \"tab:orange\", \"tab:green\", \"tab:red\"]\n", "\n", "# create alpha for color scaling\n", "Tend = time_opts.Tend\n", "alpha = np.linspace(1.0, 0.0, Nt + 1)\n", "\n", "# loop through particles, plot all time steps\n", "for i in range(Np):\n", " ax.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], alpha=alpha)\n", "\n", "circle1 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n", "\n", "ax.add_patch(circle1)\n", "ax.set_aspect(\"equal\")\n", "ax.set_xlabel(\"x\")\n", "ax.set_ylabel(\"y\")\n", "ax.set_title(f\"{Nt - 1} time steps (full color at t=0)\");" ] }, { "cell_type": "markdown", "id": "44", "metadata": {}, "source": [ "## Part 3: Particles in a cylinder with a magnetic field\n", "\n", "Now extend the setup by adding a static magnetic background field.\n", "\n", "You do this by assigning an object from the `equils` module:" ] }, { "cell_type": "code", "execution_count": null, "id": "45", "metadata": {}, "outputs": [], "source": [ "B0x = 0.0\n", "B0y = 0.0\n", "B0z = 1.0\n", "equil = equils.HomogenSlab(B0x=B0x, B0y=B0y, B0z=B0z)" ] }, { "cell_type": "markdown", "id": "46", "metadata": {}, "source": [ "To evaluate the equilibrium efficiently in particle kernels, project it onto the spline basis. This requires creating a De Rham complex:" ] }, { "cell_type": "code", "execution_count": null, "id": "47", "metadata": {}, "outputs": [], "source": [ "bcs = ((\"free\", \"free\"), None, None)\n", "derham_opts = DerhamOptions(bcs=bcs)" ] }, { "cell_type": "markdown", "id": "48", "metadata": {}, "source": [ "Create a lightweight model instance and configure species options.\n", "\n", "Here we choose to `remove` particles that hit the boundary in radial direction $\\eta_1$:" ] }, { "cell_type": "code", "execution_count": null, "id": "49", "metadata": {}, "outputs": [], "source": [ "# light-weight model instance\n", "model = Vlasov()\n", "\n", "sim_withB = sim.spawn_sister(model=model,\n", " equil=equil,\n", " derham_opts=derham_opts,)" ] }, { "cell_type": "code", "execution_count": null, "id": "50", "metadata": {}, "outputs": [], "source": [ "loading_params = LoadingParameters(Np=20)\n", "boundary_params = BoundaryParameters(bc=(\"remove\", \"periodic\", \"periodic\"))\n", "\n", "model.kinetic_ions.set_markers(\n", " loading_params=loading_params, \n", " weights_params=weights_params, \n", " boundary_params=boundary_params,\n", " saving_params=saving_params,\n", ")\n", "\n", "# propagator options\n", "model.propagators.push_vxb.options = model.propagators.push_vxb.Options()\n", "model.propagators.push_eta.options = model.propagators.push_eta.Options()\n", "\n", "# initial conditions (background + perturbation)\n", "perturbation = None\n", "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", "\n", "model.kinetic_ions.var.add_background(background)" ] }, { "cell_type": "markdown", "id": "51", "metadata": {}, "source": [ "Run the case, post-process results, load the data, and plot trajectories to see how the magnetic field modifies particle motion:" ] }, { "cell_type": "code", "execution_count": null, "id": "52", "metadata": {}, "outputs": [], "source": [ "sim_withB.run()" ] }, { "cell_type": "code", "execution_count": null, "id": "53", "metadata": {}, "outputs": [], "source": [ "sim_withB.pproc()" ] }, { "cell_type": "code", "execution_count": null, "id": "54", "metadata": {}, "outputs": [], "source": [ "sim_withB.load_plotting_data()" ] }, { "cell_type": "code", "execution_count": null, "id": "55", "metadata": {}, "outputs": [], "source": [ "orbits = sim_withB.orbits.kinetic_ions\n", "\n", "Nt = orbits.shape[0]\n", "Np = orbits.shape[1]" ] }, { "cell_type": "code", "execution_count": null, "id": "56", "metadata": {}, "outputs": [], "source": [ "fig = plt.figure()\n", "ax = fig.gca()\n", "\n", "colors = [\"tab:blue\", \"tab:orange\", \"tab:green\", \"tab:red\"]\n", "\n", "# create alpha for color scaling\n", "Tend = time_opts.Tend\n", "alpha = np.linspace(1.0, 0.0, Nt + 1)\n", "\n", "# loop through particles, plot all time steps\n", "for i in range(Np):\n", " ax.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], alpha=alpha)\n", "\n", "circle1 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n", "\n", "ax.add_patch(circle1)\n", "ax.set_aspect(\"equal\")\n", "ax.set_xlabel(\"x\")\n", "ax.set_ylabel(\"y\")\n", "ax.set_title(f\"{int(Nt - 1)} time steps (full color at t=0)\");" ] }, { "cell_type": "markdown", "id": "57", "metadata": {}, "source": [ "## Part 4: Particles in a tokamak equilibrium\n", "\n", "As a more realistic example, use an ASDEX-Upgrade equilibrium loaded from an EQDSK file.\n", "\n", "Instantiate an `EQDSKequilibrium` with mostly default settings and adjust the density:" ] }, { "cell_type": "code", "execution_count": null, "id": "58", "metadata": {}, "outputs": [], "source": [ "n1 = 0.0\n", "n2 = 0.0\n", "na = 1.0\n", "equil = equils.EQDSKequilibrium(n1=n1, n2=n2, na=na)" ] }, { "cell_type": "markdown", "id": "59", "metadata": {}, "source": [ "`EQDSKequilibrium` inherits from `AxisymmMHDequilibrium` and `CartesianMHDequilibrium`, so in principle you could choose different mappings.\n", "\n", "For a domain that matches the equilibrium boundary, use the `Tokamak` mapping:" ] }, { "cell_type": "code", "execution_count": null, "id": "60", "metadata": {}, "outputs": [], "source": [ "num_elements = (28, 72)\n", "degree = (3, 3)\n", "psi_power = 0.6\n", "psi_shifts = (1e-6, 1.0)\n", "domain = domains.Tokamak(equilibrium=equil, num_elements=num_elements, degree=degree, psi_power=psi_power, psi_shifts=psi_shifts)" ] }, { "cell_type": "markdown", "id": "61", "metadata": {}, "source": [ "The `Tokamak` domain is a `PoloidalSplineTorus`. The coordinate relation between Cartesian $(x,y,z)$ and Tokamak $(R,Z,\\phi)$ variables is\n", "\n", "$$\n", " \\begin{aligned}\n", " x &= R \\cos(\\phi)\\,,\n", " \\\\\n", " y &= -R \\sin(\\phi)\\,,\n", " \\\\\n", " z &= Z\\,,\n", " \\end{aligned}\n", "$$\n", "\n", "where $(R,Z)$ spans the poloidal plane.\n", "\n", "Tokamak coordinates are connected to torus coordinates $(r,\\theta,\\phi)$ through a polar mapping in the poloidal plane:\n", "\n", "$$\n", " \\begin{aligned}\n", " R &= R_0 + r \\cos(\\theta)\\,,\n", " \\\\\n", " Z &= r \\sin(\\theta)\\,,\n", " \\\\\n", " \\phi &= \\phi\\,,\n", " \\end{aligned}\n", "$$\n", "\n", "and torus coordinates map to Struphy logical coordinates $\\boldsymbol\\eta=(\\eta_1,\\eta_2,\\eta_3)\\in[0,1]^3$ as\n", "\n", "$$\n", " \\begin{aligned}\n", " r &= a_1 + (a_2-a_1)\\eta_1\\,,\n", " \\\\\n", " \\theta &= 2\\pi\\eta_2\\,,\n", " \\\\\n", " \\phi &= 2\\pi\\eta_3\\,,\n", " \\end{aligned}\n", "$$\n", "\n", "with $a_2>a_1\\geq 0$ the radial bounds.\n", "\n", "You can inspect this relation in mappings such as `HollowTorus` (more advanced angle parametrizations $\\theta(\\eta_1,\\eta_2)$ also exist).\n", "\n", "As a quick sanity check, plot the magnetic-field magnitude:\n", "\n", "1. in the poloidal plane at $\\phi=0$,\n", "2. in the top view at $z=0$." ] }, { "cell_type": "code", "execution_count": null, "id": "62", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "\n", "# logical grid on the unit cube\n", "e1 = np.linspace(0.0, 1.0, 101)\n", "e2 = np.linspace(0.0, 1.0, 101)\n", "e3 = np.linspace(0.0, 1.0, 101)\n", "\n", "# move away from the singular point r = 0\n", "e1[0] += 1e-5" ] }, { "cell_type": "code", "execution_count": null, "id": "63", "metadata": {}, "outputs": [], "source": [ "# logical coordinates of the poloidal plane at phi = 0\n", "eta_poloidal = (e1, e2, 0.0)\n", "# logical coordinates of the top view at theta = 0\n", "eta_topview_1 = (e1, 0.0, e3)\n", "# logical coordinates of the top view at theta = pi\n", "eta_topview_2 = (e1, 0.5, e3)" ] }, { "cell_type": "code", "execution_count": null, "id": "64", "metadata": {}, "outputs": [], "source": [ "# Cartesian coordinates (squeezed)\n", "x_pol, y_pol, z_pol = domain(*eta_poloidal, squeeze_out=True)\n", "x_top1, y_top1, z_top1 = domain(*eta_topview_1, squeeze_out=True)\n", "x_top2, y_top2, z_top2 = domain(*eta_topview_2, squeeze_out=True)\n", "\n", "print(f\"{x_pol.shape = }\")\n", "print(f\"{x_top1.shape = }\")\n", "print(f\"{x_top2.shape = }\")" ] }, { "cell_type": "code", "execution_count": null, "id": "65", "metadata": {}, "outputs": [], "source": [ "# generate two axes\n", "fig, axs = plt.subplots(2, 1, figsize=(8, 16))\n", "ax = axs[0]\n", "ax_top = axs[1]\n", "\n", "# min/max of field strength\n", "equil.domain = domain\n", "Bmax = np.max(equil.absB0(*eta_topview_2, squeeze_out=True))\n", "Bmin = np.min(equil.absB0(*eta_topview_1, squeeze_out=True))\n", "levels = np.linspace(Bmin, Bmax, 51)\n", "\n", "# absolute magnetic field at phi = 0\n", "im = ax.contourf(x_pol, z_pol, equil.absB0(*eta_poloidal, squeeze_out=True), levels=levels)\n", "\n", "# absolute magnetic field at Z = 0\n", "im_top = ax_top.contourf(x_top1, y_top1, equil.absB0(*eta_topview_1, squeeze_out=True), levels=levels)\n", "ax_top.contourf(x_top2, y_top2, equil.absB0(*eta_topview_2, squeeze_out=True), levels=levels)\n", "\n", "# last closed flux surface, poloidal\n", "ax.plot(x_pol[-1], z_pol[-1], color=\"k\")\n", "\n", "# last closed flux surface, toroidal\n", "ax_top.plot(x_top1[-1], y_top1[-1], color=\"k\")\n", "ax_top.plot(x_top2[-1], y_top2[-1], color=\"k\")\n", "\n", "# limiter, poloidal\n", "ax.plot(equil.limiter_pts_R, equil.limiter_pts_Z, \"tab:orange\")\n", "ax.axis(\"equal\")\n", "ax.set_xlabel(\"R\")\n", "ax.set_ylabel(\"Z\")\n", "ax.set_title(\"abs(B) at $\\phi=0$\")\n", "fig.colorbar(im)\n", "# limiter, toroidal\n", "limiter_Rmax = np.max(equil.limiter_pts_R)\n", "limiter_Rmin = np.min(equil.limiter_pts_R)\n", "\n", "thetas = 2 * np.pi * e2\n", "limiter_x_max = limiter_Rmax * np.cos(thetas)\n", "limiter_y_max = -limiter_Rmax * np.sin(thetas)\n", "limiter_x_min = limiter_Rmin * np.cos(thetas)\n", "limiter_y_min = -limiter_Rmin * np.sin(thetas)\n", "\n", "ax_top.plot(limiter_x_max, limiter_y_max, \"tab:orange\")\n", "ax_top.plot(limiter_x_min, limiter_y_min, \"tab:orange\")\n", "ax_top.axis(\"equal\")\n", "ax_top.set_xlabel(\"x\")\n", "ax_top.set_ylabel(\"y\")\n", "ax_top.set_title(\"abs(B) at $Z=0$\")\n", "fig.colorbar(im_top);" ] }, { "cell_type": "markdown", "id": "66", "metadata": {}, "source": [ "As before, build a De Rham complex so the equilibrium can be projected onto the spline basis:" ] }, { "cell_type": "code", "execution_count": null, "id": "67", "metadata": {}, "outputs": [], "source": [ "num_elements = (32, 72, 1)\n", "grid = grids.TensorProductGrid(num_elements=num_elements)\n", "\n", "degree = (3, 3, 1)\n", "bcs = ((\"free\", \"free\"), None, None)\n", "derham_opts = DerhamOptions(degree=degree, bcs=bcs)" ] }, { "cell_type": "markdown", "id": "68", "metadata": {}, "source": [ "For this example we run 15000 time steps using a second-order splitting scheme:" ] }, { "cell_type": "code", "execution_count": null, "id": "69", "metadata": {}, "outputs": [], "source": [ "time_opts = Time(dt=0.2, Tend=3000, split_algo=\"Strang\")" ] }, { "cell_type": "markdown", "id": "70", "metadata": {}, "source": [ "Set up a simulation with four hand-picked particle initial conditions to study representative orbit types in this equilibrium:" ] }, { "cell_type": "code", "execution_count": null, "id": "71", "metadata": {}, "outputs": [], "source": [ "# light-weight model instance\n", "model = Vlasov()\n", "\n", "sim_asdex = Simulation(model,\n", " env=env,\n", " time_opts=time_opts,\n", " domain=domain,\n", " equil=equil,\n", " grid=grid,\n", " derham_opts=derham_opts,)" ] }, { "cell_type": "code", "execution_count": null, "id": "72", "metadata": {}, "outputs": [], "source": [ "# initial particle positions in phase space\n", "initial = (\n", " (0.501, 0.001, 0.001, 0.0, 0.0450, -0.04), # co-passing particle\n", " (0.511, 0.001, 0.001, 0.0, -0.0450, -0.04), # counter passing particle\n", " (0.521, 0.001, 0.001, 0.0, 0.0105, -0.04), # co-trapped particle\n", " (0.531, 0.001, 0.001, 0.0, -0.0155, -0.04),\n", ")\n", "\n", "loading_params = LoadingParameters(Np=4, seed=1608, specific_markers=initial)\n", "boundary_params = BoundaryParameters(bc=(\"remove\", \"periodic\", \"periodic\"))\n", "\n", "model.kinetic_ions.set_markers(\n", " loading_params=loading_params, \n", " weights_params=weights_params,\n", " boundary_params=boundary_params, \n", " saving_params=saving_params,\n", " bufsize=2.0\n", ")\n", "\n", "# propagator options\n", "model.propagators.push_vxb.options = model.propagators.push_vxb.Options()\n", "model.propagators.push_eta.options = model.propagators.push_eta.Options()\n", "\n", "# initial conditions (background + perturbation)\n", "perturbation = None\n", "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", "\n", "model.kinetic_ions.var.add_background(background)" ] }, { "cell_type": "code", "execution_count": null, "id": "73", "metadata": {}, "outputs": [], "source": [ "sim_asdex.run()" ] }, { "cell_type": "code", "execution_count": null, "id": "74", "metadata": {}, "outputs": [], "source": [ "sim_asdex.pproc()" ] }, { "cell_type": "code", "execution_count": null, "id": "75", "metadata": {}, "outputs": [], "source": [ "sim_asdex.load_plotting_data()" ] }, { "cell_type": "code", "execution_count": null, "id": "76", "metadata": {}, "outputs": [], "source": [ "orbits = sim_asdex.orbits.kinetic_ions\n", "\n", "Nt = orbits.shape[0]\n", "Np = orbits.shape[1]" ] }, { "cell_type": "code", "execution_count": null, "id": "77", "metadata": {}, "outputs": [], "source": [ "import math\n", "\n", "colors = [\"tab:blue\", \"tab:orange\", \"tab:green\", \"tab:red\"]\n", "\n", "dt = time_opts.dt\n", "Tend = time_opts.Tend\n", "\n", "for i in range(Np):\n", " r = np.sqrt(orbits[:, i, 0] ** 2 + orbits[:, i, 1] ** 2)\n", " # poloidal\n", " ax.scatter(r, orbits[:, i, 2], c=colors[i % 4], s=1)\n", " # top view\n", " ax_top.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], s=1)\n", "\n", "ax.set_title(f\"{math.ceil(Tend / dt)} time steps\")\n", "ax_top.set_title(f\"{math.ceil(Tend / dt)} time steps\")\n", "fig" ] }, { "cell_type": "markdown", "id": "78", "metadata": {}, "source": [ "## Part 5: Guiding centers in a tokamak equilibrium\n", "\n", "Finally, repeat the experiment with the guiding-center model and compare the trajectory behavior to full particle tracing:" ] }, { "cell_type": "code", "execution_count": null, "id": "79", "metadata": {}, "outputs": [], "source": [ "from struphy.models import GuidingCenter\n", "\n", "# light-weight model instance\n", "model = GuidingCenter()" ] }, { "cell_type": "code", "execution_count": null, "id": "80", "metadata": {}, "outputs": [], "source": [ "time_opts = Time(dt=0.1, Tend=100, split_algo=\"Strang\")\n", "\n", "sim_gc = sim_asdex.spawn_sister(model=model, time_opts=time_opts)" ] }, { "cell_type": "code", "execution_count": null, "id": "81", "metadata": {}, "outputs": [], "source": [ "# initial phase space coordinates\n", "initial = (\n", " (0.501, 0.001, 0.001, -1.935, 1.72), # co-passing particle\n", " (0.501, 0.001, 0.001, 1.935, 1.72), # couner-passing particle\n", " (0.501, 0.001, 0.001, -0.6665, 1.72), # co-trapped particle\n", " (0.501, 0.001, 0.001, 0.4515, 1.72),\n", ") # counter-trapped particl\n", "\n", "loading_params = LoadingParameters(Np=4, seed=1608, specific_markers=initial)\n", "boundary_params = BoundaryParameters(bc=(\"remove\", \"periodic\", \"periodic\"))\n", "\n", "model.kinetic_ions.set_markers(\n", " loading_params=loading_params, \n", " weights_params=weights_params, \n", " boundary_params=boundary_params,\n", " saving_params=saving_params,\n", " bufsize=2.0,\n", ")\n", "\n", "# propagator options\n", "model.propagators.push_bxe.options = model.propagators.push_bxe.Options(tol=1e-5)\n", "model.propagators.push_parallel.options = model.propagators.push_parallel.Options(tol=1e-5)\n", "\n", "# initial conditions (background + perturbation)\n", "perturbation = None\n", "background = maxwellians.GyroMaxwellian2D(n=(1.0, perturbation), equil=equil)\n", "\n", "model.kinetic_ions.var.add_background(background)" ] }, { "cell_type": "code", "execution_count": null, "id": "82", "metadata": {}, "outputs": [], "source": [ "# generate two axes\n", "fig, axs = plt.subplots(2, 1, figsize=(8, 16))\n", "ax = axs[0]\n", "ax_top = axs[1]\n", "\n", "# min/max of field strength\n", "equil.domain = domain\n", "Bmax = np.max(equil.absB0(*eta_topview_2, squeeze_out=True))\n", "Bmin = np.min(equil.absB0(*eta_topview_1, squeeze_out=True))\n", "levels = np.linspace(Bmin, Bmax, 51)\n", "\n", "# absolute magnetic field at phi = 0\n", "im = ax.contourf(x_pol, z_pol, equil.absB0(*eta_poloidal, squeeze_out=True), levels=levels)\n", "\n", "# absolute magnetic field at Z = 0\n", "im_top = ax_top.contourf(x_top1, y_top1, equil.absB0(*eta_topview_1, squeeze_out=True), levels=levels)\n", "ax_top.contourf(x_top2, y_top2, equil.absB0(*eta_topview_2, squeeze_out=True), levels=levels)\n", "\n", "# last closed flux surface, poloidal\n", "ax.plot(x_pol[-1], z_pol[-1], color=\"k\")\n", "\n", "# last closed flux surface, toroidal\n", "ax_top.plot(x_top1[-1], y_top1[-1], color=\"k\")\n", "ax_top.plot(x_top2[-1], y_top2[-1], color=\"k\")\n", "\n", "# limiter, poloidal\n", "ax.plot(equil.limiter_pts_R, equil.limiter_pts_Z, \"tab:orange\")\n", "ax.axis(\"equal\")\n", "ax.set_xlabel(\"R\")\n", "ax.set_ylabel(\"Z\")\n", "ax.set_title(\"abs(B) at $\\phi=0$\")\n", "fig.colorbar(im)\n", "# limiter, toroidal\n", "limiter_Rmax = np.max(equil.limiter_pts_R)\n", "limiter_Rmin = np.min(equil.limiter_pts_R)\n", "\n", "thetas = 2 * np.pi * e2\n", "limiter_x_max = limiter_Rmax * np.cos(thetas)\n", "limiter_y_max = -limiter_Rmax * np.sin(thetas)\n", "limiter_x_min = limiter_Rmin * np.cos(thetas)\n", "limiter_y_min = -limiter_Rmin * np.sin(thetas)\n", "\n", "ax_top.plot(limiter_x_max, limiter_y_max, \"tab:orange\")\n", "ax_top.plot(limiter_x_min, limiter_y_min, \"tab:orange\")\n", "ax_top.axis(\"equal\")\n", "ax_top.set_xlabel(\"x\")\n", "ax_top.set_ylabel(\"y\")\n", "ax_top.set_title(\"abs(B) at $Z=0$\")\n", "fig.colorbar(im_top);" ] }, { "cell_type": "code", "execution_count": null, "id": "83", "metadata": {}, "outputs": [], "source": [ "sim_gc.run()" ] }, { "cell_type": "code", "execution_count": null, "id": "84", "metadata": {}, "outputs": [], "source": [ "sim_gc.pproc()" ] }, { "cell_type": "code", "execution_count": null, "id": "85", "metadata": {}, "outputs": [], "source": [ "sim_gc.load_plotting_data()" ] }, { "cell_type": "code", "execution_count": null, "id": "86", "metadata": {}, "outputs": [], "source": [ "orbits = sim_gc.orbits.kinetic_ions\n", "\n", "Nt = orbits.shape[0]\n", "Np = orbits.shape[1]" ] }, { "cell_type": "code", "execution_count": null, "id": "87", "metadata": {}, "outputs": [], "source": [ "import math\n", "\n", "colors = [\"tab:blue\", \"tab:orange\", \"tab:green\", \"tab:red\"]\n", "\n", "dt = time_opts.dt\n", "Tend = time_opts.Tend\n", "\n", "for i in range(Np):\n", " r = np.sqrt(orbits[:, i, 0] ** 2 + orbits[:, i, 1] ** 2)\n", " # poloidal\n", " ax.scatter(r, orbits[:, i, 2], c=colors[i % 4], s=1)\n", " # top view\n", " ax_top.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], s=1)\n", "\n", "ax.set_title(f\"{math.ceil(Tend / dt)} time steps\")\n", "ax_top.set_title(f\"{math.ceil(Tend / dt)} time steps\")\n", "fig" ] } ], "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": 5 }