From c4a09cdac72c55ddbf4d0cfba2a65b2c6e7ae08f Mon Sep 17 00:00:00 2001 From: issra-raven Date: Fri, 30 Aug 2024 12:49:46 +0200 Subject: [PATCH] initial commit --- docs/notebooks/tutorials/single-stage.ipynb | 475 ++++++++++++++++++++ requirements.txt | 2 +- 2 files changed, 476 insertions(+), 1 deletion(-) create mode 100644 docs/notebooks/tutorials/single-stage.ipynb diff --git a/docs/notebooks/tutorials/single-stage.ipynb b/docs/notebooks/tutorials/single-stage.ipynb new file mode 100644 index 0000000000..61f06859b7 --- /dev/null +++ b/docs/notebooks/tutorials/single-stage.ipynb @@ -0,0 +1,475 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "ename": "ModuleNotFoundError", + "evalue": "No module named 'desc'", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mModuleNotFoundError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[2], line 5\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;66;03m# uncomment these two lines if using a gpu for a speedup\u001b[39;00m\n\u001b[1;32m 2\u001b[0m \u001b[38;5;66;03m#! nvidia-smi\u001b[39;00m\n\u001b[1;32m 3\u001b[0m \u001b[38;5;66;03m#os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\u001b[39;00m\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mdesc\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m set_device\n\u001b[1;32m 6\u001b[0m set_device(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mgpu\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", + "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'desc'" + ] + } + ], + "source": [ + "# uncomment these two lines if using a gpu for a speedup\n", + "#! nvidia-smi\n", + "#os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n", + "\n", + "from desc import set_device\n", + "set_device(\"gpu\") " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from desc import equilibrium\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# imports\n", + "\n", + "import numpy as np\n", + "from desc.coils import CoilSet, FourierPlanarCoil\n", + "import desc.examples\n", + "from desc.equilibrium import Equilibrium\n", + "from desc.plotting import plot_surfaces, plot_2d, plot_3d, plot_coils\n", + "from desc.grid import LinearGrid, ConcentricGrid\n", + "from desc.coils import MixedCoilSet\n", + "from desc.objectives import (\n", + " ObjectiveFunction,\n", + " #coil\n", + " CoilCurvature,\n", + " CoilLength,\n", + " CoilTorsion,\n", + " CoilSetMinDistance,\n", + " PlasmaCoilSetMinDistance,\n", + " QuadraticFlux,\n", + " ToroidalFlux,\n", + " FixCoilCurrent,\n", + " FixParameters,\n", + " # plasma\n", + " FixBoundaryR,\n", + " FixBoundaryZ,\n", + " FixPressure,\n", + " FixCurrent,\n", + " FixIota,\n", + " FixPsi,\n", + " AspectRatio,\n", + " ForceBalance,\n", + " QuasisymmetryBoozer,\n", + " QuasisymmetryTwoTerm,\n", + " QuasisymmetryTripleProduct,\n", + " VacuumBoundaryError,\n", + ")\n", + "from desc.continuation import solve_continuation_automatic\n", + "from desc.optimize import Optimizer\n", + "from desc.magnetic_fields import field_line_integrate\n", + "from desc.singularities import compute_B_plasma\n", + "import time\n", + "import plotly.express as px\n", + "import plotly.io as pio\n", + "\n", + "from desc.geometry import FourierRZToroidalSurface\n", + "\n", + "\n", + "# This ensures Plotly output works in multiple places:\n", + "# plotly_mimetype: VS Code notebook UI\n", + "# notebook: \"Jupyter: Export to HTML\" command in VS Code\n", + "# See https://plotly.com/python/renderers/#multiple-renderers\n", + "pio.renderers.default = \"plotly_mimetype+notebook\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# plotting\n", + "from desc.plotting import (\n", + " plot_grid,\n", + " plot_boozer_modes,\n", + " plot_boozer_surface,\n", + " plot_qs_error,\n", + " plot_boundaries,\n", + " plot_boundary,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "##################################\n", + "# Generating initial equilibrium #\n", + "##################################\n", + "\n", + "eq_init = desc.io.load(\"../DESC/docs/notebooks/tutorials/qs_initial_guess.h5\")\n", + "\n", + "# # create initial surface. Aspect ratio ~ 8, circular cross section with slight\n", + "# # axis torsion to make it nonplanar\n", + "# surf = FourierRZToroidalSurface(\n", + "# R_lmn=[1, 0.125, 0.1],\n", + "# Z_lmn=[-0.125, -0.1],\n", + "# modes_R=[[0, 0], [1, 0], [0, 1]],\n", + "# modes_Z=[[-1, 0], [0, -1]],\n", + "# NFP=4,\n", + "# )\n", + "# # create initial equilibrium. Psi chosen to give B ~ 1 T. Could also give profiles here,\n", + "# # default is zero pressure and zero current\n", + "# eq_init = Equilibrium(M=4, N=4, Psi=0.04, surface=surf)\n", + "\n", + "\n", + "##############################\n", + "# Generating initial coilset #\n", + "##############################\n", + "\n", + "minor_radius = eq_init.compute(\"a\")[\"a\"]\n", + "offset = 0.3\n", + "num_coils = 3 # coils per half field period\n", + "\n", + "zeta = np.linspace(0, np.pi / eq_init.NFP, num_coils, endpoint=False) + np.pi / (\n", + " 2 * eq_init.NFP * num_coils\n", + ")\n", + "grid = LinearGrid(rho=[0.0], M=0, zeta=zeta, NFP=eq_init.NFP)\n", + "data = eq_init.axis.compute([\"x\", \"x_s\"], grid=grid, basis=\"rpz\")\n", + "\n", + "centers = data[\"x\"] # center coils on axis position\n", + "normals = data[\"x_s\"] # make normal to coil align with tangent along axis\n", + "\n", + "unique_coils = []\n", + "for k in range(num_coils):\n", + " coil = FourierPlanarCoil(\n", + " current=1e6,\n", + " center=centers[k, :],\n", + " normal=normals[k, :],\n", + " r_n=minor_radius + offset,\n", + " basis=\"rpz\", # we are giving the center and normal in cylindrical coordinates\n", + " ).to_FourierXYZ(\n", + " N=10\n", + " ) # fit with 10 fourier coefficients per coil\n", + " unique_coils.append(coil)\n", + "\n", + "# We package these coils together into a CoilSet, which has efficient methods for calculating\n", + "# the total field while accounting for field period and stellarator symmetry\n", + "# Note that `CoilSet` requires all the member coils to have the same parameterization and resolution.\n", + "# if we wanted to use coils of different types or resolutions, we can use a `MixedCoilSet` (see the next section below)\n", + "coilset = CoilSet(unique_coils, NFP=eq_init.NFP, sym=eq_init.sym)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fig = plot_3d(eq_init,\"|B|\")\n", + "plot_coils(coilset,fig=fig)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# #####################################################################################\n", + "# # visualize the initial coilset\n", + "\n", + "# # we use a smaller than usual plot grid to reduce memory of the notebook file\n", + "# plot_grid = LinearGrid(M=20, N=40, NFP=1, endpoint=True)\n", + "# fig = plot_3d(eq_init, \"|B|\", grid=plot_grid)\n", + "# fig = plot_coils(coilset, fig=fig)\n", + "# fig.show()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "grid_vol = ConcentricGrid(\n", + " L=eq_init.L_grid,\n", + " M=eq_init.M_grid,\n", + " N=eq_init.N_grid,\n", + " NFP=eq_init.NFP,\n", + " sym=eq_init.sym,\n", + ")\n", + "\n", + "grid_lcfs = LinearGrid(\n", + " M=eq_init.M_grid,\n", + " N=eq_init.N_grid,\n", + " NFP=eq_init.NFP,\n", + " sym=eq_init.sym,\n", + " rho=np.array(1.0),\n", + ")\n", + "\n", + "coil_grid = LinearGrid(N=300)\n", + "plasma_grid = LinearGrid(M=100, N=100, NFP=eq_init.NFP, sym=eq_init.sym)\n", + "\n", + "coil_indices_to_fix_current = [False for c in coilset]\n", + "coil_indices_to_fix_current[0] = True\n", + "constraints = (\n", + " ForceBalance(eq=eq_init),\n", + " FixCoilCurrent(coilset, indices=coil_indices_to_fix_current),\n", + " FixPressure(eq=eq_init), # fix pressure profile\n", + " #FixIota(eq=eq_init), # fix rotational transform profile\n", + " FixPsi(eq=eq_init), # fix total toroidal magnetic flux\n", + " )\n", + "\n", + "obj = ObjectiveFunction(\n", + " (\n", + "\n", + " # QuasisymmetryBoozer(\n", + " # eq = eq_init, \n", + " # #grid = grid_vol, \n", + " # helicity = (1, 1),\n", + " # weight = 100\n", + " # ),\n", + " QuasisymmetryTwoTerm(\n", + " eq = eq_init, \n", + " helicity=(1, eq_init.NFP), \n", + " grid = grid_lcfs,\n", + " weight = 1\n", + " ),\n", + " # QuasisymmetryTripleProduct(\n", + " # eq=eq_init,\n", + " # grid=grid_vol, \n", + " # weight=200\n", + " # ),\n", + " CoilSetMinDistance(\n", + " coilset,\n", + " # in normalized units, want coil-coil distance to be at least 10% of minor radius\n", + " bounds=(0.1, np.inf),\n", + " normalize_target=False, # we're giving bounds in normalized units\n", + " grid=coil_grid,\n", + " weight=100,\n", + " ),\n", + " PlasmaCoilSetMinDistance(\n", + " eq_init,\n", + " coilset,\n", + " # in normalized units, want plasma-coil distance to be at least 25% of minor radius\n", + " bounds=(0.25, np.inf),\n", + " normalize_target=False, # we're giving bounds in normalized units\n", + " plasma_grid=plasma_grid,\n", + " coil_grid=coil_grid,\n", + " eq_fixed=False, # Fix the equilibrium. For single stage optimization, this would be False\n", + " weight=100,\n", + " ),\n", + " CoilCurvature(\n", + " coilset,\n", + " # this uses signed curvature, depending on whether it curves towards\n", + " # or away from the centroid of the curve, with a circle having positive curvature.\n", + " # We give the bounds normalized units, curvature of approx 1 means circular,\n", + " # so we allow them to be a bit more strongly shaped\n", + " bounds=(-1, 2),\n", + " normalize_target=False, # we're giving bounds in normalized units\n", + " grid=coil_grid,\n", + " weight=75,\n", + " ),\n", + " CoilLength(\n", + " coilset,\n", + " bounds=(0, 2 * np.pi * (minor_radius + offset)),\n", + " normalize_target=True, # target length is in meters, not normalized\n", + " grid=coil_grid,\n", + " weight=100,\n", + " ),\n", + " )\n", + ")\n", + "\n", + "\n", + "# def compute_average_normalized_field(field, eq_init, vacuum=False):\n", + "# grid = LinearGrid(M=80, N=80, NFP=eq_init.NFP)\n", + "# Bn, surf_coords = field.compute_Bnormal(eq_init, eval_grid=grid)\n", + "# normalizing_field_vec = field.compute_magnetic_field(surf_coords)\n", + "# if not vacuum:\n", + "# # add plasma field to the normalizing field\n", + "# normalizing_field_vec += compute_B_plasma(eq_init, eval_grid=grid)\n", + "# normalizing_field = np.mean(np.linalg.norm(normalizing_field_vec, axis=1))\n", + "# return np.mean(np.abs(Bn)) / normalizing_field\n", + "\n", + "\n", + "# def plot_field_lines(field, eq_init):\n", + "# # for starting locations we'll pick positions on flux surfaces on the outboard midplane\n", + "# grid_trace = LinearGrid(rho=np.linspace(0, 1, 9))\n", + "# r0 = eq_init.compute(\"R\", grid=grid_trace)[\"R\"]\n", + "# z0 = eq_init.compute(\"Z\", grid=grid_trace)[\"Z\"]\n", + "# fig, ax = desc.plotting.plot_surfaces(eq_init)\n", + "# fig, ax = desc.plotting.poincare_plot(\n", + "# field,\n", + "# r0,\n", + "# z0,\n", + "# NFP=eq_init.NFP,\n", + "# ax=ax,\n", + "# color=\"k\",\n", + "# size=1,\n", + "# )\n", + "# return fig, ax" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "optimizer = Optimizer(\"proximal-lsq-exact\")\n", + "\n", + "\n", + "(eq_qs_T, optimized_coilset,), result_T = optimizer.optimize(\n", + " (eq_init, coilset), \n", + " objective=obj,\n", + " constraints=constraints,\n", + " ftol=1e-8, # stopping tolerance on the function value\n", + " xtol=1e-8, # stopping tolerance on the step size\n", + " gtol=1e-8, # stopping tolerance on the gradient\n", + " maxiter=1000, # maximum number of iterations\n", + " options={\n", + " \"perturb_options\": {\"order\": 2, \"verbose\": 0}, # use 2nd-order perturbations\n", + " \"solve_options\": {\n", + " \"ftol\": 5e-5,\n", + " \"xtol\": 5e-5,\n", + " \"gtol\": 1e-6,\n", + " \"verbose\": 0,\n", + " }, # for equilibrium subproblem\n", + " },\n", + " copy=True, # copy=True to make a copy of the optimized result and dont touch the orignals\n", + " verbose=3,\n", + ")\n" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "ename": "NameError", + "evalue": "name 'plot_boozer_surface' is not defined", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[1], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mplot_boozer_surface\u001b[49m(eq_init);\n\u001b[1;32m 2\u001b[0m plot_boozer_surface(eq_qs_T)\n", + "\u001b[0;31mNameError\u001b[0m: name 'plot_boozer_surface' is not defined" + ] + } + ], + "source": [ + "plot_boozer_surface(eq_init);\n", + "plot_boozer_surface(eq_qs_T); # |B| contours at rho=1 surface" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# compare f_T & f_C before (o) vs after (x) optimization\n", + "fig, ax = plot_qs_error(\n", + " eq_init, helicity=(1, eq_init.NFP), fB=False, legend=False, rho=10\n", + ")\n", + "plot_qs_error(\n", + " eq_qs_T, helicity=(1, eq_qs_T.NFP), fB=False, ax=ax, marker=[\"x\", \"x\"], rho=10\n", + ");" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# compare f_B before (o) vs after (x) optimization\n", + "fig, ax = plot_qs_error(\n", + " eq_init, helicity=(1, eq_init.NFP), fT=False, fC=False, legend=False, rho=10\n", + ")\n", + "plot_qs_error(\n", + " eq_qs_T, helicity=(1, eq_qs_T.NFP), fT=False, fC=False, ax=ax, marker=[\"x\"], rho=10\n", + ");" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fig=plot_3d(eq_qs_T.surface,\"B*n\",field=optimized_coilset, alpha = 0.5) # pass in surface so the Bplasma is not computed, as it is vacuum\n", + "plot_coils(optimized_coilset,fig=fig)\n", + "#plot_coils(coilset,fig=fig,color=\"red\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Anaconda 2023.03", + "language": "python", + "name": "anaconda_3_2023_03" + }, + "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.10.9" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/requirements.txt b/requirements.txt index a667a2a2db..b8d2fb5707 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,5 +10,5 @@ nvgpu plotly >= 5.16, < 6.0 psutil pylatexenc >= 2.0, < 3.0 -scipy >= 1.7.0, < 2.0.0 +scipy >= 1.7.0, <= 1.13.0 termcolor