Skip to content

Commit

Permalink
Merge pull request #6 from gsnoep/tracer_bugfix
Browse files Browse the repository at this point in the history 
Large revamp of tracer routine to work with domains including coils
  • Loading branch information
@gsnoep
gsnoep authored Nov 19, 2024
2 parents 51706e3 + a31f61c commit fea977c
Show file tree
Hide file tree
Showing 3 changed files with 659 additions and 313 deletions.
72 changes: 47 additions & 25 deletions megpy/equilibrium.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -663,10 +663,10 @@ def add_fluxsurfaces(self,x=None,x_label='rho_tor',refine=None,analytic_shape=Fa
interp_method = 'normal'
if np.max(psirz[np.where(psirz!=0.0)]) <= threshold:
threshold = np.max(psirz[np.where(psirz!=0.0)])
interp_method = 'bounded_extrapolation'
# interp_method = 'bounded_extrapolation'
elif np.min(psirz[np.where(psirz!=0.0)]) >= threshold:
threshold = np.min(psirz[np.where(psirz!=0.0)])
interp_method = 'bounded_extrapolation'
# interp_method = 'bounded_extrapolation'

tracer_timing = 0.
analytic_timing = 0.
Expand All @@ -681,7 +681,7 @@ def add_fluxsurfaces(self,x=None,x_label='rho_tor',refine=None,analytic_shape=Fa
psi_fs = float(interpolate.interp1d(derived[x_label],derived['psi'])(x_fs))
q_fs = float(interpolate.interp1d(derived[x_label],derived['qpsi'])(x_fs))
fpol_fs = float(interpolate.interp1d(derived[x_label],derived['fpol'])(x_fs))

# trace the flux surface contour and relabel the tracer output
time0 = time.time()
fs = tracer.contour(R,Z,psirz,psi_fs,threshold,i_center=[i_rmaxis,i_zmaxis],tracer_diag=tracer_diag,interp_method=interp_method)
Expand All @@ -696,36 +696,54 @@ def add_fluxsurfaces(self,x=None,x_label='rho_tor',refine=None,analytic_shape=Fa
fs[_key] = fs.pop(key)
del fs['label']
del fs['level']
if not np.isfinite(fs['r']):
r_res = np.sqrt((R[i_rmaxis] - derived['rmaxis']) ** 2 + (Z[i_zmaxis] - derived['zmaxis']) ** 2)
fs['r'] = r_res * (psi_fs - derived['simag']) / (psirz[i_zmaxis,i_rmaxis] - derived['simag'])
if fs['r'] == 0.0:
fs['r'] = 1.0e-4
if not np.isfinite(fs['R0']):
fs['R0'] = derived['rmaxis']
if not np.isfinite(fs['Z0']):
fs['Z0'] = derived['zmaxis']
if not np.isfinite(fs['R_Zmax']):
fs['R_Zmax'] = derived['rmaxis']
if not np.isfinite(fs['R_Zmin']):
fs['R_Zmin'] = derived['rmaxis']
if not np.isfinite(fs['Z_max']):
fs['Z_max'] = derived['zmaxis'] + fs['r']
if not np.isfinite(fs['Z_min']):
fs['Z_min'] = derived['zmaxis'] - fs['r']

if analytic_shape:
time1 = time.time()
fs['miller_geo'] = LocalEquilibrium.extract_analytic_shape(fs)
analytic_timing += time.time()-time1


_incl_B = False
if incl_B:
if isinstance(incl_B,list):
_incl_B = incl_B[i_x_fs]
else:
_incl_B = incl_B
else:
_incl_B = False

if _incl_B:
# to speed up the Bpol interpolation generate a reduced Z,R mesh
i_R_in = find(fs['R_in'],self.derived['R'])-2
i_R_out = find(fs['R_out'],self.derived['R'])+2
i_Z_min = find(fs['Z_min'],self.derived['Z'])-2
i_Z_max = find(fs['Z_max'],self.derived['Z'])+2
R_mesh,Z_mesh = np.meshgrid(self.derived['R'][i_R_in:i_R_out],self.derived['Z'][i_Z_min:i_Z_max])
RZ_mesh = np.column_stack((Z_mesh.flatten(),R_mesh.flatten()))

# interpolate Bpol and Btor
B_pol_fs = interpolate.griddata(RZ_mesh,self.derived['B_pol_rz'][i_Z_min:i_Z_max,i_R_in:i_R_out].flatten(),(fs['Z'],fs['R']),method='cubic')
#B_pol_fs = np.array([])
#for i_R,RR in enumerate(fs['R']):
# B_pol_fs = np.append(B_pol_fs,interpolate.interp2d(self.derived['R'][i_R_in:i_R_out],self.derived['Z'][i_Z_min:i_Z_max],self.derived['B_pol_rz'][i_Z_min:i_Z_max,i_R_in:i_R_out],bounds_error=False,fill_value='extrapolate')(RR,fs['Z'][i_R]))
B_tor_fs = interpolate.interp1d(self.derived['psi'],self.derived['fpol'],bounds_error=False)(psi_fs)/fs['R']

if _incl_B:
B_pol_fs = 0.0
B_tor_fs = fs['fpol'] / fs['R0']
if len(fs['R']) > 5:
# to speed up the Bpol interpolation generate a reduced Z,R mesh
i_R_in = find(fs['R_in'],self.derived['R'])-2
i_R_out = find(fs['R_out'],self.derived['R'])+2
i_Z_min = find(fs['Z_min'],self.derived['Z'])-2
i_Z_max = find(fs['Z_max'],self.derived['Z'])+2
R_mesh,Z_mesh = np.meshgrid(self.derived['R'][i_R_in:i_R_out],self.derived['Z'][i_Z_min:i_Z_max])
RZ_mesh = np.column_stack((Z_mesh.flatten(),R_mesh.flatten()))

# interpolate Bpol and Btor
B_pol_fs = interpolate.griddata(RZ_mesh,self.derived['B_pol_rz'][i_Z_min:i_Z_max,i_R_in:i_R_out].flatten(),(fs['Z'],fs['R']),method='cubic')
#B_pol_fs = np.array([])
#for i_R,RR in enumerate(fs['R']):
# B_pol_fs = np.append(B_pol_fs,interpolate.interp2d(self.derived['R'][i_R_in:i_R_out],self.derived['Z'][i_Z_min:i_Z_max],self.derived['B_pol_rz'][i_Z_min:i_Z_max,i_R_in:i_R_out],bounds_error=False,fill_value='extrapolate')(RR,fs['Z'][i_R]))
B_tor_fs = interpolate.interp1d(self.derived['psi'],self.derived['fpol'],bounds_error=False)(np.array([psi_fs]))[0]/fs['R']
fs.update({'Bpol':B_pol_fs, 'Btor':B_tor_fs, 'B':np.sqrt(B_pol_fs**2+B_tor_fs**2)})
else:
fs.update({'Bpol':np.array([]), 'Btor':np.array([]), 'B':np.array([])})
Expand Down Expand Up @@ -785,18 +803,22 @@ def add_fluxsurfaces(self,x=None,x_label='rho_tor',refine=None,analytic_shape=Fa
fluxsurfaces[key].insert(0,derived['rmaxis'])
elif key in ['Z','Z0','Z_max','Z_min']:
fluxsurfaces[key].insert(0,derived['zmaxis'])
elif key in ['kappa','delta','zeta','s_kappa','s_delta','s_zeta']:
elif key in ['q','kappa','delta','zeta','s_kappa','s_delta','s_zeta']:
fluxsurfaces[key].insert(0,fluxsurfaces[key][0])
else:
if isinstance(fluxsurfaces[key],dict):
for _key in fluxsurfaces[key]:
fluxsurfaces[key][_key].insert(0,0.*fluxsurfaces[key][_key][-1])
if _key in ['delta_u','delta_l','delta','kappa','zeta_uo','zeta_ui','zeta_li','zeta_lo','zeta']:
fluxsurfaces[key][_key].insert(0,fluxsurfaces[key][_key][0])
else:
fluxsurfaces[key][_key].insert(0,0.*fluxsurfaces[key][_key][-1])
elif isinstance(fluxsurfaces[key],list):
fluxsurfaces[key].insert(0,0.*fluxsurfaces[key][-1])

# add the midplane average geometric flux surface quantities to derived
derived['Ro'] = np.array(fluxsurfaces['R0'])
derived['Ro'][0] = derived['Ro'][1] # clear the starting zero
if derived['Ro'][0] == 0.0:
derived['Ro'][0] = derived['Ro'][1] # clear the starting zero
derived['R0'] = derived['Ro'][-1] # midplane average major radius of the lcfs
derived['Zo'] = np.array(fluxsurfaces['Z0'])
derived['Z0'] = derived['Zo'][-1] # average elevation of the lcfs
Expand Down
29 changes: 21 additions & 8 deletions megpy/localequilibrium.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -268,7 +268,10 @@ def __init__(self,param,equilibrium,x_loc,x_label='rho_tor',n_x=9,n_theta='defau

# compute the Hessian estimate from the least squares fit Jacobian
H = self.lsq['jac'].T @ self.lsq['jac']
H_inv = np.linalg.inv(H)
try:
H_inv = np.linalg.inv(H)
except:
H_inv = np.nan # in case the Hessian is singular

# compute the covariance matrix, where N_DOF = 2*n_theta - n_shape, 2*n_theta since both R and Z are minimized on
self.shape_cov = H_inv * np.sum(self.lsq['fun'])/((2*self.n_theta)-len(self.shape))
Expand Down Expand Up @@ -1041,7 +1044,8 @@ def cost_deriv(self,params):
return cost

# extraction tools
def extract_analytic_shape(fluxsurface):
@classmethod
def extract_analytic_shape(cls, fluxsurface):
"""Extract Turnbull-Miller geometry parameters [Turnbull PoP 6 1113 (1999)] from a flux-surface contour. Adapted from 'extract_miller_from_eqdsk.py' by D. Told.
Args:
Expand All @@ -1060,12 +1064,18 @@ def extract_analytic_shape(fluxsurface):
miller_geo['kappa'] = (fluxsurface['Z_max'] - fluxsurface['Z_min'])/(2*fluxsurface['r'])

# generate theta grid and interpolate the flux-surface trace to the Miller parameterization
R_miller = fluxsurface['R0'] + fluxsurface['r']*np.cos(fluxsurface['theta_RZ']+x*np.sin(fluxsurface['theta_RZ']))
Z_miller = np.hstack((interpolate.interp1d(fluxsurface['R'][:np.argmin(fluxsurface['R'])],fluxsurface['Z'][:np.argmin(fluxsurface['R'])],bounds_error=False)(R_miller[:find(np.min(fluxsurface['R']),R_miller)]),interpolate.interp1d(fluxsurface['R'][np.argmin(fluxsurface['R']):],fluxsurface['Z'][np.argmin(fluxsurface['R']):],bounds_error=False)(R_miller[find(np.min(fluxsurface['R']),R_miller):])))
theta_zeta = np.array([0.25*np.pi,0.75*np.pi,1.25*np.pi,1.75*np.pi])
theta_miller = fluxsurface['theta_RZ'] if len(fluxsurface['theta_RZ']) > 1 else theta_zeta
R_miller = fluxsurface['R0'] + fluxsurface['r']*np.cos(theta_miller+x*np.sin(theta_miller))
Z_miller = fluxsurface['Z0'] + miller_geo['kappa']*fluxsurface['r']*np.sin(theta_miller)
if len(fluxsurface['R']) > 1 and len(fluxsurface['Z']) > 1:
Z_miller = np.hstack((
interpolate.interp1d(fluxsurface['R'][:np.argmin(fluxsurface['R'])],fluxsurface['Z'][:np.argmin(fluxsurface['R'])],bounds_error=False)(R_miller[:find(np.min(fluxsurface['R']),R_miller)]),
interpolate.interp1d(fluxsurface['R'][np.argmin(fluxsurface['R']):],fluxsurface['Z'][np.argmin(fluxsurface['R']):],bounds_error=False)(R_miller[find(np.min(fluxsurface['R']),R_miller):])
))

# derive the squareness (zeta) from the Miller parameterization
theta_zeta = np.array([0.25*np.pi,0.75*np.pi,1.25*np.pi,1.75*np.pi])
Z_zeta = interpolate.interp1d(fluxsurface['theta_RZ'],Z_miller,bounds_error=False,fill_value='extrapolate')(theta_zeta)
Z_zeta = interpolate.interp1d(theta_miller,Z_miller,bounds_error=False,fill_value='extrapolate')(theta_zeta)

# invert the Miller parameterization of Z, holding off on subtracting theta/sin(2*theta)
zeta_4q = np.arcsin((Z_zeta-fluxsurface['Z0'])/(miller_geo['kappa']*fluxsurface['r']))/np.sin(2*theta_zeta)
Expand All @@ -1081,8 +1091,11 @@ def extract_analytic_shape(fluxsurface):
# compute the average squareness of the flux-surface
miller_geo['zeta'] = 0.25*(miller_geo['zeta_uo']+miller_geo['zeta_ui']+miller_geo['zeta_li']+miller_geo['zeta_lo'])

miller_geo['R_miller'] = R_miller
miller_geo['Z_miller'] = fluxsurface['Z0']+miller_geo['kappa']*fluxsurface['r']*np.sin(fluxsurface['theta_RZ']+miller_geo['zeta']*np.sin(2*fluxsurface['theta_RZ']))
miller_geo['R_miller'] = np.array([fluxsurface['R0']])
miller_geo['Z_miller'] = np.array([fluxsurface['Z0']])
if len(fluxsurface['theta_RZ']) > 1:
miller_geo['R_miller'] = R_miller
miller_geo['Z_miller'] = fluxsurface['Z0']+miller_geo['kappa']*fluxsurface['r']*np.sin(fluxsurface['theta_RZ']+miller_geo['zeta']*np.sin(2*fluxsurface['theta_RZ']))

return miller_geo

Expand Down
Loading

0 comments on commit fea977c

Please sign in to comment.