From b0d450705a43057f1c8fdf33fb56919ee3c26603 Mon Sep 17 00:00:00 2001 From: mfmart Date: Fri, 12 Apr 2024 06:46:35 -0700 Subject: [PATCH 01/49] added objective file wrapping the TERPSICHORE linear stability code --- desc/objectives/_terpsichore.py | 405 ++++++++++++++++++++++++++++++++ 1 file changed, 405 insertions(+) create mode 100644 desc/objectives/_terpsichore.py diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py new file mode 100644 index 0000000000..509d9f1b5c --- /dev/null +++ b/desc/objectives/_terpsichore.py @@ -0,0 +1,405 @@ +import numpy as np +import subprocess +from scipy.interpolate import interp1d +import os +import time +from desc.backend import jnp,put + +from .objective_funs import ObjectiveFunction, _Objective +from .utils import ( + factorize_linear_constraints, +) +from desc.utils import Timer +from desc.compute import compute as compute_fun +from desc.compute import get_profiles, get_transforms, get_params +from desc.compute.utils import dot + +from scipy.constants import mu_0, pi +from desc.grid import LinearGrid, Grid, ConcentricGrid, QuadratureGrid +from jax import core +from jax.interpreters import ad, batching +from desc.derivatives import FiniteDiffDerivative +import netCDF4 as nc +from shutil import copyfile +from desc.compute.utils import cross, dot + + +class TERPSICHORE(_Objective): + r"""Calls the linear MHD stability code TERPSICHORE to compute the linear growth rates of the fastest growing instability + + Parameters + ---------- + eq : Equilibrium, optional + Equilibrium that will be optimized to satisfy the Objective. + target : float, ndarray, optional + COMPLETE THIS + bounds : tuple, optional + COMPLETE THIS + weight : float, ndarray, optional + COMPLETE THIS + + more stuff here + """ + + # Used by GX objective, not sure where they're used yet + _scalar = True + _linear = False + _units = "[]" # Outputting with dimensions or normalized? + _print_value_fmt = "Growth rate: {:10.3e} " + + + # Need to sure up the paths + def __init__( + self, + eq=None, + target=None, + weight=1, + grid=None, + name="TERPSICHORE", + path=os.getenv['terps_dir'], + path_in=os.getenv['terps_dir'], + bounds=None, + normalize=False, + normalize_target=False, + awall=1.5, + deltajp=4.e-2, + modelk=0, + al0=-5.e-2, + nev=1, + nfp=2, + xplo=1.e-6, + max_bozm=19, + max_bozn=14, + mode_family=0, + max_modem=55, + max_moden=8, + wout_filename="" # add something + ): + + if target is None and bounds is None: + target = 0 + self.eq = eq + self.awall = awall + self.deltajp = deltajp + self.modelk = modelk + self.al0 = al0 + self.nev = nev + self.nfp = nfp + self.xplo = xplo + self.max_bozm = max_bozm + self.max_bozn = max_bozn + self.mode_family = mode_family + self.max_modem = max_modem + self.max_moden = max_moden + self.grid = grid + + super().__init__( + things=eq, + target=target, + bounds=bounds, + weight=weight, + normalize=normalize, + normalize_target=normalize_target, + name=name, + ) + + units = "" + self._callback_fmt = "Growth rate: {:10.3e} " + units + self._print_value_fmt = "Growth rate: {:10.3e} " + units + + self.path = path + self.path_in = path_in + self.wout_file = os.path.join(self.path, wout_filename) + self.vmec2terps_app = os.path.join(self.path, "thea-vmec2terps.x") + self.terps_app = os.path.join(self.path, "tpr_ap.x") + # self.t = t # -- Replace with some stability indicator? + + self.terps_compute = core.Primitive("terps") + self.terps_compute.def_impl(self.compute_impl) + # What are the following two lines? + ad.primitive_jvps[self.terps_compute] = self.compute_terps_jvp + batching.primitive_batchers[self.terps_compute] = self.compute_terps_batch + + def build(self, use_jit=False, verbose=1): + """Build constant arrays. + + Parameters + ---------- + eq : Equilibrium, optional + Equilibrium that will be optimized to satisfy the Objective. + use_jit : bool, optional + Whether to just-in-time compile the objective and derivatives + (MUST be False to run GX). + verbose : int, optional + Level of output. + + """ + eq = self.things[0] + if self.grid is None: + self.grid_eq = QuadratureGrid( + L=eq.L_grid, + M=eq.M_grid, + N=eq.N_grid, + NFP=eq.NFP, + ) + + + self._dim_f = 1 # Presumbaly this should be 1? just a growth rate + timer = Timer() + + self._eq_keys = [ + "iota", + "iota_r", + "a", + "rho", + "psi", + ] + self._field_line_keys = [ + "|B|", "|grad(psi)|^2", "grad(|B|)", "grad(alpha)", "grad(psi)", + "B", "grad(|B|)", "kappa", "B^theta", "B^zeta", "lambda_t", "lambda_z",'p_r', + "lambda_r", "lambda", "g^rr", "g^rt", "g^rz", "g^tz", "g^tt", "g^zz", + "e^rho", "e^theta", "e^zeta" + ] + + self._args = get_params( + self._field_line_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=self.grid_eq.axis.size, + ) + + if verbose > 0: + print("Precomputing transforms") + timer.start("Precomputing transforms") + + #Need separate transforms and profiles for the equilibrium and flux-tube + self.eq = eq + self._profiles = get_profiles(self._field_line_keys, obj=eq, grid=self.grid) + self._profiles_eq = get_profiles(self._eq_keys, obj=eq, grid=self.grid_eq) + self._transforms = get_transforms(self._field_line_keys, obj=eq, grid=self.grid) + self._transforms_eq = get_transforms(self._eq_keys, obj=eq, grid=self.grid_eq) + + self._constants = { + "transforms": self._transforms_eq, + "profiles": self._profiles_eq, + } + + if self._normalize: + scales = compute_scaling_factors(eq) + # I think only needed if we're not using the TERPS normalization? + + timer.stop("Precomputing transforms") + if verbose > 1: + timer.disp("Precomputing transforms") + +# self._check_dimensions() +# self._set_dimensions(eq) + super().build(use_jit=use_jit, verbose=verbose) + + + def compute_impl(self, params, constants): + +# params, constants = self._parse_args(*args, **kwargs) + if not np.any(params["R_lmn"]): + return 0 + + if constants is None: + constants = self.constants + + self.write_vmec() # Write VMEC file from DESC + self.compute_fort18() + self.write_terps_io() + self.run_terps() + + self.terps_outfile = os.path.join(self.path_in,'fort.16') # Let's change the name of this at some point + self.parse_terps_outfile() + + + # Is something analogous needed for TERPS?? + if qflux_avg > 20: + out_fail = 'fail_' + str(self.t) + '_.nc' + copyfile(out_file,out_fail) + stdout = 'stdout.out_' + str(self.t) + stdout_fail = 'stdout_fail.out_' + str(self.t) + copyfile(stdout,stdout_fail) + return jnp.atleast_1d(qflux_avg) + + + def write_terps_io(self): + t = str(self.t) # t is a time indicator here (could change to something stability-relevant + path_in_old = self.path_in + '.in' # I'm not sure which INPUT FILE changes would be made between calls, wouldn't most be fort.18 related? + path_in_new = self.path_in + '_' + t + '.in' + self.write_input(path_in_old,path_geo_old,path_in_new,path_geo_new) + + + def compute_fort18(self): + + print("Figure out how to do this directly from DESC equilibrium quantities!!") + + stdout = 'stdout.out_' + str(self.t) # Again, replace 't' with something stability-related + stderr = 'stderr.out_' + str(self.t) + fs = open('stdout.out_' + str(self.t),'w') + path_in = self.path_in + "_" + str(self.t) + '.in' + cmd = [self.vmec2terps_app, self.wout_file] + subprocess.run(cmd,stdout=fs) + fs.close() + # Need to ensure that the output file is in the correct directory + + def run_terps(self): + stdout = 'stdout.out_' + str(self.t) # Again, replace 't' with something stability-related + stderr = 'stderr.out_' + str(self.t) + fs = open('stdout.out_' + str(self.t),'w') + path_in = self.path_in + "_" + str(self.t) + '.in' + cmd = ['srun', '-N', '1', '-t', '00:45:00', '--ntasks-per-node=1', '--mem-per-cpu=100G', self.terps_app, '<', path_in] + subprocess.run(cmd,stdout=fs) + fs.close() + + def parse_terps_outfile(self): + + # Read fort.16 and search for growth rate + f = open(self.terps_outfile, 'r') + + growth_rate = [] + for line_number, line in enumerate(f, start=1): + index = line.find('GROWTH RATE') + if index != -1: + growth_rate.append(float(line.strip().split("=")[1])) + + if (len(growth_rate) > 1): + print("Currently capable of using only one growth rate. Exiting...") + exit() + elif (len(growth_rate) == 0): + print("Growth rate not found! Exiting...") + exit() + + self.growth_rate = growth_rate[0] + + # I'm not sure what's going on with these final two functions + def compute_terps_jvp(self,values,tangents): + + R_lmn, Z_lmn, L_lmn, i_l, c_l, p_l, Psi = values + primal_out = jnp.atleast_1d(0.0) + + n = len(values) + argnum = np.arange(0,n,1) + + jvp = FiniteDiffDerivative.compute_jvp(self.compute,argnum,tangents,*values,rel_step=1e-2) + + return (primal_out, jvp) + + def compute_terps_batch(self, values, axis): + numdiff = len(values[0]) + res = jnp.array([0.0]) + + for i in range(numdiff): + R_lmn = values[0][i] + Z_lmn = values[1][i] + L_lmn = values[2][i] + i_l = values[3][i] + p_l = values[4][i] + Psi = values[5][i] + + res = jnp.vstack([res,self.compute(R_lmn,Z_lmn,L_lmn,i_l,p_l,Psi)]) + + res = res[1:] + + + return res, axis[0] + + + + def write_input(self,path_in_temp,geo_temp,path_in,geo): + copyfile(path_in_temp,path_in) # Need a standard TERPS input + + terps_infile = os.path.join(path_in, "{}_N{}_family".format(eq_identifier, mode_family)) + f = open(terps_infile,"w") + + f.write(" {}\n".format(eq_identifier)) + f.write("C\n") + f.write("C MM NMIN NMAX MMS NSMIN NSMAX NPROCS INSOL\n") + f.write(" {:>2d} {:>3d} {:>3d} 55 -8 10 1 0\n".format(self.max_bozm, -self.max_bozn, self.max_bozn)) + f.write("C\n") + f.write("C TABLE OF FOURIER COEFFIENTS FOR BOOZER COORDINATES\n") + f.write("C EQUILIBRIUM SETTINGS ARE COMPUTED FROM FIT/VMEC\n") + f.write("C\n") + + boz_str_title = "C M= 0" + boz_str_neg = " 0" + boz_str_pos = " 1" + for _im in range(1,max_bozm+1): + if (_im >= 10): + boz_str_title += " "+str(_im)[1] + else: + boz_str_title += " "+str(_im) + boz_str_neg += " 1" + boz_str_pos += " 1" + + + boz_str_title += " N\n" + f.write(boz_str_title) + for _in in range(-max_bozn,max_bozn+1): + final_str_neg = boz_str_neg+"{:>3}\n".format(_in) + final_str_pos = boz_str_pos+"{:>3}\n".format(_in) + if _in < 0.0: + f.write(final_str_neg) + else: + f.write(final_str_pos) + + f.write("C\n") + f.write(" LLAMPR LVMTPR LMETPR LFOUPR\n") + f.write(" 0 0 0 0\n") + f.write(" LLHSPR LRHSPR LEIGPR LEFCPR\n") + f.write(" 9 9 1 1\n") + f.write(" LXYZPR LIOTPL LDW2PL LEFCPL\n") + f.write(" 0 1 1 1\n") + f.write(" LCURRF LMESHP LMESHV LITERS\n") + f.write(" 1 1 2 1\n") + f.write(" LXYZPL LEFPLS LEQVPL LPRESS\n") + f.write(" 1 1 0 0\n") + f.write("C\n") + f.write("C PVAC PARFAC QONAX QN DSVAC QVAC NOWALL\n") + f.write(" 1.0001e+00 0.0000e-00 0.6500e-00 0.0000e-00 1.0000e-00 1.0001e+00 -2\n") + f.write("\n") + f.write("C AWALL EWALL DWALL GWALL DRWAL DZWAL NPWALL\n") + f.write(" {:10.4e} 1.5000e+00 -1.0000e-00 5.2000e-00 -0.0000e-00 +0.0000e-00 2\n".format(awall)) + f.write("C\n") + f.write("C RPLMIN XPLO DELTAJP WCT CURFAC\n") + f.write(" 1.0000e-05 {:10.4e} {:10.4e} 6.6667e-01 1.0000e-00\n".format(xplo, deltajp)) + f.write("C\n") + f.write("C MODELK = {}\n".format(modelk)) + f.write("C\n") + f.write("C NUMBER OF EQUILIBRIUM FIELD PERIODS PER STABILITY PERIOD: NSTA = {}\n".format(nfp)) + f.write("C\n") + f.write("C TABLE OF FOURIER COEFFIENTS FOR STABILITY DISPLACEMENTS\n") + f.write("C\n") + f.write("C M= 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 N\n") + + #for _in in range(-max_moden,max_moden+1): + for _in in range(-8,11): + + in_family = False + if (_in % 2 == mode_family): # This needs to be modified for nfp != 2,3 + in_family = True + + if _in <= 0: + mode_str = " 0" + else: + if ((np.abs(_in) <= max_moden) and (in_family)): + mode_str = " 1" + else: + mode_str = " 0" + + for _im in range(1,55+1): + if ((_im <= max_modem) and (in_family)): + mode_str += " 1" + else: + mode_str += " 0" + + mode_str += "{:>3}\n".format(_in) + f.write(mode_str) + + f.write("C\n") + f.write("C NEV NITMAX AL0 EPSPAM IGREEN MPINIT\n") + f.write(" {} 4500 {:10.3e} 1.000E-04 0 0\n".format(nev,al0)) + f.write("C\n") + + f.close() From 19c90a9248f71dc6e823ebe1b9d7273839e4693a Mon Sep 17 00:00:00 2001 From: Michael Martin Date: Tue, 9 Apr 2024 09:17:29 -0400 Subject: [PATCH 02/49] initial work on terpsichore wrapper --- desc/objectives/_terpsichore.py | 41 ++++++++------------------------- 1 file changed, 9 insertions(+), 32 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 509d9f1b5c..f8021af1de 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -39,6 +39,8 @@ class TERPSICHORE(_Objective): COMPLETE THIS more stuff here +======= + """ # Used by GX objective, not sure where they're used yet @@ -49,32 +51,8 @@ class TERPSICHORE(_Objective): # Need to sure up the paths - def __init__( - self, - eq=None, - target=None, - weight=1, - grid=None, - name="TERPSICHORE", - path=os.getenv['terps_dir'], - path_in=os.getenv['terps_dir'], - bounds=None, - normalize=False, - normalize_target=False, - awall=1.5, - deltajp=4.e-2, - modelk=0, - al0=-5.e-2, - nev=1, - nfp=2, - xplo=1.e-6, - max_bozm=19, - max_bozn=14, - mode_family=0, - max_modem=55, - max_moden=8, - wout_filename="" # add something - ): + + def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", path=os.getenv['terps_dir'],path_in=os.getenv['terps_dir'],bounds=None,normalize=False,normalize_target=False, awall=1.5, deltajp=4.e-2, modelk=0, al0=-5.e-2, nev=1, nfp=2, xplo=1.e-6, max_bozm=19, max_bozn=14, mode_family=0, max_modem=55, max_moden=8): if target is None and bounds is None: target = 0 @@ -112,8 +90,6 @@ def __init__( self.wout_file = os.path.join(self.path, wout_filename) self.vmec2terps_app = os.path.join(self.path, "thea-vmec2terps.x") self.terps_app = os.path.join(self.path, "tpr_ap.x") - # self.t = t # -- Replace with some stability indicator? - self.terps_compute = core.Primitive("terps") self.terps_compute.def_impl(self.compute_impl) # What are the following two lines? @@ -143,8 +119,8 @@ def build(self, use_jit=False, verbose=1): NFP=eq.NFP, ) - self._dim_f = 1 # Presumbaly this should be 1? just a growth rate + timer = Timer() self._eq_keys = [ @@ -183,8 +159,9 @@ def build(self, use_jit=False, verbose=1): "profiles": self._profiles_eq, } - if self._normalize: - scales = compute_scaling_factors(eq) + + #if self._normalize: + #scales = compute_scaling_factors(eq) # I think only needed if we're not using the TERPS normalization? timer.stop("Precomputing transforms") @@ -197,7 +174,7 @@ def build(self, use_jit=False, verbose=1): def compute_impl(self, params, constants): - + # params, constants = self._parse_args(*args, **kwargs) if not np.any(params["R_lmn"]): return 0 From d139eac6a26347e0082eea81623486d91297ac55 Mon Sep 17 00:00:00 2001 From: Michael Martin Date: Mon, 15 Apr 2024 14:42:01 -0400 Subject: [PATCH 03/49] added terpsichore objective to __init__. Other small modifications --- desc/objectives/__init__.py | 1 + desc/objectives/_terpsichore.py | 96 ++++++++++++++++++++------------- 2 files changed, 59 insertions(+), 38 deletions(-) diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py index 33ff6e22fd..2a978304c2 100644 --- a/desc/objectives/__init__.py +++ b/desc/objectives/__init__.py @@ -72,3 +72,4 @@ FixThetaSFL, ) from .objective_funs import ObjectiveFunction +from ._terpsichore import TERPSICHORE diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index f8021af1de..1e66c838d6 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -51,8 +51,7 @@ class TERPSICHORE(_Objective): # Need to sure up the paths - - def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", path=os.getenv['terps_dir'],path_in=os.getenv['terps_dir'],bounds=None,normalize=False,normalize_target=False, awall=1.5, deltajp=4.e-2, modelk=0, al0=-5.e-2, nev=1, nfp=2, xplo=1.e-6, max_bozm=19, max_bozn=14, mode_family=0, max_modem=55, max_moden=8): + def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", path=None, bounds=None,normalize=False,normalize_target=False, awall=1.3, deltajp=4.e-2, modelk=0, al0=-5.e-2, nev=1, nfp=2, xplo=1.e-6, max_bozm=19, max_bozn=14, mode_family=0, max_modem=55, max_moden=8): if target is None and bounds is None: target = 0 @@ -84,12 +83,14 @@ def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", p units = "" self._callback_fmt = "Growth rate: {:10.3e} " + units self._print_value_fmt = "Growth rate: {:10.3e} " + units - + + + wout_filename = "wout_C640.nc" self.path = path - self.path_in = path_in self.wout_file = os.path.join(self.path, wout_filename) self.vmec2terps_app = os.path.join(self.path, "thea-vmec2terps.x") self.terps_app = os.path.join(self.path, "tpr_ap.x") + self.terps_compute = core.Primitive("terps") self.terps_compute.def_impl(self.compute_impl) # What are the following two lines? @@ -111,7 +112,9 @@ def build(self, use_jit=False, verbose=1): """ eq = self.things[0] + if self.grid is None: + ''' self.grid_eq = QuadratureGrid( L=eq.L_grid, M=eq.M_grid, @@ -119,7 +122,26 @@ def build(self, use_jit=False, verbose=1): NFP=eq.NFP, ) - self._dim_f = 1 # Presumbaly this should be 1? just a growth rate + #Construct flux-tube geometry + data = eq.compute('iota') + rhoa = eq.compute('rho') + iotad = data['iota'] + fi = interp1d(rhoa['rho'],iotad) + + rho = np.sqrt(self.psi) + iota = fi(rho) + zeta = np.linspace((-np.pi*self.npol-self.alpha)/np.abs(iota),(np.pi*self.npol-self.alpha)/np.abs(iota),2*self.nzgrid+1) + theta_sfl = iota/np.abs(iota)*self.alpha*np.ones(len(zeta)) + iota*zeta + + zeta_center = zeta[self.nzgrid] + rhoa = rho*np.ones(len(zeta)) + c = np.vstack([rhoa,theta_sfl,zeta]).T + coords = eq.compute_theta_coords(c,tol=1e-10,maxiter=50) + self.grid = Grid(coords,sort=False) + + ''' + + self._dim_f = 1 timer = Timer() @@ -130,6 +152,7 @@ def build(self, use_jit=False, verbose=1): "rho", "psi", ] + ''' self._field_line_keys = [ "|B|", "|grad(psi)|^2", "grad(|B|)", "grad(alpha)", "grad(psi)", "B", "grad(|B|)", "kappa", "B^theta", "B^zeta", "lambda_t", "lambda_z",'p_r', @@ -142,17 +165,17 @@ def build(self, use_jit=False, verbose=1): obj="desc.equilibrium.equilibrium.Equilibrium", has_axis=self.grid_eq.axis.size, ) - + ''' if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") #Need separate transforms and profiles for the equilibrium and flux-tube self.eq = eq - self._profiles = get_profiles(self._field_line_keys, obj=eq, grid=self.grid) - self._profiles_eq = get_profiles(self._eq_keys, obj=eq, grid=self.grid_eq) - self._transforms = get_transforms(self._field_line_keys, obj=eq, grid=self.grid) - self._transforms_eq = get_transforms(self._eq_keys, obj=eq, grid=self.grid_eq) + #self._profiles = get_profiles(self._field_line_keys, obj=eq, grid=self.grid) + self._profiles_eq = get_profiles(self._eq_keys, obj=eq, grid=self.grid) + #self._transforms = get_transforms(self._field_line_keys, obj=eq, grid=self.grid) + self._transforms_eq = get_transforms(self._eq_keys, obj=eq, grid=self.grid) self._constants = { "transforms": self._transforms_eq, @@ -172,6 +195,9 @@ def build(self, use_jit=False, verbose=1): # self._set_dimensions(eq) super().build(use_jit=use_jit, verbose=verbose) + def compute(self, params, constants=None): + + return self.terps_compute.bind(params,constants) def compute_impl(self, params, constants): @@ -187,7 +213,7 @@ def compute_impl(self, params, constants): self.write_terps_io() self.run_terps() - self.terps_outfile = os.path.join(self.path_in,'fort.16') # Let's change the name of this at some point + self.terps_outfile = os.path.join(self.path,'fort.16') # Let's change the name of this at some point self.parse_terps_outfile() @@ -201,32 +227,24 @@ def compute_impl(self, params, constants): return jnp.atleast_1d(qflux_avg) - def write_terps_io(self): - t = str(self.t) # t is a time indicator here (could change to something stability-relevant - path_in_old = self.path_in + '.in' # I'm not sure which INPUT FILE changes would be made between calls, wouldn't most be fort.18 related? - path_in_new = self.path_in + '_' + t + '.in' - self.write_input(path_in_old,path_geo_old,path_in_new,path_geo_new) - - def compute_fort18(self): print("Figure out how to do this directly from DESC equilibrium quantities!!") - stdout = 'stdout.out_' + str(self.t) # Again, replace 't' with something stability-related - stderr = 'stderr.out_' + str(self.t) - fs = open('stdout.out_' + str(self.t),'w') - path_in = self.path_in + "_" + str(self.t) + '.in' + #stdout = 'stdout.out_' + str(self.t) # Again, replace 't' with something stability-related + #stderr = 'stderr.out_' + str(self.t) + #fs = open('stdout.out_' + str(self.t),'w') + fs = open('stdout.vmec2terps','w') cmd = [self.vmec2terps_app, self.wout_file] subprocess.run(cmd,stdout=fs) fs.close() # Need to ensure that the output file is in the correct directory def run_terps(self): - stdout = 'stdout.out_' + str(self.t) # Again, replace 't' with something stability-related - stderr = 'stderr.out_' + str(self.t) - fs = open('stdout.out_' + str(self.t),'w') - path_in = self.path_in + "_" + str(self.t) + '.in' - cmd = ['srun', '-N', '1', '-t', '00:45:00', '--ntasks-per-node=1', '--mem-per-cpu=100G', self.terps_app, '<', path_in] + #stdout = 'stdout.out_' + str(self.t) # Again, replace 't' with something stability-related + #stderr = 'stderr.out_' + str(self.t) + fs = open('stdout.terps','w') + cmd = ['srun', '-N', '1', '-t', '00:45:00', '--ntasks-per-node=1', '--mem-per-cpu=100G', self.terps_app, '<', self.terps_infile] subprocess.run(cmd,stdout=fs) fs.close() @@ -258,7 +276,9 @@ def compute_terps_jvp(self,values,tangents): n = len(values) argnum = np.arange(0,n,1) - + + print(argnum) + exit() jvp = FiniteDiffDerivative.compute_jvp(self.compute,argnum,tangents,*values,rel_step=1e-2) return (primal_out, jvp) @@ -283,12 +303,12 @@ def compute_terps_batch(self, values, axis): return res, axis[0] + def write_terps_io(self): - def write_input(self,path_in_temp,geo_temp,path_in,geo): - copyfile(path_in_temp,path_in) # Need a standard TERPS input - - terps_infile = os.path.join(path_in, "{}_N{}_family".format(eq_identifier, mode_family)) - f = open(terps_infile,"w") + eq_identifier = "C624" + + self.terps_infile = os.path.join(self.path, "{}_N{}_family".format(eq_identifier, self.mode_family)) + f = open(self.terps_infile,"w") f.write(" {}\n".format(eq_identifier)) f.write("C\n") @@ -374,9 +394,9 @@ def write_input(self,path_in_temp,geo_temp,path_in,geo): mode_str += "{:>3}\n".format(_in) f.write(mode_str) - f.write("C\n") - f.write("C NEV NITMAX AL0 EPSPAM IGREEN MPINIT\n") - f.write(" {} 4500 {:10.3e} 1.000E-04 0 0\n".format(nev,al0)) - f.write("C\n") + f.write("C\n") + f.write("C NEV NITMAX AL0 EPSPAM IGREEN MPINIT\n") + f.write(" {} 4500 {:10.3e} 1.000E-04 0 0\n".format(nev,al0)) + f.write("C\n") - f.close() + f.close() From ec81ac1456752d932e42bdaf97fa8e309fd9c87d Mon Sep 17 00:00:00 2001 From: Michael Martin Date: Fri, 26 Apr 2024 13:50:52 -0400 Subject: [PATCH 04/49] TERPS running successfully when fixing all parameters --- desc/objectives/_terpsichore.py | 213 ++++++++++++++++++++------------ 1 file changed, 137 insertions(+), 76 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 1e66c838d6..f6b5f997ed 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -3,6 +3,7 @@ from scipy.interpolate import interp1d import os import time +import math from desc.backend import jnp,put from .objective_funs import ObjectiveFunction, _Objective @@ -46,12 +47,12 @@ class TERPSICHORE(_Objective): # Used by GX objective, not sure where they're used yet _scalar = True _linear = False - _units = "[]" # Outputting with dimensions or normalized? + _units = "[1/s]" _print_value_fmt = "Growth rate: {:10.3e} " # Need to sure up the paths - def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", path=None, bounds=None,normalize=False,normalize_target=False, awall=1.3, deltajp=4.e-2, modelk=0, al0=-5.e-2, nev=1, nfp=2, xplo=1.e-6, max_bozm=19, max_bozn=14, mode_family=0, max_modem=55, max_moden=8): + def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", path=None, bounds=None,normalize=False, submit_script_name="terps_job.submit", normalize_target=False, awall=1.3, deltajp=5.e-1, modelk=0, al0=-5.e-2, nev=1, nfp=2, xplo=1.e-6, max_bozm=19, max_bozn=14, mode_family=0, max_modem=55, min_moden=-8, max_moden=11): if target is None and bounds is None: target = 0 @@ -67,6 +68,7 @@ def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", p self.max_bozn = max_bozn self.mode_family = mode_family self.max_modem = max_modem + self.min_moden = min_moden self.max_moden = max_moden self.grid = grid @@ -85,8 +87,10 @@ def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", p self._print_value_fmt = "Growth rate: {:10.3e} " + units - wout_filename = "wout_C640.nc" + wout_filename = "wout_C640_MPOL-7_NTOR-6_NS-32.nc" self.path = path + self.submit_script_name = submit_script_name + self.submit_script_path = os.path.join(self.path, self.submit_script_name) self.wout_file = os.path.join(self.path, wout_filename) self.vmec2terps_app = os.path.join(self.path, "thea-vmec2terps.x") self.terps_app = os.path.join(self.path, "tpr_ap.x") @@ -112,35 +116,6 @@ def build(self, use_jit=False, verbose=1): """ eq = self.things[0] - - if self.grid is None: - ''' - self.grid_eq = QuadratureGrid( - L=eq.L_grid, - M=eq.M_grid, - N=eq.N_grid, - NFP=eq.NFP, - ) - - #Construct flux-tube geometry - data = eq.compute('iota') - rhoa = eq.compute('rho') - iotad = data['iota'] - fi = interp1d(rhoa['rho'],iotad) - - rho = np.sqrt(self.psi) - iota = fi(rho) - zeta = np.linspace((-np.pi*self.npol-self.alpha)/np.abs(iota),(np.pi*self.npol-self.alpha)/np.abs(iota),2*self.nzgrid+1) - theta_sfl = iota/np.abs(iota)*self.alpha*np.ones(len(zeta)) + iota*zeta - - zeta_center = zeta[self.nzgrid] - rhoa = rho*np.ones(len(zeta)) - c = np.vstack([rhoa,theta_sfl,zeta]).T - coords = eq.compute_theta_coords(c,tol=1e-10,maxiter=50) - self.grid = Grid(coords,sort=False) - - ''' - self._dim_f = 1 timer = Timer() @@ -168,8 +143,9 @@ def build(self, use_jit=False, verbose=1): ''' if verbose > 0: print("Precomputing transforms") - timer.start("Precomputing transforms") - + #timer.start("Precomputing transforms") + + ''' #Need separate transforms and profiles for the equilibrium and flux-tube self.eq = eq #self._profiles = get_profiles(self._field_line_keys, obj=eq, grid=self.grid) @@ -183,16 +159,13 @@ def build(self, use_jit=False, verbose=1): } - #if self._normalize: - #scales = compute_scaling_factors(eq) - # I think only needed if we're not using the TERPS normalization? - timer.stop("Precomputing transforms") if verbose > 1: timer.disp("Precomputing transforms") # self._check_dimensions() # self._set_dimensions(eq) + ''' super().build(use_jit=use_jit, verbose=verbose) def compute(self, params, constants=None): @@ -208,51 +181,138 @@ def compute_impl(self, params, constants): if constants is None: constants = self.constants - self.write_vmec() # Write VMEC file from DESC + #self.write_vmec() # Write VMEC file from DESC self.compute_fort18() self.write_terps_io() self.run_terps() - self.terps_outfile = os.path.join(self.path,'fort.16') # Let's change the name of this at some point self.parse_terps_outfile() - - # Is something analogous needed for TERPS?? - if qflux_avg > 20: - out_fail = 'fail_' + str(self.t) + '_.nc' - copyfile(out_file,out_fail) - stdout = 'stdout.out_' + str(self.t) - stdout_fail = 'stdout_fail.out_' + str(self.t) - copyfile(stdout,stdout_fail) - return jnp.atleast_1d(qflux_avg) + print("Growth rate = {}".format(self.growth_rate)) + + return jnp.atleast_1d(self.growth_rate) + def write_vmec(self): + + print("Figure out how to do this directly from DESC equilibrium quantities!!") + + + def compute_fort18(self): print("Figure out how to do this directly from DESC equilibrium quantities!!") - #stdout = 'stdout.out_' + str(self.t) # Again, replace 't' with something stability-related - #stderr = 'stderr.out_' + str(self.t) - #fs = open('stdout.out_' + str(self.t),'w') fs = open('stdout.vmec2terps','w') - cmd = [self.vmec2terps_app, self.wout_file] + head, tail = os.path.split(self.wout_file) + cmd = [self.vmec2terps_app, tail] subprocess.run(cmd,stdout=fs) fs.close() - # Need to ensure that the output file is in the correct directory + + def is_terps_complete(self, slurm_file, stop_time, running, runtime): + + if not os.path.exists(slurm_file): + return False + else: + f_slurm = open(slurm_file, 'r') + + if (running): + print("Stop time = {} seconds".format(stop_time)) + print("Current runtime = {} seconds".format(math.ceil(runtime))) + if (runtime > stop_time): + print("TERPS was unable to find a growth rate. Exiting...") + exit() + + terps_out_contents = f_slurm.read() + + if 'GROWTH RATE' in terps_out_contents: + f_slurm.close() + rm_cmd = ['rm', 'tpr16_dat_wall'] # There's probably a better way to handle this + subprocess.run(rm_cmd) + return True + else: + f_slurm.close() + return False def run_terps(self): - #stdout = 'stdout.out_' + str(self.t) # Again, replace 't' with something stability-related - #stderr = 'stderr.out_' + str(self.t) + + sleep_time = 10 # seconds + stop_time = 300 # seconds (kill the infinite loop if TERPS ran into an error and won't be printing growth rate) + fs = open('stdout.terps','w') - cmd = ['srun', '-N', '1', '-t', '00:45:00', '--ntasks-per-node=1', '--mem-per-cpu=100G', self.terps_app, '<', self.terps_infile] - subprocess.run(cmd,stdout=fs) + head, tail = os.path.split(self.terps_infile) + + # This could potentially launch a number of parallel TERPS jobs (probably at least want to run parallel jobs for N=0 and N=1 family) + + if not (os.path.exists(self.submit_script_path)): + f = open(self.submit_script_path,"w") + f.write("#!/bin/bash\n") + f.write("#SBATCH --job-name=terps # Job name\n") + f.write("#SBATCH --time=00:45:00 # Time limit hrs:min:sec\n") + f.write("#SBATCH --output=terps_%j.log # Standard output and error log\n") + f.write("#SBATCH --nodes=1\n") + f.write("#SBATCH --mem-per-cpu=100G\n") + f.write("#SBATCH --ntasks-per-node=1\n") + f.write("#SBATCH --partition=stellar-debug\n") + f.write("\n") + f.write("srun {} < {}\n".format(self.terps_app,tail)) + f.close() + + print("Need a command to remove remnants of pasts TERPS runs (tpr16_dat_wall in particular) or move them to a new directory") + + if (os.path.exists(os.path.join(self.path, "tpr16_dat_wall"))): + rm_cmd = ['rm', 'tpr16_dat_wall'] # There's probably a better way to handle this + subprocess.run(rm_cmd) + + cmd = ['sbatch', self.submit_script_path] + terps_process = subprocess.run(cmd,stdout=subprocess.PIPE) + out_text = terps_process.stdout.decode('utf-8') + fs.write(out_text) + jobID = out_text.split()[-1] + slurm_file = os.path.join(self.path,"terps_{}.log".format(jobID)) + + running = False + runtime = 0.0 + tic = time.perf_counter() + while not self.is_terps_complete(slurm_file, stop_time, running, runtime): + if not running: + check_status_cmd = ['squeue', '-j', jobID, '--format="%T"'] + check_status = subprocess.run(check_status_cmd, stdout=subprocess.PIPE) + status_text = check_status.stdout.decode('utf-8') + status = status_text.split()[1].replace('"','') + if status == 'RUNNING': + running = True + tic = time.perf_counter() + print("TERPS has started running") + elif status == 'PENDING': + print("TERPS is still in the queue") + else: + print(status) + exit() + + else: + print("Growth rate not found. Checking again in {} seconds".format(sleep_time)) + + time.sleep(sleep_time) + toc = time.perf_counter() + runtime = toc-tic + + print("Found growth rate!") + fs.close() + # Need a command here to wait until all TERPS runs are complete if doing some form of parallel execution + + def parse_terps_outfile(self): # Read fort.16 and search for growth rate - f = open(self.terps_outfile, 'r') - + if (os.path.exists(self.terps_outfile)): + f = open(self.terps_outfile, 'r') + else: + print("TERPS fort.16 output file not found! Exiting...") + exit() + growth_rate = [] for line_number, line in enumerate(f, start=1): index = line.find('GROWTH RATE') @@ -277,8 +337,6 @@ def compute_terps_jvp(self,values,tangents): n = len(values) argnum = np.arange(0,n,1) - print(argnum) - exit() jvp = FiniteDiffDerivative.compute_jvp(self.compute,argnum,tangents,*values,rel_step=1e-2) return (primal_out, jvp) @@ -305,8 +363,8 @@ def compute_terps_batch(self, values, axis): def write_terps_io(self): - eq_identifier = "C624" - + eq_identifier = "C640" + self.terps_infile = os.path.join(self.path, "{}_N{}_family".format(eq_identifier, self.mode_family)) f = open(self.terps_infile,"w") @@ -322,7 +380,8 @@ def write_terps_io(self): boz_str_title = "C M= 0" boz_str_neg = " 0" boz_str_pos = " 1" - for _im in range(1,max_bozm+1): + #for _im in range(1,self.max_bozm+1): + for _im in range(1,37): if (_im >= 10): boz_str_title += " "+str(_im)[1] else: @@ -333,7 +392,7 @@ def write_terps_io(self): boz_str_title += " N\n" f.write(boz_str_title) - for _in in range(-max_bozn,max_bozn+1): + for _in in range(-self.max_bozn,self.max_bozn+1): final_str_neg = boz_str_neg+"{:>3}\n".format(_in) final_str_pos = boz_str_pos+"{:>3}\n".format(_in) if _in < 0.0: @@ -357,37 +416,39 @@ def write_terps_io(self): f.write(" 1.0001e+00 0.0000e-00 0.6500e-00 0.0000e-00 1.0000e-00 1.0001e+00 -2\n") f.write("\n") f.write("C AWALL EWALL DWALL GWALL DRWAL DZWAL NPWALL\n") - f.write(" {:10.4e} 1.5000e+00 -1.0000e-00 5.2000e-00 -0.0000e-00 +0.0000e-00 2\n".format(awall)) + f.write(" {:10.4e} 1.5000e+00 -1.0000e-00 5.2000e-00 -0.0000e-00 +0.0000e-00 2\n".format(self.awall)) f.write("C\n") f.write("C RPLMIN XPLO DELTAJP WCT CURFAC\n") - f.write(" 1.0000e-05 {:10.4e} {:10.4e} 6.6667e-01 1.0000e-00\n".format(xplo, deltajp)) + f.write(" 1.0000e-05 {:10.4e} {:10.4e} 6.6667e-01 1.0000e-00\n".format(self.xplo, self.deltajp)) f.write("C\n") - f.write("C MODELK = {}\n".format(modelk)) + f.write("C MODELK = {}\n".format(self.modelk)) f.write("C\n") - f.write("C NUMBER OF EQUILIBRIUM FIELD PERIODS PER STABILITY PERIOD: NSTA = {}\n".format(nfp)) + f.write("C NUMBER OF EQUILIBRIUM FIELD PERIODS PER STABILITY PERIOD: NSTA = {}\n".format(self.nfp)) f.write("C\n") f.write("C TABLE OF FOURIER COEFFIENTS FOR STABILITY DISPLACEMENTS\n") f.write("C\n") f.write("C M= 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 N\n") - #for _in in range(-max_moden,max_moden+1): for _in in range(-8,11): in_family = False - if (_in % 2 == mode_family): # This needs to be modified for nfp != 2,3 + if (_in % 2 == self.mode_family): # This needs to be modified for nfp != 2,3 in_family = True if _in <= 0: mode_str = " 0" else: - if ((np.abs(_in) <= max_moden) and (in_family)): + if ((_in <= self.max_moden) and (_in >= self.min_moden) and (in_family)): mode_str = " 1" else: mode_str = " 0" for _im in range(1,55+1): - if ((_im <= max_modem) and (in_family)): - mode_str += " 1" + if ((_im <= self.max_modem) and (in_family)): + if ((_in <= self.max_moden) and (_in >= self.min_moden)): + mode_str += " 1" + else: + mode_str += " 0" else: mode_str += " 0" @@ -396,7 +457,7 @@ def write_terps_io(self): f.write("C\n") f.write("C NEV NITMAX AL0 EPSPAM IGREEN MPINIT\n") - f.write(" {} 4500 {:10.3e} 1.000E-04 0 0\n".format(nev,al0)) + f.write(" {} 4500 {:10.3e} 1.000E-04 0 0\n".format(self.nev,self.al0)) f.write("C\n") f.close() From c48c240e7d20c97e797b6880cdbad063180ae46c Mon Sep 17 00:00:00 2001 From: Michael Martin Date: Fri, 26 Apr 2024 15:31:12 -0400 Subject: [PATCH 05/49] merge w/ master --- desc/objectives/_terpsichore.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index f6b5f997ed..09a7eb4e7a 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -196,7 +196,8 @@ def compute_impl(self, params, constants): def write_vmec(self): print("Figure out how to do this directly from DESC equilibrium quantities!!") - + + #VMECIO.save(eq, "path/to/wout.nc", surfs=32) def compute_fort18(self): @@ -260,7 +261,7 @@ def run_terps(self): print("Need a command to remove remnants of pasts TERPS runs (tpr16_dat_wall in particular) or move them to a new directory") - if (os.path.exists(os.path.join(self.path, "tpr16_dat_wall"))): + if (os.path.exists(os.path.join(self.path, 'tpr16_dat_wall'))): rm_cmd = ['rm', 'tpr16_dat_wall'] # There's probably a better way to handle this subprocess.run(rm_cmd) From 50f044666a62cc9f11a55eeae9f36db0f76e100e Mon Sep 17 00:00:00 2001 From: Michael Martin Date: Thu, 2 May 2024 10:45:41 -0400 Subject: [PATCH 06/49] modified run_terps() and write_terps_io() to account for new dynamically-allocated resolution input variables --- desc/objectives/_terpsichore.py | 208 +++++++++++++++++++++++--------- 1 file changed, 148 insertions(+), 60 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 09a7eb4e7a..a461e47992 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -14,6 +14,7 @@ from desc.compute import compute as compute_fun from desc.compute import get_profiles, get_transforms, get_params from desc.compute.utils import dot +from desc.vmec.VMECIO import save from scipy.constants import mu_0, pi from desc.grid import LinearGrid, Grid, ConcentricGrid, QuadratureGrid @@ -52,11 +53,12 @@ class TERPSICHORE(_Objective): # Need to sure up the paths - def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", path=None, bounds=None,normalize=False, submit_script_name="terps_job.submit", normalize_target=False, awall=1.3, deltajp=5.e-1, modelk=0, al0=-5.e-2, nev=1, nfp=2, xplo=1.e-6, max_bozm=19, max_bozn=14, mode_family=0, max_modem=55, min_moden=-8, max_moden=11): + def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", wout_filename="wout_default.nc", path=None, bounds=None,normalize=False, submit_script_name="terps_job.submit", normalize_target=False, awall=2.0, deltajp=1.e-2, modelk=0, al0=-5.e-1, nev=1, nfp=2, xplo=1.e-6, max_bozm=19, max_bozn=14, mode_family=0, max_modem=55, min_moden=-8, max_moden=11, nsurf=128): if target is None and bounds is None: target = 0 self.eq = eq + self.nsurf = nsurf self.awall = awall self.deltajp = deltajp self.modelk = modelk @@ -70,6 +72,10 @@ def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", p self.max_modem = max_modem self.min_moden = min_moden self.max_moden = max_moden + self.lssl = 1 # LSSL and LSSD depend on the specified resolution and the required result is given after running the code (once for each variable) + self.lssd = 1 + self.lssl_repeat = True + self.lssd_repeat = True self.grid = grid super().__init__( @@ -87,10 +93,11 @@ def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", p self._print_value_fmt = "Growth rate: {:10.3e} " + units - wout_filename = "wout_C640_MPOL-7_NTOR-6_NS-32.nc" self.path = path + self.terps_stdout = os.path.join(self.path,'stdout.terps') self.submit_script_name = submit_script_name self.submit_script_path = os.path.join(self.path, self.submit_script_name) + self.wout_filename = wout_filename self.wout_file = os.path.join(self.path, wout_filename) self.vmec2terps_app = os.path.join(self.path, "thea-vmec2terps.x") self.terps_app = os.path.join(self.path, "tpr_ap.x") @@ -181,9 +188,10 @@ def compute_impl(self, params, constants): if constants is None: constants = self.constants - #self.write_vmec() # Write VMEC file from DESC + self.write_vmec() # Write VMEC file from DESC equilibrium self.compute_fort18() self.write_terps_io() + exit() self.run_terps() self.terps_outfile = os.path.join(self.path,'fort.16') # Let's change the name of this at some point self.parse_terps_outfile() @@ -195,9 +203,7 @@ def compute_impl(self, params, constants): def write_vmec(self): - print("Figure out how to do this directly from DESC equilibrium quantities!!") - - #VMECIO.save(eq, "path/to/wout.nc", surfs=32) + desc.vmec.VMECIO.save(eq, self.wout_file, surfs=self.nsurf) def compute_fort18(self): @@ -210,6 +216,7 @@ def compute_fort18(self): subprocess.run(cmd,stdout=fs) fs.close() + ''' def is_terps_complete(self, slurm_file, stop_time, running, runtime): if not os.path.exists(slurm_file): @@ -234,70 +241,138 @@ def is_terps_complete(self, slurm_file, stop_time, running, runtime): else: f_slurm.close() return False + ''' + + def is_terps_complete(self, stop_time, runtime): + + if not os.path.exists(self.terps_stdout): + return False + else: + f_slurm = open(self.terps_stdout, 'r') + + print("Current runtime = {} seconds".format(math.ceil(runtime))) + if (runtime > stop_time): + print("TERPS was unable to find a growth rate. Exiting...") + exit() + + terps_out_contents = f_slurm.read() + + if 'GROWTH RATE' in terps_out_contents: + f_slurm.close() + rm_cmd = ['rm', 'tpr16_dat_wall', 'tpr16_dat_pvi'] # There's probably a better way to handle this + subprocess.run(rm_cmd) + return True + + elif 'PARAMETER LSSL' in terps_out_contents: + line = f_slurm.readline() + while line: + if 'PARAMETER LSSL' in line: + self.lssl = int(line.split("TO:")[1]) + break + line = f_slurm.readline() + write_terps_io() # Rewrite the input file with suggested value of LSSL and re-run + return True + + elif 'PARAMETER LSSD' in terps_out_contents: + line = f_slurm.readline() + while line: + if 'PARAMETER LSSD' in line: + self.lssd = int(line.split("TO:")[1]) + break + line = f_slurm.readline() + write_terps_io() # Rewrite the input file with suggested value of LSSD and re-run + return True + + else: + f_slurm.close() + return False + def run_terps(self): - sleep_time = 10 # seconds - stop_time = 300 # seconds (kill the infinite loop if TERPS ran into an error and won't be printing growth rate) + sleep_time = 1 # seconds + stop_time = 60 # seconds (kill the infinite loop if TERPS ran into an error and won't be printing growth rate) - fs = open('stdout.terps','w') head, tail = os.path.split(self.terps_infile) - # This could potentially launch a number of parallel TERPS jobs (probably at least want to run parallel jobs for N=0 and N=1 family) + rm_cmd = ['rm', 'tpr16_dat_wall', 'tpr16_dat_pvi'] # There's probably a better way to handle this + subprocess.run(rm_cmd) - if not (os.path.exists(self.submit_script_path)): - f = open(self.submit_script_path,"w") - f.write("#!/bin/bash\n") - f.write("#SBATCH --job-name=terps # Job name\n") - f.write("#SBATCH --time=00:45:00 # Time limit hrs:min:sec\n") - f.write("#SBATCH --output=terps_%j.log # Standard output and error log\n") - f.write("#SBATCH --nodes=1\n") - f.write("#SBATCH --mem-per-cpu=100G\n") - f.write("#SBATCH --ntasks-per-node=1\n") - f.write("#SBATCH --partition=stellar-debug\n") - f.write("\n") - f.write("srun {} < {}\n".format(self.terps_app,tail)) - f.close() - - print("Need a command to remove remnants of pasts TERPS runs (tpr16_dat_wall in particular) or move them to a new directory") - - if (os.path.exists(os.path.join(self.path, 'tpr16_dat_wall'))): - rm_cmd = ['rm', 'tpr16_dat_wall'] # There's probably a better way to handle this - subprocess.run(rm_cmd) - - cmd = ['sbatch', self.submit_script_path] - terps_process = subprocess.run(cmd,stdout=subprocess.PIPE) - out_text = terps_process.stdout.decode('utf-8') - fs.write(out_text) - jobID = out_text.split()[-1] - slurm_file = os.path.join(self.path,"terps_{}.log".format(jobID)) - - running = False - runtime = 0.0 - tic = time.perf_counter() - while not self.is_terps_complete(slurm_file, stop_time, running, runtime): - if not running: - check_status_cmd = ['squeue', '-j', jobID, '--format="%T"'] - check_status = subprocess.run(check_status_cmd, stdout=subprocess.PIPE) - status_text = check_status.stdout.decode('utf-8') - status = status_text.split()[1].replace('"','') - if status == 'RUNNING': - running = True - tic = time.perf_counter() - print("TERPS has started running") - elif status == 'PENDING': - print("TERPS is still in the queue") - else: - print(status) - exit() - - else: - print("Growth rate not found. Checking again in {} seconds".format(sleep_time)) + cmd = ['srun', self.terps_app, '<', tail, '>', self.terps_stdout] + terps_process = subprocess.run(cmd, stdout=subprocess.PIPE) + runtime = 0.0 + while not self.is_terps_complete(self.terps_stdout, stop_time, runtime): + tic = time.perf_counter() time.sleep(sleep_time) toc = time.perf_counter() runtime = toc-tic + if (self.lssl_repeat): + self.lssl_repeat = False + run_terps() + + elif (self.lssd_repeat): + self.lssd_repeat = False + run_terps() + + exit() + + ''' + if (submit_scipt): + # This could potentially launch a number of parallel TERPS jobs (probably at least want to run parallel jobs for N=0 and N=1 family) + fs = open('stdout.terps','w') + + if not (os.path.exists(self.submit_script_path)): + f = open(self.submit_script_path,"w") + f.write("#!/bin/bash\n") + f.write("#SBATCH --job-name=terps # Job name\n") + f.write("#SBATCH --time=00:45:00 # Time limit hrs:min:sec\n") + f.write("#SBATCH --output=terps_%j.log # Standard output and error log\n") + f.write("#SBATCH --nodes=1\n") + f.write("#SBATCH --mem-per-cpu=100G\n") + f.write("#SBATCH --ntasks-per-node=1\n") + f.write("#SBATCH --partition=stellar-debug\n") + f.write("\n") + f.write("srun {} < {}\n".format(self.terps_app,tail)) + f.close() + + print("Need a command to remove remnants of pasts TERPS runs (tpr16_dat_wall in particular) or move them to a new directory") + + + cmd = ['sbatch', self.submit_script_path] + terps_process = subprocess.run(cmd,stdout=subprocess.PIPE) + out_text = terps_process.stdout.decode('utf-8') + fs.write(out_text) + jobID = out_text.split()[-1] + slurm_file = os.path.join(self.path,"terps_{}.log".format(jobID)) + + running = False + runtime = 0.0 + tic = time.perf_counter() + while not self.is_terps_complete(slurm_file, stop_time, running, runtime): + if not running: + check_status_cmd = ['squeue', '-j', jobID, '--format="%T"'] + check_status = subprocess.run(check_status_cmd, stdout=subprocess.PIPE) + status_text = check_status.stdout.decode('utf-8') + status = status_text.split()[1].replace('"','') + if status == 'RUNNING': + running = True + tic = time.perf_counter() + print("TERPS has started running") + elif status == 'PENDING': + print("TERPS is still in the queue") + else: + print(status) + exit() + + else: + print("Growth rate not found. Checking again in {} seconds".format(sleep_time)) + + time.sleep(sleep_time) + toc = time.perf_counter() + runtime = toc-tic + ''' print("Found growth rate!") fs.close() @@ -363,9 +438,19 @@ def compute_terps_batch(self, values, axis): def write_terps_io(self): - + eq_identifier = "C640" - + ivac = self.nsurf // 4 + if (max_moden > 8) or (max_modem > 16): + nj = 150 + nk = 150 + elif (max_moden > 4) or (max_modem > 8): + nj = 100 + nk = 100 + else: + nj = 50 + nk = 50 + self.terps_infile = os.path.join(self.path, "{}_N{}_family".format(eq_identifier, self.mode_family)) f = open(self.terps_infile,"w") @@ -374,6 +459,9 @@ def write_terps_io(self): f.write("C MM NMIN NMAX MMS NSMIN NSMAX NPROCS INSOL\n") f.write(" {:>2d} {:>3d} {:>3d} 55 -8 10 1 0\n".format(self.max_bozm, -self.max_bozn, self.max_bozn)) f.write("C\n") + f.write("C NJ NK IVAC LSSL LSSD MMAXDF NMAXDF\n") + f.write(" {:>3d} {:>3d} {:>3d} {:>4d} {:>4d} 120 64\n".format(nj, nk, ivac, self.lssl, self.lssd)) # Not exactly sure how to set LSSL and LSSD..usually it tells you to increase if not high enough (does this affect runtime?) + f.write("C\n") f.write("C TABLE OF FOURIER COEFFIENTS FOR BOOZER COORDINATES\n") f.write("C EQUILIBRIUM SETTINGS ARE COMPUTED FROM FIT/VMEC\n") f.write("C\n") From 6611b24f51b4f64eb5674b65ffe02c00c9cb646e Mon Sep 17 00:00:00 2001 From: Michael Martin Date: Fri, 3 May 2024 17:31:54 -0400 Subject: [PATCH 07/49] terps finite differences working with old version of terps (need to compile for resolution changes) --- desc/objectives/_terpsichore.py | 228 +++++++++----------------------- 1 file changed, 66 insertions(+), 162 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index a461e47992..7716e1c8d0 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -14,7 +14,6 @@ from desc.compute import compute as compute_fun from desc.compute import get_profiles, get_transforms, get_params from desc.compute.utils import dot -from desc.vmec.VMECIO import save from scipy.constants import mu_0, pi from desc.grid import LinearGrid, Grid, ConcentricGrid, QuadratureGrid @@ -58,7 +57,7 @@ def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", w if target is None and bounds is None: target = 0 self.eq = eq - self.nsurf = nsurf + self.nsurf = nsurf # No functionality unless using dynamic allocation TERPS self.awall = awall self.deltajp = deltajp self.modelk = modelk @@ -72,8 +71,9 @@ def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", w self.max_modem = max_modem self.min_moden = min_moden self.max_moden = max_moden - self.lssl = 1 # LSSL and LSSD depend on the specified resolution and the required result is given after running the code (once for each variable) - self.lssd = 1 + # lssl, lssd, lssl_repeat, and lssd_repeat are only used with dynamic allocation TERPS + self.lssl = 200 # LSSL and LSSD depend on the specified resolution and the required result is given after running the code (once for each variable) + self.lssd = 100 self.lssl_repeat = True self.lssd_repeat = True self.grid = grid @@ -127,58 +127,14 @@ def build(self, use_jit=False, verbose=1): timer = Timer() - self._eq_keys = [ - "iota", - "iota_r", - "a", - "rho", - "psi", - ] - ''' - self._field_line_keys = [ - "|B|", "|grad(psi)|^2", "grad(|B|)", "grad(alpha)", "grad(psi)", - "B", "grad(|B|)", "kappa", "B^theta", "B^zeta", "lambda_t", "lambda_z",'p_r', - "lambda_r", "lambda", "g^rr", "g^rt", "g^rz", "g^tz", "g^tt", "g^zz", - "e^rho", "e^theta", "e^zeta" - ] - - self._args = get_params( - self._field_line_keys, - obj="desc.equilibrium.equilibrium.Equilibrium", - has_axis=self.grid_eq.axis.size, - ) - ''' - if verbose > 0: - print("Precomputing transforms") - #timer.start("Precomputing transforms") - - ''' - #Need separate transforms and profiles for the equilibrium and flux-tube - self.eq = eq - #self._profiles = get_profiles(self._field_line_keys, obj=eq, grid=self.grid) - self._profiles_eq = get_profiles(self._eq_keys, obj=eq, grid=self.grid) - #self._transforms = get_transforms(self._field_line_keys, obj=eq, grid=self.grid) - self._transforms_eq = get_transforms(self._eq_keys, obj=eq, grid=self.grid) - - self._constants = { - "transforms": self._transforms_eq, - "profiles": self._profiles_eq, - } - - - timer.stop("Precomputing transforms") - if verbose > 1: - timer.disp("Precomputing transforms") - -# self._check_dimensions() -# self._set_dimensions(eq) - ''' super().build(use_jit=use_jit, verbose=verbose) + def compute(self, params, constants=None): return self.terps_compute.bind(params,constants) - + + def compute_impl(self, params, constants): # params, constants = self._parse_args(*args, **kwargs) @@ -191,7 +147,6 @@ def compute_impl(self, params, constants): self.write_vmec() # Write VMEC file from DESC equilibrium self.compute_fort18() self.write_terps_io() - exit() self.run_terps() self.terps_outfile = os.path.join(self.path,'fort.16') # Let's change the name of this at some point self.parse_terps_outfile() @@ -203,52 +158,35 @@ def compute_impl(self, params, constants): def write_vmec(self): - desc.vmec.VMECIO.save(eq, self.wout_file, surfs=self.nsurf) - - + from desc.vmec import VMECIO + VMECIO.save(self.eq, self.wout_file, surfs=self.nsurf) + def compute_fort18(self): print("Figure out how to do this directly from DESC equilibrium quantities!!") - fs = open('stdout.vmec2terps','w') head, tail = os.path.split(self.wout_file) cmd = [self.vmec2terps_app, tail] - subprocess.run(cmd,stdout=fs) - fs.close() + subprocess.run(cmd,stdout=subprocess.PIPE, check=True) - ''' - def is_terps_complete(self, slurm_file, stop_time, running, runtime): - if not os.path.exists(slurm_file): - return False - else: - f_slurm = open(slurm_file, 'r') - - if (running): - print("Stop time = {} seconds".format(stop_time)) - print("Current runtime = {} seconds".format(math.ceil(runtime))) - if (runtime > stop_time): - print("TERPS was unable to find a growth rate. Exiting...") - exit() - - terps_out_contents = f_slurm.read() + def remove_terps_files(self): - if 'GROWTH RATE' in terps_out_contents: - f_slurm.close() + if (os.path.exists(os.path.join(self.path, 'tpr16_dat_wall'))): rm_cmd = ['rm', 'tpr16_dat_wall'] # There's probably a better way to handle this subprocess.run(rm_cmd) - return True - else: - f_slurm.close() - return False - ''' - def is_terps_complete(self, stop_time, runtime): + if (os.path.exists(os.path.join(self.path, 'tpr16_dat_pvi'))): + rm_cmd = ['rm', 'tpr16_dat_pvi'] # There's probably a better way to handle this + subprocess.run(rm_cmd) + + + def is_terps_complete(self, stop_time, runtime, terps_subprocess): - if not os.path.exists(self.terps_stdout): - return False - else: + if (os.path.exists(self.terps_stdout)): f_slurm = open(self.terps_stdout, 'r') + else: + return False print("Current runtime = {} seconds".format(math.ceil(runtime))) if (runtime > stop_time): @@ -256,36 +194,47 @@ def is_terps_complete(self, stop_time, runtime): exit() terps_out_contents = f_slurm.read() - + if 'GROWTH RATE' in terps_out_contents: + if terps_subprocess.returncode == None: + terps_subprocess.terminate() f_slurm.close() - rm_cmd = ['rm', 'tpr16_dat_wall', 'tpr16_dat_pvi'] # There's probably a better way to handle this - subprocess.run(rm_cmd) + self.remove_terps_files() return True - + + else: + f_slurm.close() + return False + + # Uncomment the following block when dynamic allocation is working correctly + ''' elif 'PARAMETER LSSL' in terps_out_contents: + if terps_subprocess.returncode == None: + terps_subprocess.terminate() + f_slurm.seek(0) line = f_slurm.readline() while line: if 'PARAMETER LSSL' in line: self.lssl = int(line.split("TO:")[1]) break line = f_slurm.readline() - write_terps_io() # Rewrite the input file with suggested value of LSSL and re-run + self.write_terps_io() # Rewrite the input file with suggested value of LSSL and re-run return True elif 'PARAMETER LSSD' in terps_out_contents: + if terps_subprocess.returncode == None: + terps_subprocess.terminate() + f_slurm.seek(0) line = f_slurm.readline() while line: if 'PARAMETER LSSD' in line: self.lssd = int(line.split("TO:")[1]) break line = f_slurm.readline() - write_terps_io() # Rewrite the input file with suggested value of LSSD and re-run + self.write_terps_io() # Rewrite the input file with suggested value of LSSD and re-run return True + ''' - else: - f_slurm.close() - return False def run_terps(self): @@ -293,86 +242,33 @@ def run_terps(self): sleep_time = 1 # seconds stop_time = 60 # seconds (kill the infinite loop if TERPS ran into an error and won't be printing growth rate) + fs_error = open('error.terps','w') + fs = open('stdout.terps','w') head, tail = os.path.split(self.terps_infile) + + self.remove_terps_files() - rm_cmd = ['rm', 'tpr16_dat_wall', 'tpr16_dat_pvi'] # There's probably a better way to handle this - subprocess.run(rm_cmd) - - cmd = ['srun', self.terps_app, '<', tail, '>', self.terps_stdout] - terps_process = subprocess.run(cmd, stdout=subprocess.PIPE) + cmd = ['srun', '-n', '1', '-t', '00:01:00', self.terps_app] + terps_subprocess = subprocess.run(cmd, stdin=open(self.terps_infile,'r'), stdout=fs, stderr=fs_error) runtime = 0.0 - while not self.is_terps_complete(self.terps_stdout, stop_time, runtime): - tic = time.perf_counter() + tic = time.perf_counter() + while not self.is_terps_complete(stop_time, runtime, terps_subprocess): time.sleep(sleep_time) toc = time.perf_counter() runtime = toc-tic + # Uncomment the following block when dynamic allocation TERPS is working + ''' if (self.lssl_repeat): self.lssl_repeat = False - run_terps() + self.run_terps() elif (self.lssd_repeat): self.lssd_repeat = False - run_terps() + self.run_terps() + ''' - exit() - - ''' - if (submit_scipt): - # This could potentially launch a number of parallel TERPS jobs (probably at least want to run parallel jobs for N=0 and N=1 family) - fs = open('stdout.terps','w') - - if not (os.path.exists(self.submit_script_path)): - f = open(self.submit_script_path,"w") - f.write("#!/bin/bash\n") - f.write("#SBATCH --job-name=terps # Job name\n") - f.write("#SBATCH --time=00:45:00 # Time limit hrs:min:sec\n") - f.write("#SBATCH --output=terps_%j.log # Standard output and error log\n") - f.write("#SBATCH --nodes=1\n") - f.write("#SBATCH --mem-per-cpu=100G\n") - f.write("#SBATCH --ntasks-per-node=1\n") - f.write("#SBATCH --partition=stellar-debug\n") - f.write("\n") - f.write("srun {} < {}\n".format(self.terps_app,tail)) - f.close() - - print("Need a command to remove remnants of pasts TERPS runs (tpr16_dat_wall in particular) or move them to a new directory") - - - cmd = ['sbatch', self.submit_script_path] - terps_process = subprocess.run(cmd,stdout=subprocess.PIPE) - out_text = terps_process.stdout.decode('utf-8') - fs.write(out_text) - jobID = out_text.split()[-1] - slurm_file = os.path.join(self.path,"terps_{}.log".format(jobID)) - - running = False - runtime = 0.0 - tic = time.perf_counter() - while not self.is_terps_complete(slurm_file, stop_time, running, runtime): - if not running: - check_status_cmd = ['squeue', '-j', jobID, '--format="%T"'] - check_status = subprocess.run(check_status_cmd, stdout=subprocess.PIPE) - status_text = check_status.stdout.decode('utf-8') - status = status_text.split()[1].replace('"','') - if status == 'RUNNING': - running = True - tic = time.perf_counter() - print("TERPS has started running") - elif status == 'PENDING': - print("TERPS is still in the queue") - else: - print(status) - exit() - - else: - print("Growth rate not found. Checking again in {} seconds".format(sleep_time)) - - time.sleep(sleep_time) - toc = time.perf_counter() - runtime = toc-tic - ''' print("Found growth rate!") fs.close() @@ -440,17 +336,21 @@ def compute_terps_batch(self, values, axis): def write_terps_io(self): eq_identifier = "C640" + + # Uncomment the following block when dynamic allocation TERPS is working + ''' ivac = self.nsurf // 4 - if (max_moden > 8) or (max_modem > 16): + if (self.max_moden > 8) or (self.max_modem > 16): nj = 150 nk = 150 - elif (max_moden > 4) or (max_modem > 8): + elif (self.max_moden > 4) or (self.max_modem > 8): nj = 100 nk = 100 else: nj = 50 nk = 50 - + ''' + self.terps_infile = os.path.join(self.path, "{}_N{}_family".format(eq_identifier, self.mode_family)) f = open(self.terps_infile,"w") @@ -458,9 +358,13 @@ def write_terps_io(self): f.write("C\n") f.write("C MM NMIN NMAX MMS NSMIN NSMAX NPROCS INSOL\n") f.write(" {:>2d} {:>3d} {:>3d} 55 -8 10 1 0\n".format(self.max_bozm, -self.max_bozn, self.max_bozn)) + + # Uncomment the following block when dynamic allocation TERPS is working + ''' f.write("C\n") f.write("C NJ NK IVAC LSSL LSSD MMAXDF NMAXDF\n") f.write(" {:>3d} {:>3d} {:>3d} {:>4d} {:>4d} 120 64\n".format(nj, nk, ivac, self.lssl, self.lssd)) # Not exactly sure how to set LSSL and LSSD..usually it tells you to increase if not high enough (does this affect runtime?) + ''' f.write("C\n") f.write("C TABLE OF FOURIER COEFFIENTS FOR BOOZER COORDINATES\n") f.write("C EQUILIBRIUM SETTINGS ARE COMPUTED FROM FIT/VMEC\n") From ae872145fd9b260cf298e6b78b025395c1db9692 Mon Sep 17 00:00:00 2001 From: Michael Martin Date: Mon, 6 May 2024 16:52:45 -0400 Subject: [PATCH 08/49] not currently recognizing 'seff' in subprocess... --- desc/objectives/_terpsichore.py | 43 ++++++++++++++++++++++++++++----- 1 file changed, 37 insertions(+), 6 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 7716e1c8d0..da047a25ea 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -4,6 +4,7 @@ import os import time import math +import multiprocessing from desc.backend import jnp,put from .objective_funs import ObjectiveFunction, _Objective @@ -52,11 +53,12 @@ class TERPSICHORE(_Objective): # Need to sure up the paths - def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", wout_filename="wout_default.nc", path=None, bounds=None,normalize=False, submit_script_name="terps_job.submit", normalize_target=False, awall=2.0, deltajp=1.e-2, modelk=0, al0=-5.e-1, nev=1, nfp=2, xplo=1.e-6, max_bozm=19, max_bozn=14, mode_family=0, max_modem=55, min_moden=-8, max_moden=11, nsurf=128): + def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", wout_filename="wout_default.nc", try_parallel=True, path=None, bounds=None,normalize=False, submit_script_name="terps_job.submit", normalize_target=False, awall=2.0, deltajp=1.e-2, modelk=0, al0=-5.e-1, nev=1, nfp=2, xplo=1.e-6, max_bozm=19, max_bozn=14, mode_family=0, max_modem=55, min_moden=-8, max_moden=11, nsurf=128): if target is None and bounds is None: target = 0 self.eq = eq + self.try_parallel = try_parallel self.nsurf = nsurf # No functionality unless using dynamic allocation TERPS self.awall = awall self.deltajp = deltajp @@ -144,9 +146,9 @@ def compute_impl(self, params, constants): if constants is None: constants = self.constants - self.write_vmec() # Write VMEC file from DESC equilibrium - self.compute_fort18() - self.write_terps_io() + #self.write_vmec() # Write VMEC file from DESC equilibrium + #self.compute_fort18() + #self.write_terps_io() self.run_terps() self.terps_outfile = os.path.join(self.path,'fort.16') # Let's change the name of this at some point self.parse_terps_outfile() @@ -213,7 +215,7 @@ def is_terps_complete(self, stop_time, runtime, terps_subprocess): terps_subprocess.terminate() f_slurm.seek(0) line = f_slurm.readline() - while line: +v while line: if 'PARAMETER LSSL' in line: self.lssl = int(line.split("TO:")[1]) break @@ -235,10 +237,39 @@ def is_terps_complete(self, stop_time, runtime, terps_subprocess): return True ''' - +# def processor_work(self, task_id): +# # stuff + def run_terps(self): + # Get jobID + job_id = os.environ.get('SLURM_JOB_ID') + print(job_id) + + #time.sleep(60) + # Get number processors + #cmd = ['squeue', '-l', '-j', job_id] + #cmd = ['squeue', '-h', '-j', job_id, '-o', '"%.18i %.9P %.8j %.8u %.2t %.10M %.6D %.20R %C"']#, '|', 'awk', "'{sum+=$NF} END {print sum}'"] + #cmd = ['squeue'] + #cmd = ['/usr/bin/seff', job_id] + #cmd = ['seff', job_id] + job_check = subprocess.run(['seff'], stdout=subprocess.PIPE) + out_text = job_check.stdout.decode('utf-8') + print(out_text.split("\n")) + #jobID = out_text.split()[-1] + #slurm_file = os.path.join(self.path,"terps_{}.log".format(jobID)) + exit() + + # Figure out number of parallel jobs to run + + pool = multiprocessing.Pool(nproc) + pool.map(processor_task, range(nproc)) + + pool.close() + pool.join() + + sleep_time = 1 # seconds stop_time = 60 # seconds (kill the infinite loop if TERPS ran into an error and won't be printing growth rate) From 2f7f9b127c7e8950f041cb05d2cc54f6133874ba Mon Sep 17 00:00:00 2001 From: Michael Martin Date: Tue, 4 Jun 2024 16:36:44 -0400 Subject: [PATCH 09/49] commented block with some initial parallelization debugging --- desc/objectives/_terpsichore.py | 23 ++--------------------- 1 file changed, 2 insertions(+), 21 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index da047a25ea..8e37f5c4d4 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -208,8 +208,6 @@ def is_terps_complete(self, stop_time, runtime, terps_subprocess): f_slurm.close() return False - # Uncomment the following block when dynamic allocation is working correctly - ''' elif 'PARAMETER LSSL' in terps_out_contents: if terps_subprocess.returncode == None: terps_subprocess.terminate() @@ -235,14 +233,11 @@ def is_terps_complete(self, stop_time, runtime, terps_subprocess): line = f_slurm.readline() self.write_terps_io() # Rewrite the input file with suggested value of LSSD and re-run return True - ''' - -# def processor_work(self, task_id): -# # stuff def run_terps(self): + ''' # Get jobID job_id = os.environ.get('SLURM_JOB_ID') print(job_id) @@ -268,7 +263,7 @@ def run_terps(self): pool.close() pool.join() - + ''' sleep_time = 1 # seconds stop_time = 60 # seconds (kill the infinite loop if TERPS ran into an error and won't be printing growth rate) @@ -289,8 +284,6 @@ def run_terps(self): toc = time.perf_counter() runtime = toc-tic - # Uncomment the following block when dynamic allocation TERPS is working - ''' if (self.lssl_repeat): self.lssl_repeat = False self.run_terps() @@ -298,7 +291,6 @@ def run_terps(self): elif (self.lssd_repeat): self.lssd_repeat = False self.run_terps() - ''' print("Found growth rate!") @@ -368,8 +360,6 @@ def write_terps_io(self): eq_identifier = "C640" - # Uncomment the following block when dynamic allocation TERPS is working - ''' ivac = self.nsurf // 4 if (self.max_moden > 8) or (self.max_modem > 16): nj = 150 @@ -380,7 +370,6 @@ def write_terps_io(self): else: nj = 50 nk = 50 - ''' self.terps_infile = os.path.join(self.path, "{}_N{}_family".format(eq_identifier, self.mode_family)) f = open(self.terps_infile,"w") @@ -389,13 +378,6 @@ def write_terps_io(self): f.write("C\n") f.write("C MM NMIN NMAX MMS NSMIN NSMAX NPROCS INSOL\n") f.write(" {:>2d} {:>3d} {:>3d} 55 -8 10 1 0\n".format(self.max_bozm, -self.max_bozn, self.max_bozn)) - - # Uncomment the following block when dynamic allocation TERPS is working - ''' - f.write("C\n") - f.write("C NJ NK IVAC LSSL LSSD MMAXDF NMAXDF\n") - f.write(" {:>3d} {:>3d} {:>3d} {:>4d} {:>4d} 120 64\n".format(nj, nk, ivac, self.lssl, self.lssd)) # Not exactly sure how to set LSSL and LSSD..usually it tells you to increase if not high enough (does this affect runtime?) - ''' f.write("C\n") f.write("C TABLE OF FOURIER COEFFIENTS FOR BOOZER COORDINATES\n") f.write("C EQUILIBRIUM SETTINGS ARE COMPUTED FROM FIT/VMEC\n") @@ -404,7 +386,6 @@ def write_terps_io(self): boz_str_title = "C M= 0" boz_str_neg = " 0" boz_str_pos = " 1" - #for _im in range(1,self.max_bozm+1): for _im in range(1,37): if (_im >= 10): boz_str_title += " "+str(_im)[1] From e9f68483c97e298a8df505044e37c66662425604 Mon Sep 17 00:00:00 2001 From: Michael Martin Date: Tue, 4 Jun 2024 17:13:56 -0400 Subject: [PATCH 10/49] small modifications for dynamic terps --- desc/objectives/_terpsichore.py | 41 +++++++-------------------------- 1 file changed, 8 insertions(+), 33 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 8e37f5c4d4..48faae7dc1 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -146,9 +146,9 @@ def compute_impl(self, params, constants): if constants is None: constants = self.constants - #self.write_vmec() # Write VMEC file from DESC equilibrium - #self.compute_fort18() - #self.write_terps_io() + self.write_vmec() # Write VMEC file from DESC equilibrium + #self.compute_fort18() # No longer needed + self.write_terps_io() self.run_terps() self.terps_outfile = os.path.join(self.path,'fort.16') # Let's change the name of this at some point self.parse_terps_outfile() @@ -162,7 +162,10 @@ def write_vmec(self): from desc.vmec import VMECIO VMECIO.save(self.eq, self.wout_file, surfs=self.nsurf) - + +''' + I think this function is defunct now + def compute_fort18(self): print("Figure out how to do this directly from DESC equilibrium quantities!!") @@ -170,7 +173,7 @@ def compute_fort18(self): head, tail = os.path.split(self.wout_file) cmd = [self.vmec2terps_app, tail] subprocess.run(cmd,stdout=subprocess.PIPE, check=True) - +''' def remove_terps_files(self): @@ -236,34 +239,6 @@ def is_terps_complete(self, stop_time, runtime, terps_subprocess): def run_terps(self): - - ''' - # Get jobID - job_id = os.environ.get('SLURM_JOB_ID') - print(job_id) - - #time.sleep(60) - # Get number processors - #cmd = ['squeue', '-l', '-j', job_id] - #cmd = ['squeue', '-h', '-j', job_id, '-o', '"%.18i %.9P %.8j %.8u %.2t %.10M %.6D %.20R %C"']#, '|', 'awk', "'{sum+=$NF} END {print sum}'"] - #cmd = ['squeue'] - #cmd = ['/usr/bin/seff', job_id] - #cmd = ['seff', job_id] - job_check = subprocess.run(['seff'], stdout=subprocess.PIPE) - out_text = job_check.stdout.decode('utf-8') - print(out_text.split("\n")) - #jobID = out_text.split()[-1] - #slurm_file = os.path.join(self.path,"terps_{}.log".format(jobID)) - exit() - - # Figure out number of parallel jobs to run - - pool = multiprocessing.Pool(nproc) - pool.map(processor_task, range(nproc)) - - pool.close() - pool.join() - ''' sleep_time = 1 # seconds stop_time = 60 # seconds (kill the infinite loop if TERPS ran into an error and won't be printing growth rate) From 54c9f44ad87c5e67984c6112c4f06fc0aad7cd0b Mon Sep 17 00:00:00 2001 From: Michael Martin Date: Wed, 5 Jun 2024 15:49:39 -0400 Subject: [PATCH 11/49] modifications for dynamic TERPS and assuming the tpr16* files are no longer output during TERPS runs --- desc/objectives/_terpsichore.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 48faae7dc1..12f71d47fa 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -177,12 +177,14 @@ def compute_fort18(self): def remove_terps_files(self): + # Keeping this function for now in case there's a reason we might want these files in the future + if (os.path.exists(os.path.join(self.path, 'tpr16_dat_wall'))): - rm_cmd = ['rm', 'tpr16_dat_wall'] # There's probably a better way to handle this + rm_cmd = ['rm', 'tpr16_dat_wall'] subprocess.run(rm_cmd) if (os.path.exists(os.path.join(self.path, 'tpr16_dat_pvi'))): - rm_cmd = ['rm', 'tpr16_dat_pvi'] # There's probably a better way to handle this + rm_cmd = ['rm', 'tpr16_dat_pvi'] subprocess.run(rm_cmd) @@ -204,7 +206,7 @@ def is_terps_complete(self, stop_time, runtime, terps_subprocess): if terps_subprocess.returncode == None: terps_subprocess.terminate() f_slurm.close() - self.remove_terps_files() + #self.remove_terps_files() return True else: @@ -247,9 +249,9 @@ def run_terps(self): fs = open('stdout.terps','w') head, tail = os.path.split(self.terps_infile) - self.remove_terps_files() + # self.remove_terps_files() # Newest TERPS version does not write out problematic files - cmd = ['srun', '-n', '1', '-t', '00:01:00', self.terps_app] + cmd = ['srun', '-n', '1', '-t', '00:01:00', '<', self.terps_app, self.wout_filename] terps_subprocess = subprocess.run(cmd, stdin=open(self.terps_infile,'r'), stdout=fs, stderr=fs_error) runtime = 0.0 @@ -259,6 +261,7 @@ def run_terps(self): toc = time.perf_counter() runtime = toc-tic + # !!!!! This check and setting of LSSL and LSSD should be done in a TERPS run BEFORE sending out a bunch of parallel TERPS runs (since they should all use the same LSSL and LSSD values) !!!!! if (self.lssl_repeat): self.lssl_repeat = False self.run_terps() @@ -266,11 +269,11 @@ def run_terps(self): elif (self.lssd_repeat): self.lssd_repeat = False self.run_terps() - - print("Found growth rate!") - + fs.close() + # The parallel jobs could start here + # Need a command here to wait until all TERPS runs are complete if doing some form of parallel execution From 29bfc8c048824d16cbb9cc9c28e07c6a1196e6fe Mon Sep 17 00:00:00 2001 From: Michael Martin Date: Fri, 7 Jun 2024 11:06:44 -0400 Subject: [PATCH 12/49] Changed modelk input parameter default to modelk=1 --- desc/objectives/_terpsichore.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 12f71d47fa..fc09cd6e16 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -53,7 +53,7 @@ class TERPSICHORE(_Objective): # Need to sure up the paths - def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", wout_filename="wout_default.nc", try_parallel=True, path=None, bounds=None,normalize=False, submit_script_name="terps_job.submit", normalize_target=False, awall=2.0, deltajp=1.e-2, modelk=0, al0=-5.e-1, nev=1, nfp=2, xplo=1.e-6, max_bozm=19, max_bozn=14, mode_family=0, max_modem=55, min_moden=-8, max_moden=11, nsurf=128): + def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", wout_filename="wout_default.nc", try_parallel=True, path=None, bounds=None,normalize=False, submit_script_name="terps_job.submit", normalize_target=False, awall=2.0, deltajp=1.e-2, modelk=1, al0=-5.e-1, nev=1, nfp=2, xplo=1.e-6, max_bozm=19, max_bozn=14, mode_family=0, max_modem=55, min_moden=-8, max_moden=11, nsurf=128): if target is None and bounds is None: target = 0 From fea999fa22a4edabbb0da60fa60820aa78a1b6a6 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Fri, 7 Jun 2024 16:33:17 -0600 Subject: [PATCH 13/49] clean up TERPSICHORE to work with _ExternalObjective --- desc/objectives/_terpsichore.py | 793 +++++++++++++++----------------- 1 file changed, 379 insertions(+), 414 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index fc09cd6e16..dc6928741e 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -1,446 +1,411 @@ -import numpy as np -import subprocess -from scipy.interpolate import interp1d +import multiprocessing import os +import subprocess import time -import math -import multiprocessing -from desc.backend import jnp,put - -from .objective_funs import ObjectiveFunction, _Objective -from .utils import ( - factorize_linear_constraints, -) -from desc.utils import Timer -from desc.compute import compute as compute_fun -from desc.compute import get_profiles, get_transforms, get_params -from desc.compute.utils import dot - -from scipy.constants import mu_0, pi -from desc.grid import LinearGrid, Grid, ConcentricGrid, QuadratureGrid -from jax import core -from jax.interpreters import ad, batching -from desc.derivatives import FiniteDiffDerivative -import netCDF4 as nc -from shutil import copyfile -from desc.compute.utils import cross, dot - - -class TERPSICHORE(_Objective): - r"""Calls the linear MHD stability code TERPSICHORE to compute the linear growth rates of the fastest growing instability - - Parameters - ---------- - eq : Equilibrium, optional - Equilibrium that will be optimized to satisfy the Objective. - target : float, ndarray, optional - COMPLETE THIS - bounds : tuple, optional - COMPLETE THIS - weight : float, ndarray, optional - COMPLETE THIS - more stuff here -======= +import numpy as np +from desc.utils import errorif + +from ._generic import _ExternalObjective + + +def terps_fun( + eq, + path="", + exec="", + eq_id="", + mode_family=0, + surfs=16, + M_max=2, + N_min=-2, + N_max=2, + M_booz_max=4, + N_booz_max=4, + awall=2.0, + xplo=1e-6, + deltajp=1e-2, + modelk=1, + nev=1, + al0=-5e-1, + sleep_time=1, + stop_time=60, +): + """TERPSICHORE driver function.""" + process = multiprocessing.current_process() + pid = str(process.pid) + + exec_path = os.path.join(path, exec) + input_path = os.path.join(path, "{}_N{}_{}".format(eq_id, mode_family, pid)) + wout_path = os.path.join(path, "wout_{}.nc".format(pid)) + fort16_path = os.path.join(path, "fort_{}.16".format(pid)) + + _write_wout(eq=eq, path=wout_path, surfs=surfs) + _write_terps_input( + path=input_path, + eq_id=eq_id, + mode_family=mode_family, + surfs=surfs, + M_max=M_max, + N_min=N_min, + N_max=N_max, + M_booz_max=M_booz_max, + N_booz_max=N_booz_max, + awall=awall, + xplo=xplo, + deltajp=deltajp, + modelk=modelk, + nfp=eq.NFP, + nev=nev, + al0=al0, + ) + _run_terps( + exec=exec_path, + input=input_path, + wout=wout_path, + sleep_time=sleep_time, + stop_time=stop_time, + ) + growth_rate = _read_terps_output(path=fort16_path) + + return np.atleast_1d(growth_rate) + + +def _write_wout(eq, path, surfs): + """Write the wout NetCDF file from the equilibrium.""" + from desc.vmec import VMECIO + + VMECIO.save(eq=eq, path=path, surfs=surfs, verbose=0) + + +def _write_terps_input( + path, + eq_id, + mode_family, + surfs, + M_max, + N_min, + N_max, + M_booz_max, + N_booz_max, + awall, + xplo, + deltajp, + modelk, + nfp, + nev, + al0, +): + """Write TERPSICHORE input file.""" + """ + ivac = surfs // 4 + if (N_max > 8) or (M_max > 16): + nj = 150 + nk = 150 + elif (N_max > 4) or (M_max > 8): + nj = 100 + nk = 100 + else: + nj = 50 + nk = 50 """ - # Used by GX objective, not sure where they're used yet - _scalar = True - _linear = False - _units = "[1/s]" - _print_value_fmt = "Growth rate: {:10.3e} " - - - # Need to sure up the paths - def __init__(self, eq=None, target=0, weight=1, grid=None, name="TERPSICHORE", wout_filename="wout_default.nc", try_parallel=True, path=None, bounds=None,normalize=False, submit_script_name="terps_job.submit", normalize_target=False, awall=2.0, deltajp=1.e-2, modelk=1, al0=-5.e-1, nev=1, nfp=2, xplo=1.e-6, max_bozm=19, max_bozn=14, mode_family=0, max_modem=55, min_moden=-8, max_moden=11, nsurf=128): - - if target is None and bounds is None: - target = 0 - self.eq = eq - self.try_parallel = try_parallel - self.nsurf = nsurf # No functionality unless using dynamic allocation TERPS - self.awall = awall - self.deltajp = deltajp - self.modelk = modelk - self.al0 = al0 - self.nev = nev - self.nfp = nfp - self.xplo = xplo - self.max_bozm = max_bozm - self.max_bozn = max_bozn - self.mode_family = mode_family - self.max_modem = max_modem - self.min_moden = min_moden - self.max_moden = max_moden - # lssl, lssd, lssl_repeat, and lssd_repeat are only used with dynamic allocation TERPS - self.lssl = 200 # LSSL and LSSD depend on the specified resolution and the required result is given after running the code (once for each variable) - self.lssd = 100 - self.lssl_repeat = True - self.lssd_repeat = True - self.grid = grid - - super().__init__( - things=eq, - target=target, - bounds=bounds, - weight=weight, - normalize=normalize, - normalize_target=normalize_target, - name=name, - ) - - units = "" - self._callback_fmt = "Growth rate: {:10.3e} " + units - self._print_value_fmt = "Growth rate: {:10.3e} " + units - + f = open(path, "w") - self.path = path - self.terps_stdout = os.path.join(self.path,'stdout.terps') - self.submit_script_name = submit_script_name - self.submit_script_path = os.path.join(self.path, self.submit_script_name) - self.wout_filename = wout_filename - self.wout_file = os.path.join(self.path, wout_filename) - self.vmec2terps_app = os.path.join(self.path, "thea-vmec2terps.x") - self.terps_app = os.path.join(self.path, "tpr_ap.x") + f.write(" {}\n".format(eq_id)) + f.write("C\n") + f.write("C MM NMIN NMAX MMS NSMIN NSMAX NPROCS INSOL\n") + f.write( + " {:>2d} {:>3d} {:>3d} 55 -8 10 1 0\n".format( + M_booz_max, -N_booz_max, N_booz_max + ) + ) + f.write("C\n") + f.write("C TABLE OF FOURIER COEFFIENTS FOR BOOZER COORDINATES\n") + f.write("C EQUILIBRIUM SETTINGS ARE COMPUTED FROM FIT/VMEC\n") + f.write("C\n") + + boz_str_title = "C M= 0" + boz_str_neg = " 0" + boz_str_pos = " 1" + for _im in range(1, 37): + if _im >= 10: + boz_str_title += " " + str(_im)[1] + else: + boz_str_title += " " + str(_im) + boz_str_neg += " 1" + boz_str_pos += " 1" + + boz_str_title += " N\n" + f.write(boz_str_title) + for _in in range(-N_booz_max, N_booz_max + 1): + final_str_neg = boz_str_neg + "{:>3}\n".format(_in) + final_str_pos = boz_str_pos + "{:>3}\n".format(_in) + if _in < 0.0: + f.write(final_str_neg) + else: + f.write(final_str_pos) + + f.write("C\n") + f.write(" LLAMPR LVMTPR LMETPR LFOUPR\n") + f.write(" 0 0 0 0\n") + f.write(" LLHSPR LRHSPR LEIGPR LEFCPR\n") + f.write(" 9 9 1 1\n") + f.write(" LXYZPR LIOTPL LDW2PL LEFCPL\n") + f.write(" 0 1 1 1\n") + f.write(" LCURRF LMESHP LMESHV LITERS\n") + f.write(" 1 1 2 1\n") + f.write(" LXYZPL LEFPLS LEQVPL LPRESS\n") + f.write(" 1 1 0 0\n") + f.write("C\n") + f.write( + "C PVAC PARFAC QONAX QN DSVAC " + + "QVAC NOWALL\n" + ) + f.write( + " 1.0001e+00 0.0000e-00 0.6500e-00 0.0000e-00 1.0000e-00 " + + "1.0001e+00 -2\n" + ) + f.write("\n") + f.write( + "C AWALL EWALL DWALL GWALL DRWAL DZWAL " + + "NPWALL\n" + ) + f.write( + " {:10.4e} ".format(awall) + + "1.5000e+00 -1.0000e-00 5.2000e-00 -0.0000e-00 +0.0000e-00 2\n" + ) + f.write("C\n") + f.write("C RPLMIN XPLO DELTAJP WCT CURFAC\n") + f.write( + " 1.0000e-05 {:10.4e} {:10.4e} 6.6667e-01 1.0000e-00\n".format( + xplo, deltajp + ) + ) + f.write("C\n") + f.write( + "C " + + "MODELK = {}\n".format(modelk) + ) + f.write("C\n") + f.write( + "C NUMBER OF EQUILIBRIUM FIELD PERIODS PER STABILITY PERIOD: " + + "NSTA = {}\n".format(nfp) + ) + f.write("C\n") + f.write("C TABLE OF FOURIER COEFFIENTS FOR STABILITY DISPLACEMENTS\n") + f.write("C\n") + f.write( + "C M= 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 " + + "0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 N\n" + ) + + for n in range(-8, 11): + + in_family = False + if n % 2 == mode_family: # FIXME: must be modified for nfp != 2,3 + in_family = True + + if n <= 0: + mode_str = " 0" + else: + if (n <= N_max) and (n >= N_min) and (in_family): + mode_str = " 1" + else: + mode_str = " 0" - self.terps_compute = core.Primitive("terps") - self.terps_compute.def_impl(self.compute_impl) - # What are the following two lines? - ad.primitive_jvps[self.terps_compute] = self.compute_terps_jvp - batching.primitive_batchers[self.terps_compute] = self.compute_terps_batch + for m in range(1, 55 + 1): + if (m <= M_max) and (in_family): + if (n <= N_max) and (n >= N_min): + mode_str += " 1" + else: + mode_str += " 0" + else: + mode_str += " 0" - def build(self, use_jit=False, verbose=1): - """Build constant arrays. + mode_str += "{:>3}\n".format(_in) + f.write(mode_str) - Parameters - ---------- - eq : Equilibrium, optional - Equilibrium that will be optimized to satisfy the Objective. - use_jit : bool, optional - Whether to just-in-time compile the objective and derivatives - (MUST be False to run GX). - verbose : int, optional - Level of output. + f.write("C\n") + f.write("C NEV NITMAX AL0 EPSPAM IGREEN MPINIT\n") + f.write(" {} 4500 {:10.3e} 1.000E-04 0 0\n".format(nev, al0)) + f.write("C\n") - """ - eq = self.things[0] - self._dim_f = 1 + f.close() - timer = Timer() - super().build(use_jit=use_jit, verbose=verbose) - +def _run_terps(exec, input, wout, sleep_time, stop_time): + """Run TERPSICHORE.""" + stdout_path = "stdout.terps" + stderr_path = "error.terps" - def compute(self, params, constants=None): + fout = open(stdout_path, "w") + ferr = open(stderr_path, "w") - return self.terps_compute.bind(params,constants) + cmd = ["srun", "-n", "1", "-t", "00:01:00", "<", exec, wout] + terps_subprocess = subprocess.run(cmd, stdin=open(input), stdout=fout, stderr=ferr) - - def compute_impl(self, params, constants): + run_time = 0.0 + tic = time.perf_counter() + while not _is_terps_complete( + path=stdout_path, + run_time=run_time, + stop_time=stop_time, + terps_subprocess=terps_subprocess, + ): + time.sleep(sleep_time) + toc = time.perf_counter() + run_time = toc - tic -# params, constants = self._parse_args(*args, **kwargs) - if not np.any(params["R_lmn"]): - return 0 + fout.close() + ferr.close() - if constants is None: - constants = self.constants - self.write_vmec() # Write VMEC file from DESC equilibrium - #self.compute_fort18() # No longer needed - self.write_terps_io() - self.run_terps() - self.terps_outfile = os.path.join(self.path,'fort.16') # Let's change the name of this at some point - self.parse_terps_outfile() +def _is_terps_complete(path, run_time, stop_time, terps_subprocess): + """Check if TERPSICHORE is finished running.""" + if not os.path.exists(path): + return False - print("Growth rate = {}".format(self.growth_rate)) - - return jnp.atleast_1d(self.growth_rate) - + errorif( + run_time > stop_time, RuntimeError, "TERPS was unable to find a growth rate!" + ) - def write_vmec(self): + file = open(path) + terps_output = file.read() + if "GROWTH RATE" in terps_output: + if terps_subprocess.returncode is None: + terps_subprocess.terminate() + file.close() + return True + else: + file.close() + return False - from desc.vmec import VMECIO - VMECIO.save(self.eq, self.wout_file, surfs=self.nsurf) -''' - I think this function is defunct now - - def compute_fort18(self): +def _read_terps_output(path): + """Read TERPSICHORE output file and return the growth rate.""" + errorif( + not os.path.exists(path), RuntimeError, "TERPS fort.16 output file not found!" + ) - print("Figure out how to do this directly from DESC equilibrium quantities!!") - - head, tail = os.path.split(self.wout_file) - cmd = [self.vmec2terps_app, tail] - subprocess.run(cmd,stdout=subprocess.PIPE, check=True) -''' + file = open(path) + growth_rates = [] + for line in file: + index = line.find("GROWTH RATE") + if index != -1: + growth_rates.append(float(line.strip().split("=")[1])) + file.close() - def remove_terps_files(self): + errorif(len(growth_rates) == 0, RuntimeError, "Growth rate not found!") + errorif( + len(growth_rates) > 1, + NotImplementedError, + "Not capable of handling multiple growth rates.", + ) - # Keeping this function for now in case there's a reason we might want these files in the future - - if (os.path.exists(os.path.join(self.path, 'tpr16_dat_wall'))): - rm_cmd = ['rm', 'tpr16_dat_wall'] - subprocess.run(rm_cmd) + return growth_rates[0] - if (os.path.exists(os.path.join(self.path, 'tpr16_dat_pvi'))): - rm_cmd = ['rm', 'tpr16_dat_pvi'] - subprocess.run(rm_cmd) - - def is_terps_complete(self, stop_time, runtime, terps_subprocess): +class TERPSICHORE(_ExternalObjective): + """Computes linear MHD stability from calls to the code TERPSICHORE. - if (os.path.exists(self.terps_stdout)): - f_slurm = open(self.terps_stdout, 'r') - else: - return False - - print("Current runtime = {} seconds".format(math.ceil(runtime))) - if (runtime > stop_time): - print("TERPS was unable to find a growth rate. Exiting...") - exit() - - terps_out_contents = f_slurm.read() - - if 'GROWTH RATE' in terps_out_contents: - if terps_subprocess.returncode == None: - terps_subprocess.terminate() - f_slurm.close() - #self.remove_terps_files() - return True + Returns the linear growth rate of the fastest growing instability. - else: - f_slurm.close() - return False - - elif 'PARAMETER LSSL' in terps_out_contents: - if terps_subprocess.returncode == None: - terps_subprocess.terminate() - f_slurm.seek(0) - line = f_slurm.readline() -v while line: - if 'PARAMETER LSSL' in line: - self.lssl = int(line.split("TO:")[1]) - break - line = f_slurm.readline() - self.write_terps_io() # Rewrite the input file with suggested value of LSSL and re-run - return True - - elif 'PARAMETER LSSD' in terps_out_contents: - if terps_subprocess.returncode == None: - terps_subprocess.terminate() - f_slurm.seek(0) - line = f_slurm.readline() - while line: - if 'PARAMETER LSSD' in line: - self.lssd = int(line.split("TO:")[1]) - break - line = f_slurm.readline() - self.write_terps_io() # Rewrite the input file with suggested value of LSSD and re-run - return True - - - def run_terps(self): - - sleep_time = 1 # seconds - stop_time = 60 # seconds (kill the infinite loop if TERPS ran into an error and won't be printing growth rate) - - fs_error = open('error.terps','w') - fs = open('stdout.terps','w') - head, tail = os.path.split(self.terps_infile) - - # self.remove_terps_files() # Newest TERPS version does not write out problematic files - - cmd = ['srun', '-n', '1', '-t', '00:01:00', '<', self.terps_app, self.wout_filename] - terps_subprocess = subprocess.run(cmd, stdin=open(self.terps_infile,'r'), stdout=fs, stderr=fs_error) - - runtime = 0.0 - tic = time.perf_counter() - while not self.is_terps_complete(stop_time, runtime, terps_subprocess): - time.sleep(sleep_time) - toc = time.perf_counter() - runtime = toc-tic - - # !!!!! This check and setting of LSSL and LSSD should be done in a TERPS run BEFORE sending out a bunch of parallel TERPS runs (since they should all use the same LSSL and LSSD values) !!!!! - if (self.lssl_repeat): - self.lssl_repeat = False - self.run_terps() - - elif (self.lssd_repeat): - self.lssd_repeat = False - self.run_terps() - - fs.close() - - # The parallel jobs could start here - - # Need a command here to wait until all TERPS runs are complete if doing some form of parallel execution - - - def parse_terps_outfile(self): - - # Read fort.16 and search for growth rate - if (os.path.exists(self.terps_outfile)): - f = open(self.terps_outfile, 'r') - else: - print("TERPS fort.16 output file not found! Exiting...") - exit() - - growth_rate = [] - for line_number, line in enumerate(f, start=1): - index = line.find('GROWTH RATE') - if index != -1: - growth_rate.append(float(line.strip().split("=")[1])) - - if (len(growth_rate) > 1): - print("Currently capable of using only one growth rate. Exiting...") - exit() - elif (len(growth_rate) == 0): - print("Growth rate not found! Exiting...") - exit() - - self.growth_rate = growth_rate[0] - - # I'm not sure what's going on with these final two functions - def compute_terps_jvp(self,values,tangents): - - R_lmn, Z_lmn, L_lmn, i_l, c_l, p_l, Psi = values - primal_out = jnp.atleast_1d(0.0) - - n = len(values) - argnum = np.arange(0,n,1) - - jvp = FiniteDiffDerivative.compute_jvp(self.compute,argnum,tangents,*values,rel_step=1e-2) - - return (primal_out, jvp) - - def compute_terps_batch(self, values, axis): - numdiff = len(values[0]) - res = jnp.array([0.0]) - - for i in range(numdiff): - R_lmn = values[0][i] - Z_lmn = values[1][i] - L_lmn = values[2][i] - i_l = values[3][i] - p_l = values[4][i] - Psi = values[5][i] - - res = jnp.vstack([res,self.compute(R_lmn,Z_lmn,L_lmn,i_l,p_l,Psi)]) - - res = res[1:] - - - return res, axis[0] - - - def write_terps_io(self): - - eq_identifier = "C640" - - ivac = self.nsurf // 4 - if (self.max_moden > 8) or (self.max_modem > 16): - nj = 150 - nk = 150 - elif (self.max_moden > 4) or (self.max_modem > 8): - nj = 100 - nk = 100 - else: - nj = 50 - nk = 50 - - self.terps_infile = os.path.join(self.path, "{}_N{}_family".format(eq_identifier, self.mode_family)) - f = open(self.terps_infile,"w") - - f.write(" {}\n".format(eq_identifier)) - f.write("C\n") - f.write("C MM NMIN NMAX MMS NSMIN NSMAX NPROCS INSOL\n") - f.write(" {:>2d} {:>3d} {:>3d} 55 -8 10 1 0\n".format(self.max_bozm, -self.max_bozn, self.max_bozn)) - f.write("C\n") - f.write("C TABLE OF FOURIER COEFFIENTS FOR BOOZER COORDINATES\n") - f.write("C EQUILIBRIUM SETTINGS ARE COMPUTED FROM FIT/VMEC\n") - f.write("C\n") - - boz_str_title = "C M= 0" - boz_str_neg = " 0" - boz_str_pos = " 1" - for _im in range(1,37): - if (_im >= 10): - boz_str_title += " "+str(_im)[1] - else: - boz_str_title += " "+str(_im) - boz_str_neg += " 1" - boz_str_pos += " 1" - - - boz_str_title += " N\n" - f.write(boz_str_title) - for _in in range(-self.max_bozn,self.max_bozn+1): - final_str_neg = boz_str_neg+"{:>3}\n".format(_in) - final_str_pos = boz_str_pos+"{:>3}\n".format(_in) - if _in < 0.0: - f.write(final_str_neg) - else: - f.write(final_str_pos) - - f.write("C\n") - f.write(" LLAMPR LVMTPR LMETPR LFOUPR\n") - f.write(" 0 0 0 0\n") - f.write(" LLHSPR LRHSPR LEIGPR LEFCPR\n") - f.write(" 9 9 1 1\n") - f.write(" LXYZPR LIOTPL LDW2PL LEFCPL\n") - f.write(" 0 1 1 1\n") - f.write(" LCURRF LMESHP LMESHV LITERS\n") - f.write(" 1 1 2 1\n") - f.write(" LXYZPL LEFPLS LEQVPL LPRESS\n") - f.write(" 1 1 0 0\n") - f.write("C\n") - f.write("C PVAC PARFAC QONAX QN DSVAC QVAC NOWALL\n") - f.write(" 1.0001e+00 0.0000e-00 0.6500e-00 0.0000e-00 1.0000e-00 1.0001e+00 -2\n") - f.write("\n") - f.write("C AWALL EWALL DWALL GWALL DRWAL DZWAL NPWALL\n") - f.write(" {:10.4e} 1.5000e+00 -1.0000e-00 5.2000e-00 -0.0000e-00 +0.0000e-00 2\n".format(self.awall)) - f.write("C\n") - f.write("C RPLMIN XPLO DELTAJP WCT CURFAC\n") - f.write(" 1.0000e-05 {:10.4e} {:10.4e} 6.6667e-01 1.0000e-00\n".format(self.xplo, self.deltajp)) - f.write("C\n") - f.write("C MODELK = {}\n".format(self.modelk)) - f.write("C\n") - f.write("C NUMBER OF EQUILIBRIUM FIELD PERIODS PER STABILITY PERIOD: NSTA = {}\n".format(self.nfp)) - f.write("C\n") - f.write("C TABLE OF FOURIER COEFFIENTS FOR STABILITY DISPLACEMENTS\n") - f.write("C\n") - f.write("C M= 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 N\n") - - for _in in range(-8,11): - - in_family = False - if (_in % 2 == self.mode_family): # This needs to be modified for nfp != 2,3 - in_family = True - - if _in <= 0: - mode_str = " 0" - else: - if ((_in <= self.max_moden) and (_in >= self.min_moden) and (in_family)): - mode_str = " 1" - else: - mode_str = " 0" - - for _im in range(1,55+1): - if ((_im <= self.max_modem) and (in_family)): - if ((_in <= self.max_moden) and (_in >= self.min_moden)): - mode_str += " 1" - else: - mode_str += " 0" - else: - mode_str += " 0" + Parameters + ---------- + eq : Equilibrium + Equilibrium that will be optimized to satisfy the Objective. + target : {float, ndarray}, optional + Target value(s) of the objective. Only used if bounds is None. + Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. + bounds : tuple of {float, ndarray}, optional + Lower and upper bounds on the objective. Overrides target. + Both bounds must be broadcastable to to Objective.dim_f. + Defaults to ``target=0``. + weight : {float, ndarray}, optional + Weighting to apply to the Objective, relative to other Objectives. + Must be broadcastable to to Objective.dim_f + normalize : bool, optional + Whether to compute the error in physical units or non-dimensionalize. + Has no effect for this objective. + normalize_target : bool, optional + Whether target and bounds should be normalized before comparing to computed + values. If `normalize` is `True` and the target is in physical units, + this should also be set to True. + loss_function : {None, 'mean', 'min', 'max'}, optional + Loss function to apply to the objective values once computed. This loss function + is called on the raw compute value, before any shifting, scaling, or + normalization. + name : str, optional + Name of the objective function. + + # TODO: update Parameters docs - mode_str += "{:>3}\n".format(_in) - f.write(mode_str) + """ - f.write("C\n") - f.write("C NEV NITMAX AL0 EPSPAM IGREEN MPINIT\n") - f.write(" {} 4500 {:10.3e} 1.000E-04 0 0\n".format(self.nev,self.al0)) - f.write("C\n") - - f.close() + _units = "(???)" + _print_value_fmt = "TERPSICHORE growth rate: {:10.3e}" + + def __init__( + self, + eq=None, + target=None, + bounds=None, + weight=1, + normalize=False, + normalize_target=False, + loss_function=None, + fd_step=1e-4, + vectorized=True, + path="", + exec="", + eq_id="", + mode_family=0, + surfs=16, + M_max=2, + N_min=-2, + N_max=2, + M_booz_max=4, + N_booz_max=4, + awall=2.0, + xplo=1e-6, + deltajp=1e-2, + modelk=1, + nev=1, + al0=-5e-1, + sleep_time=1, + stop_time=60, + name="terpsichore", + ): + super().__init__( + eq=eq, + fun=terps_fun, + dim_f=1, + target=target, + bounds=bounds, + weight=weight, + normalize=normalize, + normalize_target=normalize_target, + loss_function=loss_function, + fd_step=fd_step, + vectorized=vectorized, + name=name, + path=path, + exec=exec, + eq_id=eq_id, + mode_family=mode_family, + surfs=surfs, + M_max=M_max, + N_min=N_min, + N_max=N_max, + M_booz_max=M_booz_max, + N_booz_max=N_booz_max, + awall=awall, + xplo=xplo, + deltajp=deltajp, + modelk=modelk, + nev=nev, + al0=al0, + sleep_time=sleep_time, + stop_time=stop_time, + ) From 199c657bc93f05fbb4294047ec675cc2305c3caf Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Fri, 7 Jun 2024 16:35:11 -0600 Subject: [PATCH 14/49] fix TERPS input order --- desc/objectives/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py index 2349642e86..1e774d016a 100644 --- a/desc/objectives/__init__.py +++ b/desc/objectives/__init__.py @@ -38,6 +38,7 @@ ) from ._profiles import Pressure, RotationalTransform, Shear, ToroidalCurrent from ._stability import MagneticWell, MercierStability +from ._terpsichore import TERPSICHORE from .getters import ( get_equilibrium_objective, get_fixed_axis_constraints, @@ -81,4 +82,3 @@ FixThetaSFL, ) from .objective_funs import ObjectiveFunction -from ._terpsichore import TERPSICHORE From 4156f4089b2bb241543ee2678796c0f6e6ba45ac Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Fri, 7 Jun 2024 16:49:45 -0600 Subject: [PATCH 15/49] adding lssl and lssd --- desc/objectives/_terpsichore.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index dc6928741e..cd29f04f2c 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -17,6 +17,8 @@ def terps_fun( eq_id="", mode_family=0, surfs=16, + lssl=1, + lssd=1, M_max=2, N_min=-2, N_max=2, @@ -46,6 +48,8 @@ def terps_fun( eq_id=eq_id, mode_family=mode_family, surfs=surfs, + lssl=lssl, + lssd=lssd, M_max=M_max, N_min=N_min, N_max=N_max, @@ -83,6 +87,8 @@ def _write_terps_input( eq_id, mode_family, surfs, + lssl, + lssd, M_max, N_min, N_max, @@ -348,7 +354,7 @@ class TERPSICHORE(_ExternalObjective): def __init__( self, - eq=None, + eq, target=None, bounds=None, weight=1, @@ -362,6 +368,8 @@ def __init__( eq_id="", mode_family=0, surfs=16, + lssl=1, + lssd=1, M_max=2, N_min=-2, N_max=2, @@ -395,6 +403,8 @@ def __init__( eq_id=eq_id, mode_family=mode_family, surfs=surfs, + lssl=lssl, + lssd=lssd, M_max=M_max, N_min=N_min, N_max=N_max, From 9c24b732943af4906ff98b018a2986d4f3ec3819 Mon Sep 17 00:00:00 2001 From: Michael Martin Date: Mon, 10 Jun 2024 10:30:46 -0400 Subject: [PATCH 16/49] debugging dynamic terps --- desc/objectives/_terpsichore.py | 29 +++++++++-------------------- 1 file changed, 9 insertions(+), 20 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index fc09cd6e16..9acbdebc13 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -147,7 +147,6 @@ def compute_impl(self, params, constants): constants = self.constants self.write_vmec() # Write VMEC file from DESC equilibrium - #self.compute_fort18() # No longer needed self.write_terps_io() self.run_terps() self.terps_outfile = os.path.join(self.path,'fort.16') # Let's change the name of this at some point @@ -163,18 +162,6 @@ def write_vmec(self): from desc.vmec import VMECIO VMECIO.save(self.eq, self.wout_file, surfs=self.nsurf) -''' - I think this function is defunct now - - def compute_fort18(self): - - print("Figure out how to do this directly from DESC equilibrium quantities!!") - - head, tail = os.path.split(self.wout_file) - cmd = [self.vmec2terps_app, tail] - subprocess.run(cmd,stdout=subprocess.PIPE, check=True) -''' - def remove_terps_files(self): # Keeping this function for now in case there's a reason we might want these files in the future @@ -208,17 +195,13 @@ def is_terps_complete(self, stop_time, runtime, terps_subprocess): f_slurm.close() #self.remove_terps_files() return True - - else: - f_slurm.close() - return False - + elif 'PARAMETER LSSL' in terps_out_contents: if terps_subprocess.returncode == None: terps_subprocess.terminate() f_slurm.seek(0) line = f_slurm.readline() -v while line: + while line: if 'PARAMETER LSSL' in line: self.lssl = int(line.split("TO:")[1]) break @@ -238,7 +221,11 @@ def is_terps_complete(self, stop_time, runtime, terps_subprocess): line = f_slurm.readline() self.write_terps_io() # Rewrite the input file with suggested value of LSSD and re-run return True - + + else: + f_slurm.close() + return False + def run_terps(self): @@ -357,6 +344,8 @@ def write_terps_io(self): f.write("C MM NMIN NMAX MMS NSMIN NSMAX NPROCS INSOL\n") f.write(" {:>2d} {:>3d} {:>3d} 55 -8 10 1 0\n".format(self.max_bozm, -self.max_bozn, self.max_bozn)) f.write("C\n") + f.write("C NJ NK IVAC LSSL LSSD MMAXDF NMAXDF\n") + f.write(" {:>3d} {:>3d} {:>3d} {:>4d} {:>4d} 120 64\n".format(nj, nk, ivac, self.lssl, self.lssd)) f.write("C TABLE OF FOURIER COEFFIENTS FOR BOOZER COORDINATES\n") f.write("C EQUILIBRIUM SETTINGS ARE COMPUTED FROM FIT/VMEC\n") f.write("C\n") From 8929534b7b5bf2ec370115701b4d78c67f42ef38 Mon Sep 17 00:00:00 2001 From: Michael Martin Date: Mon, 10 Jun 2024 12:53:26 -0400 Subject: [PATCH 17/49] write input additions for dynamic TERPS --- desc/objectives/_terpsichore.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index cd29f04f2c..91fd1e8772 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -17,8 +17,8 @@ def terps_fun( eq_id="", mode_family=0, surfs=16, - lssl=1, - lssd=1, + lssl=200, + lssd=100, M_max=2, N_min=-2, N_max=2, @@ -42,6 +42,12 @@ def terps_fun( wout_path = os.path.join(path, "wout_{}.nc".format(pid)) fort16_path = os.path.join(path, "fort_{}.16".format(pid)) + #if (process.name == 'Process-1'): + # # Need to run a couple times to get LSSL and LSSD + # _determine_lssl_lssd() + + #barrier.wait() + _write_wout(eq=eq, path=wout_path, surfs=surfs) _write_terps_input( path=input_path, @@ -103,7 +109,7 @@ def _write_terps_input( al0, ): """Write TERPSICHORE input file.""" - """ + ivac = surfs // 4 if (N_max > 8) or (M_max > 16): nj = 150 @@ -114,7 +120,6 @@ def _write_terps_input( else: nj = 50 nk = 50 - """ f = open(path, "w") @@ -127,6 +132,9 @@ def _write_terps_input( ) ) f.write("C\n") + f.write("C NJ NK IVAC LSSL LSSD MMAXDF NMAXDF\n") + f.write(" {:>3d} {:>3d} {:>3d} {:>4d} {:>4d} 120 64\n".format(nj, nk, ivac, lssl, lssd)) + f.write("C\n") f.write("C TABLE OF FOURIER COEFFIENTS FOR BOOZER COORDINATES\n") f.write("C EQUILIBRIUM SETTINGS ARE COMPUTED FROM FIT/VMEC\n") f.write("C\n") From 89ade9bbd1f76a958b3673dccb80e32c4b0c8215 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Mon, 10 Jun 2024 11:18:53 -0600 Subject: [PATCH 18/49] edit srun command --- desc/objectives/_terpsichore.py | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 91fd1e8772..fea629fd34 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -10,7 +10,7 @@ from ._generic import _ExternalObjective -def terps_fun( +def terpsichore( eq, path="", exec="", @@ -42,12 +42,6 @@ def terps_fun( wout_path = os.path.join(path, "wout_{}.nc".format(pid)) fort16_path = os.path.join(path, "fort_{}.16".format(pid)) - #if (process.name == 'Process-1'): - # # Need to run a couple times to get LSSL and LSSD - # _determine_lssl_lssd() - - #barrier.wait() - _write_wout(eq=eq, path=wout_path, surfs=surfs) _write_terps_input( path=input_path, @@ -109,7 +103,6 @@ def _write_terps_input( al0, ): """Write TERPSICHORE input file.""" - ivac = surfs // 4 if (N_max > 8) or (M_max > 16): nj = 150 @@ -133,7 +126,11 @@ def _write_terps_input( ) f.write("C\n") f.write("C NJ NK IVAC LSSL LSSD MMAXDF NMAXDF\n") - f.write(" {:>3d} {:>3d} {:>3d} {:>4d} {:>4d} 120 64\n".format(nj, nk, ivac, lssl, lssd)) + f.write( + " {:>3d} {:>3d} {:>3d} {:>4d} {:>4d} 120 64\n".format( + nj, nk, ivac, lssl, lssd + ) + ) f.write("C\n") f.write("C TABLE OF FOURIER COEFFIENTS FOR BOOZER COORDINATES\n") f.write("C EQUILIBRIUM SETTINGS ARE COMPUTED FROM FIT/VMEC\n") @@ -256,8 +253,8 @@ def _run_terps(exec, input, wout, sleep_time, stop_time): fout = open(stdout_path, "w") ferr = open(stderr_path, "w") - cmd = ["srun", "-n", "1", "-t", "00:01:00", "<", exec, wout] - terps_subprocess = subprocess.run(cmd, stdin=open(input), stdout=fout, stderr=ferr) + cmd = ["srun", "-n", "1", "-t", "00:01:00", exec, "<", input, wout] + terps_subprocess = subprocess.run(cmd, stdout=fout, stderr=ferr) run_time = 0.0 tic = time.perf_counter() @@ -395,7 +392,7 @@ def __init__( ): super().__init__( eq=eq, - fun=terps_fun, + fun=terpsichore, dim_f=1, target=target, bounds=bounds, From a340f9573f35ba94675fd991d04793650785da66 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Mon, 10 Jun 2024 14:12:51 -0600 Subject: [PATCH 19/49] subprocess.run with single command --- desc/objectives/_terpsichore.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index fea629fd34..89048bfe10 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -248,13 +248,13 @@ def _write_terps_input( def _run_terps(exec, input, wout, sleep_time, stop_time): """Run TERPSICHORE.""" stdout_path = "stdout.terps" - stderr_path = "error.terps" + stderr_path = "stderr.terps" fout = open(stdout_path, "w") ferr = open(stderr_path, "w") - cmd = ["srun", "-n", "1", "-t", "00:01:00", exec, "<", input, wout] - terps_subprocess = subprocess.run(cmd, stdout=fout, stderr=ferr) + cmd = "srun -n 1 -t 00:01:00 " + exec + " < " + input + " " + wout + terps_subprocess = subprocess.run(cmd, stdout=fout, stderr=ferr, shell=True) run_time = 0.0 tic = time.perf_counter() From 4e03bca2e5ef766ec881d97beea261be5a3289a9 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Mon, 10 Jun 2024 16:23:54 -0600 Subject: [PATCH 20/49] temporary dir for external fun calls --- desc/objectives/_generic.py | 1 + desc/objectives/_terpsichore.py | 32 ++++++++++++++++++++++++-------- 2 files changed, 25 insertions(+), 8 deletions(-) diff --git a/desc/objectives/_generic.py b/desc/objectives/_generic.py index 0114a0d0b8..4d22d5b12c 100644 --- a/desc/objectives/_generic.py +++ b/desc/objectives/_generic.py @@ -143,6 +143,7 @@ def fun_wrapped(params): results = pool.map( functools.partial(self._fun, **self._kwargs), eqs ) + pool.join() return jnp.vstack(results, dtype=float) else: # no vectorization # update equilibrium with new params diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 89048bfe10..6ec2029432 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -1,5 +1,6 @@ import multiprocessing import os +import shutil import subprocess import time @@ -36,11 +37,19 @@ def terpsichore( """TERPSICHORE driver function.""" process = multiprocessing.current_process() pid = str(process.pid) + print(pid) - exec_path = os.path.join(path, exec) - input_path = os.path.join(path, "{}_N{}_{}".format(eq_id, mode_family, pid)) - wout_path = os.path.join(path, "wout_{}.nc".format(pid)) - fort16_path = os.path.join(path, "fort_{}.16".format(pid)) + # create temporary directory to store I/O files + pid_path = os.path.join(path, pid) + os.mkdir(pid_path) + + exec_path = os.path.join(pid_path, exec) + input_path = os.path.join(pid_path, "input") + wout_path = os.path.join(pid_path, "wout.nc") + fort16_path = os.path.join(pid_path, "fort.16") + + # copy executable to temporary directory + shutil.copy(os.path.join(path, exec), exec_path) _write_wout(eq=eq, path=wout_path, surfs=surfs) _write_terps_input( @@ -64,6 +73,7 @@ def terpsichore( al0=al0, ) _run_terps( + dir=pid_path, exec=exec_path, input=input_path, wout=wout_path, @@ -72,6 +82,10 @@ def terpsichore( ) growth_rate = _read_terps_output(path=fort16_path) + # remove temporary directory + shutil.rmtree(pid_path) + + print(growth_rate) return np.atleast_1d(growth_rate) @@ -245,16 +259,18 @@ def _write_terps_input( f.close() -def _run_terps(exec, input, wout, sleep_time, stop_time): +def _run_terps(dir, exec, input, wout, sleep_time, stop_time): """Run TERPSICHORE.""" - stdout_path = "stdout.terps" - stderr_path = "stderr.terps" + stdout_path = os.path.join(dir, "stdout.terps") + stderr_path = os.path.join(dir, "stderr.terps") fout = open(stdout_path, "w") ferr = open(stderr_path, "w") cmd = "srun -n 1 -t 00:01:00 " + exec + " < " + input + " " + wout - terps_subprocess = subprocess.run(cmd, stdout=fout, stderr=ferr, shell=True) + terps_subprocess = subprocess.run( + cmd, cwd=dir, shell=True, stdout=fout, stderr=ferr + ) run_time = 0.0 tic = time.perf_counter() From c4bbd7973a593e0e38e13af9eacf68372882819c Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Mon, 10 Jun 2024 22:46:57 -0600 Subject: [PATCH 21/49] remove debugging print statements --- desc/objectives/_generic.py | 5 ++++- desc/objectives/_terpsichore.py | 2 -- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/desc/objectives/_generic.py b/desc/objectives/_generic.py index 4d22d5b12c..8d9fa6761b 100644 --- a/desc/objectives/_generic.py +++ b/desc/objectives/_generic.py @@ -3,6 +3,7 @@ import functools import inspect import multiprocessing +import os import re from abc import ABC @@ -139,7 +140,9 @@ def fun_wrapped(params): if len(param_value): setattr(eq, param_key, param_value) # parallelize calls to external function - with multiprocessing.Pool(processes=num_eq) as pool: + with multiprocessing.Pool( + processes=min(os.cpu_count(), num_eq) + ) as pool: results = pool.map( functools.partial(self._fun, **self._kwargs), eqs ) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 6ec2029432..3b472c2d18 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -37,7 +37,6 @@ def terpsichore( """TERPSICHORE driver function.""" process = multiprocessing.current_process() pid = str(process.pid) - print(pid) # create temporary directory to store I/O files pid_path = os.path.join(path, pid) @@ -85,7 +84,6 @@ def terpsichore( # remove temporary directory shutil.rmtree(pid_path) - print(growth_rate) return np.atleast_1d(growth_rate) From 0d76cc55072f7ae64ec44d4b5f629d364691985a Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 11 Jun 2024 14:24:00 -0600 Subject: [PATCH 22/49] remove srun --- desc/objectives/_terpsichore.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 3b472c2d18..6d2292d713 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -1,4 +1,4 @@ -import multiprocessing +import multiprocessing as mp import os import shutil import subprocess @@ -35,7 +35,7 @@ def terpsichore( stop_time=60, ): """TERPSICHORE driver function.""" - process = multiprocessing.current_process() + process = mp.current_process() pid = str(process.pid) # create temporary directory to store I/O files @@ -265,7 +265,7 @@ def _run_terps(dir, exec, input, wout, sleep_time, stop_time): fout = open(stdout_path, "w") ferr = open(stderr_path, "w") - cmd = "srun -n 1 -t 00:01:00 " + exec + " < " + input + " " + wout + cmd = exec + " < " + input + " " + wout terps_subprocess = subprocess.run( cmd, cwd=dir, shell=True, stdout=fout, stderr=ferr ) From 35e4859dbd7ebc0e0f2eac524185f28da08ccc7e Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 11 Jun 2024 16:13:10 -0600 Subject: [PATCH 23/49] adding comments --- desc/objectives/_generic.py | 3 +++ desc/objectives/_terpsichore.py | 1 + 2 files changed, 4 insertions(+) diff --git a/desc/objectives/_generic.py b/desc/objectives/_generic.py index 937cd2d88d..fc16bc6fb1 100644 --- a/desc/objectives/_generic.py +++ b/desc/objectives/_generic.py @@ -28,6 +28,8 @@ class _ExternalObjective(_Objective, ABC): The user supplied function must take an Equilibrium as its only positional argument, but can take additional keyword arguments. + # TODO: document that driver script has to include ``if __name__ == "__main__":`` + # TODO: add Parameters documentation Parameters @@ -146,6 +148,7 @@ def fun_wrapped(params): if isinstance(self._vectorized, int) else os.cpu_count() ) + # TODO: re-use same pool with mp.Pool(processes=min(max_processes, num_eq)) as pool: results = pool.map( functools.partial(self._fun, **self._kwargs), eqs diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 6d2292d713..90d0187ef6 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -50,6 +50,7 @@ def terpsichore( # copy executable to temporary directory shutil.copy(os.path.join(path, exec), exec_path) + # FIXME: remove this from the vectorized part? _write_wout(eq=eq, path=wout_path, surfs=surfs) _write_terps_input( path=input_path, From e9575ae842019bed664ecdf78a819123c1adc8ec Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 11 Jun 2024 17:33:55 -0600 Subject: [PATCH 24/49] remove eq_id --- desc/objectives/_terpsichore.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 90d0187ef6..e0e3c6826d 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -15,7 +15,6 @@ def terpsichore( eq, path="", exec="", - eq_id="", mode_family=0, surfs=16, lssl=200, @@ -50,11 +49,9 @@ def terpsichore( # copy executable to temporary directory shutil.copy(os.path.join(path, exec), exec_path) - # FIXME: remove this from the vectorized part? _write_wout(eq=eq, path=wout_path, surfs=surfs) _write_terps_input( path=input_path, - eq_id=eq_id, mode_family=mode_family, surfs=surfs, lssl=lssl, @@ -97,7 +94,6 @@ def _write_wout(eq, path, surfs): def _write_terps_input( path, - eq_id, mode_family, surfs, lssl, @@ -129,7 +125,7 @@ def _write_terps_input( f = open(path, "w") - f.write(" {}\n".format(eq_id)) + f.write(" TERPSICHORE\n") f.write("C\n") f.write("C MM NMIN NMAX MMS NSMIN NSMAX NPROCS INSOL\n") f.write( @@ -385,7 +381,6 @@ def __init__( vectorized=True, path="", exec="", - eq_id="", mode_family=0, surfs=16, lssl=1, @@ -420,7 +415,6 @@ def __init__( name=name, path=path, exec=exec, - eq_id=eq_id, mode_family=mode_family, surfs=surfs, lssl=lssl, From 4e482039d1ccd66d5539371aa208157fe7448acc Mon Sep 17 00:00:00 2001 From: mfmart Date: Thu, 13 Jun 2024 13:37:05 -0700 Subject: [PATCH 25/49] set mode_family < 0 to include all mode families --- desc/objectives/_terpsichore.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index e0e3c6826d..aad0f380d0 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -15,7 +15,7 @@ def terpsichore( eq, path="", exec="", - mode_family=0, + mode_family=-1, surfs=16, lssl=200, lssd=100, @@ -222,6 +222,8 @@ def _write_terps_input( for n in range(-8, 11): + # If mode_family < 0, includes all modes in desired range + in_family = False if n % 2 == mode_family: # FIXME: must be modified for nfp != 2,3 in_family = True @@ -229,17 +231,19 @@ def _write_terps_input( if n <= 0: mode_str = " 0" else: - if (n <= N_max) and (n >= N_min) and (in_family): - mode_str = " 1" + if (n <= N_max) and (n >= N_min): + if (mode_family < 0) or (in_family): + mode_str = " 1" else: mode_str = " 0" for m in range(1, 55 + 1): - if (m <= M_max) and (in_family): - if (n <= N_max) and (n >= N_min): - mode_str += " 1" - else: - mode_str += " 0" + if (m <= M_max): + if (mode_family < 0) or (in_family): + if (n <= N_max) and (n >= N_min): + mode_str += " 1" + else: + mode_str += " 0" else: mode_str += " 0" From 95d26c7867fbf8a364bec6f1376a4a390958db7a Mon Sep 17 00:00:00 2001 From: mfmart Date: Fri, 14 Jun 2024 08:23:37 -0700 Subject: [PATCH 26/49] set mode_family < 0 to use all mode families within active mode range --- desc/objectives/_terpsichore.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index aad0f380d0..d350ff5790 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -69,6 +69,7 @@ def terpsichore( nev=nev, al0=al0, ) + exit() _run_terps( dir=pid_path, exec=exec_path, @@ -247,7 +248,7 @@ def _write_terps_input( else: mode_str += " 0" - mode_str += "{:>3}\n".format(_in) + mode_str += "{:>3}\n".format(n) f.write(mode_str) f.write("C\n") From bc1f6fd140a10d9bafa2d1f21c67da952f58317c Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Mon, 17 Jun 2024 15:40:48 -0400 Subject: [PATCH 27/49] black formatting --- desc/objectives/_terpsichore.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index d350ff5790..89adde167b 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -69,7 +69,6 @@ def terpsichore( nev=nev, al0=al0, ) - exit() _run_terps( dir=pid_path, exec=exec_path, @@ -93,7 +92,7 @@ def _write_wout(eq, path, surfs): VMECIO.save(eq=eq, path=path, surfs=surfs, verbose=0) -def _write_terps_input( +def _write_terps_input( # noqa: C901 path, mode_family, surfs, @@ -223,8 +222,6 @@ def _write_terps_input( for n in range(-8, 11): - # If mode_family < 0, includes all modes in desired range - in_family = False if n % 2 == mode_family: # FIXME: must be modified for nfp != 2,3 in_family = True @@ -239,7 +236,7 @@ def _write_terps_input( mode_str = " 0" for m in range(1, 55 + 1): - if (m <= M_max): + if m <= M_max: if (mode_family < 0) or (in_family): if (n <= N_max) and (n >= N_min): mode_str += " 1" @@ -367,6 +364,7 @@ class TERPSICHORE(_ExternalObjective): Name of the objective function. # TODO: update Parameters docs + # If mode_family < 0, includes all modes in desired range """ @@ -386,10 +384,10 @@ def __init__( vectorized=True, path="", exec="", - mode_family=0, + mode_family=-1, surfs=16, - lssl=1, - lssd=1, + lssl=200, + lssd=100, M_max=2, N_min=-2, N_max=2, From 32ece01a2e03fe988d4cb7997e1b772e58b48e04 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Mon, 17 Jun 2024 17:53:55 -0400 Subject: [PATCH 28/49] streamlined write_wout function --- desc/objectives/_terpsichore.py | 276 +++++++++++++++++++++++++++++++- 1 file changed, 273 insertions(+), 3 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 89adde167b..4f8e914b6c 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -5,8 +5,13 @@ import time import numpy as np +from netCDF4 import Dataset, stringtochar +from desc.basis import DoubleFourierSeries +from desc.grid import LinearGrid +from desc.transform import Transform from desc.utils import errorif +from desc.vmec_utils import ptolemy_identity_rev, zernike_to_fourier from ._generic import _ExternalObjective @@ -79,17 +84,282 @@ def terpsichore( ) growth_rate = _read_terps_output(path=fort16_path) + # TODO: default value for when growth rate is not found (al0?) + # remove temporary directory shutil.rmtree(pid_path) return np.atleast_1d(growth_rate) -def _write_wout(eq, path, surfs): +def _write_wout(eq, path, surfs): # noqa: C901 """Write the wout NetCDF file from the equilibrium.""" - from desc.vmec import VMECIO + # this is a lightweight version of VMECIO.save + file = Dataset(path, mode="w", format="NETCDF3_64BIT_OFFSET") + + Psi = eq.Psi + NFP = eq.NFP + M = eq.M + N = eq.N + M_nyq = M + 4 + N_nyq = N + 2 if N > 0 else 0 + + # VMEC radial coordinate: s = rho^2 = Psi / Psi(LCFS) + s_full = np.linspace(0, 1, surfs) + s_half = s_full[0:-1] + 0.5 / (surfs - 1) + r_full = np.sqrt(s_full) + r_half = np.sqrt(s_half) + + # dimensions + file.createDimension("radius", surfs) # number of flux surfaces + file.createDimension("mn_mode", (2 * N + 1) * M + N + 1) # number of Fourier modes + file.createDimension( + "mn_mode_nyq", (2 * N_nyq + 1) * M_nyq + N_nyq + 1 + ) # number of Nyquist Fourier modes + file.createDimension("n_tor", N + 1) # number of toroidal Fourier modes + file.createDimension("preset", 21) # dimension of profile inputs + file.createDimension("ndfmax", 101) # used for am_aux & ai_aux + file.createDimension("time", 100) # used for fsq* & wdot + file.createDimension("dim_00001", 1) + file.createDimension("dim_00020", 20) + file.createDimension("dim_00100", 100) + file.createDimension("dim_00200", 200) + + grid_lcfs = LinearGrid(M=M_nyq, N=N_nyq, rho=np.array([1.0]), NFP=NFP) + grid_half = LinearGrid(M=M_nyq, N=N_nyq, NFP=NFP, rho=r_half) + data_half = eq.compute( + [ + "B_rho", + "B_theta", + "B_zeta", + "G", + "I", + "J", + "iota", + "p", + "sqrt(g)", + "V_r(r)", + "|B|", + "<|B|^2>", + ], + grid=grid_half, + ) + + mgrid_file = file.createVariable("mgrid_file", "S1", ("dim_00200",)) + mgrid_file[:] = stringtochar( + np.array(["none" + " " * 196], "S" + str(file.dimensions["dim_00200"].size)) + ) + + lasym = file.createVariable("lasym__logical__", np.int32) + lasym.long_name = "asymmetry logical (0 = stellarator symmetry)" + lasym[:] = int(not eq.sym) + + nfp = file.createVariable("nfp", np.int32) + nfp.long_name = "number of field periods" + nfp[:] = NFP + + ns = file.createVariable("ns", np.int32) + ns.long_name = "number of flux surfaces" + ns[:] = surfs + + mpol = file.createVariable("mpol", np.int32) + mpol.long_name = "number of poloidal Fourier modes" + mpol[:] = M + 1 + + ntor = file.createVariable("ntor", np.int32) + ntor.long_name = "number of positive toroidal Fourier modes" + ntor[:] = N + + mnmax = file.createVariable("mnmax", np.int32) + mnmax.long_name = "total number of Fourier modes" + mnmax[:] = file.dimensions["mn_mode"].size + + xm = file.createVariable("xm", np.float64, ("mn_mode",)) + xm.long_name = "poloidal mode numbers" + xm[:] = np.tile(np.linspace(0, M, M + 1), (2 * N + 1, 1)).T.flatten()[ + -file.dimensions["mn_mode"].size : + ] + + xn = file.createVariable("xn", np.float64, ("mn_mode",)) + xn.long_name = "toroidal mode numbers" + xn[:] = np.tile(np.linspace(-N, N, 2 * N + 1) * NFP, M + 1)[ + -file.dimensions["mn_mode"].size : + ] + + mnmax_nyq = file.createVariable("mnmax_nyq", np.int32) + mnmax_nyq.long_name = "total number of Nyquist Fourier modes" + mnmax_nyq[:] = file.dimensions["mn_mode_nyq"].size + + xm_nyq = file.createVariable("xm_nyq", np.float64, ("mn_mode_nyq",)) + xm_nyq.long_name = "poloidal Nyquist mode numbers" + xm_nyq[:] = np.tile( + np.linspace(0, M_nyq, M_nyq + 1), (2 * N_nyq + 1, 1) + ).T.flatten()[-file.dimensions["mn_mode_nyq"].size :] + + xn_nyq = file.createVariable("xn_nyq", np.float64, ("mn_mode_nyq",)) + xn_nyq.long_name = "toroidal Nyquist mode numbers" + xn_nyq[:] = np.tile(np.linspace(-N_nyq, N_nyq, 2 * N_nyq + 1) * NFP, M_nyq + 1)[ + -file.dimensions["mn_mode_nyq"].size : + ] + + gamma = file.createVariable("gamma", np.float64) + gamma.long_name = "compressibility index (0 = pressure prescribed)" + gamma[:] = 0 + + # half mesh quantities + + pres = file.createVariable("pres", np.float64, ("radius",)) + pres.long_name = "pressure on half mesh" + pres.units = "Pa" + pres[0] = 0 + pres[1:] = grid_half.compress(data_half["p"]) + + mass = file.createVariable("mass", np.float64, ("radius",)) + mass.long_name = "mass on half mesh" + mass.units = "Pa" + mass[:] = pres[:] + + iotas = file.createVariable("iotas", np.float64, ("radius",)) + iotas.long_name = "rotational transform on half mesh" + iotas.units = "None" + iotas[0] = 0 + iotas[1:] = -grid_half.compress(data_half["iota"]) # - for negative Jacobian + + phips = file.createVariable("phips", np.float64, ("radius",)) + phips.long_name = "d(phi)/ds * -1/2pi: toroidal flux derivative, on half mesh" + phips[0] = 0 + phips[1:] = -Psi * np.ones((surfs - 1,)) / (2 * np.pi) + + vp = file.createVariable("vp", np.float64, ("radius",)) + vp.long_name = "dV/ds normalized by 4*pi^2, on half mesh" + vp.units = "m^3" + vp[1:] = grid_half.compress(data_half["V_r(r)"]) / ( + 8 * np.pi**2 * grid_half.compress(data_half["rho"]) + ) + vp[0] = 0 + + # R + rmnc = file.createVariable("rmnc", np.float64, ("radius", "mn_mode")) + rmnc.long_name = "cos(m*t-n*p) component of cylindrical R, on full mesh" + rmnc.units = "m" + if not eq.sym: + rmns = file.createVariable("rmns", np.float64, ("radius", "mn_mode")) + rmns.long_name = "sin(m*t-n*p) component of cylindrical R, on full mesh" + rmns.units = "m" + r1 = np.ones_like(eq.R_lmn) + r1[eq.R_basis.modes[:, 1] < 0] *= -1 + m, n, x_mn = zernike_to_fourier(r1 * eq.R_lmn, basis=eq.R_basis, rho=r_full) + xm, xn, s, c = ptolemy_identity_rev(m, n, x_mn) + rmnc[:] = c + if not eq.sym: + rmns[:] = s + + # Z + zmns = file.createVariable("zmns", np.float64, ("radius", "mn_mode")) + zmns.long_name = "sin(m*t-n*p) component of cylindrical Z, on full mesh" + zmns.units = "m" + if not eq.sym: + zmnc = file.createVariable("zmnc", np.float64, ("radius", "mn_mode")) + zmnc.long_name = "cos(m*t-n*p) component of cylindrical Z, on full mesh" + zmnc.units = "m" + z1 = np.ones_like(eq.Z_lmn) + z1[eq.Z_basis.modes[:, 1] < 0] *= -1 + m, n, x_mn = zernike_to_fourier(z1 * eq.Z_lmn, basis=eq.Z_basis, rho=r_full) + xm, xn, s, c = ptolemy_identity_rev(m, n, x_mn) + zmns[:] = s + if not eq.sym: + zmnc[:] = c + + # derived quantities (approximate conversion) + + if eq.sym: + cos_basis = DoubleFourierSeries(M=M_nyq, N=N_nyq, NFP=NFP, sym="cos") + cos_transform = Transform( + grid=grid_lcfs, basis=cos_basis, build=False, build_pinv=True + ) + + def cosfit(x): + y = cos_transform.fit(x) + return np.where(cos_transform.basis.modes[:, 1] < 0, -y, y) - VMECIO.save(eq=eq, path=path, surfs=surfs, verbose=0) + else: + full_basis = DoubleFourierSeries(M=M_nyq, N=N_nyq, NFP=NFP, sym=None) + full_transform = Transform( + grid=grid_lcfs, basis=full_basis, build=False, build_pinv=True + ) + + def fullfit(x): + y = full_transform.fit(x) + return np.where(full_transform.basis.modes[:, 1] < 0, -y, y) + + # Jacobian + gmnc = file.createVariable("gmnc", np.float64, ("radius", "mn_mode_nyq")) + gmnc.long_name = "cos(m*t-n*p) component of Jacobian, on half mesh" + gmnc.units = "m" + m = cos_basis.modes[:, 1] + n = cos_basis.modes[:, 2] + if not eq.sym: + gmns = file.createVariable("gmns", np.float64, ("radius", "mn_mode_nyq")) + gmns.long_name = "sin(m*t-n*p) component of Jacobian, on half mesh" + gmns.units = "m" + m = full_basis.modes[:, 1] + n = full_basis.modes[:, 2] + # d(rho)/d(s) = 1/(2*rho) + data = ( + (data_half["sqrt(g)"] / (2 * data_half["rho"])) + .reshape( + (grid_half.num_theta, grid_half.num_rho, grid_half.num_zeta), order="F" + ) + .transpose((1, 0, 2)) + .reshape((grid_half.num_rho, -1), order="F") + ) + x_mn = np.zeros((surfs - 1, m.size)) + for i in range(surfs - 1): + if eq.sym: + x_mn[i, :] = cosfit(data[i, :]) + else: + x_mn[i, :] = fullfit(data[i, :]) + xm, xn, s, c = ptolemy_identity_rev(m, n, x_mn) + gmnc[0, :] = 0 + gmnc[1:, :] = -c # negative sign for negative Jacobian + if not eq.sym: + gmns[0, :] = 0 + gmns[1:, :] = -s + + # |B| + bmnc = file.createVariable("bmnc", np.float64, ("radius", "mn_mode_nyq")) + bmnc.long_name = "cos(m*t-n*p) component of |B|, on half mesh" + bmnc.units = "T" + m = cos_basis.modes[:, 1] + n = cos_basis.modes[:, 2] + if not eq.sym: + bmns = file.createVariable("bmns", np.float64, ("radius", "mn_mode_nyq")) + bmns.long_name = "sin(m*t-n*p) component of |B|, on half mesh" + bmns.units = "T" + m = full_basis.modes[:, 1] + n = full_basis.modes[:, 2] + data = ( + data_half["|B|"] + .reshape( + (grid_half.num_theta, grid_half.num_rho, grid_half.num_zeta), order="F" + ) + .transpose((1, 0, 2)) + .reshape((grid_half.num_rho, -1), order="F") + ) + x_mn = np.zeros((surfs - 1, m.size)) + for i in range(surfs - 1): + if eq.sym: + x_mn[i, :] = cosfit(data[i, :]) + else: + x_mn[i, :] = full_transform.fit(data[i, :]) + xm, xn, s, c = ptolemy_identity_rev(m, n, x_mn) + bmnc[0, :] = 0 + bmnc[1:, :] = c + if not eq.sym: + bmns[0, :] = 0 + bmns[1:, :] = s + + file.close() def _write_terps_input( # noqa: C901 From e941eab697a26e36d7e1c7fe82d4b2e70cf5d2cf Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Mon, 15 Jul 2024 17:25:37 -0600 Subject: [PATCH 29/49] set default growth rate to al0 --- desc/objectives/_terpsichore.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 4f8e914b6c..6595505275 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -54,6 +54,7 @@ def terpsichore( # copy executable to temporary directory shutil.copy(os.path.join(path, exec), exec_path) + # write input files _write_wout(eq=eq, path=wout_path, surfs=surfs) _write_terps_input( path=input_path, @@ -74,17 +75,20 @@ def terpsichore( nev=nev, al0=al0, ) - _run_terps( - dir=pid_path, - exec=exec_path, - input=input_path, - wout=wout_path, - sleep_time=sleep_time, - stop_time=stop_time, - ) - growth_rate = _read_terps_output(path=fort16_path) - # TODO: default value for when growth rate is not found (al0?) + # run TERPSICHORE + try: + _run_terps( + dir=pid_path, + exec=exec_path, + input=input_path, + wout=wout_path, + sleep_time=sleep_time, + stop_time=stop_time, + ) + growth_rate = _read_terps_output(path=fort16_path) + except RuntimeError: + growth_rate = al0 # default growth rate if it was unable to find one # remove temporary directory shutil.rmtree(pid_path) From 1836369a6e97541b05bddcd83f0f135f28d83129 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Mon, 22 Jul 2024 14:24:38 -0600 Subject: [PATCH 30/49] fix vmec save typo --- desc/vmec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/vmec.py b/desc/vmec.py index 7a68f52de9..a74deccb07 100644 --- a/desc/vmec.py +++ b/desc/vmec.py @@ -932,7 +932,7 @@ def fullfit(x): if eq.sym: x_mn[i, :] = cosfit(data[i, :]) else: - x_mn[i, :] = full_transform.fit(data[i, :]) + x_mn[i, :] = fullfit(data[i, :]) xm, xn, s, c = ptolemy_identity_rev(m, n, x_mn) bmnc[0, :] = 0 bmnc[1:, :] = c From 06bc1033ca4f410dcd1f773bac3dc24ef4e17169 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Mon, 22 Jul 2024 14:25:24 -0600 Subject: [PATCH 31/49] precompute vmec transforms --- desc/objectives/_terpsichore.py | 196 ++++++++++++++++++-------------- 1 file changed, 111 insertions(+), 85 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 6595505275..6891d38b32 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -8,12 +8,13 @@ from netCDF4 import Dataset, stringtochar from desc.basis import DoubleFourierSeries +from desc.compute.utils import get_transforms from desc.grid import LinearGrid from desc.transform import Transform from desc.utils import errorif from desc.vmec_utils import ptolemy_identity_rev, zernike_to_fourier -from ._generic import _ExternalObjective +from ._generic import ExternalObjective def terpsichore( @@ -37,6 +38,8 @@ def terpsichore( al0=-5e-1, sleep_time=1, stop_time=60, + data_transforms=None, + fit_transform=None, ): """TERPSICHORE driver function.""" process = mp.current_process() @@ -55,7 +58,13 @@ def terpsichore( shutil.copy(os.path.join(path, exec), exec_path) # write input files - _write_wout(eq=eq, path=wout_path, surfs=surfs) + _write_wout( + eq=eq, + path=wout_path, + surfs=surfs, + data_transforms=data_transforms, + fit_transform=fit_transform, + ) _write_terps_input( path=input_path, mode_family=mode_family, @@ -88,7 +97,7 @@ def terpsichore( ) growth_rate = _read_terps_output(path=fort16_path) except RuntimeError: - growth_rate = al0 # default growth rate if it was unable to find one + growth_rate = abs(al0) # default growth rate if it was unable to find one # remove temporary directory shutil.rmtree(pid_path) @@ -96,7 +105,7 @@ def terpsichore( return np.atleast_1d(growth_rate) -def _write_wout(eq, path, surfs): # noqa: C901 +def _write_wout(eq, path, surfs, data_transforms, fit_transform): # noqa: C901 """Write the wout NetCDF file from the equilibrium.""" # this is a lightweight version of VMECIO.save file = Dataset(path, mode="w", format="NETCDF3_64BIT_OFFSET") @@ -129,7 +138,6 @@ def _write_wout(eq, path, surfs): # noqa: C901 file.createDimension("dim_00100", 100) file.createDimension("dim_00200", 200) - grid_lcfs = LinearGrid(M=M_nyq, N=N_nyq, rho=np.array([1.0]), NFP=NFP) grid_half = LinearGrid(M=M_nyq, N=N_nyq, NFP=NFP, rho=r_half) data_half = eq.compute( [ @@ -147,6 +155,7 @@ def _write_wout(eq, path, surfs): # noqa: C901 "<|B|^2>", ], grid=grid_half, + transforms=data_transforms, ) mgrid_file = file.createVariable("mgrid_file", "S1", ("dim_00200",)) @@ -276,38 +285,20 @@ def _write_wout(eq, path, surfs): # noqa: C901 # derived quantities (approximate conversion) - if eq.sym: - cos_basis = DoubleFourierSeries(M=M_nyq, N=N_nyq, NFP=NFP, sym="cos") - cos_transform = Transform( - grid=grid_lcfs, basis=cos_basis, build=False, build_pinv=True - ) - - def cosfit(x): - y = cos_transform.fit(x) - return np.where(cos_transform.basis.modes[:, 1] < 0, -y, y) - - else: - full_basis = DoubleFourierSeries(M=M_nyq, N=N_nyq, NFP=NFP, sym=None) - full_transform = Transform( - grid=grid_lcfs, basis=full_basis, build=False, build_pinv=True - ) - - def fullfit(x): - y = full_transform.fit(x) - return np.where(full_transform.basis.modes[:, 1] < 0, -y, y) + def fit(x): + y = fit_transform.fit(x) + return np.where(fit_transform.basis.modes[:, 1] < 0, -y, y) # Jacobian gmnc = file.createVariable("gmnc", np.float64, ("radius", "mn_mode_nyq")) gmnc.long_name = "cos(m*t-n*p) component of Jacobian, on half mesh" gmnc.units = "m" - m = cos_basis.modes[:, 1] - n = cos_basis.modes[:, 2] if not eq.sym: gmns = file.createVariable("gmns", np.float64, ("radius", "mn_mode_nyq")) gmns.long_name = "sin(m*t-n*p) component of Jacobian, on half mesh" gmns.units = "m" - m = full_basis.modes[:, 1] - n = full_basis.modes[:, 2] + m = fit_transform.basis.modes[:, 1] + n = fit_transform.basis.modes[:, 2] # d(rho)/d(s) = 1/(2*rho) data = ( (data_half["sqrt(g)"] / (2 * data_half["rho"])) @@ -319,10 +310,7 @@ def fullfit(x): ) x_mn = np.zeros((surfs - 1, m.size)) for i in range(surfs - 1): - if eq.sym: - x_mn[i, :] = cosfit(data[i, :]) - else: - x_mn[i, :] = fullfit(data[i, :]) + x_mn[i, :] = fit(data[i, :]) xm, xn, s, c = ptolemy_identity_rev(m, n, x_mn) gmnc[0, :] = 0 gmnc[1:, :] = -c # negative sign for negative Jacobian @@ -334,14 +322,12 @@ def fullfit(x): bmnc = file.createVariable("bmnc", np.float64, ("radius", "mn_mode_nyq")) bmnc.long_name = "cos(m*t-n*p) component of |B|, on half mesh" bmnc.units = "T" - m = cos_basis.modes[:, 1] - n = cos_basis.modes[:, 2] if not eq.sym: bmns = file.createVariable("bmns", np.float64, ("radius", "mn_mode_nyq")) bmns.long_name = "sin(m*t-n*p) component of |B|, on half mesh" bmns.units = "T" - m = full_basis.modes[:, 1] - n = full_basis.modes[:, 2] + m = fit_transform.basis.modes[:, 1] + n = fit_transform.basis.modes[:, 2] data = ( data_half["|B|"] .reshape( @@ -352,10 +338,7 @@ def fullfit(x): ) x_mn = np.zeros((surfs - 1, m.size)) for i in range(surfs - 1): - if eq.sym: - x_mn[i, :] = cosfit(data[i, :]) - else: - x_mn[i, :] = full_transform.fit(data[i, :]) + x_mn[i, :] = fit(data[i, :]) xm, xn, s, c = ptolemy_identity_rev(m, n, x_mn) bmnc[0, :] = 0 bmnc[1:, :] = c @@ -422,23 +405,21 @@ def _write_terps_input( # noqa: C901 boz_str_title = "C M= 0" boz_str_neg = " 0" boz_str_pos = " 1" - for _im in range(1, 37): - if _im >= 10: - boz_str_title += " " + str(_im)[1] + for m in range(1, 37): + if m >= 10: + boz_str_title += " " + str(m)[1] else: - boz_str_title += " " + str(_im) + boz_str_title += " " + str(m) boz_str_neg += " 1" boz_str_pos += " 1" - boz_str_title += " N\n" f.write(boz_str_title) - for _in in range(-N_booz_max, N_booz_max + 1): - final_str_neg = boz_str_neg + "{:>3}\n".format(_in) - final_str_pos = boz_str_pos + "{:>3}\n".format(_in) - if _in < 0.0: - f.write(final_str_neg) + + for n in range(-N_booz_max, N_booz_max + 1): + if n < 0: + f.write(boz_str_neg + "{:>3}\n".format(n)) else: - f.write(final_str_pos) + f.write(boz_str_pos + "{:>3}\n".format(n)) f.write("C\n") f.write(" LLAMPR LVMTPR LMETPR LFOUPR\n") @@ -494,28 +475,20 @@ def _write_terps_input( # noqa: C901 + "0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 N\n" ) - for n in range(-8, 11): - - in_family = False - if n % 2 == mode_family: # FIXME: must be modified for nfp != 2,3 - in_family = True - - if n <= 0: - mode_str = " 0" - else: - if (n <= N_max) and (n >= N_min): - if (mode_family < 0) or (in_family): - mode_str = " 1" - else: - mode_str = " 0" - - for m in range(1, 55 + 1): - if m <= M_max: - if (mode_family < 0) or (in_family): - if (n <= N_max) and (n >= N_min): - mode_str += " 1" + for n in range(-9, 10): + mode_str = " " + for m in range(56): + if mode_family < 0 or n % 2 == mode_family: # FIXME: modify for nfp != 2,3 + if (m <= M_max) and (n >= N_min) and (n <= N_max): + if m == 0: + if n > 0: + mode_str += " 1" + else: + mode_str += " 0" else: - mode_str += " 0" + mode_str += " 1" + else: + mode_str += " 0" else: mode_str += " 0" @@ -604,7 +577,7 @@ def _read_terps_output(path): return growth_rates[0] -class TERPSICHORE(_ExternalObjective): +class TERPSICHORE(ExternalObjective): """Computes linear MHD stability from calls to the code TERPSICHORE. Returns the linear growth rate of the fastest growing instability. @@ -654,23 +627,24 @@ def __init__( normalize=False, normalize_target=False, loss_function=None, - fd_step=1e-4, - vectorized=True, + vectorized=False, + abs_step=1e-4, + rel_step=0, path="", exec="", mode_family=-1, surfs=16, lssl=200, lssd=100, - M_max=2, - N_min=-2, - N_max=2, - M_booz_max=4, - N_booz_max=4, - awall=2.0, + M_max=8, + N_min=-4, + N_max=4, + M_booz_max=19, + N_booz_max=18, + awall=1.3, xplo=1e-6, - deltajp=1e-2, - modelk=1, + deltajp=5e-1, + modelk=0, nev=1, al0=-5e-1, sleep_time=1, @@ -687,9 +661,9 @@ def __init__( normalize=normalize, normalize_target=normalize_target, loss_function=loss_function, - fd_step=fd_step, vectorized=vectorized, - name=name, + abs_step=abs_step, + rel_step=rel_step, path=path, exec=exec, mode_family=mode_family, @@ -709,4 +683,56 @@ def __init__( al0=al0, sleep_time=sleep_time, stop_time=stop_time, + name=name, + ) + + def build(self, use_jit=True, verbose=1): + """Build constant arrays. + + Parameters + ---------- + use_jit : bool, optional + Whether to just-in-time compile the objective and derivatives. + verbose : int, optional + Level of output. + + """ + # transforms for _write_wout + surfs = self._kwargs.get("surfs") + NFP = self._eq.NFP + M = self._eq.M + N = self._eq.N + M_nyq = M + 4 + N_nyq = N + 2 if N > 0 else 0 + s_full = np.linspace(0, 1, surfs) + s_half = s_full[0:-1] + 0.5 / (surfs - 1) + r_half = np.sqrt(s_half) + grid_lcfs = LinearGrid(M=M_nyq, N=N_nyq, rho=np.array([1.0]), NFP=NFP) + grid_half = LinearGrid(M=M_nyq, N=N_nyq, NFP=NFP, rho=r_half) + self._kwargs["data_transforms"] = get_transforms( + keys=[ + "B_rho", + "B_theta", + "B_zeta", + "G", + "I", + "J", + "iota", + "p", + "sqrt(g)", + "V_r(r)", + "|B|", + "<|B|^2>", + ], + obj=self._eq, + grid=grid_half, ) + if self._eq.sym: + fit_basis = DoubleFourierSeries(M=M_nyq, N=N_nyq, NFP=NFP, sym="cos") + else: + fit_basis = DoubleFourierSeries(M=M_nyq, N=N_nyq, NFP=NFP, sym=None) + self._kwargs["fit_transform"] = Transform( + grid=grid_lcfs, basis=fit_basis, build=False, build_pinv=True + ) + + super().build(use_jit=use_jit, verbose=verbose) From 685cf007654b762a70bf0b2e4c52eefe89ff0b27 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Mon, 22 Jul 2024 15:59:43 -0600 Subject: [PATCH 32/49] multiprocessing (at least on CPU, maybe GPU?) --- desc/objectives/_generic.py | 2 +- desc/objectives/_terpsichore.py | 159 ++++++++++++++++++++------------ 2 files changed, 101 insertions(+), 60 deletions(-) diff --git a/desc/objectives/_generic.py b/desc/objectives/_generic.py index b52b8e3fa2..ecd371b57f 100644 --- a/desc/objectives/_generic.py +++ b/desc/objectives/_generic.py @@ -77,7 +77,7 @@ class ExternalObjective(_Objective, ABC): """ _units = "(Unknown)" - _print_value_fmt = "External objective value: {:10.3e}" + _print_value_fmt = "External objective value: {:10.3e} " def __init__( self, diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 6891d38b32..6a75a4f093 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -1,3 +1,4 @@ +import functools import multiprocessing as mp import os import shutil @@ -7,6 +8,7 @@ import numpy as np from netCDF4 import Dataset, stringtochar +from desc.backend import jnp from desc.basis import DoubleFourierSeries from desc.compute.utils import get_transforms from desc.grid import LinearGrid @@ -19,6 +21,7 @@ def terpsichore( eq, + processes=1, path="", exec="", mode_family=-1, @@ -42,67 +45,80 @@ def terpsichore( fit_transform=None, ): """TERPSICHORE driver function.""" - process = mp.current_process() - pid = str(process.pid) + idxs = list(range(len(eq))) # equilibrium indices # create temporary directory to store I/O files - pid_path = os.path.join(path, pid) - os.mkdir(pid_path) - - exec_path = os.path.join(pid_path, exec) - input_path = os.path.join(pid_path, "input") - wout_path = os.path.join(pid_path, "wout.nc") - fort16_path = os.path.join(pid_path, "fort.16") - - # copy executable to temporary directory - shutil.copy(os.path.join(path, exec), exec_path) - - # write input files - _write_wout( - eq=eq, - path=wout_path, - surfs=surfs, - data_transforms=data_transforms, - fit_transform=fit_transform, - ) - _write_terps_input( - path=input_path, - mode_family=mode_family, - surfs=surfs, - lssl=lssl, - lssd=lssd, - M_max=M_max, - N_min=N_min, - N_max=N_max, - M_booz_max=M_booz_max, - N_booz_max=N_booz_max, - awall=awall, - xplo=xplo, - deltajp=deltajp, - modelk=modelk, - nfp=eq.NFP, - nev=nev, - al0=al0, - ) + tmp_path = os.path.join(path, "tmp-TERPS") + os.mkdir(tmp_path) + + # write input files for each equilibrium in serial + for k in idxs: + idx_path = os.path.join(tmp_path, str(k)) + os.mkdir(idx_path) + exec_path = os.path.join(idx_path, exec) + input_path = os.path.join(idx_path, "input") + wout_path = os.path.join(idx_path, "wout.nc") + shutil.copy(os.path.join(path, exec), exec_path) + _write_wout( + eq=eq[k], + path=wout_path, + surfs=surfs, + data_transforms=data_transforms, + fit_transform=fit_transform, + ) + _write_terps_input( + path=input_path, + mode_family=mode_family, + surfs=surfs, + lssl=lssl, + lssd=lssd, + M_max=M_max, + N_min=N_min, + N_max=N_max, + M_booz_max=M_booz_max, + N_booz_max=N_booz_max, + awall=awall, + xplo=xplo, + deltajp=deltajp, + modelk=modelk, + nfp=eq[k].NFP, + nev=nev, + al0=al0, + ) - # run TERPSICHORE - try: - _run_terps( - dir=pid_path, - exec=exec_path, - input=input_path, - wout=wout_path, - sleep_time=sleep_time, - stop_time=stop_time, + # run TERPSICHORE on list of equilibria in parallel + if len(eq) == 1: # no multiprocessing if only one equilibrium + result = jnp.atleast_1d( + _pool_fun( + 0, + path=tmp_path, + exec=exec, + al0=al0, + sleep_time=sleep_time, + stop_time=stop_time, + ) ) - growth_rate = _read_terps_output(path=fort16_path) - except RuntimeError: - growth_rate = abs(al0) # default growth rate if it was unable to find one + else: + with mp.Pool(processes=min(processes, len(eq))) as pool: + results = pool.map( + functools.partial( + _pool_fun, + path=tmp_path, + exec=exec, + al0=al0, + sleep_time=sleep_time, + stop_time=stop_time, + ), + idxs, + ) + pool.close() + pool.join() + result = jnp.vstack(results, dtype=float) # remove temporary directory - shutil.rmtree(pid_path) + shutil.rmtree(tmp_path) - return np.atleast_1d(growth_rate) + return result def _write_wout(eq, path, surfs, data_transforms, fit_transform): # noqa: C901 @@ -503,6 +519,30 @@ def _write_terps_input( # noqa: C901 f.close() +def _pool_fun(k, path, exec, al0, sleep_time, stop_time): + """Run TERPSICHORE and read growth rate for equilibrium with index k.""" + idx_path = os.path.join(path, str(k)) + exec_path = os.path.join(idx_path, exec) + fort16_path = os.path.join(idx_path, "fort.16") + input_path = os.path.join(idx_path, "input") + wout_path = os.path.join(idx_path, "wout.nc") + + try: + _run_terps( + dir=idx_path, + exec=exec_path, + input=input_path, + wout=wout_path, + sleep_time=sleep_time, + stop_time=stop_time, + ) + growth_rate = _read_terps_output(path=fort16_path) + except RuntimeError: + growth_rate = abs(al0) # default growth rate if it was unable to find one + + return np.atleast_1d(growth_rate) + + def _run_terps(dir, exec, input, wout, sleep_time, stop_time): """Run TERPSICHORE.""" stdout_path = os.path.join(dir, "stdout.terps") @@ -615,8 +655,8 @@ class TERPSICHORE(ExternalObjective): """ - _units = "(???)" - _print_value_fmt = "TERPSICHORE growth rate: {:10.3e}" + _units = "(?)" + _print_value_fmt = "TERPSICHORE growth rate: {:10.3e} " def __init__( self, @@ -627,9 +667,9 @@ def __init__( normalize=False, normalize_target=False, loss_function=None, - vectorized=False, abs_step=1e-4, rel_step=0, + processes=1, path="", exec="", mode_family=-1, @@ -661,9 +701,10 @@ def __init__( normalize=normalize, normalize_target=normalize_target, loss_function=loss_function, - vectorized=vectorized, + vectorized=True, abs_step=abs_step, rel_step=rel_step, + processes=processes, path=path, exec=exec, mode_family=mode_family, @@ -697,7 +738,7 @@ def build(self, use_jit=True, verbose=1): Level of output. """ - # transforms for _write_wout + # transforms for writing the wout file surfs = self._kwargs.get("surfs") NFP = self._eq.NFP M = self._eq.M From 557637f731f9b4f821966220154015d990e68c06 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 23 Jul 2024 13:26:08 -0600 Subject: [PATCH 33/49] set terpsichore fun to run on cpu --- desc/objectives/_terpsichore.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 6a75a4f093..ba47e1b743 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -8,7 +8,7 @@ import numpy as np from netCDF4 import Dataset, stringtochar -from desc.backend import jnp +from desc.backend import jnp, set_default_cpu from desc.basis import DoubleFourierSeries from desc.compute.utils import get_transforms from desc.grid import LinearGrid @@ -19,6 +19,7 @@ from ._generic import ExternalObjective +@set_default_cpu def terpsichore( eq, processes=1, From 6bb920e220bb9e965804c7006425f94c3d3b6cfc Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 23 Jul 2024 20:42:03 -0600 Subject: [PATCH 34/49] update defaults --- desc/objectives/_terpsichore.py | 28 +++++++++++++++------------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index ba47e1b743..aa5523ac40 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -26,14 +26,14 @@ def terpsichore( path="", exec="", mode_family=-1, - surfs=16, - lssl=200, - lssd=100, - M_max=2, - N_min=-2, - N_max=2, - M_booz_max=4, - N_booz_max=4, + surfs=32, + lssl=1000, + lssd=1000, + M_max=8, + N_min=-4, + N_max=4, + M_booz_max=19, + N_booz_max=18, awall=2.0, xplo=1e-6, deltajp=1e-2, @@ -656,7 +656,7 @@ class TERPSICHORE(ExternalObjective): """ - _units = "(?)" + _units = "(?)" # FIXME: normalized by Alfven frequency _print_value_fmt = "TERPSICHORE growth rate: {:10.3e} " def __init__( @@ -674,15 +674,15 @@ def __init__( path="", exec="", mode_family=-1, - surfs=16, - lssl=200, - lssd=100, + surfs=32, + lssl=1000, + lssd=1000, M_max=8, N_min=-4, N_max=4, M_booz_max=19, N_booz_max=18, - awall=1.3, + awall=2.0, xplo=1e-6, deltajp=5e-1, modelk=0, @@ -692,6 +692,8 @@ def __init__( stop_time=60, name="terpsichore", ): + if target is None and bounds is None: + bounds = (-np.inf, 0) super().__init__( eq=eq, fun=terpsichore, From c49c3e05ec17a183ff5556eebf1a0155458a5cb0 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Wed, 24 Jul 2024 10:40:17 -0600 Subject: [PATCH 35/49] execute_on_cpu --- desc/objectives/_terpsichore.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index aa5523ac40..5604733a8e 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -8,7 +8,7 @@ import numpy as np from netCDF4 import Dataset, stringtochar -from desc.backend import jnp, set_default_cpu +from desc.backend import execute_on_cpu, jnp from desc.basis import DoubleFourierSeries from desc.compute.utils import get_transforms from desc.grid import LinearGrid @@ -19,7 +19,7 @@ from ._generic import ExternalObjective -@set_default_cpu +@execute_on_cpu def terpsichore( eq, processes=1, From 782b23229a315d3951769659b4a7544923dce43b Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Fri, 9 Aug 2024 13:05:59 -0600 Subject: [PATCH 36/49] fix bug in write_terps_input --- desc/objectives/_terpsichore.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 5604733a8e..81295322b6 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -13,7 +13,7 @@ from desc.compute.utils import get_transforms from desc.grid import LinearGrid from desc.transform import Transform -from desc.utils import errorif +from desc.utils import errorif, warnif from desc.vmec_utils import ptolemy_identity_rev, zernike_to_fourier from ._generic import ExternalObjective @@ -403,8 +403,8 @@ def _write_terps_input( # noqa: C901 f.write("C\n") f.write("C MM NMIN NMAX MMS NSMIN NSMAX NPROCS INSOL\n") f.write( - " {:>2d} {:>3d} {:>3d} 55 -8 10 1 0\n".format( - M_booz_max, -N_booz_max, N_booz_max + " {:>2d} {:>3d} {:>3d} 55 {:>3d} {:>3d} 1 0\n".format( + M_booz_max, -N_booz_max, N_booz_max, N_min, N_max ) ) f.write("C\n") @@ -492,7 +492,7 @@ def _write_terps_input( # noqa: C901 + "0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 N\n" ) - for n in range(-9, 10): + for n in range(N_min, N_max + 1): mode_str = " " for m in range(56): if mode_family < 0 or n % 2 == mode_family: # FIXME: modify for nfp != 2,3 @@ -539,6 +539,9 @@ def _pool_fun(k, path, exec, al0, sleep_time, stop_time): ) growth_rate = _read_terps_output(path=fort16_path) except RuntimeError: + warnif( + True, UserWarning, "TERPSICHORE growth rate not found, using default value." + ) growth_rate = abs(al0) # default growth rate if it was unable to find one return np.atleast_1d(growth_rate) From c9d42e061fb4860ca3e942a021341301f455e605 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Mon, 26 Aug 2024 17:24:21 -0600 Subject: [PATCH 37/49] update units and print format --- desc/objectives/_terpsichore.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 81295322b6..bc0ec934e8 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -659,8 +659,8 @@ class TERPSICHORE(ExternalObjective): """ - _units = "(?)" # FIXME: normalized by Alfven frequency - _print_value_fmt = "TERPSICHORE growth rate: {:10.3e} " + _units = "(dimensionless)" # normalized by Alfven frequency + _print_value_fmt = "TERPSICHORE growth rate: " def __init__( self, From 90c26559c07f2e99c5b5d12de43232ce4251356c Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Wed, 28 Aug 2024 11:19:08 -0600 Subject: [PATCH 38/49] update notes --- desc/objectives/_terpsichore.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index bc0ec934e8..e1ad019502 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -656,6 +656,8 @@ class TERPSICHORE(ExternalObjective): # TODO: update Parameters docs # If mode_family < 0, includes all modes in desired range + # NOTE: that TERPSICHORE assumes theta=0 is defined on the outboard midplane! + # TODO: include documentation of example script for multiprocessing """ From 549105ad4338e9fd45b6adb1323fdc1c3ebf488e Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Thu, 7 Nov 2024 16:26:48 -0700 Subject: [PATCH 39/49] add scalar option for dW instead of growth rate --- desc/objectives/_terpsichore.py | 60 +++++++++++++++++++++++---------- 1 file changed, 43 insertions(+), 17 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index e1ad019502..5052a00172 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -25,6 +25,7 @@ def terpsichore( processes=1, path="", exec="", + scalar=True, mode_family=-1, surfs=32, lssl=1000, @@ -94,6 +95,8 @@ def terpsichore( 0, path=tmp_path, exec=exec, + scalar=scalar, + surfs=surfs, al0=al0, sleep_time=sleep_time, stop_time=stop_time, @@ -520,8 +523,8 @@ def _write_terps_input( # noqa: C901 f.close() -def _pool_fun(k, path, exec, al0, sleep_time, stop_time): - """Run TERPSICHORE and read growth rate for equilibrium with index k.""" +def _pool_fun(k, path, exec, scalar, surfs, al0, sleep_time, stop_time): + """Run TERPSICHORE and read output for equilibrium with index k.""" idx_path = os.path.join(path, str(k)) exec_path = os.path.join(idx_path, exec) fort16_path = os.path.join(idx_path, "fort.16") @@ -537,14 +540,15 @@ def _pool_fun(k, path, exec, al0, sleep_time, stop_time): sleep_time=sleep_time, stop_time=stop_time, ) - growth_rate = _read_terps_output(path=fort16_path) + output = _read_terps_output(path=fort16_path, scalar=scalar, surfs=surfs) except RuntimeError: warnif( True, UserWarning, "TERPSICHORE growth rate not found, using default value." ) - growth_rate = abs(al0) # default growth rate if it was unable to find one + # default values if it was unable to find a growth rate + output = abs(al0) if scalar else np.ones(surfs - 1) * al0**2 - return np.atleast_1d(growth_rate) + return np.atleast_1d(output) def _run_terps(dir, exec, input, wout, sleep_time, stop_time): @@ -597,34 +601,52 @@ def _is_terps_complete(path, run_time, stop_time, terps_subprocess): return False -def _read_terps_output(path): - """Read TERPSICHORE output file and return the growth rate.""" +def _read_terps_output(path, scalar, surfs): + """Read TERPSICHORE output file and return the growth rate / dW values.""" errorif( - not os.path.exists(path), RuntimeError, "TERPS fort.16 output file not found!" + not os.path.exists(path), + RuntimeError, + "TERPSICHORE fort.16 output file not found!", ) file = open(path) - growth_rates = [] - for line in file: + line_num = 0 + growth_rate = [] + delta_W = [] + for k, line in enumerate(file): index = line.find("GROWTH RATE") if index != -1: - growth_rates.append(float(line.strip().split("=")[1])) + growth_rate.append(float(line.strip().split("=")[1])) + continue + index = line.find("DELTAW") + if index != -1: + line_num = k + 4 # table values start 4 lines below heading + continue + if line_num > 0 and k == line_num: + row = line.strip().split() + delta_W.append(float(row[2])) + surf_idx = int(row[0]) + line_num = line_num + 1 + if surf_idx >= surfs - 1: + break # NOTE: not including vacuum surfaces file.close() - errorif(len(growth_rates) == 0, RuntimeError, "Growth rate not found!") + errorif(len(growth_rate) == 0, RuntimeError, "Growth rate not found!") errorif( - len(growth_rates) > 1, + len(growth_rate) > 1, NotImplementedError, "Not capable of handling multiple growth rates.", ) + errorif(len(delta_W) != surfs - 1, RuntimeError, "Incorrect number of dW values!") - return growth_rates[0] + return growth_rate[0] if scalar else delta_W class TERPSICHORE(ExternalObjective): """Computes linear MHD stability from calls to the code TERPSICHORE. - Returns the linear growth rate of the fastest growing instability. + Returns the linear growth rate of the fastest growing instability, + or the ΔW values at each flux surface. Parameters ---------- @@ -656,13 +678,14 @@ class TERPSICHORE(ExternalObjective): # TODO: update Parameters docs # If mode_family < 0, includes all modes in desired range + # If scalar = True return growth rate, if scalar = False return dW at surfaces # NOTE: that TERPSICHORE assumes theta=0 is defined on the outboard midplane! # TODO: include documentation of example script for multiprocessing """ _units = "(dimensionless)" # normalized by Alfven frequency - _print_value_fmt = "TERPSICHORE growth rate: " + _print_value_fmt = "TERPSICHORE " def __init__( self, @@ -678,6 +701,7 @@ def __init__( processes=1, path="", exec="", + scalar=True, mode_family=-1, surfs=32, lssl=1000, @@ -702,7 +726,7 @@ def __init__( super().__init__( eq=eq, fun=terpsichore, - dim_f=1, + dim_f=1 if scalar else surfs - 1, target=target, bounds=bounds, weight=weight, @@ -715,6 +739,7 @@ def __init__( processes=processes, path=path, exec=exec, + scalar=scalar, mode_family=mode_family, surfs=surfs, lssl=lssl, @@ -734,6 +759,7 @@ def __init__( stop_time=stop_time, name=name, ) + self._print_value_fmt += "growth rate: " if scalar else "delta W: " def build(self, use_jit=True, verbose=1): """Build constant arrays. From ced733690212cafa33b2adb6d2b0b1751d9e9559 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 12 Nov 2024 11:25:15 -0700 Subject: [PATCH 40/49] add new args to multiprocessing call --- desc/objectives/_terpsichore.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 5052a00172..02c19f61be 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -109,6 +109,8 @@ def terpsichore( _pool_fun, path=tmp_path, exec=exec, + scalar=scalar, + surfs=surfs, al0=al0, sleep_time=sleep_time, stop_time=stop_time, From 5030d7419e59f795e1bfb014d2525067f9977d7f Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Thu, 14 Nov 2024 15:54:36 -0600 Subject: [PATCH 41/49] better default values for delta W --- desc/objectives/_terpsichore.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index 02c19f61be..f82cca86d6 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -548,7 +548,7 @@ def _pool_fun(k, path, exec, scalar, surfs, al0, sleep_time, stop_time): True, UserWarning, "TERPSICHORE growth rate not found, using default value." ) # default values if it was unable to find a growth rate - output = abs(al0) if scalar else np.ones(surfs - 1) * al0**2 + output = abs(al0) if scalar else -np.ones(surfs - 1) * al0**2 return np.atleast_1d(output) @@ -724,7 +724,8 @@ def __init__( name="terpsichore", ): if target is None and bounds is None: - bounds = (-np.inf, 0) + # want to minimize growth rate or maximize delta W + bounds = (-np.inf, 0) if scalar else (0, np.inf) super().__init__( eq=eq, fun=terpsichore, From ea093e90024cf0422336a1a232e8ccecf5d2aedd Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 19 Nov 2024 14:46:08 -0700 Subject: [PATCH 42/49] comment notes about thresholds --- desc/objectives/_terpsichore.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/desc/objectives/_terpsichore.py b/desc/objectives/_terpsichore.py index f82cca86d6..11f534fb0c 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/objectives/_terpsichore.py @@ -725,6 +725,8 @@ def __init__( ): if target is None and bounds is None: # want to minimize growth rate or maximize delta W + # good threshold for growth rate is < -1e-3 + # good threshold for delta W is > +1e-4 bounds = (-np.inf, 0) if scalar else (0, np.inf) super().__init__( eq=eq, From 4933cd5566d7f2c78ccc9bf1514d7276a744c6e3 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Wed, 18 Dec 2024 12:25:56 -0700 Subject: [PATCH 43/49] experimental folder --- desc/experimental/README.rst | 6 ++++++ desc/experimental/__init__.py | 3 +++ desc/{objectives => experimental}/_terpsichore.py | 3 +-- desc/objectives/__init__.py | 1 - 4 files changed, 10 insertions(+), 3 deletions(-) create mode 100644 desc/experimental/README.rst create mode 100644 desc/experimental/__init__.py rename desc/{objectives => experimental}/_terpsichore.py (99%) diff --git a/desc/experimental/README.rst b/desc/experimental/README.rst new file mode 100644 index 0000000000..f660aa4c44 --- /dev/null +++ b/desc/experimental/README.rst @@ -0,0 +1,6 @@ +**Warning:** + +These objectives depend on other codes external to DESC. They are not routinely tested +and may not be compatible with all versions of the external codes. The DESC team is not +responsible for maintaining or documenting those external codes, and we do not guarantee +to regularly maintain these objectives. Use these at your own risk! diff --git a/desc/experimental/__init__.py b/desc/experimental/__init__.py new file mode 100644 index 0000000000..ce774f513c --- /dev/null +++ b/desc/experimental/__init__.py @@ -0,0 +1,3 @@ +"""Classes defining objectives that wrap external codes.""" + +from ._terpsichore import TERPSICHORE diff --git a/desc/objectives/_terpsichore.py b/desc/experimental/_terpsichore.py similarity index 99% rename from desc/objectives/_terpsichore.py rename to desc/experimental/_terpsichore.py index 11f534fb0c..f4026a2416 100644 --- a/desc/objectives/_terpsichore.py +++ b/desc/experimental/_terpsichore.py @@ -12,12 +12,11 @@ from desc.basis import DoubleFourierSeries from desc.compute.utils import get_transforms from desc.grid import LinearGrid +from desc.objectives import ExternalObjective from desc.transform import Transform from desc.utils import errorif, warnif from desc.vmec_utils import ptolemy_identity_rev, zernike_to_fourier -from ._generic import ExternalObjective - @execute_on_cpu def terpsichore( diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py index 13fdb17c80..8cbd4b4d77 100644 --- a/desc/objectives/__init__.py +++ b/desc/objectives/__init__.py @@ -53,7 +53,6 @@ from ._power_balance import FusionPower, HeatingPowerISS04 from ._profiles import Pressure, RotationalTransform, Shear, ToroidalCurrent from ._stability import BallooningStability, MagneticWell, MercierStability -from ._terpsichore import TERPSICHORE from .getters import ( get_equilibrium_objective, get_fixed_axis_constraints, From 26c8bb264206653245a2dd8e52246649f8f3a1d8 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Wed, 18 Dec 2024 13:22:27 -0700 Subject: [PATCH 44/49] reorder vectorized arg --- desc/experimental/_terpsichore.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/desc/experimental/_terpsichore.py b/desc/experimental/_terpsichore.py index f4026a2416..57f3a1d7cb 100644 --- a/desc/experimental/_terpsichore.py +++ b/desc/experimental/_terpsichore.py @@ -704,7 +704,7 @@ def __init__( exec="", scalar=True, mode_family=-1, - surfs=32, + surfs=101, lssl=1000, lssd=1000, M_max=8, @@ -731,13 +731,13 @@ def __init__( eq=eq, fun=terpsichore, dim_f=1 if scalar else surfs - 1, + vectorized=True, target=target, bounds=bounds, weight=weight, normalize=normalize, normalize_target=normalize_target, loss_function=loss_function, - vectorized=True, abs_step=abs_step, rel_step=rel_step, processes=processes, From 58ff25313dc7df8f42b5ddf33a3153d43807c933 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Wed, 18 Dec 2024 17:13:17 -0700 Subject: [PATCH 45/49] tutorial for using external codes --- desc/experimental/_terpsichore.py | 2 +- docs/external_objectives.rst | 320 ++++++++++++++++++++++++++++++ docs/index.rst | 5 +- 3 files changed, 322 insertions(+), 5 deletions(-) create mode 100644 docs/external_objectives.rst diff --git a/desc/experimental/_terpsichore.py b/desc/experimental/_terpsichore.py index 57f3a1d7cb..5b2efd90a9 100644 --- a/desc/experimental/_terpsichore.py +++ b/desc/experimental/_terpsichore.py @@ -26,7 +26,7 @@ def terpsichore( exec="", scalar=True, mode_family=-1, - surfs=32, + surfs=101, lssl=1000, lssd=1000, M_max=8, diff --git a/docs/external_objectives.rst b/docs/external_objectives.rst new file mode 100644 index 0000000000..b0a1f7db5c --- /dev/null +++ b/docs/external_objectives.rst @@ -0,0 +1,320 @@ +Using external codes +-------------------- + +The ``ExternalObjective`` class allows external codes to be used within a DESC +optimization. This takes a user-supplied function that does not need not be JAX +transformable, and wraps it with forward finite differences. This is most useful for +cases when the function needs to call external codes that are written in another +language besides Python, or cannot be rewritten with JAX for some other reason. If the +user function can be made JAX transformable, it is recommended to use the +``ObjectiveFromUser`` class instead. + +This guide begins with a simple example of how to use a custom function with the +``ExternalObjective``. It then uses the ``TERPSICHORE`` objective as a more realistic +example, which can be used as a template for wrapping other codes. Finally, an example +script is shown for how to run the TERPSICHORE optimization with multiprocessing. + +Simple example +-------------- + +The following is a simple example of a custom function that is not JAX transformable. +It was adapted from the test ``test_external_vs_generic_objectives`` in +``tests/test_examples.py``. + +The function must take a single positional argument, which can be either a single +Equilibrium or a list of Equilibria. Additional inputs can be passed as keyword +arguments. In this example, the function returns the three scalar quatities +:math:`\langle\beta\rangle_{vol}`, :math:`\langle\beta\rangle_{pol}`, and +:math:`\langle\beta\rangle_{tor}`. It writes these quantities to a NetCDF file and then +reads them back from the file. This mimics I/O from the common VMEC "wout" format, and +also demonstrates an operation that is not automatically differentiable. + +:: + + def beta_from_file(eq, path=""): + # write data + file = Dataset(path, mode="w", format="NETCDF3_64BIT_OFFSET") + data = eq.compute(["_vol", "_vol", "_vol"]) + betatotal = file.createVariable("betatotal", np.float64) + betatotal[:] = data["_vol"] + betapol = file.createVariable("betapol", np.float64) + betapol[:] = data["_vol"] + betator = file.createVariable("betator", np.float64) + betator[:] = data["_vol"] + file.close() + # read data + file = Dataset(path, mode="r") + betatot = float(file.variables["betatotal"][0]) + betapol = float(file.variables["betapol"][0]) + betator = float(file.variables["betator"][0]) + file.close() + return np.atleast_1d([betatot, betapol, betator]) + +This function can then be used in an optimization with the ``ExternalObjective`` class. +This objective must be initialized with the Equilibrium to be optimized, ``eq``, and the +following other arguments: + +* ``fun``, the external function. In this example it is the function ``beta_from_file``. +* ``dim_f``, the dimension of the output of the function ``fun``. In this example + ``dim_f=3`` because ``beta_from_file`` returns an array of size 3. +* ``vectorized``, whether ``fun`` expects a single Equilibrium or a list of + Equilibria. Since the function ``beta_from_file`` takes a single Equilibrium as its + only positional argument, ``vectorized=False``. + +All other input parameters that the external function requires can be passed as keyword +arguments to ``ExternalObjective``. In this example, ``path`` specifies the file name +of the NetCDF file that the function writes/reads. Since ``vectorized=False``, the +function ``beta_from_file`` will be evaluated sequentially when computing the finite +differences (rather than in parallel), so the same file will be overwritten each time +for the different Equilibria. + +:: + + from desc.objectives import ExternalObjective, ObjectiveFunction + + objective = ObjectiveFunction( + ExternalObjective( + eq=eq, + fun=beta_from_file, + dim_f=3, + vectorized=False, + path="path/to/file.nc", + ) + ) + +Note that for this example, the same optimization objective could be accomplished +naitively in DESC with automatic differentiation instead of finite differences as: + +:: + + from desc.objectives import GenericObjective, ObjectiveFunction + + objective = ObjectiveFunction( + ( + GenericObjective("_vol", thing=eq), + GenericObjective("_vol", thing=eq), + GenericObjective("_vol", thing=eq), + ) + ) + + +TERPSICHORE objective +--------------------- + +This section walks through the implementation of wrapping the ideal MHD linear stability +code TERPSICHORE, which is written in FORTRAN. It summarizes how the external function +was written to call TERPSICHORE from DESC, and can be used as a template for wrapping +other codes. The code is abbreviated for brevity, but the full source code can be found +in ``desc/experimental/_terpsichore.py``. + +The external function in this case is named ``terpsichore``, and takes the following +arguments: + +* ``eq``, a list of Equilibria. +* ``processes``, the maximum number of processes to use for multiprocessing. +* ``path``, the path to the directory where the optimization script is being executed. +* ``exec``, the file name of the TERPSICHORE executable. + +The ``TERPSICHORE`` objective takes other arguments that have been omitted for +simplicity. The outline of this function is: + +1. Create a temporary directory where I/O files will be written. +2. Write the equilibria data to files in the format that TERPSICHORE expects from a VMEC + input. +3. Write the TERPSICHORE input files for each equilibrium. +4. Run TERPSICHORE with the inputs from steps 2 and 3. +5. Read the instability growth rates from the TERPSICHORE output files for each + equilibrium. +6. Return the growth rates. + +Running TERPSICHORE is relatively slow compared to other computations in DESC. The +bottleneck of an optimization is computing the Jacobian matrix with finite differences, +which scales with the number of optimization degrees of freedom. Evaluating the +TERPSICHORE growth rates for each Equilibrium can be performed in parallel on different +CPU threads using Python multiprocessing. Note that writing the equilibria data in step +2 cannot be easily parallelized, since it involves computations using JAX that has +issues with multiprocessing. + +:: + + # TERPSICHORE only runs on a CPU, but DESC is optimized to run on a GPU. + # This decorator will run this function on a CPU, even if other functions are being + # run on a GPU. + @execute_on_cpu + def terpsichore(eq, processes=1, path="", exec=""): + """TERPSICHORE driver function.""" + # these indices are used to give each equilibrium's I/O files unique file names + idxs = list(range(len(eq))) # equilibrium indices + + # create temporary directory to store I/O files + tmp_path = os.path.join(path, "tmp-TERPS") + os.mkdir(tmp_path) + + # write input files for each equilibrium in serial + for k in idxs: + # create a sub-directory for each equilibrium + idx_path = os.path.join(tmp_path, str(k)) + os.mkdir(idx_path) + exec_path = os.path.join(idx_path, exec) + input_path = os.path.join(idx_path, "input") + wout_path = os.path.join(idx_path, "wout.nc") + shutil.copy(os.path.join(path, exec), exec_path) + _write_wout(eq=eq[k], path=wout_path) # write equilibrium input data + _write_terps_input(path=input_path) # write TERPSICHORE input file + + # run TERPSICHORE on list of equilibria in parallel + if len(eq) == 1: # no multiprocessing if only one equilibrium + result = jnp.atleast_1d(_pool_fun(0, path=tmp_path, exec=exec)) + else: # use multiprocessing if there are multiple equilibria + with mp.Pool(processes=min(processes, len(eq))) as pool: + results = pool.map( + functools.partial(_pool_fun, path=tmp_path, exec=exec), + idxs, + ) + pool.close() + pool.join() + result = jnp.vstack(results, dtype=float) + + # remove temporary directory and all sub-directories + shutil.rmtree(tmp_path) + + return result + +The function ``_write_wout`` is a simplified version of ``VMECIO.save`` that only saves +the output quantities that TERPSICHORE needs. Avoiding computation of the unnecessary +quantities greatly reduces the overall run time. The function +``_write_terps_input`` writes the TERPSICHORE input file, which is a text file with a +specific format. The details of these two functions are not important for the scope of +this guide. + +``_pool_fun`` is the function that is run in parallel for each Equilibrium. It calls +``_run_terps`` (also shown below) to execute the TERPSICHORE Fortran code through a +Python subprocess call, and ``_read_terps_output`` (not shown) to parse the output file +and extract the instability growth rate. + +:: + + def _pool_fun(k, path, exec): + """Run TERPSICHORE and read output for equilibrium with index k.""" + idx_path = os.path.join(path, str(k)) + exec_path = os.path.join(idx_path, exec) + fort16_path = os.path.join(idx_path, "fort.16") + input_path = os.path.join(idx_path, "input") + wout_path = os.path.join(idx_path, "wout.nc") + + try: # try to run TERPSICHORE and read the growth rate from the output file + _run_terps(dir=idx_path, exec=exec_path, input=input_path, wout=wout_path) + output = _read_terps_output(path=fort16_path) + except RuntimeError: + output = 1.0 # default value if TERPSICHORE failed to run + + return np.atleast_1d(output) + + + def _run_terps(dir, exec, input, wout): + """Run TERPSICHORE.""" + stdout_path = os.path.join(dir, "stdout.terps") + stderr_path = os.path.join(dir, "stderr.terps") + + fout = open(stdout_path, "w") + ferr = open(stderr_path, "w") + + # execute TERPSICHORE + cmd = exec + " < " + input + " " + wout + terps_subprocess = subprocess.run( + cmd, cwd=dir, shell=True, stdout=fout, stderr=ferr + ) + + # not shown: a delay to wait until TERPSICHORE finishes running + terps_subprocess.terminate() + + fout.close() + ferr.close() + +Finally, the ``TERPSICHORE`` objective function simply inherits from +``ExternalObjective`` and passes ``fun=terpsichore`` as the external function. +``dim_f=1`` because TERPSICHORE is returning a scalar growth rate in this example, and +``vectorized=True`` because the function ``terpsichore`` expects a list of Equilibria +as its only positional argument. (Parts of the full class definition have been ommitted +here for simplicity.) + +:: + + class TERPSICHORE(ExternalObjective): + """Computes linear MHD stability from calls to the code TERPSICHORE.""" + + def __init__(self, eq, processes=1, path="", exec=""): + super().__init__( + eq=eq, + fun=terpsichore, + dim_f=1, + vectorized=True, + processes=processes, + path=path, + exec=exec, + ) + +Multiprocessing +--------------- + +Due to complexities of Python multiprocessing, one must guard against spawning unwanted +child processes. The following is a simple example script for performing an optimization +with the ``TERPSICHORE`` objective function. Note that the step size used in the finite +differencing of ``ExternalObjective`` can be controlled with the arguments ``abs_step`` +and ``rel_step``. ``processes=os.cpu_count()`` will use the maximum number of CPU +threads that are available. + +:: + + import os + import sys + + import multiprocessing as mp + import numpy as np + + # this ensures that this driver code only runs once, for the main process + if mp.current_process().name == "MainProcess": + from desc import set_device + + set_device("gpu") + + import numpy as np + + from desc.examples import get + from desc.experimental import TERPSICHORE + from desc.objectives import ( + ForceBalance, + FixBoundaryR, + FixCurrent, + FixPressure, + FixPsi, + ObjectiveFunction, + ) + from desc.optimize import Optimizer + + eq = get("precise_QA") + optimizer = Optimizer("proximal-lsq-exact") + objective = ObjectiveFunction( + ( + TERPSICHORE( + eq=eq, + abs_step=1e-2, + rel_step=0, + processes=os.cpu_count(), + path=os.getcwd(), + exec="terps_exec.x", + ), + ), + ) + constraints = ( + FixBoundaryR(eq=eq, modes=np.array([[0, 0, 0]])), + FixCurrent(eq=eq), + FixPressure(eq=eq), + FixPsi(eq=eq), + ForceBalance(eq=eq), + ) + [eq], _ = optimizer.optimize( + things=eq, + objective=objective, + constraints=constraints, + ) diff --git a/docs/index.rst b/docs/index.rst index 4d8bda4bf1..7e0ac2e843 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -49,7 +49,6 @@ api_fields api - .. toctree:: :maxdepth: 1 :caption: Command Line Interface and I/O @@ -58,7 +57,6 @@ input output - .. toctree:: :maxdepth: 1 :caption: Developer guides @@ -67,10 +65,9 @@ notebooks/dev_guide/grid.ipynb adding_compute_funs adding_objectives + external_objectives adding_optimizers - - Indices and tables ================== From 16cfa89491321bfd60dd32fc5e0944277cc95681 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Wed, 18 Dec 2024 17:38:57 -0700 Subject: [PATCH 46/49] minor edits --- desc/experimental/_terpsichore.py | 2 +- docs/external_objectives.rst | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/desc/experimental/_terpsichore.py b/desc/experimental/_terpsichore.py index 5b2efd90a9..e791979375 100644 --- a/desc/experimental/_terpsichore.py +++ b/desc/experimental/_terpsichore.py @@ -644,7 +644,7 @@ def _read_terps_output(path, scalar, surfs): class TERPSICHORE(ExternalObjective): - """Computes linear MHD stability from calls to the code TERPSICHORE. + """Computes ideal MHD linear stability from calls to the code TERPSICHORE. Returns the linear growth rate of the fastest growing instability, or the ΔW values at each flux surface. diff --git a/docs/external_objectives.rst b/docs/external_objectives.rst index b0a1f7db5c..1810f7e2e2 100644 --- a/docs/external_objectives.rst +++ b/docs/external_objectives.rst @@ -241,7 +241,7 @@ here for simplicity.) :: class TERPSICHORE(ExternalObjective): - """Computes linear MHD stability from calls to the code TERPSICHORE.""" + """Computes ideal MHD linear stability from calls to the code TERPSICHORE.""" def __init__(self, eq, processes=1, path="", exec=""): super().__init__( @@ -285,14 +285,14 @@ threads that are available. from desc.objectives import ( ForceBalance, FixBoundaryR, - FixCurrent, + FixIota, FixPressure, FixPsi, ObjectiveFunction, ) from desc.optimize import Optimizer - eq = get("precise_QA") + eq = get("W7-X") optimizer = Optimizer("proximal-lsq-exact") objective = ObjectiveFunction( ( @@ -308,7 +308,7 @@ threads that are available. ) constraints = ( FixBoundaryR(eq=eq, modes=np.array([[0, 0, 0]])), - FixCurrent(eq=eq), + FixIota(eq=eq), FixPressure(eq=eq), FixPsi(eq=eq), ForceBalance(eq=eq), From 1ef649249fe774fddeee11960a935e56a6a55a43 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Thu, 19 Dec 2024 14:08:49 -0700 Subject: [PATCH 47/49] update objective args docs --- desc/experimental/_terpsichore.py | 145 +++++++++++++++++++++--------- docs/external_objectives.rst | 12 +-- 2 files changed, 107 insertions(+), 50 deletions(-) diff --git a/desc/experimental/_terpsichore.py b/desc/experimental/_terpsichore.py index e791979375..47ace2a4eb 100644 --- a/desc/experimental/_terpsichore.py +++ b/desc/experimental/_terpsichore.py @@ -10,9 +10,11 @@ from desc.backend import execute_on_cpu, jnp from desc.basis import DoubleFourierSeries +from desc.compat import flip_theta from desc.compute.utils import get_transforms from desc.grid import LinearGrid from desc.objectives import ExternalObjective +from desc.objectives.objective_funs import collect_docs from desc.transform import Transform from desc.utils import errorif, warnif from desc.vmec_utils import ptolemy_identity_rev, zernike_to_fourier @@ -35,8 +37,7 @@ def terpsichore( M_booz_max=19, N_booz_max=18, awall=2.0, - xplo=1e-6, - deltajp=1e-2, + deltaJp=1e-2, modelk=1, nev=1, al0=-5e-1, @@ -44,15 +45,19 @@ def terpsichore( stop_time=60, data_transforms=None, fit_transform=None, + theta0_outboard=True, ): """TERPSICHORE driver function.""" - idxs = list(range(len(eq))) # equilibrium indices + # TERPSICHORE assumes theta=0 is on the outboard midplane + if not theta0_outboard: + eq = flip_theta(eq) # create temporary directory to store I/O files tmp_path = os.path.join(path, "tmp-TERPS") os.mkdir(tmp_path) # write input files for each equilibrium in serial + idxs = list(range(len(eq))) # equilibrium indices for k in idxs: idx_path = os.path.join(tmp_path, str(k)) os.mkdir(idx_path) @@ -79,8 +84,7 @@ def terpsichore( M_booz_max=M_booz_max, N_booz_max=N_booz_max, awall=awall, - xplo=xplo, - deltajp=deltajp, + deltaJp=deltaJp, modelk=modelk, nfp=eq[k].NFP, nev=nev, @@ -382,8 +386,7 @@ def _write_terps_input( # noqa: C901 M_booz_max, N_booz_max, awall, - xplo, - deltajp, + deltaJp, modelk, nfp, nev, @@ -474,9 +477,7 @@ def _write_terps_input( # noqa: C901 f.write("C\n") f.write("C RPLMIN XPLO DELTAJP WCT CURFAC\n") f.write( - " 1.0000e-05 {:10.4e} {:10.4e} 6.6667e-01 1.0000e-00\n".format( - xplo, deltajp - ) + " 1.0000e-05 1.0000e-06 {:10.4e} 6.6667e-01 1.0000e-00\n".format(deltaJp) ) f.write("C\n") f.write( @@ -653,39 +654,87 @@ class TERPSICHORE(ExternalObjective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - Has no effect for this objective. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - name : str, optional - Name of the objective function. - - # TODO: update Parameters docs - # If mode_family < 0, includes all modes in desired range - # If scalar = True return growth rate, if scalar = False return dW at surfaces - # NOTE: that TERPSICHORE assumes theta=0 is defined on the outboard midplane! - # TODO: include documentation of example script for multiprocessing + abs_step : float, optional + Absolute finite difference step size. Default = 1e-4. + Total step size is ``abs_step + rel_step * mean(abs(x))``. + rel_step : float, optional + Relative finite difference step size. Default = 0. + Total step size is ``abs_step + rel_step * mean(abs(x))``. + processes : int, optional + Maximum number of CPU threads to use for multiprocessing. Default = 1. + path : str, optional + Path to the directory where temporary files will be stored. + exec : str, optional + File name of the TERPSICHORE executable. Must be located in the directory + specified by ``path``. + scalar : bool, optional + If Ture, returns the growth rate of the fastest growing instability (default). + If False, returns the ΔW values at each flux surface. + mode_family : int, optional + The mode family of the instabilities to consider. Only implemented for NFP = + 2 or 3, which only have two mode families: N=0 and N=1. If ``mode_family < 0``, + both mode families are considered. Default = -1. + surfs : int, optional + Number of surfaces to include in the equilibrium input. More surfaces provides + more accuracy at the cost of longer compute times. Default = 101. + lssl : int, optional + Minimum number of possible permutations of Boozer mode combinations + (determined by ``M_booz_max`` and ``N_booz_max``). If TERPSICHORE fails to run, + try increasing this parameter. Default = 1000. + lssd : int, optional + Minimum number of possible permutations of stability mode combinations + (determined by ``M_max`` and ``N_max``). If TERPSICHORE fails to run, + try increasing this parameter. Default = 1000. + M_max : int, optional + Maximum poloidal mode number of stability modes to consider. Will include modes + with m ∈ [0,M_max] (if ``mode_family < 0``). Default = 8. + N_min : int, optional + Minimum toroidal mode number of stability modes to consider. Will include modes + with n ∈ [N_min,N_max] (if ``mode_family < 0``). Default = -4. + N_max : int, optional + Maximum toroidal mode number of stability modes to consider. Will include modes + with n ∈ [N_min,N_max] (if ``mode_family < 0``). Default = 4. + M_booz_max : int, optional + Maximum poloidal mode number of Boozer spectrum. Will include modes with + m ∈ [0,M_booz_max]. Default = 19. + N_booz_max : int, optional + Maximum poloidal mode number of Boozer spectrum. Will include modes with + n ∈ [-N_booz_max,N_booz_max]. Default = 18. + awall : float, optional + Ratio of the radius of the conformal conducting wall to the plasma minor radius. + The conducting wall is obtained by scaling the m ≠ 0 Fourier components of the + plasma boundary by ``awall``. A shorter wall offset will help stabilize the + plasma. If TERPSICHORE fails to run, try decreasing this parameter. Default = 2. + deltaJp : float, optional + Resonance detuning parameter to resolve parallel current density singularities. + Default = 0.5. + modelk : int, optional + 0 = Noninteracting anisotropic fast particle stability model with reduced + kintetic energy. 1 = Kruskal-Oberman anisotropic energy principle model with + reduced kinetic energy. 2 = Noninteracting anisotropic fast particle stability + model with physical kinetic energy. 3 = Kruskal-Oberman anisotropic energy + principle model with physical kinetic energy. Default = 0. + nev : int, optional + Number of eigenvalue computations. Default = 1. + al0 : float, optional + Initial guess of the eigenvalue. Use a sufficiently negative value to find the + most unstable growth rate. If TERPSICHORE fails to run, the objective will + return a growth rate of ``abs(al0)``. Default = -0.5. + sleep_time : float, optional + Time in seconds to wait between checks to determine if TERPSICHORE has finished + executing. Default = 1. + stop_time : float, optional + Time in seconds to wait for TERPSICHORE to execute before terminating its run. + Default = 60. """ - _units = "(dimensionless)" # normalized by Alfven frequency + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``bounds=(-np.inf, 0)``.", + bounds_default="``bounds=(-np.inf, 0)``.", + ) + + _units = "(dimensionless)" _print_value_fmt = "TERPSICHORE " def __init__( @@ -713,8 +762,7 @@ def __init__( M_booz_max=19, N_booz_max=18, awall=2.0, - xplo=1e-6, - deltajp=5e-1, + deltaJp=5e-1, modelk=0, nev=1, al0=-5e-1, @@ -754,8 +802,7 @@ def __init__( M_booz_max=M_booz_max, N_booz_max=N_booz_max, awall=awall, - xplo=xplo, - deltajp=deltajp, + deltaJp=deltaJp, modelk=modelk, nev=nev, al0=al0, @@ -763,6 +810,7 @@ def __init__( stop_time=stop_time, name=name, ) + self._scalar = scalar self._print_value_fmt += "growth rate: " if scalar else "delta W: " def build(self, use_jit=True, verbose=1): @@ -776,6 +824,15 @@ def build(self, use_jit=True, verbose=1): Level of output. """ + # check if theta needs to be flipped so that theta=0 is on the outboard midplane + R0 = self._eq.Rb_lmn[self._eq.surface.R_basis.get_idx(M=0, N=0)] + R1 = np.sum( + self._eq.Rb_lmn[ + np.where((self._eq.surface.R_basis.modes[:, 1:] >= 0).all(axis=1))[0] + ] + ) + self._kwargs["theta0_outboard"] = R1 > R0 # R(rho=1,theta=0,phi=0) > R(rho=0) + # transforms for writing the wout file surfs = self._kwargs.get("surfs") NFP = self._eq.NFP diff --git a/docs/external_objectives.rst b/docs/external_objectives.rst index 1810f7e2e2..4d9249639a 100644 --- a/docs/external_objectives.rst +++ b/docs/external_objectives.rst @@ -104,8 +104,8 @@ TERPSICHORE objective This section walks through the implementation of wrapping the ideal MHD linear stability code TERPSICHORE, which is written in FORTRAN. It summarizes how the external function was written to call TERPSICHORE from DESC, and can be used as a template for wrapping -other codes. The code is abbreviated for brevity, but the full source code can be found -in ``desc/experimental/_terpsichore.py``. +other codes. The code shown here is abbreviated and slightly modified for brevity, but +the full source code can be found in ``desc/experimental/_terpsichore.py``. The external function in this case is named ``terpsichore``, and takes the following arguments: @@ -143,14 +143,13 @@ issues with multiprocessing. @execute_on_cpu def terpsichore(eq, processes=1, path="", exec=""): """TERPSICHORE driver function.""" - # these indices are used to give each equilibrium's I/O files unique file names - idxs = list(range(len(eq))) # equilibrium indices - # create temporary directory to store I/O files tmp_path = os.path.join(path, "tmp-TERPS") os.mkdir(tmp_path) # write input files for each equilibrium in serial + # these indices are used to give each equilibrium's I/O files unique file names + idxs = list(range(len(eq))) # equilibrium indices for k in idxs: # create a sub-directory for each equilibrium idx_path = os.path.join(tmp_path, str(k)) @@ -190,7 +189,8 @@ this guide. ``_pool_fun`` is the function that is run in parallel for each Equilibrium. It calls ``_run_terps`` (also shown below) to execute the TERPSICHORE Fortran code through a Python subprocess call, and ``_read_terps_output`` (not shown) to parse the output file -and extract the instability growth rate. +and extract the instability growth rate. If TERPSICHORE fails to execute for any reason +or takes too long to run, a large unstable growth rate is returned. :: From d3f99a6dd237c281eb55592b83cf59f9b2a43387 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Thu, 19 Dec 2024 14:20:49 -0700 Subject: [PATCH 48/49] make TERPSICHORE args keyword only --- desc/experimental/_terpsichore.py | 1 + 1 file changed, 1 insertion(+) diff --git a/desc/experimental/_terpsichore.py b/desc/experimental/_terpsichore.py index 47ace2a4eb..07698b7aa1 100644 --- a/desc/experimental/_terpsichore.py +++ b/desc/experimental/_terpsichore.py @@ -740,6 +740,7 @@ class TERPSICHORE(ExternalObjective): def __init__( self, eq, + *, target=None, bounds=None, weight=1, From 3691c7314ba5dff47840244ca8a53746deaf431f Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Fri, 20 Dec 2024 14:42:27 -0700 Subject: [PATCH 49/49] update bounds docs --- desc/experimental/_terpsichore.py | 13 ++++++------- docs/external_objectives.rst | 2 -- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/desc/experimental/_terpsichore.py b/desc/experimental/_terpsichore.py index 07698b7aa1..c7ee949bad 100644 --- a/desc/experimental/_terpsichore.py +++ b/desc/experimental/_terpsichore.py @@ -647,8 +647,8 @@ def _read_terps_output(path, scalar, surfs): class TERPSICHORE(ExternalObjective): """Computes ideal MHD linear stability from calls to the code TERPSICHORE. - Returns the linear growth rate of the fastest growing instability, - or the ΔW values at each flux surface. + Returns the linear growth rate of the fastest growing instability, or the ΔW values + at each flux surface. A positive growth rate or negative ΔW denotes instability. Parameters ---------- @@ -730,8 +730,10 @@ class TERPSICHORE(ExternalObjective): """ __doc__ = __doc__.rstrip() + collect_docs( - target_default="``bounds=(-np.inf, 0)``.", - bounds_default="``bounds=(-np.inf, 0)``.", + target_default="``bounds=(-np.inf, 0)`` if ``scalar=True`` " + + "else ``bounds=(0, np.inf)``.", + bounds_default="``bounds=(-np.inf, 0)`` if ``scalar=True`` " + + "else ``bounds=(0, np.inf)``.", ) _units = "(dimensionless)" @@ -772,9 +774,6 @@ def __init__( name="terpsichore", ): if target is None and bounds is None: - # want to minimize growth rate or maximize delta W - # good threshold for growth rate is < -1e-3 - # good threshold for delta W is > +1e-4 bounds = (-np.inf, 0) if scalar else (0, np.inf) super().__init__( eq=eq, diff --git a/docs/external_objectives.rst b/docs/external_objectives.rst index 4d9249639a..61cde86805 100644 --- a/docs/external_objectives.rst +++ b/docs/external_objectives.rst @@ -278,8 +278,6 @@ threads that are available. set_device("gpu") - import numpy as np - from desc.examples import get from desc.experimental import TERPSICHORE from desc.objectives import (