Automatic continuation method for free boundary

[f0uriest]

solve_continuation_automatic_free_bdry where we give the field we want to find free bdry for.

[dpanici]

@dpanici add your script

[dpanici]

929 related

[dpanici]

this is loading from VMEC but a similar structure would work (should not have that first no sheet current solve commented out tho)

"""Run free-bdry DESC on a finite beta eq given from a VMEC .nc file
First cmd line arg: path to the VMEC wout file
Second cmd line arg: path to the MAKEGRID coils file for this equilibrium
Third cmd line arg : path to the VMEC input file (so DESC can get the pressure profile)
"""

import os
import pathlib
import pickle
import sys

import jax
import numpy as np

sys.path.insert(0, os.path.abspath("."))
sys.path.append(os.path.abspath("../"))
from desc import set_device

set_device("gpu")

import desc.examples
from desc.continuation import solve_continuation_automatic
from desc.coils import CoilSet, MixedCoilSet
from desc.equilibrium import EquilibriaFamily, Equilibrium
from desc.geometry import FourierRZToroidalSurface
from desc.magnetic_fields import FourierCurrentPotentialField, SplineMagneticField
from desc.objectives import (
    BoundaryError,
    FixBoundaryR,
    FixBoundaryZ,
    FixCurrent,
    FixPressure,
    FixPsi,
    ForceBalance,
    ObjectiveFunction,
)
from desc.profiles import PowerSeriesProfile
from desc.vmec import VMECIO
from desc.input_reader import InputReader
from desc.equilibrium.utils import parse_profile
optimizer = "proximal-lsq-exact"
nn = 7
maxiter = 100
wout_filename = sys.argv[1] # first input is the name of the file
subfoldername = wout_filename.split("/")[0]
name = wout_filename.split("/")[1]
name = name.strip("wout_").strip(".nc")

path_init_fixed_solve = str(pathlib.Path(__file__).parent.resolve()) + "/" + subfoldername + "/" + f"desc_initial_fixed_bdry_solve_{name}.h5"
path_with_K = str(pathlib.Path(__file__).parent.resolve()) + "/" + subfoldername + "/" + f"desc_with_sheet_current_{name}.h5"
path_no_K = str(pathlib.Path(__file__).parent.resolve()) + "/" + subfoldername + "/"+  f"desc_no_sheet_current_{name}.h5"

print("SAVING NO SHEET CURRENT SOLVE TO ", path_no_K)
print("SAVING  SHEET CURRENT SOLVE TO ", path_with_K)

def get_constraints(eq, i):

    R_modes = eq.surface.R_basis.modes[
        np.max(np.abs(eq.surface.R_basis.modes), 1) > int(i / nn * eq.M), :
    ]
    Z_modes = eq.surface.Z_basis.modes[
        np.max(np.abs(eq.surface.Z_basis.modes), 1) > int(i / nn * eq.M), :
    ]

    return (
        ForceBalance(eq=eq),
        FixBoundaryR(eq=eq, modes=R_modes),
        FixBoundaryZ(eq=eq, modes=Z_modes),
        FixPressure(eq=eq),
        FixCurrent(eq=eq),
        FixPsi(eq=eq),
    )

def get_objective(ext_field, eq):
    return ObjectiveFunction(BoundaryError(eq, ext_field, field_fixed=True))

os.getcwd()

ext_field = CoilSet.from_makegrid_coilfile(
    sys.argv[2],
).to_FourierXYZ(N=40)
if isinstance(ext_field, MixedCoilSet):
    # make a CoilSet if it is a MixedCoilSet
    ext_field = CoilSet(*ext_field)

eq = VMECIO.load(wout_filename,profile="current")

# get the pressure, which is a power series, from the input file
input_info = InputReader.parse_vmec_inputs(sys.argv[3])
pressure = parse_profile(input_info[0]["pressure"])
pressure.change_resolution(eq.L)
eq.pressure = pressure
eq.change_resolution(10, 10, 10, 20, 20, 20)
eq.solve()
eq.save(path_init_fixed_solve)
# eqf = EquilibriaFamily(eq)

###### First solve with no sheet current ######
# because these coils were made for these equilibria, sheet
# current should be small, so it is better to initially solve without

# for i in range(1, nn + 1):

#     jax.clear_caches()
#     print("\n==================================")
#     print("Optimizing boundary modes M,N <= {}".format(int(i / nn * eq.M)))
#     print("====================================")
#     eqf[-1].solve(ftol=1e-2, verbose=3)

#     eq_new, out = eqf[-1].optimize(
#         get_objective(ext_field, eqf[-1]),
#         get_constraints(eqf[-1], i),
#         optimizer=optimizer,
#         maxiter=maxiter,
#         ftol=1e-4,
#         xtol=0,
#         verbose=3,
#         copy=True,
#         options={},
#     )
#     eqf.append(eq_new)
#     eqf.save(path_no_K)

###### Then solve with sheet current included ######

eq.surface = FourierCurrentPotentialField.from_surface(eq.surface)
eqf = EquilibriaFamily(eq)

for i in range(1, nn + 1):

    jax.clear_caches()
    print("\n==================================")
    print("Optimizing boundary modes M,N <= {}".format(int(i / nn * eq.M)))
    print("====================================")
    eq.surface.change_Phi_resolution(int(i / nn * eq.M), int(i / nn * eq.N))
    eqf[-1].solve(ftol=1e-2, verbose=3)

    eq_new, out = eqf[-1].optimize(
        get_objective(ext_field, eqf[-1]),
        get_constraints(eqf[-1], i),
        optimizer=optimizer,
        maxiter=maxiter,
        ftol=1e-4,
        xtol=0,
        verbose=3,
        copy=True,
        options={},
    )
    eqf.append(eq_new)
    eqf.save(path_with_K)