|
| 1 | +#!/usr/bin/env python3 |
| 2 | + |
| 3 | +""" |
| 4 | +This example demonstrates how to use SIMSOPT to compute Poincare plots. |
| 5 | +
|
| 6 | +This example also illustrates how the coil shapes resulting from a |
| 7 | +simsopt stage-2 optimization can be loaded in to another script for |
| 8 | +analysis. The coil shape data used in this script, |
| 9 | +``inputs/biot_savart_opt.json``, can be re-generated by running the |
| 10 | +example 2_Intermediate/stage_two_optimization.py. |
| 11 | +""" |
| 12 | + |
| 13 | +import time |
| 14 | +import os |
| 15 | +import logging |
| 16 | +import sys |
| 17 | +from pathlib import Path |
| 18 | +import numpy as np |
| 19 | +try: |
| 20 | + from mpi4py import MPI |
| 21 | + comm = MPI.COMM_WORLD |
| 22 | +except ImportError: |
| 23 | + comm = None |
| 24 | + |
| 25 | +import simsopt |
| 26 | +from simsopt.field.biotsavart import BiotSavart |
| 27 | +from simsopt.field.magneticfieldclasses import InterpolatedField, UniformInterpolationRule |
| 28 | +from simsopt.geo.surfacexyztensorfourier import SurfaceRZFourier |
| 29 | +from simsopt.field.coil import coils_via_symmetries |
| 30 | +from simsopt.field.tracing import SurfaceClassifier, \ |
| 31 | + particles_to_vtk, compute_fieldlines, LevelsetStoppingCriterion, plot_poincare_data |
| 32 | +from simsopt.geo.curve import curves_to_vtk |
| 33 | + |
| 34 | +print("Running 1_Simple/tracing_fieldlines_QA.py") |
| 35 | +print("=========================================") |
| 36 | + |
| 37 | +logging.basicConfig() |
| 38 | +logger = logging.getLogger('simsopt.field.tracing') |
| 39 | +logger.setLevel(1) |
| 40 | + |
| 41 | +# check whether we're in CI, in that case we make the run a bit cheaper |
| 42 | +ci = "CI" in os.environ and os.environ['CI'].lower() in ['1', 'true'] |
| 43 | +nfieldlines = 3 if ci else 10 |
| 44 | +tmax_fl = 10000 if ci else 20000 |
| 45 | +degree = 2 if ci else 4 |
| 46 | + |
| 47 | +# Directory for output |
| 48 | +OUT_DIR = "./output/" |
| 49 | +os.makedirs(OUT_DIR, exist_ok=True) |
| 50 | + |
| 51 | +# Load in the boundary surface: |
| 52 | +TEST_DIR = (Path(__file__).parent / ".." / ".." / "tests" / "test_files").resolve() |
| 53 | +filename = TEST_DIR / 'input.LandremanPaul2021_QA' |
| 54 | +# Note that the range must be "full torus"! |
| 55 | +surf = SurfaceRZFourier.from_vmec_input(filename, nphi=200, ntheta=30, range="full torus") |
| 56 | +nfp = surf.nfp |
| 57 | + |
| 58 | +# Load in coils from stage_two_optimization.py: |
| 59 | +bs = simsopt.load_simsopt("inputs/biot_savart_opt.json") |
| 60 | + |
| 61 | +surf.to_vtk(OUT_DIR + 'surface') |
| 62 | +sc_fieldline = SurfaceClassifier(surf, h=0.03, p=2) |
| 63 | +sc_fieldline.to_vtk(OUT_DIR + 'levelset', h=0.02) |
| 64 | + |
| 65 | + |
| 66 | +def trace_fieldlines(bfield, label): |
| 67 | + t1 = time.time() |
| 68 | + # Set initial grid of points for field line tracing, going from |
| 69 | + # the magnetic axis to the surface. The actual plasma boundary is |
| 70 | + # at R=1.300425, but the outermost initial point is a bit inward |
| 71 | + # from that, R = 1.295, so the SurfaceClassifier does not think we |
| 72 | + # have exited the surface |
| 73 | + R0 = np.linspace(1.2125346, 1.295, nfieldlines) |
| 74 | + Z0 = np.zeros(nfieldlines) |
| 75 | + phis = [(i/4)*(2*np.pi/nfp) for i in range(4)] |
| 76 | + fieldlines_tys, fieldlines_phi_hits = compute_fieldlines( |
| 77 | + bfield, R0, Z0, tmax=tmax_fl, tol=1e-16, comm=comm, |
| 78 | + phis=phis, stopping_criteria=[LevelsetStoppingCriterion(sc_fieldline.dist)]) |
| 79 | + t2 = time.time() |
| 80 | + print(f"Time for fieldline tracing={t2-t1:.3f}s. Num steps={sum([len(l) for l in fieldlines_tys])//nfieldlines}", flush=True) |
| 81 | + if comm is None or comm.rank == 0: |
| 82 | + particles_to_vtk(fieldlines_tys, OUT_DIR + f'fieldlines_{label}') |
| 83 | + plot_poincare_data(fieldlines_phi_hits, phis, OUT_DIR + f'poincare_fieldline_{label}.png', dpi=150) |
| 84 | + |
| 85 | + |
| 86 | +# uncomment this to run tracing using the biot savart field (very slow!) |
| 87 | +# trace_fieldlines(bs, 'bs') |
| 88 | + |
| 89 | + |
| 90 | +# Bounds for the interpolated magnetic field chosen so that the surface is |
| 91 | +# entirely contained in it |
| 92 | +n = 20 |
| 93 | +rs = np.linalg.norm(surf.gamma()[:, :, 0:2], axis=2) |
| 94 | +zs = surf.gamma()[:, :, 2] |
| 95 | +rrange = (np.min(rs), np.max(rs), n) |
| 96 | +phirange = (0, 2*np.pi/nfp, n*2) |
| 97 | +# exploit stellarator symmetry and only consider positive z values: |
| 98 | +zrange = (0, np.max(zs), n//2) |
| 99 | + |
| 100 | + |
| 101 | +def skip(rs, phis, zs): |
| 102 | + # The RegularGrindInterpolant3D class allows us to specify a function that |
| 103 | + # is used in order to figure out which cells to be skipped. Internally, |
| 104 | + # the class will evaluate this function on the nodes of the regular mesh, |
| 105 | + # and if *all* of the eight corners are outside the domain, then the cell |
| 106 | + # is skipped. Since the surface may be curved in a way that for some |
| 107 | + # cells, all mesh nodes are outside the surface, but the surface still |
| 108 | + # intersects with a cell, we need to have a bit of buffer in the signed |
| 109 | + # distance (essentially blowing up the surface a bit), to avoid ignoring |
| 110 | + # cells that shouldn't be ignored |
| 111 | + rphiz = np.asarray([rs, phis, zs]).T.copy() |
| 112 | + dists = sc_fieldline.evaluate_rphiz(rphiz) |
| 113 | + skip = list((dists < -0.05).flatten()) |
| 114 | + print("Skip", sum(skip), "cells out of", len(skip), flush=True) |
| 115 | + return skip |
| 116 | + |
| 117 | + |
| 118 | +print('Initializing InterpolatedField') |
| 119 | +bsh = InterpolatedField( |
| 120 | + bs, degree, rrange, phirange, zrange, True, nfp=nfp, stellsym=True, skip=skip |
| 121 | +) |
| 122 | +print('Done initializing InterpolatedField.') |
| 123 | + |
| 124 | +bsh.set_points(surf.gamma().reshape((-1, 3))) |
| 125 | +bs.set_points(surf.gamma().reshape((-1, 3))) |
| 126 | +Bh = bsh.B() |
| 127 | +B = bs.B() |
| 128 | +print("Mean(|B|) on plasma surface =", np.mean(bs.AbsB())) |
| 129 | + |
| 130 | +print("|B-Bh| on surface:", np.sort(np.abs(B-Bh).flatten())) |
| 131 | + |
| 132 | +print('Beginning field line tracing') |
| 133 | +trace_fieldlines(bsh, 'bsh') |
| 134 | + |
| 135 | +print("End of 1_Simple/tracing_fieldlines_QA.py") |
| 136 | +print("========================================") |
0 commit comments