From 1f8d33be4cd1dc1dafe452e8d88c4e9ff13a5961 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Tue, 31 Mar 2026 16:12:32 -0400 Subject: [PATCH 1/8] add gitignore --- .gitignore | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..8e202ef --- /dev/null +++ b/.gitignore @@ -0,0 +1,32 @@ +__pycache__/ +build/ +data/ +dist/ +docs/_api +docs/_build +htmlcov/ +local/ +matlab/ +tests/benchmarks/.benchmarks/ +.coverage +.ipynb_checkpoints/ +.spyproject/ +*.egg-info +*.gif +*.h5 +*.ini +*.mat +*.out +*.output +*.pdf +*.pyc +*.swp +*.swo +.pymon + +# UV +uv.lock + +# Environments +.env +.venv \ No newline at end of file From 7c6710391fa108574f3689006977926298f8dd92 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Tue, 31 Mar 2026 16:12:55 -0400 Subject: [PATCH 2/8] make all plotting optional --- efit2desc/efit2desc.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/efit2desc/efit2desc.py b/efit2desc/efit2desc.py index 949815f..6d412e2 100644 --- a/efit2desc/efit2desc.py +++ b/efit2desc/efit2desc.py @@ -268,10 +268,11 @@ def convert_EFIT_to_DESC( efit_Psi = efit["AuxQuantities"]["PHI"] # get bdry - plt.figure() - for k in range(0, len(fluxsurf["flux"]))[::-10]: - plt.plot(fluxsurf["flux"][k]["R"], fluxsurf["flux"][k]["Z"]) - plt.axis("equal") + if plot: + plt.figure() + for k in range(0, len(fluxsurf["flux"]))[::-10]: + plt.plot(fluxsurf["flux"][k]["R"], fluxsurf["flux"][k]["Z"]) + plt.axis("equal") # choose the LCFS as the bdry # TODO: use spectral condensation (ideally when implemented in DESC) to choose a better angle @@ -292,10 +293,11 @@ def convert_EFIT_to_DESC( M=20, N=0, ) - plt.figure() - for k in range(0, len(fluxsurf["flux"]))[::-10]: - plt.plot(fluxsurf["flux"][k]["R"], fluxsurf["flux"][k]["Z"]) - plt.axis("equal") + if plot: + plt.figure() + for k in range(0, len(fluxsurf["flux"]))[::-10]: + plt.plot(fluxsurf["flux"][k]["R"], fluxsurf["flux"][k]["Z"]) + plt.axis("equal") # choose the LCFS as the bdry lastind = len(fluxsurf["flux"]) - 1 @@ -330,8 +332,9 @@ def convert_EFIT_to_DESC( N=0, ) data_surf = surface.compute(["R", "Z"], grid=LinearGrid(M=50, rho=1.0)) - plt.plot(data_surf["R"], data_surf["Z"], "k--") - plt.savefig(savefolder + "/" + f"initial_surfs_and_bdry_{efitfile}_{name}.png") + if plot: + plt.plot(data_surf["R"], data_surf["Z"], "k--") + plt.savefig(savefolder + "/" + f"initial_surfs_and_bdry_{efitfile}_{name}.png") efit_rho = efit["RHOVN"] # so if one were to integrate it over chi, we would get psi_T(chi) From b1e3a6c89f96a19f3cdec1ec813114168fee8693 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Tue, 31 Mar 2026 16:29:02 -0400 Subject: [PATCH 3/8] fix the relative file path issue in file saving --- efit2desc/efit2desc.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/efit2desc/efit2desc.py b/efit2desc/efit2desc.py index 6d412e2..7ca0e56 100644 --- a/efit2desc/efit2desc.py +++ b/efit2desc/efit2desc.py @@ -1,6 +1,7 @@ from omfit_classes import omfit_eqdsk import matplotlib.pyplot as plt import numpy as np +import os from scipy import integrate from desc.profiles import SplineProfile, PowerSeriesProfile from desc.equilibrium import Equilibrium @@ -261,6 +262,7 @@ def convert_EFIT_to_DESC( "arclength", "polar", ], "poloidal_angle must be one of polar or arclength" + efitname = os.path.basename(efitfile) name = f"{current_or_iota}_{profile_type}_M_{M}_prof_L_{profile_L}_psimax_{psiN_cutoff}" # _surfind_{fluxsurfind}" efit = read_EFIT_and_get_fluxsurfs(efitfile, psiN_cutoff) fluxsurf = efit["fluxSurfaces"] @@ -334,7 +336,7 @@ def convert_EFIT_to_DESC( data_surf = surface.compute(["R", "Z"], grid=LinearGrid(M=50, rho=1.0)) if plot: plt.plot(data_surf["R"], data_surf["Z"], "k--") - plt.savefig(savefolder + "/" + f"initial_surfs_and_bdry_{efitfile}_{name}.png") + plt.savefig(savefolder + "/" + f"initial_surfs_and_bdry_{efitname}_{name}.png") efit_rho = efit["RHOVN"] # so if one were to integrate it over chi, we would get psi_T(chi) @@ -407,25 +409,25 @@ def convert_EFIT_to_DESC( ), ) if save: - eq.save(savefolder + "/" + f"DESC_eq_{efitfile}_{name}.h5") + eq.save(savefolder + "/" + f"DESC_eq_{efitname}_{name}.h5") if plot: plot_1d(eq, "iota", label="DESC", lw=3) plt.plot(efit_rho, efit_iota, "r--", label="EFIT", lw=3) plt.legend() if save: - plt.savefig(savefolder + "/" + f"iota_comp_{efitfile}_{name}.png") + plt.savefig(savefolder + "/" + f"iota_comp_{efitname}_{name}.png") plot_1d(eq, "p", label="DESC", lw=3) plt.plot(efit_rho, p, "r--", label="EFIT", lw=3) plt.legend() if save: - plt.savefig(savefolder + "/" + f"pressure_comp_{efitfile}_{name}.png") + plt.savefig(savefolder + "/" + f"pressure_comp_{efitname}_{name}.png") plot_1d(eq, "current", label="DESC", lw=3) plt.plot(efit_rho, current, "r--", label="EFIT", lw=3) plt.legend() if save: - plt.savefig(savefolder + "/" + f"current_comp_{efitfile}_{name}.png") + plt.savefig(savefolder + "/" + f"current_comp_{efitname}_{name}.png") plt.figure() inds = np.arange(len(fluxsurf["flux"]))[::-10] @@ -454,14 +456,14 @@ def convert_EFIT_to_DESC( plt.legend() if save: plt.savefig( - savefolder + "/" + f"final_surfs_and_bdry_{efitfile}_{name}.png" + savefolder + "/" + f"final_surfs_and_bdry_{efitname}_{name}.png" ) plot_surfaces(eq, figsize=(8, 8), rho_lw=3, rho=rho_to_plot) if save: plt.savefig( savefolder + "/" - + f"final_surfs_with_sfl_theta_and_bdry_{efitfile}_{name}.png" + + f"final_surfs_with_sfl_theta_and_bdry_{efitname}_{name}.png" ) return eq, efit @@ -611,9 +613,9 @@ def S2(R_star): # eq 14b # compute same defs as EFIT circum = lcfs_data["perimeter(z)"][-1] vol = lcfs_data["V"] - r_0 = lcfs_data[ - "R0" - ] # TODO: this is actually vac center?? https://github.com/gafusion/OMFIT-source/blob/ebfff46939e1e8a56ff6add2fd617ccbf80eee1f/omfit/omfit_classes/fluxSurface.py#L1746 + # TODO: this is actually vac center?? + # https://github.com/gafusion/OMFIT-source/blob/ebfff46939e1e8a56ff6add2fd617ccbf80eee1f/omfit/omfit_classes/fluxSurface.py#L1746 + r_0 = lcfs_data["R0"] r_axis = r_0 ip = lcfs_data["current"][-1] Bp2_vol = vol_int(Bsq_P_vol) From 3f47dc7d1a6539753a8138dfa32111b1a8d4a7f2 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Tue, 31 Mar 2026 17:29:27 -0400 Subject: [PATCH 4/8] add solve options, add missing sqrt, remove redundant surface creation, current[0] is 0 by default --- efit2desc/efit2desc.py | 84 +++++++++++++++++++----------------------- 1 file changed, 37 insertions(+), 47 deletions(-) diff --git a/efit2desc/efit2desc.py b/efit2desc/efit2desc.py index 7ca0e56..bd91642 100644 --- a/efit2desc/efit2desc.py +++ b/efit2desc/efit2desc.py @@ -197,6 +197,7 @@ def convert_EFIT_to_DESC( profile_L=24, psiN_cutoff=0.99, solve=True, + solve_options=None, plot=True, save=True, savefolder=".", @@ -240,7 +241,11 @@ def convert_EFIT_to_DESC( Which normalized poloidal flux to cut the EFIT equilibrium off at and consider as the LCFS for the DESC Equilibrium solve : bool, Whether or not to solve the DESC Equilibrium before returning. if False, will return an unsolved DESC Equilibrium object, - which will not be in equilibrium and will not have the correct interior flux surfaces. + which will not be in equilibrium and will not have the correct interior flux surfaces. + solve_options : dict, optional + Keyword arguments forwarded to ``eq.solve()``. Any key not provided defaults to: + ``ftol=1e-12``, ``gtol=1e-14``, ``xtol=1e-14``, ``maxiter=150``, ``verbose=3``, + ``objective=ObjectiveFunction(ForceBalance(eq, grid=QuadratureGrid(L=eq.L_grid, M=eq.M_grid, N=0)))``. plot : bool, Whether or not to create the plots showing the initial and final flux surfaces and profiles compared to EFIT. save : bool, @@ -253,11 +258,14 @@ def convert_EFIT_to_DESC( Returns ------- - eq : desc.equilibrium.Equilibrium, the DESC ``Equilibrium`` object. - efit : omfit_classes.omfit_eqdsk.OMFITgeqdsk, the ``OMFITgeqdsk`` object containing the read-in and post-processed EFIT data from the gfile. - + eq : desc.equilibrium.Equilibrium + the DESC ``Equilibrium`` object. + efit : omfit_classes.omfit_eqdsk.OMFITgeqdsk + the ``OMFITgeqdsk`` object containing the read-in and post-processed EFIT data from the gfile. """ + if solve_options is None: + solve_options = {} assert poloidal_angle in [ "arclength", "polar", @@ -285,37 +293,14 @@ def convert_EFIT_to_DESC( Zaxis = np.mean(fluxsurf["flux"][0]["Z"]) x1 = Zbdry - Zaxis x2 = Rbdry - Raxis - thetas = np.arctan2(x1, x2) - - surface = FourierRZToroidalSurface.from_values( - coords=np.vstack([Rbdry, np.zeros_like(Rbdry), Zbdry]).T, - theta=thetas, - sym=False, - NFP=1, - M=20, - N=0, - ) - if plot: - plt.figure() - for k in range(0, len(fluxsurf["flux"]))[::-10]: - plt.plot(fluxsurf["flux"][k]["R"], fluxsurf["flux"][k]["Z"]) - plt.axis("equal") - - # choose the LCFS as the bdry - lastind = len(fluxsurf["flux"]) - 1 - Rbdry = fluxsurf["flux"][lastind]["R"] - Zbdry = fluxsurf["flux"][lastind]["Z"] - Raxis = np.mean(fluxsurf["flux"][0]["R"]) - Zaxis = np.mean(fluxsurf["flux"][0]["Z"]) - x1 = Zbdry - Zaxis - x2 = Rbdry - Raxis # use arclength as the angle if poloidal_angle == "arclength": arclengths = np.sqrt( (Rbdry[1:] - Rbdry[0:-1]) ** 2 + (Zbdry[1:] - Zbdry[0:-1]) ** 2 ) arclengths = np.append( - arclengths, (Rbdry[0] - Rbdry[-1]) ** 2 + (Zbdry[0] - Zbdry[-1]) ** 2 + arclengths, + np.sqrt((Rbdry[0] - Rbdry[-1]) ** 2 + (Zbdry[0] - Zbdry[-1]) ** 2), ) theta_norm_arclength = integrate.cumulative_trapezoid(y=arclengths, initial=0) theta_norm_arclength = ( @@ -334,6 +319,11 @@ def convert_EFIT_to_DESC( N=0, ) data_surf = surface.compute(["R", "Z"], grid=LinearGrid(M=50, rho=1.0)) + if plot: + plt.figure() + for k in range(0, len(fluxsurf["flux"]))[::-10]: + plt.plot(fluxsurf["flux"][k]["R"], fluxsurf["flux"][k]["Z"]) + plt.axis("equal") if plot: plt.plot(data_surf["R"], data_surf["Z"], "k--") plt.savefig(savefolder + "/" + f"initial_surfs_and_bdry_{efitname}_{name}.png") @@ -349,19 +339,17 @@ def convert_EFIT_to_DESC( psi_T = np.insert(psi_T, 0, 0) * 2 * np.pi * -1 # need this factor apparently efit_rho = np.sqrt(abs(psi_T / np.max(abs(psi_T)))) + # current[0] is always 0 current = integrate.cumtrapz( fluxsurf["avg"]["dip/dpsi"], fluxsurf["geo"]["psi"], initial=0 ) - current_shifted = current - current[0] # make current[0]=0 for spline fit - current_spline = SplineProfile( - knots=efit_rho, values=current_shifted, method="cubic2" - ) + current_spline = SplineProfile(knots=efit_rho, values=current, method="cubic2") current_poly = PowerSeriesProfile.from_values( efit_rho, current, order=profile_L, sym="even" ) - current_poly.params[0] = ( - 0.0 # make current[0]=0 for poly to enforce zero on-axis net toroidal current - ) + # make current.params[0]=0 strictly to enforce zero on-axis net toroidal current + # can be nonzero due to small profile_L + current_poly.params[0] = 0.0 p = fluxsurf["avg"]["P"] p_spline = SplineProfile(knots=efit_rho, values=p) @@ -377,13 +365,13 @@ def convert_EFIT_to_DESC( # make axis initial guess from the Raxis, Zaxis earlier axis = FourierRZCurve(R_n=Raxis, Z_n=Zaxis, sym=False, modes_R=[0], modes_Z=[0]) + pprof = p_poly if profile_type == "power_series" else p_spline iprof = i_poly if profile_type == "power_series" else i_spline - iprof = None if current_or_iota == "current" else iprof - currprof = current_poly if profile_type == "power_series" else current_spline - currprof = None if current_or_iota == "iota" else currprof - pprof = p_poly if profile_type == "power_series" else p_spline + # assign only choosen profile + iprof = None if current_or_iota == "current" else iprof + currprof = None if current_or_iota == "iota" else currprof efit_Psi = psi_T[-1] eq = Equilibrium( @@ -398,16 +386,18 @@ def convert_EFIT_to_DESC( L=L, ) if solve: - eq.solve( - ftol=1e-12, - maxiter=150, - gtol=0, - verbose=3, - xtol=0, - objective=ObjectiveFunction( + solve_options.setdefault("ftol", 1e-8) + solve_options.setdefault("gtol", 0) + solve_options.setdefault("xtol", 0) + solve_options.setdefault("maxiter", 100) + solve_options.setdefault("verbose", 3) + solve_options.setdefault( + "objective", + ObjectiveFunction( ForceBalance(eq, grid=QuadratureGrid(L=eq.L_grid, M=eq.M_grid, N=0)) ), ) + eq.solve(**solve_options) if save: eq.save(savefolder + "/" + f"DESC_eq_{efitname}_{name}.h5") From ec219a2426631662a340bb10f5bc5d302a148195 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Tue, 31 Mar 2026 17:46:50 -0400 Subject: [PATCH 5/8] minor formatting --- efit2desc/efit2desc.py | 124 ++++++++++++++++++++++------------------- 1 file changed, 67 insertions(+), 57 deletions(-) diff --git a/efit2desc/efit2desc.py b/efit2desc/efit2desc.py index bd91642..956aaf1 100644 --- a/efit2desc/efit2desc.py +++ b/efit2desc/efit2desc.py @@ -1,8 +1,12 @@ -from omfit_classes import omfit_eqdsk +import os + import matplotlib.pyplot as plt import numpy as np -import os +from scipy.constants import mu_0 from scipy import integrate + +from omfit_classes import omfit_eqdsk + from desc.profiles import SplineProfile, PowerSeriesProfile from desc.equilibrium import Equilibrium from desc.plotting import * @@ -203,65 +207,74 @@ def convert_EFIT_to_DESC( savefolder=".", poloidal_angle="polar", ): - """Read the EFIT file and generate a solved DESC equilibrium, as well as return the OMFITgeqdsk class object. + """Read the EFIT file and generate a solved DESC equilibrium. + + Also returns the OMFITgeqdsk class object. This function: - - Read the EFIT equilibrium information from the gfile using the ``omfit_eqdsk.OMFITgeqdsk`` class - - Use the ``omfit_eqdsk.OMFITgeqdsk.addAuxQuantities()`` and ``omfit_eqdsk.OMFITgeqdsk.addFluxSurfaces()`` functions to - populate the object with flux-surface quantities such as the flux surface geometries and the safety factor (q = 1/iota) and toroidal current density - - Find the last-closed-flux-surface based on the passed-in desired psiN_cutoff - and parametrize this surface with a Fourier series based off of a geometric poloidal angle - - integrates the q profile over the poloidal flux ``psi`` in order to get the toroidal flux ``psi_T(chi)`` - and uses it to define the DESC radial variable ``rho = sqrt(psi_T/psi_T(bdry))`` - - integrates the toroidal current density in order to the net toroidal current as a flux function - and shifts the toroidal current such that at the magnetic axis, there is zero net enclosed toroidal current - - fits the pressure, 1/q profile and the net toroidal current profiles as functions of ``rho``, with either a power series - or a spline, in order to form the necessary inputs for a DESC fixed-boundary Equilibrium. - - creates a DESC ``Equilibrium`` object using the LCFS, profiles, and net enclosed toroidal flux calculated from EFIT - - Optionally, solves this equilibrium in DESC using by default the grid ``QuadratureGrid(L=eq.L_grid, M=eq.M, N=0)`` - - Optionally, plots the results against the EFIT profiles and the EFIT flux surfaces - - returns the DESC ``Equilibrium`` object as well as the ``omfit_eqdsk.OMFITgeqdsk`` object - - NOTE: up-down asymmetry is assumed by default. - - Paramters - --------- - - efitfile: str, - Path to eqdsk file - current_or_iota: {"current", "iota"} - Whether to fix the iota or current profile - profile_type: {"power_series", "spline"}, - What type of Profile to use for pressure and iota/current - L,M : int, - Radial/poloidal spectral resolution to use for DESC equilibrium - profile_L : int, - Radial resolution to use for the profile fits (if using a power series) - psiN_cutoff : int, - Which normalized poloidal flux to cut the EFIT equilibrium off at and consider as the LCFS for the DESC Equilibrium - solve : bool, - Whether or not to solve the DESC Equilibrium before returning. if False, will return an unsolved DESC Equilibrium object, - which will not be in equilibrium and will not have the correct interior flux surfaces. + + - Reads the EFIT equilibrium from the gfile using ``omfit_eqdsk.OMFITgeqdsk``. + - Calls ``addAuxQuantities()`` and ``addFluxSurfaces()`` to populate flux-surface + quantities (geometry, safety factor q = 1/iota, toroidal current density). + - Finds the LCFS at ``psiN_cutoff`` and parametrizes it with a Fourier series + using the chosen poloidal angle. + - Integrates the q profile over poloidal flux ``psi`` to obtain toroidal flux + ``psi_T(psi)`` and defines the DESC radial variable + ``rho = sqrt(psi_T / psi_T(bdry))``. + - Integrates the toroidal current density to get net enclosed toroidal current + as a flux function (zero at the magnetic axis by construction). + - Fits pressure, iota = 1/q, and current profiles as functions of ``rho`` + using either a power series or spline. + - Constructs a DESC ``Equilibrium`` with the LCFS, profiles, and total + toroidal flux from EFIT. + - Optionally solves the equilibrium, plots results against EFIT, and saves + outputs to disk. + + NOTE: up-down asymmetry is assumed by default (``sym=False``). + + Parameters + ---------- + efitfile : str + Path to the eqdsk file. + current_or_iota : {"current", "iota"} + Whether to fix the current or iota profile in the DESC equilibrium. + profile_type : {"power_series", "spline"} + Profile type to use for pressure and iota/current. + L : int + Radial spectral resolution for the DESC equilibrium. + M : int + Poloidal spectral resolution for the DESC equilibrium. + profile_L : int + Radial order for power-series profile fits. + psiN_cutoff : float + Normalized poloidal flux value to treat as the LCFS. + solve : bool + If False, return the Equilibrium object without solving it. The + returned equilibrium will not satisfy force balance and will not have + correct interior flux surfaces. solve_options : dict, optional - Keyword arguments forwarded to ``eq.solve()``. Any key not provided defaults to: - ``ftol=1e-12``, ``gtol=1e-14``, ``xtol=1e-14``, ``maxiter=150``, ``verbose=3``, - ``objective=ObjectiveFunction(ForceBalance(eq, grid=QuadratureGrid(L=eq.L_grid, M=eq.M_grid, N=0)))``. - plot : bool, - Whether or not to create the plots showing the initial and final flux surfaces and profiles compared to EFIT. - save : bool, - Whether or not to save the Equilibrium and the plots (plots are saved only if ``plot=True``) - savefolder : str, - What folder to save the Equilibrium and figures to (if save=True) - poloidal angle : {"arclength", "polar"} - which poloidal angle to use when fitting the LCFS in DESC from EFIT. - Defaults to polar + Keyword arguments forwarded to ``eq.solve()``. Keys not present are + filled with defaults: ``ftol=1e-8``, ``gtol=0``, ``xtol=0``, + ``maxiter=100``, ``verbose=3``, ``objective=ObjectiveFunction( + ForceBalance(eq, grid=QuadratureGrid(L=eq.L_grid, M=eq.M_grid, N=0)))``. + plot : bool + Whether to produce comparison plots of flux surfaces and profiles + against EFIT. + save : bool + Whether to save the Equilibrium and figures (figures only if + ``plot=True``). + savefolder : str + Directory to write output files to. + poloidal_angle : {"arclength", "polar"} + Poloidal angle convention used when fitting the LCFS boundary. Returns ------- eq : desc.equilibrium.Equilibrium - the DESC ``Equilibrium`` object. + The DESC ``Equilibrium`` object. efit : omfit_classes.omfit_eqdsk.OMFITgeqdsk - the ``OMFITgeqdsk`` object containing the read-in and post-processed EFIT data from the gfile. + The ``OMFITgeqdsk`` object with the read-in and post-processed EFIT + data from the gfile. """ if solve_options is None: @@ -271,7 +284,7 @@ def convert_EFIT_to_DESC( "polar", ], "poloidal_angle must be one of polar or arclength" efitname = os.path.basename(efitfile) - name = f"{current_or_iota}_{profile_type}_M_{M}_prof_L_{profile_L}_psimax_{psiN_cutoff}" # _surfind_{fluxsurfind}" + name = f"{current_or_iota}_{profile_type}_M_{M}_prof_L_{profile_L}_psimax_{psiN_cutoff}" efit = read_EFIT_and_get_fluxsurfs(efitfile, psiN_cutoff) fluxsurf = efit["fluxSurfaces"] # this is the toroidal flux enclosed by the bdry, as calc by EFIT @@ -459,9 +472,6 @@ def convert_EFIT_to_DESC( return eq, efit -from scipy.constants import mu_0 - - def compute_betap_li_shaf_integrals(eq, efit=None): """Given a DESC eq, compute some common volumetrics. From afcc768d143be6732e05df140c9785d93265f54b Mon Sep 17 00:00:00 2001 From: YigitElma Date: Wed, 1 Apr 2026 19:34:20 -0400 Subject: [PATCH 6/8] small updates, extra explanations and clean-up --- efit2desc/efit2desc.py | 78 +++++++++++++++++++++++++++++------------- 1 file changed, 55 insertions(+), 23 deletions(-) diff --git a/efit2desc/efit2desc.py b/efit2desc/efit2desc.py index 956aaf1..b47eb15 100644 --- a/efit2desc/efit2desc.py +++ b/efit2desc/efit2desc.py @@ -24,9 +24,6 @@ def read_EFIT_and_get_fluxsurfs(efitfile, psiN_cutoff=1.0): efit.addAuxQuantities() efit.addFluxSurfaces(levels=list(np.linspace(0, psiN_cutoff, 129))) fluxsurf = efit["fluxSurfaces"] - Jt = efit["AuxQuantities"]["Jt"] - # this is the toroidal flux enclosed by the bdry, as calc by EFIT - efit_Psi = efit["AuxQuantities"]["PHI"] # [fluxsurfind] # this method obtains the iota, etc on the flux surfaces fluxsurf.surfAvg() return efit @@ -153,23 +150,30 @@ def plot_eq_iota_against_efit( fluxsurf = efit["fluxSurfaces"] # this is the toroidal flux enclosed by the bdry, as calc by EFIT - efit_rho = efit["RHOVN"] - # so if one were to integrate it over chi, we would get psi_T(chi) - + # Compute rho = sqrt(Phi_tor / Phi_tor_LCFS) on the flux surface grid. + # + # We do NOT use efit["RHOVN"] here even though it is the same quantity in + # principle, for two reasons: + # 1. RHOVN lives on the gfile 1D psiN grid (0→1, NPTS points) and is + # computed from the 1D QPSI profile. The flux surface profiles + # (fluxsurf["avg"]["q/p/dip...]) live on the flux surface grid + # (0→psiN_cutoff, 129 points, from 2D field-line tracing). Knots and + # values must be on the same grid when building the SplineProfile / + # PowerSeriesProfile objects below. + # 2. RHOVN is normalised to 1 at psiN=1 (the true separatrix). Our LCFS + # is at psiN_cutoff, so RHOVN at that surface is slightly < 1 and the + # profiles would not span rho=[0,1] as DESC requires. + # + # By integrating q from the same flux surface grid we get rho values that + # are on the same 129-point grid as the profiles and are normalised to + # exactly 1 at the chosen LCFS. psi_T = integrate.cumtrapz( fluxsurf["avg"]["q"], abs(fluxsurf["geo"]["psi"] - np.max(fluxsurf["geo"]["psi"])), ) - psi_T = np.insert(psi_T, 0, 0) * 2 * np.pi * -1 # need this factor apparently - efit_rho = np.sqrt(abs(psi_T / np.max(abs(psi_T)))) - - current = integrate.cumtrapz( - fluxsurf["avg"]["dip/dpsi"], fluxsurf["geo"]["psi"], initial=0 - ) - current_shifted = current - current[0] # make current[0]=0 for spline fit - - efit_iota = 1 / fluxsurf["avg"]["q"] + psi_T = np.insert(psi_T, 0, 0) + efit_rho = np.sqrt(psi_T / psi_T[-1]) fig, ax = plt.subplots(dpi=1000) @@ -287,8 +291,10 @@ def convert_EFIT_to_DESC( name = f"{current_or_iota}_{profile_type}_M_{M}_prof_L_{profile_L}_psimax_{psiN_cutoff}" efit = read_EFIT_and_get_fluxsurfs(efitfile, psiN_cutoff) fluxsurf = efit["fluxSurfaces"] - # this is the toroidal flux enclosed by the bdry, as calc by EFIT - efit_Psi = efit["AuxQuantities"]["PHI"] + print(efit.keys()) + print(efit["_desc"].keys()) + print(efit["_desc"]) + print(fluxsurf.keys()) # get bdry if plot: @@ -341,16 +347,21 @@ def convert_EFIT_to_DESC( plt.plot(data_surf["R"], data_surf["Z"], "k--") plt.savefig(savefolder + "/" + f"initial_surfs_and_bdry_{efitname}_{name}.png") - efit_rho = efit["RHOVN"] - # so if one were to integrate it over chi, we would get psi_T(chi) - + # Compute rho = sqrt(Phi_tor / Phi_tor_LCFS) on the flux surface grid. + # See comment in plot_eq_iota_against_efit for why RHOVN is not used here. + # Additional reason specific to convert_EFIT_to_DESC: efit_Psi is derived + # from psi_T[-1] (see below), so using the same psi_T for both rho and Psi + # keeps the two consistent with each other and with the LCFS choice. psi_T = integrate.cumtrapz( fluxsurf["avg"]["q"], abs(fluxsurf["geo"]["psi"] - np.max(fluxsurf["geo"]["psi"])), ) - psi_T = np.insert(psi_T, 0, 0) * 2 * np.pi * -1 # need this factor apparently - efit_rho = np.sqrt(abs(psi_T / np.max(abs(psi_T)))) + # flux surfaces run psiN=0 (axis) → psiN_cutoff (boundary), so + # abs(psi - max(psi)) is monotonically increasing axis→boundary, + # making psi_T strictly non-negative — no sign or 2π correction needed. + psi_T = np.insert(psi_T, 0, 0) + efit_rho = np.sqrt(psi_T / psi_T[-1]) # current[0] is always 0 current = integrate.cumtrapz( @@ -386,7 +397,28 @@ def convert_EFIT_to_DESC( iprof = None if current_or_iota == "current" else iprof currprof = None if current_or_iota == "iota" else currprof - efit_Psi = psi_T[-1] + # psi_T[-1] = ∫ q d|ψ| where ψ = Ψ_pol/(2π) is in Wb/rad (gfile convention). + # Therefore psi_T[-1] = |φ_tor| = |Φ_tor|/(2π) = |DESC's psi (lowercase)|. + # DESC's Psi (uppercase) = Φ_tor in Webers = 2π * psi_T[-1] * sign. + # + # psi_T is always positive (abs() forces integration variable to increase + # axis→boundary regardless of EFIT's ψ sign convention), so the sign of + # Φ_tor must come from EFIT's own PHI array (reflects actual B_tor direction). + # + # We use psi_T for magnitude rather than PHI directly because psi_T is + # computed from the same 2D flux surfaces used to define the LCFS, making + # the two internally consistent. OMFIT's PHI uses the 1D QPSI gfile profile + # which can differ from the 2D surface-averaged q. + _phi_sign = float( + np.sign( + np.interp( + psiN_cutoff, + np.linspace(0, 1, len(efit["AuxQuantities"]["PHI"])), + efit["AuxQuantities"]["PHI"], + ) + ) + ) + efit_Psi = _phi_sign * 2 * np.pi * psi_T[-1] eq = Equilibrium( surface=surface, axis=axis, From 0cd48a80e9ac9b047d26a75b940b89bdbafc4750 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Wed, 1 Apr 2026 19:54:49 -0400 Subject: [PATCH 7/8] format code, remove redundant comments/code etc, no change to actual computation --- efit2desc/efit2desc.py | 173 ++++++++++++++++------------------------- 1 file changed, 65 insertions(+), 108 deletions(-) diff --git a/efit2desc/efit2desc.py b/efit2desc/efit2desc.py index b47eb15..919c7a5 100644 --- a/efit2desc/efit2desc.py +++ b/efit2desc/efit2desc.py @@ -7,36 +7,30 @@ from omfit_classes import omfit_eqdsk -from desc.profiles import SplineProfile, PowerSeriesProfile +from desc.profiles import PowerSeriesProfile, SplineProfile from desc.equilibrium import Equilibrium from desc.plotting import * -from desc.grid import QuadratureGrid +from desc.grid import LinearGrid, QuadratureGrid from desc.objectives import ForceBalance, ObjectiveFunction -from desc.geometry import FourierRZToroidalSurface, FourierRZCurve -from desc.grid import LinearGrid +from desc.geometry import FourierRZCurve, FourierRZToroidalSurface def read_EFIT_and_get_fluxsurfs(efitfile, psiN_cutoff=1.0): - efit = omfit_eqdsk.OMFITgeqdsk(efitfile) - # run the methods of the OMFITgeqdsk class to get - # aux and flux surface quantities for the EFIT + # populate aux and flux-surface quantities efit.addAuxQuantities() efit.addFluxSurfaces(levels=list(np.linspace(0, psiN_cutoff, 129))) - fluxsurf = efit["fluxSurfaces"] - # this method obtains the iota, etc on the flux surfaces - fluxsurf.surfAvg() + # surfAvg computes safety factor, pressure, etc. on each flux surface + efit["fluxSurfaces"].surfAvg() return efit def plot_eq_surfaces_against_efit(efitfile, desc_eq, levels=20): - efit = read_EFIT_and_get_fluxsurfs(efitfile, 1.0) fluxsurf = efit["fluxSurfaces"] inds = np.arange(len(fluxsurf["flux"]))[::-10] - efit_rho = efit["RHOVN"] - rho_to_plot = efit_rho[inds] + rho_to_plot = efit["RHOVN"][inds] fig, ax = plot_surfaces( desc_eq, figsize=(8, 8), @@ -46,23 +40,15 @@ def plot_eq_surfaces_against_efit(efitfile, desc_eq, levels=20): ) is_labelled = False for k in inds: - if not is_labelled: - plt.plot( - fluxsurf["flux"][k]["R"], - fluxsurf["flux"][k]["Z"], - "k--", - label="EFIT", - lw=3, - ) - is_labelled = True - else: - plt.plot( - fluxsurf["flux"][k]["R"], - fluxsurf["flux"][k]["Z"], - "k--", - lw=3, - label="EFIT", - ) + label = "EFIT" if not is_labelled else None + plt.plot( + fluxsurf["flux"][k]["R"], + fluxsurf["flux"][k]["Z"], + "k--", + lw=3, + label=label, + ) + is_labelled = True # also want to add a couple contours outside plt.contour( efit["AuxQuantities"]["R"], @@ -72,17 +58,12 @@ def plot_eq_surfaces_against_efit(efitfile, desc_eq, levels=20): linewidths=[3], linestyles=["--"], ) - # plt.colorbar() plt.axis("equal") - plt.axis("equal") - # plt.scatter(Raxis, Zaxis, marker="x", label="EFIT axis", c="k") - # desc_axis = desc_eq.axis.compute(["R", "Z"]) - # plt.scatter(desc_axis["R"][0], desc_axis["Z"][0], label="DESC Axis") plt.legend() return fig, ax -### taken from vmeclauncher.py truncateEFIT.py by Wingen ### +# taken from vmeclauncher.py truncateEFIT.py by Wingen # find psi with q(psi) = m/n for all m in [n*qmin, n*qmax] def scan_q(q, n=3): rhos = np.linspace(0, 1.0, 1000) @@ -96,18 +77,12 @@ def scan_q(q, n=3): for m in range(int(n * qmin) + 1, N): psi = bisec(lambda x: q(x) - float(m) / n, a=0, b=1) - # print (m, psi, q(psi)) psia.append(psi) qa.append(q(psi)) - # m = int(n * qmax) - 1 - # qbest = (m + 0.2) / n - # psi = bisec(lambda x: q(x) - qbest, a=0, b=1) - - return psia, qa # , psi + return psia, qa -# ---------------------------------------------------------------------------------------- # find root through bisection def bisec(funct, a=0, b=1.5): eps = 1e-14 @@ -133,9 +108,6 @@ def bisec(funct, a=0, b=1.5): return x -################################ - - def plot_eq_iota_against_efit( efitfile, desc_eq, @@ -145,10 +117,8 @@ def plot_eq_iota_against_efit( max_n=6, method="cubic", ): - efit = read_EFIT_and_get_fluxsurfs(efitfile, psiN_cutoff) fluxsurf = efit["fluxSurfaces"] - # this is the toroidal flux enclosed by the bdry, as calc by EFIT # Compute rho = sqrt(Phi_tor / Phi_tor_LCFS) on the flux surface grid. # @@ -171,7 +141,6 @@ def plot_eq_iota_against_efit( fluxsurf["avg"]["q"], abs(fluxsurf["geo"]["psi"] - np.max(fluxsurf["geo"]["psi"])), ) - psi_T = np.insert(psi_T, 0, 0) efit_rho = np.sqrt(psi_T / psi_T[-1]) @@ -179,7 +148,7 @@ def plot_eq_iota_against_efit( if show_rationals: qprof = SplineProfile( - knots=abs(psi_T / np.max(abs(psi_T))), + knots=efit_rho, values=np.abs(fluxsurf["avg"]["q"]), method=method, ) @@ -288,20 +257,12 @@ def convert_EFIT_to_DESC( "polar", ], "poloidal_angle must be one of polar or arclength" efitname = os.path.basename(efitfile) - name = f"{current_or_iota}_{profile_type}_M_{M}_prof_L_{profile_L}_psimax_{psiN_cutoff}" + name = ( + f"{current_or_iota}_{profile_type}_M_{M}" + f"_prof_L_{profile_L}_psimax_{psiN_cutoff}" + ) efit = read_EFIT_and_get_fluxsurfs(efitfile, psiN_cutoff) fluxsurf = efit["fluxSurfaces"] - print(efit.keys()) - print(efit["_desc"].keys()) - print(efit["_desc"]) - print(fluxsurf.keys()) - - # get bdry - if plot: - plt.figure() - for k in range(0, len(fluxsurf["flux"]))[::-10]: - plt.plot(fluxsurf["flux"][k]["R"], fluxsurf["flux"][k]["Z"]) - plt.axis("equal") # choose the LCFS as the bdry # TODO: use spectral condensation (ideally when implemented in DESC) to choose a better angle @@ -312,7 +273,7 @@ def convert_EFIT_to_DESC( Zaxis = np.mean(fluxsurf["flux"][0]["Z"]) x1 = Zbdry - Zaxis x2 = Rbdry - Raxis - # use arclength as the angle + if poloidal_angle == "arclength": arclengths = np.sqrt( (Rbdry[1:] - Rbdry[0:-1]) ** 2 + (Zbdry[1:] - Zbdry[0:-1]) ** 2 @@ -322,10 +283,7 @@ def convert_EFIT_to_DESC( np.sqrt((Rbdry[0] - Rbdry[-1]) ** 2 + (Zbdry[0] - Zbdry[-1]) ** 2), ) theta_norm_arclength = integrate.cumulative_trapezoid(y=arclengths, initial=0) - theta_norm_arclength = ( - theta_norm_arclength / np.max(theta_norm_arclength) * 2 * np.pi - ) - thetas = theta_norm_arclength + thetas = theta_norm_arclength / np.max(theta_norm_arclength) * 2 * np.pi elif poloidal_angle == "polar": thetas = np.arctan2(x1, x2) @@ -337,13 +295,12 @@ def convert_EFIT_to_DESC( M=20, N=0, ) - data_surf = surface.compute(["R", "Z"], grid=LinearGrid(M=50, rho=1.0)) if plot: + data_surf = surface.compute(["R", "Z"], grid=LinearGrid(M=50, rho=1.0)) plt.figure() for k in range(0, len(fluxsurf["flux"]))[::-10]: plt.plot(fluxsurf["flux"][k]["R"], fluxsurf["flux"][k]["Z"]) plt.axis("equal") - if plot: plt.plot(data_surf["R"], data_surf["Z"], "k--") plt.savefig(savefolder + "/" + f"initial_surfs_and_bdry_{efitname}_{name}.png") @@ -356,14 +313,13 @@ def convert_EFIT_to_DESC( fluxsurf["avg"]["q"], abs(fluxsurf["geo"]["psi"] - np.max(fluxsurf["geo"]["psi"])), ) - # flux surfaces run psiN=0 (axis) → psiN_cutoff (boundary), so # abs(psi - max(psi)) is monotonically increasing axis→boundary, # making psi_T strictly non-negative — no sign or 2π correction needed. psi_T = np.insert(psi_T, 0, 0) efit_rho = np.sqrt(psi_T / psi_T[-1]) - # current[0] is always 0 + # current[0] is always 0 (cumtrapz initial=0) current = integrate.cumtrapz( fluxsurf["avg"]["dip/dpsi"], fluxsurf["geo"]["psi"], initial=0 ) @@ -371,8 +327,7 @@ def convert_EFIT_to_DESC( current_poly = PowerSeriesProfile.from_values( efit_rho, current, order=profile_L, sym="even" ) - # make current.params[0]=0 strictly to enforce zero on-axis net toroidal current - # can be nonzero due to small profile_L + # enforce zero on-axis net toroidal current (can be nonzero due to finite profile_L) current_poly.params[0] = 0.0 p = fluxsurf["avg"]["P"] @@ -380,20 +335,18 @@ def convert_EFIT_to_DESC( p_poly = PowerSeriesProfile.from_values(efit_rho, p, order=profile_L, sym="even") efit_iota = 1 / fluxsurf["avg"]["q"] - i_spline = SplineProfile(knots=efit_rho, values=efit_iota) i_poly = PowerSeriesProfile.from_values( efit_rho, efit_iota, order=profile_L, sym="even" ) - # make axis initial guess from the Raxis, Zaxis earlier axis = FourierRZCurve(R_n=Raxis, Z_n=Zaxis, sym=False, modes_R=[0], modes_Z=[0]) pprof = p_poly if profile_type == "power_series" else p_spline iprof = i_poly if profile_type == "power_series" else i_spline currprof = current_poly if profile_type == "power_series" else current_spline - # assign only choosen profile + # assign only the chosen profile type iprof = None if current_or_iota == "current" else iprof currprof = None if current_or_iota == "iota" else currprof @@ -419,6 +372,7 @@ def convert_EFIT_to_DESC( ) ) efit_Psi = _phi_sign * 2 * np.pi * psi_T[-1] + eq = Equilibrium( surface=surface, axis=axis, @@ -452,6 +406,7 @@ def convert_EFIT_to_DESC( plt.legend() if save: plt.savefig(savefolder + "/" + f"iota_comp_{efitname}_{name}.png") + plot_1d(eq, "p", label="DESC", lw=3) plt.plot(efit_rho, p, "r--", label="EFIT", lw=3) plt.legend() @@ -464,26 +419,22 @@ def convert_EFIT_to_DESC( if save: plt.savefig(savefolder + "/" + f"current_comp_{efitname}_{name}.png") - plt.figure() inds = np.arange(len(fluxsurf["flux"]))[::-10] rho_to_plot = efit_rho[inds] + + plt.figure() plot_surfaces(eq, figsize=(8, 8), theta=0, rho_lw=3, rho=rho_to_plot) is_labelled = False for k in inds: - if not is_labelled: - plt.plot( - fluxsurf["flux"][k]["R"], - fluxsurf["flux"][k]["Z"], - "k--", - label="EFIT", - lw=3, - ) - is_labelled = True - else: - plt.plot( - fluxsurf["flux"][k]["R"], fluxsurf["flux"][k]["Z"], "k--", lw=3 - ) - + label = "EFIT" if not is_labelled else None + plt.plot( + fluxsurf["flux"][k]["R"], + fluxsurf["flux"][k]["Z"], + "k--", + lw=3, + label=label, + ) + is_labelled = True plt.axis("equal") plt.scatter(Raxis, Zaxis, marker="x", label="EFIT axis", c="k") desc_axis = eq.axis.compute(["R", "Z"]) @@ -493,6 +444,7 @@ def convert_EFIT_to_DESC( plt.savefig( savefolder + "/" + f"final_surfs_and_bdry_{efitname}_{name}.png" ) + plot_surfaces(eq, figsize=(8, 8), rho_lw=3, rho=rho_to_plot) if save: plt.savefig( @@ -524,10 +476,19 @@ def compute_betap_li_shaf_integrals(eq, efit=None): _type_ _description_ """ - vol_grid = QuadratureGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N, NFP=eq.NFP) vol_data = eq.compute( - ["p", "B_theta", "sqrt(g)", "B_R", "B_Z", "V", "B_phi", "R", "_vol"], + [ + "p", + "B_theta", + "sqrt(g)", + "B_R", + "B_Z", + "V", + "B_phi", + "R", + "_vol", + ], grid=vol_grid, ) @@ -552,18 +513,15 @@ def compute_betap_li_shaf_integrals(eq, efit=None): grid=lcfs_grid, ) - def vol_int(q): # , grid,vol_data): + def vol_int(q): return np.sum(vol_grid.weights * q * vol_data["sqrt(g)"]) - # return np.sum(grid.weights * q ) - def lcfs_int(q): # ,grid, lcfs_data): # _V from Hirshman 1993 paper def eq 8 + def lcfs_int(q): # _V from Hirshman 1993 paper def eq 8 return ( np.sum(lcfs_grid.weights * q * lcfs_data["|e_theta x e_zeta|"]) * lcfs_data["V"] / lcfs_data["S"] ) - # return np.sum(grid.weights * q ) * lcfs_data["R0"] * 2 * np.pi - # return np.sum(grid.weights * q ) * lcfs_data["V"]/lcfs_data["A"] Bsq_P_lcfs = lcfs_data["B_R"] ** 2 + lcfs_data["B_Z"] ** 2 Bsq_P_vol = vol_data["B_R"] ** 2 + vol_data["B_Z"] ** 2 @@ -599,12 +557,10 @@ def lcfs_int(q): # ,grid, lcfs_data): # _V from Hirshman 1993 paper def eq 8 sig_hat_R_alpha_1 = A * one_over_RT # eq 10a but with alpha=1 def S1(R_star): # eq 14a - return ( - sig_hat_R + sig_hat_Z - R_star * sig_hat_R_alpha_1 - ) # this last one should be assuming alpha=1... + return sig_hat_R + sig_hat_Z - R_star * sig_hat_R_alpha_1 def S2(R_star): # eq 14b - return R_star * sig_hat_R_alpha_1 # this last one should be assuming alpha=1... + return R_star * sig_hat_R_alpha_1 S3 = C # = sig_hat_Z(Z) eq 14c @@ -616,23 +572,26 @@ def S2(R_star): # eq 14b flao = Rshaf / Rlao # shafranov integrals for different choices of Rstar - # for some reason, hirshman multiplies the RG and RL by fgeo=Rshaf/Rgeo and flao=Rshaf/Rgeo... dont ask me why + # for some reason, hirshman multiplies the RG and RL by fgeo=Rshaf/Rgeo and + # flao=Rshaf/Rgeo... dont ask me why s1 = S1(Rlao) / 2 / flao s2 = S2(Rlao) / 2 / flao s3 = S3 / 2 print("#" * 10) print("DESC") print("#" * 10) - print(f"s3 = {0.5*(betai-lsubi-musubi+2*lsubR)}") # #S3/2, eq 14c + print(f"s3 = {0.5*(betai-lsubi-musubi+2*lsubR)}") # S3/2, eq 14c print(f"{lsubi=}") print(f"{lsubi_current_normalized=}") print(f"{musubi=}") print(f"{betai=}") print( - f"s1 = S1/2 = {S1(1/one_over_RT)/2} (RT) , {S1(Rgeo)/2/fgeo} (RG?) {S1(Rlao)/2/flao} (RL)" + f"s1 = S1/2 = {S1(1/one_over_RT)/2} (RT) , " + f"{S1(Rgeo)/2/fgeo} (RG?) {S1(Rlao)/2/flao} (RL)" ) print( - f"s2 = S2/2 = {S2(1/one_over_RT)/2} (RT) , {S2(Rgeo)/2/fgeo} (RG?) {S2(Rlao)/2/flao} (RL)" + f"s2 = S2/2 = {S2(1/one_over_RT)/2} (RT) , " + f"{S2(Rgeo)/2/fgeo} (RG?) {S2(Rlao)/2/flao} (RL)" ) print(f"DESC poloidal beta calc: {vol_data['_vol']}") # FIXME: does this actually work? @@ -679,8 +638,6 @@ def S2(R_star): # eq 14b ) print(f"{line}" + " " * len(line) + efit_line) - # print(f"EFIT li: {efit["fluxSurfaces"]['info']['internal_inductance']}") - return { "betai": betai, "_vol": vol_data["_vol"], From cb47b545f27248a5560b1ebcaff88e0a1a43fa47 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Wed, 1 Apr 2026 20:25:46 -0400 Subject: [PATCH 8/8] clean-up, reorder, format file --- efit2desc/get_coilset_for_shot.py | 66 +++++++++++++------------------ 1 file changed, 27 insertions(+), 39 deletions(-) diff --git a/efit2desc/get_coilset_for_shot.py b/efit2desc/get_coilset_for_shot.py index 9a06273..24b4ce9 100644 --- a/efit2desc/get_coilset_for_shot.py +++ b/efit2desc/get_coilset_for_shot.py @@ -1,22 +1,11 @@ +import warnings + +import numpy as np from scipy.interpolate import interp1d from scipy.constants import mu_0 -from desc.coils import MixedCoilSet, CoilSet, FourierRZCoil, SplineXYZCoil -from desc.magnetic_fields import ToroidalMagneticField, SplineMagneticField -from desc.utils import dot -import numpy as np -from desc.io import load -import io -import jax -import matplotlib.pyplot as plt -import warnings -### Get coil currents from ptdata ### -try: - from ptdata import fetch_ptdata, PtDataFetcher -except Exception as e: - print("could not import ptdata, got exception") - print(e) -from matplotlib.backends.backend_pdf import PdfPages +from desc.coils import MixedCoilSet +from desc.magnetic_fields import SplineMagneticField, ToroidalMagneticField def get_coilset_for_shot( @@ -48,6 +37,9 @@ def get_coilset_for_shot( whether to use an idealized toroidal field for the TF coilset, in which case the BCOIL coil current will be set to zero. False by default. + base_coils_file : str, optional + Path to the base coils file. If None, defaults to + ``"coils.d3d_efbic_kp48_from_andreas"``. Returns ------- @@ -83,7 +75,7 @@ def get_coilset_for_shot( ptnames.append(f"IL{ang}") def update_F_coils(t, coils, interpolable_ptnames): - # first 18 coils in Andreas' arrays is are the Fcoils + # first 18 coils in Andreas' arrays are the Fcoils Fturns = np.array( [ 58, @@ -108,8 +100,8 @@ def update_F_coils(t, coils, interpolable_ptnames): ) for i in range(18): assert "f" in coils[i].name - total_curr = interpolable_ptnames[coils[i].name[2:].upper().strip()](t) - current = total_curr # the current measured is per turn, and they actually have the number of turns given by Fturns + # current is per-turn; number of turns per coil is given by Fturns + current = interpolable_ptnames[coils[i].name[2:].upper().strip()](t) for j in range(len(coils[i])): coils[i][j].current = current * np.sign(coils[i][j].current) @@ -148,17 +140,13 @@ def update_Iu_coils(t, coils, interpolable_ptnames): CCoil199 , ind = 35 """ - import re - nums = [str(i) for i in range(10)] def update_C_coils(t, coils, interpolable_ptnames): for i in range(3): ind = i + 33 - - for coil in coils[ - ind - ]: # go into the coilset so can set correct currents per coil + # go into the coilset so we can set correct currents per coil + for coil in coils[ind]: angle_str = "" for c in coil.name[2:]: angle_str += c if c in nums else "" @@ -194,18 +182,20 @@ def update_all_coil_currents( # iL coils are total_coils[27:32] update_IL_coils(t, total_coils, interpolable_ptnames) - # coilsfile I got from Andreas Wingen - coilset_name = base_coils_file # "coils.d3d_efbic_kp48_from_andreas" - # get interp1d representations of the currents interpolable_ptnames = {} - PCS_SYS_D3 = ":/fusion/projects/codes/pcs/data/ptdata:/fusion/projects/codes/pcs/data/ptdata/uncomp:" + PCS_SYS_D3 = ( + ":/fusion/projects/codes/pcs/data/ptdata" + ":/fusion/projects/codes/pcs/data/ptdata/uncomp:" + ) + # get coil currents from ptdata + try: + from ptdata import fetch_ptdata, PtDataFetcher + except Exception as e: + print(f"could not import ptdata, got exception: {e}") for ptname in ptnames: - # print(ptname) try: fetcher = PtDataFetcher(ptname, round(shot), sys_d3=PCS_SYS_D3) - - header = fetcher.header results = fetcher.fetch() interpolable_ptnames[ptname] = interp1d( x=results["times"], y=results["data"], kind=interp_method @@ -223,7 +213,8 @@ def update_all_coil_currents( else: rename_coils = False warnings.warn( - "Make sure the coilset you load has the correct coil names", UserWarning + "Make sure the coilset you load has the correct coil names", + UserWarning, ) coilset_name = base_coils_file @@ -232,11 +223,8 @@ def update_all_coil_currents( ) if rename_coils: - # I rename his C coils as he has it setup so - # 79 and 259 are in the same group - # 139 and 319 - # 199 and 19 - # but we want to have them named acc. to what they individually are + # rename C coils: Andreas groups 79+259, 139+319, 199+19 together, + # but we want them named individually full_nominal_coils[33][1].name = "34 CCoil259" full_nominal_coils[34][1].name = "35 CCoil319" full_nominal_coils[35][1].name = "36 CCoil19" @@ -281,7 +269,7 @@ def update_all_coil_currents( EXTCUR(34) = 1.0540000E+03 EXTCUR(35) = -1.4100000E+02 EXTCUR(36) = -1.2510000E+03 - + !----- Final comments --------------------------------------------------- ! mapcode boundary fraction is 0.9936 ! file input.d3d.166439.04400_lcfs9936_fixed_ns97 generated by