From d8fa4004c82696c8e824668c76c3bf3b77b311bf Mon Sep 17 00:00:00 2001 From: aaronkho Date: Thu, 19 Sep 2024 19:59:51 -0400 Subject: [PATCH 1/9] Large revamp of tracer routine to work with domains including coils --- megpy/equilibrium.py | 68 ++- megpy/localequilibrium.py | 24 +- megpy/tracer.py | 856 +++++++++++++++++++++++--------------- 3 files changed, 592 insertions(+), 356 deletions(-) diff --git a/megpy/equilibrium.py b/megpy/equilibrium.py index cf2eb43..ca21702 100755 --- a/megpy/equilibrium.py +++ b/megpy/equilibrium.py @@ -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) @@ -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([])}) @@ -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 diff --git a/megpy/localequilibrium.py b/megpy/localequilibrium.py index e5f8da7..624f058 100755 --- a/megpy/localequilibrium.py +++ b/megpy/localequilibrium.py @@ -1041,7 +1041,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: @@ -1060,12 +1061,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) @@ -1081,8 +1088,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 diff --git a/megpy/tracer.py b/megpy/tracer.py index a56b1ff..e94fd67 100755 --- a/megpy/tracer.py +++ b/megpy/tracer.py @@ -9,6 +9,79 @@ from operator import itemgetter from .utils import * +def check_less(one, two): + return np.all(one < np.nanmax(two)) + +def check_more(one, two): + return np.all(one > np.nanmin(two)) + +def f_bnd(data, sign='>', early=0, verbose=0): + found = (len(data) <= 0) + val = np.nanmin(data) if sign == '>' else np.nanmax(data) + ii = 0 + jj = 0 + while (not found): + if sign == '>': + if data[ii] >= val: + val = data[ii] + jj = 0 + else: + jj += 1 + elif sign == '<': + if data[ii] <= val: + val = data[ii] + jj = 0 + else: + jj += 1 + ii += 1 + if ii >= len(data): + found = True + if early > 0 and jj == early: + found = True + if verbose > 0: + print(data, ii, jj) + return ii - jj + +def find_dual_value(x, y, split_index, level, threshold, method='normal', sign='>', tol=0, tracer_diag='none'): + + lower_index = split_index - f_bnd(y[:split_index+1][::-1], sign=sign, early=tol) + if lower_index <= 0: + lower_index = None + upper_index = split_index + 1 + f_bnd(y[split_index:], sign=sign, early=tol) + if upper_index >= len(y): + upper_index = None + + # chop x,y in two parts to separate two solutions + x_lower = x[lower_index:split_index+1] + x_upper = x[split_index:upper_index] + y_lower = y[lower_index:split_index+1] + y_upper = y[split_index:upper_index] + + # if the normal method provides spikey contour traces, bound the interpolation domain and extrapolate the intersection + if method == 'bounded_extrapolation': + lower_mask = (y_lower >= threshold) if sign == '>' else (y_lower <= threshold) + upper_mask = (y_upper >= threshold) if sign == '<' else (y_upper <= threshold) + x_lower = x_lower[lower_mask] + x_upper = x_upper[upper_mask] + y_lower = y_lower[lower_mask] + y_upper = y_upper[upper_mask] + + # interpolate the x coordinate to either side + #print(y_lower[0], y_upper[-1]) + xl_out = interpolate.interp1d(y_lower,x_lower,bounds_error=False,fill_value='extrapolate')(level)[0] + xu_out = interpolate.interp1d(y_upper,x_upper,bounds_error=False,fill_value='extrapolate')(level)[0] + + # diagnostic plot to check interpolation issues for rows + if tracer_diag == 'interp': + plt.figure() + #plt.title('Y:{}'.format(third_axis)) + plt.plot(y_lower,x_lower,'b-',label='lower') + plt.plot(y_upper,x_upper,'r-',label='upper') + plt.axvline(level,ls='dashed',color='black') + plt.legend() + + return xl_out, xu_out + def contour(X,Y,Z=None,level=None,threshold=None,i_center=None,interp_method='normal',tracer_diag='none',symmetrise=False): """Find the X,Y trace of a closed(!) contour in Z. @@ -27,6 +100,12 @@ def contour(X,Y,Z=None,level=None,threshold=None,i_center=None,interp_method='no dict: a dict containing the X, Y, level (and center + extrema) values of the contour trace in Z. """ with np.errstate(divide='ignore'): + + lvl = np.array([level]) if not isinstance(level, np.ndarray) else level + thr = np.array([threshold]) if not isinstance(threshold, np.ndarray) else threshold + dX = np.mean(np.abs(np.diff(X))) + dY = np.mean(np.abs(np.diff(Y))) + # define storage for contour coordinates XY_contour = {'left':{'top':[],'bottom':[]},'right':{'top':[],'bottom':[]}} @@ -37,26 +116,37 @@ def contour(X,Y,Z=None,level=None,threshold=None,i_center=None,interp_method='no else: i_xcenter = int(len(X)/2) i_ycenter = int(len(Y)/2) + i_xcenter_pad = int(0.3 * len(X)) + i_ycenter_pad = int(0.3 * len(Y)) center = Z[i_ycenter,i_xcenter] # find the vertical extrema of the threshold contour at the xcenter + search_tol = 3 if threshold > center: - f_threshold = np.min + #f_threshold = np.min f_split = np.argmin - f_bnd = np.argmax + #f_bnd = np.argmax f_sgn = '>' + f_check = check_more elif threshold < center: - f_threshold = np.max + #f_threshold = np.max f_split = np.argmax - f_bnd = np.argmin + #f_bnd = np.argmin f_sgn = '<' - i_ZY_xcenter_split = f_split(Z[:,i_xcenter]) - i_ZY_xcenter_bottom_max = f_bnd(Z[:i_ZY_xcenter_split,i_xcenter]) - i_ZY_xcenter_top_max = i_ZY_xcenter_split+f_bnd(Z[i_ZY_xcenter_split:,i_xcenter]) - - ZY_xcenter_bottom = Z[i_ZY_xcenter_bottom_max:i_ZY_xcenter_split,i_xcenter] - ZY_xcenter_top = Z[i_ZY_xcenter_split:i_ZY_xcenter_top_max,i_xcenter] + f_check = check_less + i_ZY_xcenter_split = f_split(Z[i_ycenter-i_ycenter_pad:i_ycenter+i_ycenter_pad+1,i_xcenter].flatten()) + i_ycenter - i_ycenter_pad + i_ZY_xcenter_bottom_max = i_ZY_xcenter_split - f_bnd(Z[:i_ZY_xcenter_split+1,i_xcenter].flatten()[::-1], sign=f_sgn, early=3) + i_ZY_xcenter_top_max = i_ZY_xcenter_split + f_bnd(Z[i_ZY_xcenter_split:,i_xcenter].flatten(), sign=f_sgn, early=3) + 1 + if i_ZY_xcenter_bottom_max <= 0: + i_ZY_xcenter_bottom_max = None + if i_ZY_xcenter_top_max >= len(Y): + i_ZY_xcenter_top_max = None + + Y_xcenter_bottom = Y[i_ZY_xcenter_bottom_max:i_ZY_xcenter_split+1].flatten() + Y_xcenter_top = Y[i_ZY_xcenter_split:i_ZY_xcenter_top_max].flatten() + ZY_xcenter_bottom = Z[i_ZY_xcenter_bottom_max:i_ZY_xcenter_split+1,i_xcenter].flatten() + ZY_xcenter_top = Z[i_ZY_xcenter_split:i_ZY_xcenter_top_max,i_xcenter].flatten() '''# patch for rare cases where the Z map becomes flat below the threshold values (e.g. some ESCO EQDSKs) if not np.any(ZY_xcenter_bottom >= threshold): @@ -74,15 +164,21 @@ def contour(X,Y,Z=None,level=None,threshold=None,i_center=None,interp_method='no i_ZY_xcenter_top_max = i_ZY_xcenter_split+i_dZY_xcenter_top_dr ZY_xcenter_top = ZY_xcenter_top[:i_dZY_xcenter_top_dr]''' - Y_max = interpolate.interp1d(ZY_xcenter_top,Y[i_ZY_xcenter_split:i_ZY_xcenter_top_max],bounds_error=False,fill_value='extrapolate')(threshold) - Y_min = interpolate.interp1d(ZY_xcenter_bottom,Y[i_ZY_xcenter_bottom_max:i_ZY_xcenter_split],bounds_error=False,fill_value='extrapolate')(threshold) + Y_search_max = interpolate.interp1d(ZY_xcenter_top,Y_xcenter_top,bounds_error=False,fill_value='extrapolate')(thr)[0] + 3.0 * dY + Y_search_min = interpolate.interp1d(ZY_xcenter_bottom,Y_xcenter_bottom,bounds_error=False,fill_value='extrapolate')(thr)[0] - 3.0 * dY - i_ZX_ycenter_split = f_split(Z[i_ycenter,:]) - i_ZX_ycenter_bottom_max = f_bnd(Z[i_ycenter,:i_ZX_ycenter_split]) - i_ZX_ycenter_top_max = i_ZY_xcenter_split+f_bnd(Z[i_ycenter,i_ZX_ycenter_split:]) + i_ZX_ycenter_split = f_split(Z[i_ycenter,i_xcenter-i_xcenter_pad:i_xcenter+i_xcenter_pad+1].flatten()) + i_xcenter - i_xcenter_pad + i_ZX_ycenter_left_max = i_ZX_ycenter_split - f_bnd(Z[i_ycenter,:i_ZX_ycenter_split+1].flatten()[::-1], sign=f_sgn, early=3) + i_ZX_ycenter_right_max = i_ZX_ycenter_split + f_bnd(Z[i_ycenter,i_ZX_ycenter_split:].flatten(), sign=f_sgn, early=3) + 1 + if i_ZX_ycenter_left_max <= 0: + i_ZX_ycenter_left_max = None + if i_ZX_ycenter_right_max >= len(X): + i_ZX_ycenter_right_max = None - ZX_ycenter_bottom = Z[i_ycenter,i_ZX_ycenter_bottom_max:i_ZX_ycenter_split] - ZX_ycenter_top = Z[i_ycenter,i_ZX_ycenter_split:i_ZX_ycenter_top_max] + X_ycenter_left = X[i_ZX_ycenter_left_max:i_ZX_ycenter_split+1].flatten() + X_ycenter_right = X[i_ZX_ycenter_split:i_ZX_ycenter_right_max].flatten() + ZX_ycenter_left = Z[i_ycenter,i_ZX_ycenter_left_max:i_ZX_ycenter_split+1].flatten() + ZX_ycenter_right = Z[i_ycenter,i_ZX_ycenter_split:i_ZX_ycenter_right_max].flatten() '''# patch for rare cases where the Z map becomes flat below the threshold values (e.g. some ESCO EQDSKs) if not np.any(ZX_ycenter_bottom >= threshold): @@ -100,8 +196,8 @@ def contour(X,Y,Z=None,level=None,threshold=None,i_center=None,interp_method='no i_ZX_ycenter_top_max = i_ZX_ycenter_split+i_dZX_ycenter_top_dr ZX_ycenter_top = ZX_ycenter_top[:i_dZX_ycenter_top_dr]''' - X_max = interpolate.interp1d(ZX_ycenter_top,X[i_ZX_ycenter_split:i_ZX_ycenter_top_max],bounds_error=False,fill_value='extrapolate')(threshold) - X_min = interpolate.interp1d(ZX_ycenter_bottom,X[i_ZX_ycenter_bottom_max:i_ZX_ycenter_split],bounds_error=False,fill_value='extrapolate')(threshold) + X_search_max = interpolate.interp1d(ZX_ycenter_right,X_ycenter_right,bounds_error=False,fill_value='extrapolate')(thr)[0] + 3.0 * dY + X_search_min = interpolate.interp1d(ZX_ycenter_left,X_ycenter_left,bounds_error=False,fill_value='extrapolate')(thr)[0] - 3.0 * dY # diagnostic plot to check interpolation issues for the vertical bounds if tracer_diag == 'bounds': @@ -114,313 +210,405 @@ def contour(X,Y,Z=None,level=None,threshold=None,i_center=None,interp_method='no plt.show() # set the starting coordinates for the contour tracing algorithm - i, j = i_ycenter, i_ycenter - k, l = i_xcenter, i_xcenter + i, j = i_ycenter-1, i_ycenter + k, l = i_xcenter-1, i_xcenter # while the level intersects with the current Z slice gather the intersection coordinates - while (eval('level' + f_sgn + 'f_threshold(Z[i])') and i < Z.shape[0]-1): - if Y[i] <= Y_max: + do_j = True + while (do_j and j < Z.shape[0] and f_check(lvl, Z[j,:])): + if Y[j] <= Y_search_max: try: # find the split in the Y slice of Z - i_ZX_slice_split = f_split(Z[i]) - - # chop the Z and X slices in two parts to separate the left and right - ZX_slice_top_left = Z[i,:i_ZX_slice_split] - ZX_slice_top_right = Z[i,i_ZX_slice_split:] - - # interpolate the X coordinate of the top half left and right - if interp_method == 'normal': - X_top_left = float(interpolate.interp1d(ZX_slice_top_left,X[i_ZX_slice_split-len(ZX_slice_top_left):i_ZX_slice_split],bounds_error=False)(level)) - X_top_right = float(interpolate.interp1d(ZX_slice_top_right,X[i_ZX_slice_split:i_ZX_slice_split+len(ZX_slice_top_right)],bounds_error=False)(level)) - # if the normal method provides spikey contour traces, bound the interpolation domain and extrapolate the intersection - elif interp_method == 'bounded_extrapolation': - ZX_slice_top_left = ZX_slice_top_left[ZX_slice_top_left<=threshold] - ZX_slice_top_right = ZX_slice_top_right[ZX_slice_top_right<=threshold] - - X_top_left = float(interpolate.interp1d(ZX_slice_top_left,X[i_ZX_slice_split-len(ZX_slice_top_left):i_ZX_slice_split][ZX_slice_top_left<=threshold],bounds_error=False,fill_value='extrapolate')(level)) - X_top_right = float(interpolate.interp1d(ZX_slice_top_right,X[i_ZX_slice_split:i_ZX_slice_split+len(ZX_slice_top_right)][ZX_slice_top_right<=threshold],bounds_error=False,fill_value='extrapolate')(level)) - - # diagnostic plot to check interpolation issues for rows - if tracer_diag == 'interp': - plt.figure() - plt.title('Y:{}'.format(Y[i])) - plt.plot(ZX_slice_top_left,X[i_ZX_slice_split-len(ZX_slice_top_left):i_ZX_slice_split],'b-',label='left, top') - plt.plot(ZX_slice_top_right,X[i_ZX_slice_split:i_ZX_slice_split+len(ZX_slice_top_right)],'r-',label='right, top') - plt.axvline(level,ls='dashed',color='black') - plt.legend() - - # insert the coordinate pairs into the contour trace dict if not nan (bounds error) and order properly for merging later - if not np.isnan(X_top_left): - XY_contour['left']['top'].append((X_top_left,Y[i])) - if X_top_left < X_min: - X_min = X_top_left - if not np.isnan(X_top_right): - XY_contour['right']['top'].append((X_top_right,Y[i])) - if X_top_right > X_max: - X_max = X_top_right + j_ZX_slice_split = f_split(Z[j,i_xcenter-i_xcenter_pad:i_xcenter+i_xcenter_pad+1].flatten()) + i_xcenter - i_xcenter_pad + X_top_left, X_top_right = find_dual_value(X.flatten(), Z[j,:].flatten(), j_ZX_slice_split, lvl, thr, interp_method, f_sgn, search_tol) + + # insert the coordinates into the contour trace dict if not nan (bounds error) and order properly for merging later + if np.isfinite(X_top_left): + theta_X_top_left = np.arctan2(Y[j] - Y[i_ycenter], X_top_left - X[i_xcenter]) + if theta_X_top_left < 0.0: + theta_X_top_left = 2.0 * np.pi + theta_X_top_left + XY_contour['left']['top'].append((X_top_left,Y[j],theta_X_top_left)) + if X_top_left < (X_search_min + 3.0 * dX): + X_search_min = X_top_left - 3.0 * dX + if np.isfinite(X_top_right): + theta_X_top_right = np.arctan2(Y[j] - Y[i_ycenter], X_top_right - X[i_xcenter]) + if theta_X_top_right < 0.0: + theta_X_top_right = 2.0 * np.pi + theta_X_top_right + XY_contour['right']['top'].append((X_top_right,Y[j],theta_X_top_right)) + if X_top_right > (X_search_max - 3.0 * dX): + X_search_max = X_top_right + 3.0 * dX + #print('1', Y[j], X_top_left, X_top_right, theta_X_top_left, theta_X_top_right) except: pass + # update the slice coordinates + j += 1 + else: + do_j = False - # update the slice coordinates - if i < Z.shape[0]-1: - i+=1 + # while the level intersects with the current Z slice gather the intersection coordinates + do_k = True + while (do_k and k >= 0 and f_check(lvl, Z[:,k])): + # interpolate the X coordinate of the bottom half of the contour on both the right and left + if X[k] >= X_search_min: + try: + # split the current column of the Z map in top and bottom + # find the split and the extrema of the Y slice of Z + k_ZY_slice_split = f_split(Z[i_ycenter-i_ycenter_pad:i_ycenter+i_ycenter_pad+1,k].flatten()) + i_ycenter - i_ycenter_pad + Y_bottom_left, Y_top_left = find_dual_value(Y.flatten(), Z[:,k].flatten(), k_ZY_slice_split, lvl, thr, interp_method, f_sgn, search_tol) + + # insert the coordinates into the contour trace dict if not nan (bounds error) and order properly for merging later + if np.isfinite(Y_bottom_left): + theta_Y_bottom_left = np.arctan2(Y_bottom_left - Y[i_ycenter], X[k] - X[i_xcenter]) + if theta_Y_bottom_left < 0.0: + theta_Y_bottom_left = 2.0 * np.pi + theta_Y_bottom_left + XY_contour['left']['bottom'].append((X[k],Y_bottom_left,theta_Y_bottom_left)) + # update the higher vertical bound if it is higher than the vertical bound found at the magnetic axis + if Y_bottom_left < (Y_search_min + 3.0 * dY): + Y_search_min = Y_bottom_left - 3.0 * dY + if np.isfinite(Y_top_left): + theta_Y_top_left = np.arctan2(Y_top_left - Y[i_ycenter], X[k] - X[i_xcenter]) + if theta_Y_top_left < 0.0: + theta_Y_top_left = 2.0 * np.pi + theta_Y_top_left + XY_contour['left']['top'].append((X[k],Y_top_left,theta_Y_top_left)) + # update the lower vertical bound if it is lower than the vertical bound found at the magnetic axis + if Y_top_left > (Y_search_max - 3.0 * dY): + Y_search_max = Y_top_left + 3.0 * dY + #print('1', X[k], Y_bottom_left, Y_top_left, theta_Y_bottom_left, theta_Y_top_left) + except: + pass + # update the slice coordinates + k -= 1 + else: + do_k = False # while the level intersects with the current Z slice gather the intersection coordinates - while (eval('level' + f_sgn + 'f_threshold(Z[:,k])') and k < Z.shape[1]-1): - if X[k] <= X_max: + do_i = True + while (do_i and i >= 0 and f_check(lvl, Z[i,:])): + # interpolate the R coordinate of the bottom half of the contour on both the right and left + if Y[i] >= Y_search_min: + try: + # find the split and the extrema of the Y slice of Z + i_ZX_slice_split = f_split(Z[i,i_xcenter-i_xcenter_pad:i_xcenter+i_xcenter_pad+1].flatten()) + i_xcenter - i_xcenter_pad + X_bottom_left, X_bottom_right = find_dual_value(X.flatten(), Z[i,:].flatten(), i_ZX_slice_split, lvl, thr, interp_method, f_sgn, search_tol) + + if np.isfinite(X_bottom_left): + theta_X_bottom_left = np.arctan2(Y[i] - Y[i_ycenter], X_bottom_left - X[i_xcenter]) + if theta_X_bottom_left < 0.0: + theta_X_bottom_left = 2.0 * np.pi + theta_X_bottom_left + XY_contour['left']['bottom'].append((X_bottom_left,Y[i],theta_X_bottom_left)) + if X_bottom_left < (X_search_min + 3.0 * dX): + X_search_min = X_bottom_left - 3.0 * dX + if np.isfinite(X_bottom_right): + theta_X_bottom_right = np.arctan2(Y[i] - Y[i_ycenter], X_bottom_right - X[i_xcenter]) + if theta_X_bottom_right < 0.0: + theta_X_bottom_right = 2.0 * np.pi + theta_X_bottom_right + XY_contour['right']['bottom'].append((X_bottom_right,Y[i],theta_X_bottom_right)) + if X_bottom_right > (X_search_max - 3.0 * dX): + X_search_max = X_bottom_right + 3.0 * dX + #print('1', Y[i], X_bottom_left, X_bottom_right, theta_X_bottom_left, theta_X_bottom_right) + except: + pass + # update the slice coordinates + i -= 1 + else: + do_i = False + + # while the level intersects with the current Z slice gather the intersection coordinates + do_l = True + while (do_l and l < Z.shape[1] and f_check(lvl, Z[:,l])): + if X[l] <= X_search_max: try: # find the split and the extrema of the X slice of Z - k_ZY_slice_split = f_split(Z[:,k]) - k_ZY_slice_bottom_max = f_bnd(Z[:k_ZY_slice_split,k]) - k_ZY_slice_top_max = k_ZY_slice_split+f_bnd(Z[k_ZY_slice_split:,k]) - - # chop the Z slices in two parts to separate the top and bottom halves - ZY_slice_bottom_right = Z[k_ZY_slice_bottom_max:k_ZY_slice_split+1,k] - ZY_slice_top_right = Z[k_ZY_slice_split:k_ZY_slice_top_max,k] - - # interpolate the Y coordinate for the right top and bottom - if interp_method == 'normal': - Y_bottom_right = float(interpolate.interp1d(ZY_slice_bottom_right,Y[k_ZY_slice_bottom_max:k_ZY_slice_split+1],bounds_error=False)(level)) - Y_top_right = float(interpolate.interp1d(ZY_slice_top_right,Y[k_ZY_slice_split:k_ZY_slice_top_max],bounds_error=False)(level)) - elif interp_method == 'bounded_extrapolation': - Y_slice_bottom_right = Y[k_ZY_slice_bottom_max:k_ZY_slice_split+1][ZY_slice_bottom_right<=threshold] - Y_slice_top_right = Y[k_ZY_slice_split:k_ZY_slice_top_max][ZY_slice_top_right<=threshold] - ZY_slice_bottom_right = ZY_slice_bottom_right[ZY_slice_bottom_right<=threshold] - ZY_slice_top_right = ZY_slice_top_right[ZY_slice_top_right<=threshold] - - Y_bottom_right = float(interpolate.interp1d(ZY_slice_bottom_right,Y_slice_bottom_right,bounds_error=False,fill_value='extrapolate')(level)) - Y_top_right = float(interpolate.interp1d(ZY_slice_top_right,Y_slice_top_right,bounds_error=False,fill_value='extrapolate')(level)) - - # diagnostic plot to check interpolation issues for columns - if tracer_diag == 'interp': - plt.figure() - plt.title('X:{}'.format(X[k])) - plt.plot(Z[k_ZY_slice_bottom_max:k_ZY_slice_split+1,k],Y[k_ZY_slice_bottom_max:k_ZY_slice_split+1],'b--',label='right, bottom') - plt.plot(ZY_slice_top_right,Y[k_ZY_slice_split:k_ZY_slice_top_max],'r--',label='right, top') - plt.axvline(level,ls='dashed',color='black') - plt.legend() - - if not np.isnan(Y_bottom_right): - XY_contour['right']['bottom'].append((X[k],Y_bottom_right)) + l_ZY_slice_split = f_split(Z[i_ycenter-i_ycenter_pad:i_ycenter+i_ycenter_pad+1,l].flatten()) + i_ycenter - i_ycenter_pad + Y_bottom_right, Y_top_right = find_dual_value(Y.flatten(), Z[:,l].flatten(), l_ZY_slice_split, lvl, thr, interp_method, f_sgn, search_tol) + + # insert the coordinates into the contour trace dict if not nan (bounds error) and order properly for merging later + if np.isfinite(Y_bottom_right): + theta_Y_bottom_right = np.arctan2(Y_bottom_right - Y[i_ycenter], X[l] - X[i_xcenter]) + if theta_Y_bottom_right < 0.0: + theta_Y_bottom_right = 2.0 * np.pi + theta_Y_bottom_right + XY_contour['right']['bottom'].append((X[l],Y_bottom_right,theta_Y_bottom_right)) # update the higher vertical bound if it is higher than the vertical bound found at the magnetic axis - if Y_bottom_right < Y_min: - Y_min = Y_bottom_right - if not np.isnan(Y_top_right): - XY_contour['right']['top'].append((X[k],Y_top_right)) + if Y_bottom_right < (Y_search_min + 3.0 * dY): + Y_search_min = Y_bottom_right - 3.0 * dY + if np.isfinite(Y_top_right): + theta_Y_top_right = np.arctan2(Y_top_right - Y[i_ycenter], X[l] - X[i_xcenter]) + if theta_Y_top_right < 0.0: + theta_Y_top_right = 2.0 * np.pi + theta_Y_top_right + XY_contour['right']['top'].append((X[l],Y_top_right,theta_Y_top_right)) # update the lower vertical bound if it is lower than the vertical bound found at the magnetic axis - if Y_top_right > Y_max: - Y_max = Y_top_right + if Y_top_right > (Y_search_max - 3.0 * dY): + Y_search_max = Y_top_right + 3.0 * dY + #print('1', X[l], Y_bottom_right, Y_top_right, theta_Y_bottom_right, theta_Y_top_right) except: pass + # update the slice coordinates + l += 1 + else: + do_l = False - # update the slice coordinates - if k < Z.shape[1]-1: - k+=1 + ### GO THROUGH LOOP AGAIN IN ANOTHER ORDER TO CATCH MISSING POINTS DUE TO INADEQUATE SEARCH RANGES! # while the level intersects with the current Z slice gather the intersection coordinates - while (eval('level' + f_sgn + 'f_threshold(Z[j])') and j > 0): + do_i = True + while (do_i and i >= 0 and f_check(lvl, Z[i,:])): # interpolate the R coordinate of the bottom half of the contour on both the right and left - if Y[j] >= Y_min: + if Y[i] >= Y_search_min: try: # find the split and the extrema of the Y slice of Z - j_ZX_slice_split = f_split(Z[j]) - - ZX_slice_bottom_left = Z[j,:j_ZX_slice_split] - ZX_slice_bottom_right = Z[j,j_ZX_slice_split:] - - # interpolate the X coordinate of the bottom half left and right - if interp_method == 'normal': - X_bottom_left = float(interpolate.interp1d(ZX_slice_bottom_left,X[j_ZX_slice_split-len(ZX_slice_bottom_left):j_ZX_slice_split],bounds_error=False)(level)) - X_bottom_right = float(interpolate.interp1d(ZX_slice_bottom_right,X[j_ZX_slice_split:j_ZX_slice_split+len(ZX_slice_bottom_right)],bounds_error=False)(level)) - elif interp_method == 'bounded_extrapolation': - ZX_slice_bottom_left = ZX_slice_bottom_left[ZX_slice_bottom_left<=threshold] - ZX_slice_bottom_right = ZX_slice_bottom_right[ZX_slice_bottom_right<=threshold] - X_bottom_left = float(interpolate.interp1d(ZX_slice_bottom_left,X[j_ZX_slice_split-len(ZX_slice_bottom_left):j_ZX_slice_split][ZX_slice_bottom_left<=threshold],bounds_error=False,fill_value='extrapolate')(level)) - X_bottom_right = float(interpolate.interp1d(ZX_slice_bottom_right,X[j_ZX_slice_split:j_ZX_slice_split+len(ZX_slice_bottom_right)][ZX_slice_bottom_right<=threshold],bounds_error=False,fill_value='extrapolate')(level)) - - # diagnostic plot to check interpolation issues for rows - if tracer_diag == 'interp': - plt.figure() - plt.title('Y:{}'.format(Y[j])) - plt.plot(ZX_slice_bottom_left,X[j_ZX_slice_split-len(ZX_slice_bottom_left):j_ZX_slice_split],'b-',label='left, bottom') - plt.plot(ZX_slice_bottom_right,X[j_ZX_slice_split:j_ZX_slice_split+len(ZX_slice_bottom_right)],'r-',label='right, bottom') - plt.axvline(level,ls='dashed',color='black') - plt.legend() - - if not np.isnan(X_bottom_left): - XY_contour['left']['bottom'].append((X_bottom_left,Y[j])) - if X_bottom_left < X_min: - X_min = X_bottom_left - if not np.isnan(X_bottom_right): - XY_contour['right']['bottom'].append((X_bottom_right,Y[j])) - if X_bottom_right > X_max: - X_max = X_bottom_right + i_ZX_slice_split = f_split(Z[i,i_xcenter-i_xcenter_pad:i_xcenter+i_xcenter_pad+1].flatten()) + i_xcenter - i_xcenter_pad + X_bottom_left, X_bottom_right = find_dual_value(X.flatten(), Z[i,:].flatten(), i_ZX_slice_split, lvl, thr, interp_method, f_sgn, search_tol) + + if np.isfinite(X_bottom_left): + theta_X_bottom_left = np.arctan2(Y[i] - Y[i_ycenter], X_bottom_left - X[i_xcenter]) + if theta_X_bottom_left < 0.0: + theta_X_bottom_left = 2.0 * np.pi + theta_X_bottom_left + XY_contour['left']['bottom'].append((X_bottom_left,Y[i],theta_X_bottom_left)) + if X_bottom_left < (X_search_min + 3.0 * dX): + X_search_min = X_bottom_left - 3.0 * dX + if np.isfinite(X_bottom_right): + theta_X_bottom_right = np.arctan2(Y[i] - Y[i_ycenter], X_bottom_right - X[i_xcenter]) + if theta_X_bottom_right < 0.0: + theta_X_bottom_right = 2.0 * np.pi + theta_X_bottom_right + XY_contour['right']['bottom'].append((X_bottom_right,Y[i],theta_X_bottom_right)) + if X_bottom_right > (X_search_max - 3.0 * dX): + X_search_max = X_bottom_right + 3.0 * dX + #print('2', Y[i], X_bottom_left, X_bottom_right, theta_X_bottom_left, theta_X_bottom_right) except: pass - # update the slice coordinates - if j > 0: - j-=1 + # update the slice coordinates + i -= 1 + else: + do_i = False # while the level intersects with the current Z slice gather the intersection coordinates - while (eval('level' + f_sgn + 'f_threshold(Z[:,l])') and l > 0): + do_k = True + while (do_k and k >= 0 and f_check(lvl, Z[:,k])): # interpolate the X coordinate of the bottom half of the contour on both the right and left - if X[l] >= X_min: + if X[k] >= X_search_min: try: # split the current column of the Z map in top and bottom # find the split and the extrema of the Y slice of Z - l_ZY_slice_split = f_split(Z[:,l]) - l_ZY_slice_bottom_max = f_bnd(Z[:l_ZY_slice_split,l]) - l_ZY_slice_top_max = l_ZY_slice_split+f_bnd(Z[l_ZY_slice_split:,l]) - - ZY_slice_bottom_left = Z[l_ZY_slice_bottom_max:l_ZY_slice_split+1,l] - ZY_slice_top_left = Z[l_ZY_slice_split:l_ZY_slice_top_max,l] - - # interpolate the Y coordinate for the left top and bottom - if interp_method == 'normal': - Y_bottom_left = float(interpolate.interp1d(ZY_slice_bottom_left,Y[l_ZY_slice_bottom_max:l_ZY_slice_split+1],bounds_error=False)(level)) - Y_top_left = float(interpolate.interp1d(ZY_slice_top_left,Y[l_ZY_slice_split:l_ZY_slice_top_max],bounds_error=False)(level)) - elif interp_method == 'bounded_extrapolation': - Y_slice_bottom_left = Y[l_ZY_slice_bottom_max:l_ZY_slice_split+1][ZY_slice_bottom_left<=threshold] - Y_slice_top_left = Y[l_ZY_slice_split:l_ZY_slice_top_max][ZY_slice_top_left<=threshold] - ZY_slice_bottom_left = ZY_slice_bottom_left[ZY_slice_bottom_left<=threshold] - ZY_slice_top_left = ZY_slice_top_left[ZY_slice_top_left<=threshold] - - Y_bottom_left = float(interpolate.interp1d(ZY_slice_bottom_left,Y_slice_bottom_left,bounds_error=False,fill_value='extrapolate')(level)) - Y_top_left = float(interpolate.interp1d(ZY_slice_top_left,Y_slice_top_left,bounds_error=False,fill_value='extrapolate')(level)) - - # diagnostic plot to check interpolation issues for columns - if tracer_diag == 'interp': - plt.figure() - plt.title('X:{}'.format(X[l])) - plt.plot(Z[l_ZY_slice_bottom_max:l_ZY_slice_split+1,l],Y[l_ZY_slice_bottom_max:l_ZY_slice_split+1],'b--',label='left, bottom') - plt.plot(ZY_slice_top_left,Y[l_ZY_slice_split:l_ZY_slice_top_max],'r--',label='left, top') - plt.axvline(level,ls='dashed',color='black') - plt.legend() - - if not np.isnan(Y_bottom_left): - XY_contour['left']['bottom'].append((X[l],Y_bottom_left)) + k_ZY_slice_split = f_split(Z[i_ycenter-i_ycenter_pad:i_ycenter+i_ycenter_pad+1,k].flatten()) + i_ycenter - i_ycenter_pad + Y_bottom_left, Y_top_left = find_dual_value(Y.flatten(), Z[:,k].flatten(), k_ZY_slice_split, lvl, thr, interp_method, f_sgn, search_tol) + + # insert the coordinates into the contour trace dict if not nan (bounds error) and order properly for merging later + if np.isfinite(Y_bottom_left): + theta_Y_bottom_left = np.arctan2(Y_bottom_left - Y[i_ycenter], X[k] - X[i_xcenter]) + if theta_Y_bottom_left < 0.0: + theta_Y_bottom_left = 2.0 * np.pi + theta_Y_bottom_left + XY_contour['left']['bottom'].append((X[k],Y_bottom_left,theta_Y_bottom_left)) # update the higher vertical bound if it is higher than the vertical bound found at the magnetic axis - if Y_bottom_left < Y_min: - Y_min = Y_bottom_left - if not np.isnan(Y_top_left): - XY_contour['left']['top'].append((X[l],Y_top_left)) + if Y_bottom_left < (Y_search_min + 3.0 * dY): + Y_search_min = Y_bottom_left - 3.0 * dY + if np.isfinite(Y_top_left): + theta_Y_top_left = np.arctan2(Y_top_left - Y[i_ycenter], X[k] - X[i_xcenter]) + if theta_Y_top_left < 0.0: + theta_Y_top_left = 2.0 * np.pi + theta_Y_top_left + XY_contour['left']['top'].append((X[k],Y_top_left,theta_Y_top_left)) # update the lower vertical bound if it is lower than the vertical bound found at the magnetic axis - if Y_top_left > Y_max: - Y_max = Y_top_left + if Y_top_left > (Y_search_max - 3.0 * dY): + Y_search_max = Y_top_left + 3.0 * dY + #print('2', X[k], Y_bottom_left, Y_top_left, theta_Y_bottom_left, theta_Y_top_left) except: pass + # update the slice coordinates + k -= 1 + else: + do_k = False - # update the slice coordinates - if l > 0: - l-=1 + # while the level intersects with the current Z slice gather the intersection coordinates + do_j = True + while (do_j and j < Z.shape[0] and f_check(lvl, Z[j,:])): + if Y[j] <= Y_search_max: + try: + # find the split in the Y slice of Z + j_ZX_slice_split = f_split(Z[j,i_xcenter-i_xcenter_pad:i_xcenter+i_xcenter_pad+1].flatten()) + i_xcenter - i_xcenter_pad + X_top_left, X_top_right = find_dual_value(X.flatten(), Z[j,:].flatten(), j_ZX_slice_split, lvl, thr, interp_method, f_sgn, search_tol) + + # insert the coordinates into the contour trace dict if not nan (bounds error) and order properly for merging later + if np.isfinite(X_top_left): + theta_X_top_left = np.arctan2(Y[j] - Y[i_ycenter], X_top_left - X[i_xcenter]) + if theta_X_top_left < 0.0: + theta_X_top_left = 2.0 * np.pi + theta_X_top_left + XY_contour['left']['top'].append((X_top_left,Y[j],theta_X_top_left)) + if X_top_left < (X_search_min + 3.0 * dX): + X_search_min = X_top_left - 3.0 * dX + if np.isfinite(X_top_right): + theta_X_top_right = np.arctan2(Y[j] - Y[i_ycenter], X_top_right - X[i_xcenter]) + if theta_X_top_right < 0.0: + theta_X_top_right = 2.0 * np.pi + theta_X_top_right + XY_contour['right']['top'].append((X_top_right,Y[j],theta_X_top_right)) + if X_top_right > (X_search_max - 3.0 * dX): + X_search_max = X_top_right + 3.0 * dX + #print('2', Y[j], X_top_left, X_top_right, theta_X_top_left, theta_X_top_right) + except: + pass + # update the slice coordinates + j += 1 + else: + do_j = False - # collate all the traced coordinates of the contour and sort by increasing X - XY_contour = sorted(XY_contour['right']['top']+XY_contour['left']['top']+XY_contour['left']['bottom']+XY_contour['right']['bottom']) + # while the level intersects with the current Z slice gather the intersection coordinates + do_l = True + while (do_l and l < Z.shape[1] and f_check(lvl, Z[:,l])): + if X[l] <= X_search_max: + try: + # find the split and the extrema of the X slice of Z + l_ZY_slice_split = f_split(Z[i_ycenter-i_ycenter_pad:i_ycenter+i_ycenter_pad+1,l].flatten()) + i_ycenter - i_ycenter_pad + Y_bottom_right, Y_top_right = find_dual_value(Y.flatten(), Z[:,l].flatten(), l_ZY_slice_split, lvl, thr, interp_method, f_sgn, search_tol) + + # insert the coordinates into the contour trace dict if not nan (bounds error) and order properly for merging later + if np.isfinite(Y_bottom_right): + theta_Y_bottom_right = np.arctan2(Y_bottom_right - Y[i_ycenter], X[l] - X[i_xcenter]) + if theta_Y_bottom_right < 0.0: + theta_Y_bottom_right = 2.0 * np.pi + theta_Y_bottom_right + XY_contour['right']['bottom'].append((X[l],Y_bottom_right,theta_Y_bottom_right)) + # update the higher vertical bound if it is higher than the vertical bound found at the magnetic axis + if Y_bottom_right < (Y_search_min + 3.0 * dY): + Y_search_min = Y_bottom_right - 3.0 * dY + if np.isfinite(Y_top_right): + theta_Y_top_right = np.arctan2(Y_top_right - Y[i_ycenter], X[l] - X[i_xcenter]) + if theta_Y_top_right < 0.0: + theta_Y_top_right = 2.0 * np.pi + theta_Y_top_right + XY_contour['right']['top'].append((X[l],Y_top_right,theta_Y_top_right)) + # update the lower vertical bound if it is lower than the vertical bound found at the magnetic axis + if Y_top_right > (Y_search_max - 3.0 * dY): + Y_search_max = Y_top_right + 3.0 * dY + #print('2', X[l], Y_bottom_right, Y_top_right, theta_Y_bottom_right, theta_Y_top_right) + except: + pass + # update the slice coordinates + l += 1 + else: + do_l = False + + # collate all the traced coordinates of the contour and sort by increasing X - why did you do this Garud? + # collate in order of theta, counterclockwise from outer midplane + XY_contour = sorted(XY_contour['right']['top']+XY_contour['left']['top']+XY_contour['left']['bottom']+XY_contour['right']['bottom'], key=itemgetter(2)) + max_dist = 1.2 * np.sqrt(dX ** 2 + dY ** 2) + X_contour = [x for x,y,t in XY_contour] + Y_contour = [y for x,y,t in XY_contour] + T_contour = [t for x,y,t in XY_contour] + d_contour = np.sqrt(np.diff(X_contour) ** 2 + np.diff(Y_contour) ** 2) + itr = 0 + while itr < len(d_contour): + if d_contour[itr] > max_dist: + X_contour.pop(itr+1) + Y_contour.pop(itr+1) + T_contour.pop(itr+1) + d_contour = np.sqrt(np.diff(X_contour) ** 2 + np.diff(Y_contour) ** 2) + else: + itr += 1 + XY_contour = zip(X_contour, Y_contour, T_contour) #print(XY_contour) XY_contour = list(dict.fromkeys(XY_contour)) - try: + #try: # find the extrema of the traced coordinates - Y_X_min = min(XY_contour) # coordinates for min(X) - Y_X_max = max(XY_contour) # coordinates for max(X) - X_Y_min = min(XY_contour,key=itemgetter(1)) # coordinates for min(Y) - X_Y_max = max(XY_contour,key=itemgetter(1)) # coordinates for max(Y) - - # split all the contour coordinates in four sequential quadrants, all sorted by increasing X - XY_right_top = [(x,y) for x,y in XY_contour if x>=X_Y_max[0] and y>Y_X_max[1]] - XY_left_top = [(x,y) for x,y in XY_contour if x=Y_X_min[1]] - XY_right_bottom = [(x,y) for x,y in XY_contour if x>X_Y_min[0] and y<=Y_X_max[1]] - XY_left_bottom = [(x,y) for x,y in XY_contour if x<=X_Y_min[0] and y threshold) and (level > threshold)): - plt.contour(X,Y,Z,[level],linestyles='dashed',colors='purple') - plt.axis('scaled') - plt.xlabel('X [m]') - plt.ylabel('Y [m]') - plt.show() - - if symmetrise: - R_sym = (contour_['X']+contour_['X'][::-1])/2 - Z_sym = (contour_['Y']-contour_['Y'][::-1])/2+integrate.trapz(contour_['X']*contour_['Y'],contour_['Y'])/integrate.trapz(contour_['X'],contour_['Y']) - contour_['X_sym']=R_sym - contour_['Y_sym']=Z_sym - - # compute a normalised level label for the contour level - contour_['label'] = np.sqrt((level - center)/(threshold - center)) - # find the contour center quantities and add them to the contour dict - contour_.update(contour_center(contour_,tracer_diag=tracer_diag)) - - # zipsort the contour from 0 - 2 pi - contour_['theta_XY'] = arctan2pi(contour_['Y']-contour_['Y0'],contour_['X']-contour_['X0']) - contour_['theta_XY'], contour_['X'], contour_['Y'] = zipsort(contour_['theta_XY'], contour_['X'], contour_['Y']) - - # close the contour - contour_['theta_XY'] = np.append(contour_['theta_XY'],contour_['theta_XY'][0]) - contour_['X'] = np.append(contour_['X'],contour_['X'][0]) - contour_['Y'] = np.append(contour_['Y'],contour_['Y'][0]) - - #plt.show() - return contour_ - except ValueError: - print('tracer.contour: No contour could be traced at the requested level!') + #Y_X_min = min(XY_contour,key=itemgetter(0)) # coordinates for min(X) + #Y_X_max = max(XY_contour,key=itemgetter(0)) # coordinates for max(X) + #X_Y_min = min(XY_contour,key=itemgetter(1)) # coordinates for min(Y) + #X_Y_max = max(XY_contour,key=itemgetter(1)) # coordinates for max(Y) + + # split all the contour coordinates in four sequential quadrants, all sorted by increasing X + XY_right_top = [(x,y) for x,y,t in XY_contour if (t >= 0.0) and (t < 0.5 * np.pi)] #if x>=X_Y_max[0] and y>Y_X_max[1]] + XY_left_top = [(x,y) for x,y,t in XY_contour if (t >= 0.5 * np.pi) and (t < np.pi)] #if x=Y_X_min[1]] + XY_right_bottom = [(x,y) for x,y,t in XY_contour if (t >= np.pi) and (t < 1.5 * np.pi)] #if x>X_Y_min[0] and y<=Y_X_max[1]] + XY_left_bottom = [(x,y) for x,y,t in XY_contour if (t >= 1.5 * np.pi) and (t < 2.0 * np.pi)] #if x<=X_Y_min[0] and y thr) and (lvl > thr)): + plt.contour(X,Y,Z,lvl,linestyles='dashed',colors='purple') + plt.axis('scaled') + plt.xlabel('X [m]') + plt.ylabel('Y [m]') + plt.show() + + if symmetrise: + R_sym = (contour_['X']+contour_['X'][::-1])/2 + Z_sym = (contour_['Y']-contour_['Y'][::-1])/2+integrate.trapz(contour_['X']*contour_['Y'],contour_['Y'])/integrate.trapz(contour_['X'],contour_['Y']) + contour_['X_sym']=R_sym + contour_['Y_sym']=Z_sym + + # compute a normalised level label for the contour level + contour_['label'] = np.sqrt((lvl - center)/(thr - center))[0] + # find the contour center quantities and add them to the contour dict + contour_.update(contour_center(contour_,tracer_diag=tracer_diag)) + + # zipsort the contour from 0 - 2 pi + contour_['theta_XY'] = arctan2pi(contour_['Y'] - contour_['Y0'], contour_['X'] - contour_['X0']) + contour_['theta_XY'], contour_['X'], contour_['Y'] = zipsort(contour_['theta_XY'], contour_['X'], contour_['Y']) + + # close the contour + contour_['theta_XY'] = np.append(contour_['theta_XY'],contour_['theta_XY'][0]) + contour_['X'] = np.append(contour_['X'],contour_['X'][0]) + contour_['Y'] = np.append(contour_['Y'],contour_['Y'][0]) + + #plt.show() + return contour_ + #except ValueError: + # print('tracer.contour: No contour could be traced at the requested level!') def contour_center(c,tracer_diag='none'): - """Find the geometric center of a contour trace c given by c['X'], c['Y']. + """Find the geometric center of a contour trace c given by c['X'], c['Y']. - Args: - `c` (dict): description of the contour, c={['X'],['Y'],['label'],...}, where: - - X,Y: the contour coordinates, - - label is the normalised contour level label np.sqrt((level - center)/(threshold - center)). + Args: + `c` (dict): description of the contour, c={['X'],['Y'],['label'],...}, where: + - X,Y: the contour coordinates, + - label is the normalised contour level label np.sqrt((level - center)/(threshold - center)). - Returns: - (dict): the contour with the extrema information added - """ + Returns: + (dict): the contour with the extrema information added + """ - # close the contour if not closed - if c['X'][-1] != c['X'][0] or c['Y'][-1] != c['Y'][0]: - c['X'] = np.append(c['X'],c['X'][0]) - c['Y'] = np.append(c['Y'],c['Y'][0]) + # close the contour if not closed + if c['X'][-1] != c['X'][0] or c['Y'][-1] != c['Y'][0]: + c['X'] = np.append(c['X'],c['X'][0]) + c['Y'] = np.append(c['Y'],c['Y'][0]) - # find the average elevation (midplane) of the contour by computing the vertical centroid [Candy PPCF 51 (2009) 105009] - c['Y0'] = integrate.trapz(c['X']*c['Y'],c['Y'])/integrate.trapz(c['X'],c['Y']) + # find the average elevation (midplane) of the contour by computing the vertical centroid [Candy PPCF 51 (2009) 105009] + moment = float(integrate.trapz(c['X']*c['Y'],c['Y'])) + area = float(integrate.trapz(c['X'],c['Y'])) + c['Y0'] = moment / area + #print('Y0', c['level'], moment, area) - # find the extrema of the contour in the radial direction at the average elevation - c = contour_extrema(c,tracer_diag=tracer_diag) + # find the extrema of the contour in the radial direction at the average elevation + c = contour_extrema(c,tracer_diag=tracer_diag) - # compute the minor and major radii of the contour at the average elevation - c['r'] = (c['X_out']-c['X_in'])/2 - c['X0'] = (c['X_out']+c['X_in'])/2 - #c['X0'] = integrate.trapz(c['X']*c['Y'],c['X'])/integrate.trapz(c['Y'],c['X']) + # compute the minor and major radii of the contour at the average elevation + c['r'] = (c['X_out']-c['X_in'])/2 + c['X0'] = (c['X_out']+c['X_in'])/2 + #c['X0'] = integrate.trapz(c['X']*c['Y'],c['X'])/integrate.trapz(c['Y'],c['X']) - return c + return c def contour_extrema(c,tracer_diag='none'): """Find the (true) extrema in X and Y of a contour trace c given by c['X'], c['Y']. @@ -434,90 +622,106 @@ def contour_extrema(c,tracer_diag='none'): Returns: (dict): the contour with the extrema information added """ + n = 16 + zpad = 2 if len(c['Y']) <= n else 3 + zmid = int(len(c['Y']) / 2) + # restack R_fs and Z_fs to get a continuous midplane outboard trace - X_out = np.hstack((c['X'][int(0.9*len(c['Y'])):],c['X'][:int(0.1*len(c['Y']))])) - Y_out = np.hstack((c['Y'][int(0.9*len(c['Y'])):],c['Y'][:int(0.1*len(c['Y']))])) + X_out = np.hstack((c['X'][-zpad:],c['X'][1:zpad+1])) + Y_out = np.hstack((c['Y'][-zpad:],c['Y'][1:zpad+1])) - X_in = c['X'][int(len(c['Y'])/2)-int(0.1*len(c['Y'])):int(len(c['Y'])/2)+int(0.1*len(c['Y']))] - Y_in = c['Y'][int(len(c['Y'])/2)-int(0.1*len(c['Y'])):int(len(c['Y'])/2)+int(0.1*len(c['Y']))] + X_in = c['X'][zmid-zpad:zmid+zpad+1][::-1] + Y_in = c['Y'][zmid-zpad:zmid+zpad+1][::-1] # find the approximate(!) extrema in Y of the contour - Y_max = np.max(c['Y']) - Y_min = np.min(c['Y']) + i_Y_max = np.argmax(c['Y']) + i_Y_min = np.argmin(c['Y']) # check if the midplane of the contour is provided if 'Y0' not in c: - c['Y0'] = Y_min+((Y_max-Y_min)/2) + Y_max = c['Y'][i_Y_max] + Y_min = c['Y'][i_Y_min] + c['Y0'] = float(Y_min+((Y_max-Y_min)/2)) + Y0 = np.array([c['Y0']]) # find the extrema in X of the contour at the midplane - c['X_out'] = float(interpolate.interp1d(Y_out,X_out,bounds_error=False)(c['Y0'])) - c['X_in'] = float(interpolate.interp1d(Y_in,X_in,bounds_error=False)(c['Y0'])) + X_out = interpolate.interp1d(Y_out,X_out,bounds_error=False)(Y0)[0] + X_in = interpolate.interp1d(Y_in,X_in,bounds_error=False)(Y0)[0] - # in case level is out of bounds in these interpolations - if np.isnan(c['X_out']) or np.isinf(c['X_out']): + # in case Y0 is out of bounds in these interpolations + if not np.isfinite(X_out): # restack X to get continuous trace on right side - X_ = np.hstack((c['X'][np.argmin(c['X']):],c['X'][:np.argmin(c['X'])])) + X_ = np.hstack((c['X'][np.argmin(c['X']):],c['X'][1:np.argmin(c['X'])])) # take the derivative of X_ dX_ = np.gradient(X_,edge_order=2) # find X_out by interpolating the derivative of X to 0. - dX_out = dX_[np.argmax(dX_):np.argmin(dX_)] - X_out = X_[np.argmax(dX_):np.argmin(dX_)] - c['X_out'] = float(interpolate.interp1d(dX_out,X_out,bounds_error=False)(0.)) - if np.isnan(c['X_in']) or np.isinf(c['X_in']): + dX_vec_out = dX_[np.argmax(dX_):np.argmin(dX_)] + X_vec_out = X_[np.argmax(dX_):np.argmin(dX_)] + X_out = interpolate.interp1d(dX_vec_out,X_vec_out,bounds_error=False)(np.array([0.0]))[0] + if not np.isfinite(X_in): dX = np.gradient(c['X'],edge_order=2) - dX_in = dX[np.argmin(dX):np.argmax(dX)] - X_in = c['X'][np.argmin(dX):np.argmax(dX)] - c['X_in'] = float(interpolate.interp1d(dX_in,X_in,bounds_error=False)(0.)) + dX_vec_in = dX[np.argmin(dX):np.argmax(dX)] + X_vec_in = c['X'][np.argmin(dX):np.argmax(dX)] + X_in = interpolate.interp1d(dX_vec_in,X_vec_in,bounds_error=False)(np.array([0.0]))[0] # generate filter lists that take a representative slice of the max and min of the contour coordinates around the approximate Y_max and Y_min - alpha = (0.9+0.075*c['label']**2) # magic to ensure just enough points are - max_filter = [z > alpha*(Y_max-c['Y0']) for z in c['Y']-c['Y0']] - min_filter = [z < alpha*(Y_min-c['Y0']) for z in c['Y']-c['Y0']] + #alpha = (0.9+0.075*c['label']**2) # magic to ensure just enough points are + #max_filter = [z > alpha*(Y_max-c['Y0']) for z in c['Y']-c['Y0']] + #min_filter = [z < alpha*(Y_min-c['Y0']) for z in c['Y']-c['Y0']] + #max_filter = c['Y'] > (alpha * Y_max + Y0 * (1.0 - alpha)) + #min_filter = c['Y'] < (alpha * Y_min + Y0 * (1.0 - alpha)) # patch for the filter lists in case the filter criteria results in < 7 points (minimum of required for 5th order fit + 1) - i_Y_max = np.argmax(c['Y']) - i_Y_min = np.argmin(c['Y']) - - if np.array(max_filter).sum() < 7: - for i in range(i_Y_max-3,i_Y_max+4): - if c['Y'][i] >= (np.min(c['Y'])+np.max(c['Y']))/2: - max_filter[i] = True - - if np.array(min_filter).sum() < 7: - for i in range(i_Y_min-3,i_Y_min+4): - if c['Y'][i] <= (np.min(c['Y'])+np.max(c['Y']))/2: - min_filter[i] = True + order = 5 if len(c['Y']) > n else 3 + mpad = int((order + 1) / 2) + l = 2 * n + while (l < len(c['Y'])): + mpad += 1 + l += n + X_max_vec = c['X'][i_Y_max-mpad:i_Y_max+mpad+1] + Y_max_vec = c['Y'][i_Y_max-mpad:i_Y_max+mpad+1] + X_min_vec = c['X'][i_Y_min-mpad:i_Y_min+mpad+1] + Y_min_vec = c['Y'][i_Y_min-mpad:i_Y_min+mpad+1] + #if np.count_nonzero(max_filter) < (order + 2): + # for i in range(i_Y_max-mpad,i_Y_max+mpad+1): + # if c['Y'][i] >= (np.min(c['Y'])+np.max(c['Y']))/2: + # max_filter[i] = True + #if np.count_nonzero(min_filter) < (order + 2): + # for i in range(i_Y_min-mpad,i_Y_min+mpad+1): + # if c['Y'][i] <= (np.min(c['Y'])+np.max(c['Y']))/2: + # min_filter[i] = True # fit the max and min slices of the contour, compute the gradient of these fits and interpolate to zero to find X_Ymax, Y_max, X_Ymin and Y_min - X_Ymax_fit = np.linspace(c['X'][max_filter][-1],c['X'][max_filter][0],5000) + X_Ymax_fit = np.linspace(np.nanmin(X_max_vec),np.nanmax(X_max_vec),5000) try: - Y_max_fit = interpolate.UnivariateSpline(c['X'][max_filter][::-1],c['Y'][max_filter][::-1],k=5)(X_Ymax_fit) + Y_max_fit = interpolate.UnivariateSpline(X_max_vec[::-1],Y_max_vec[::-1],k=order)(X_Ymax_fit) except: - Y_max_fit = np.poly1d(np.polyfit(c['X'][max_filter][::-1],c['Y'][max_filter][::-1],5))(X_Ymax_fit) + Y_max_fit = np.poly1d(np.polyfit(X_max_vec[::-1],Y_max_vec[::-1],order))(X_Ymax_fit) Y_max_fit_grad = np.gradient(Y_max_fit,X_Ymax_fit) - X_Ymax = interpolate.interp1d(Y_max_fit_grad,X_Ymax_fit,bounds_error=False)(0.) - Y_max = interpolate.interp1d(X_Ymax_fit,Y_max_fit,bounds_error=False)(X_Ymax) + X_Ymax = interpolate.interp1d(Y_max_fit_grad,X_Ymax_fit,bounds_error=False)(np.array([0.0]))[0] + Y_max = interpolate.interp1d(X_Ymax_fit,Y_max_fit,bounds_error=False)(np.array([X_Ymax]))[0] - X_Ymin_fit = np.linspace(c['X'][min_filter][-1],c['X'][min_filter][0],5000) + X_Ymin_fit = np.linspace(np.nanmin(X_min_vec),np.nanmax(X_min_vec),5000) try: - Y_min_fit = interpolate.UnivariateSpline(c['X'][min_filter],c['Y'][min_filter],k=5)(X_Ymin_fit) + Y_min_fit = interpolate.UnivariateSpline(X_min_vec,Y_min_vec,k=order)(X_Ymin_fit) except: - Y_min_fit = np.poly1d(np.polyfit(c['X'][min_filter][::-1],c['Y'][min_filter][::-1],5))(X_Ymin_fit) + Y_min_fit = np.poly1d(np.polyfit(X_min_vec,Y_min_vec,order))(X_Ymin_fit) Y_min_fit_grad = np.gradient(Y_min_fit,X_Ymin_fit) - X_Ymin = interpolate.interp1d(Y_min_fit_grad,X_Ymin_fit,bounds_error=False)(0.) - Y_min = interpolate.interp1d(X_Ymin_fit,Y_min_fit,bounds_error=False)(X_Ymin) + X_Ymin = interpolate.interp1d(Y_min_fit_grad,X_Ymin_fit,bounds_error=False)(np.array([0.0]))[0] + Y_min = interpolate.interp1d(X_Ymin_fit,Y_min_fit,bounds_error=False)(np.array([X_Ymin]))[0] if tracer_diag =='fs': # diagnostic plots - plt.plot(c['X'][max_filter],c['Y'][max_filter],'r.') - plt.plot(c['X'][min_filter],c['Y'][min_filter],'r.') + plt.plot(X_max_vec,Y_max_vec,'r.') + plt.plot(X_min_vec,Y_min_vec,'r.') plt.plot(X_Ymax_fit,Y_max_fit,'g-') plt.plot(X_Ymin_fit,Y_min_fit,'g-') #plt.axis('equal') #plt.show() - c.update({'X_Ymax':float(X_Ymax),'Y_max':float(Y_max),'X_Ymin':float(X_Ymin),'Y_min':float(Y_min)}) + c.update({'X_out': X_out, 'X_in': X_in, 'X_Ymax': X_Ymax, 'Y_max': Y_max, 'X_Ymin': X_Ymin, 'Y_min': Y_min}) + #print(c['level'], X_out, X_in, X_Ymax, X_Ymin, Y_max, Y_min) - return c \ No newline at end of file + return c From 04ed21c2a70997a1f2ea0bd9b073a542b1bdb033 Mon Sep 17 00:00:00 2001 From: Garud Snoep Date: Thu, 17 Oct 2024 18:18:59 +0200 Subject: [PATCH 2/9] tracer: contour closing bugfix for contour_center() --- megpy/tracer.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/megpy/tracer.py b/megpy/tracer.py index e94fd67..aeb5259 100755 --- a/megpy/tracer.py +++ b/megpy/tracer.py @@ -590,15 +590,13 @@ def contour_center(c,tracer_diag='none'): """ # close the contour if not closed - if c['X'][-1] != c['X'][0] or c['Y'][-1] != c['Y'][0]: - c['X'] = np.append(c['X'],c['X'][0]) - c['Y'] = np.append(c['Y'],c['Y'][0]) + c_ = copy.deepcopy(c) + if c_['X'][-1] != c_['X'][0] or c_['Y'][-1] != c_['Y'][0]: + c_['X'] = np.append(c_['X'],c_['X'][0]) + c_['Y'] = np.append(c_['Y'],c_['Y'][0]) # find the average elevation (midplane) of the contour by computing the vertical centroid [Candy PPCF 51 (2009) 105009] - moment = float(integrate.trapz(c['X']*c['Y'],c['Y'])) - area = float(integrate.trapz(c['X'],c['Y'])) - c['Y0'] = moment / area - #print('Y0', c['level'], moment, area) + c['Y0'] = integrate.trapezoid(c_['X']*c_['Y'],c_['Y'])/integrate.trapezoid(c_['X'],c_['Y']) # find the extrema of the contour in the radial direction at the average elevation c = contour_extrema(c,tracer_diag=tracer_diag) @@ -606,7 +604,7 @@ def contour_center(c,tracer_diag='none'): # compute the minor and major radii of the contour at the average elevation c['r'] = (c['X_out']-c['X_in'])/2 c['X0'] = (c['X_out']+c['X_in'])/2 - #c['X0'] = integrate.trapz(c['X']*c['Y'],c['X'])/integrate.trapz(c['Y'],c['X']) + #c['X0'] = integrate.trapezoid(c['X']*c['Y'],c['X'])/integrate.trapezoid(c['Y'],c['X']) return c From 4b9a8ee4c5663b6e6c0b88902e766db053150921 Mon Sep 17 00:00:00 2001 From: aaronkho Date: Mon, 21 Oct 2024 19:43:12 -0400 Subject: [PATCH 3/9] Reverse direction of level finder to avoid finding wrong edge --- megpy/tracer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/megpy/tracer.py b/megpy/tracer.py index aeb5259..95f5f4e 100755 --- a/megpy/tracer.py +++ b/megpy/tracer.py @@ -52,9 +52,9 @@ def find_dual_value(x, y, split_index, level, threshold, method='normal', sign=' upper_index = None # chop x,y in two parts to separate two solutions - x_lower = x[lower_index:split_index+1] + x_lower = x[lower_index:split_index+1][::-1] x_upper = x[split_index:upper_index] - y_lower = y[lower_index:split_index+1] + y_lower = y[lower_index:split_index+1][::-1] y_upper = y[split_index:upper_index] # if the normal method provides spikey contour traces, bound the interpolation domain and extrapolate the intersection From 7fbfa01a473cd040933c383bbe05b9875715308d Mon Sep 17 00:00:00 2001 From: aaronkho Date: Mon, 21 Oct 2024 19:44:57 -0400 Subject: [PATCH 4/9] Added contour closing point immediately after tracing --- megpy/tracer.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/megpy/tracer.py b/megpy/tracer.py index 95f5f4e..0edd1f1 100755 --- a/megpy/tracer.py +++ b/megpy/tracer.py @@ -490,6 +490,13 @@ def contour(X,Y,Z=None,level=None,threshold=None,i_center=None,interp_method='no X_contour = [x for x,y,t in XY_contour] Y_contour = [y for x,y,t in XY_contour] T_contour = [t for x,y,t in XY_contour] + if X_contour[0] != X_contour[-1]: + X_contour.append(X_contour[0]) + if Y_contour[0] != Y_contour[-1]: + Y_contour.append(Y_contour[0]) + if T_contour[0] != T_contour[-1]: + T_contour.append(T_contour[0]) + d_contour = np.sqrt(np.diff(X_contour) ** 2 + np.diff(Y_contour) ** 2) itr = 0 while itr < len(d_contour): From 19b11724aa911201a08e039b1851cd1e38482f1a Mon Sep 17 00:00:00 2001 From: aaronkho Date: Mon, 21 Oct 2024 19:48:30 -0400 Subject: [PATCH 5/9] Small readability changes to tracer routine --- megpy/tracer.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/megpy/tracer.py b/megpy/tracer.py index 0edd1f1..f0c8618 100755 --- a/megpy/tracer.py +++ b/megpy/tracer.py @@ -483,10 +483,9 @@ def contour(X,Y,Z=None,level=None,threshold=None,i_center=None,interp_method='no else: do_l = False - # collate all the traced coordinates of the contour and sort by increasing X - why did you do this Garud? # collate in order of theta, counterclockwise from outer midplane XY_contour = sorted(XY_contour['right']['top']+XY_contour['left']['top']+XY_contour['left']['bottom']+XY_contour['right']['bottom'], key=itemgetter(2)) - max_dist = 1.2 * np.sqrt(dX ** 2 + dY ** 2) + dist_limit = 1.2 * np.sqrt(dX ** 2 + dY ** 2) X_contour = [x for x,y,t in XY_contour] Y_contour = [y for x,y,t in XY_contour] T_contour = [t for x,y,t in XY_contour] @@ -500,7 +499,7 @@ def contour(X,Y,Z=None,level=None,threshold=None,i_center=None,interp_method='no d_contour = np.sqrt(np.diff(X_contour) ** 2 + np.diff(Y_contour) ** 2) itr = 0 while itr < len(d_contour): - if d_contour[itr] > max_dist: + if d_contour[itr] > dist_limit: X_contour.pop(itr+1) Y_contour.pop(itr+1) T_contour.pop(itr+1) @@ -549,7 +548,7 @@ def contour(X,Y,Z=None,level=None,threshold=None,i_center=None,interp_method='no contour_ = {'X':X_,'Y':Y_,'level':lvl[0]} # diagnostic plot for checking the complete contour trace, e.g. in combination with the contour extrema fitting - if tracer_diag == 'contour' and lvl < -4.8: + if tracer_diag == 'contour': # and np.abs(lvl - center) > (0.7 * np.abs(thr - center)): #plt.figure() plt.plot(contour_['X'],contour_['Y'],'b-') if ((center < thr) and (lvl < thr)) or ((center > thr) and (lvl > thr)): @@ -557,6 +556,7 @@ def contour(X,Y,Z=None,level=None,threshold=None,i_center=None,interp_method='no plt.axis('scaled') plt.xlabel('X [m]') plt.ylabel('Y [m]') + plt.title(f'psi = {lvl[0]:.4f}') plt.show() if symmetrise: From e14ea0b7d65c8a3524cd09756dac5ef62276c0dc Mon Sep 17 00:00:00 2001 From: aaronkho Date: Mon, 21 Oct 2024 19:57:33 -0400 Subject: [PATCH 6/9] Cosmetic change to bring variables closer to their usage --- megpy/tracer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megpy/tracer.py b/megpy/tracer.py index f0c8618..ce5db83 100755 --- a/megpy/tracer.py +++ b/megpy/tracer.py @@ -485,7 +485,6 @@ def contour(X,Y,Z=None,level=None,threshold=None,i_center=None,interp_method='no # collate in order of theta, counterclockwise from outer midplane XY_contour = sorted(XY_contour['right']['top']+XY_contour['left']['top']+XY_contour['left']['bottom']+XY_contour['right']['bottom'], key=itemgetter(2)) - dist_limit = 1.2 * np.sqrt(dX ** 2 + dY ** 2) X_contour = [x for x,y,t in XY_contour] Y_contour = [y for x,y,t in XY_contour] T_contour = [t for x,y,t in XY_contour] @@ -496,6 +495,7 @@ def contour(X,Y,Z=None,level=None,threshold=None,i_center=None,interp_method='no if T_contour[0] != T_contour[-1]: T_contour.append(T_contour[0]) + dist_limit = 1.2 * np.sqrt(dX ** 2 + dY ** 2) d_contour = np.sqrt(np.diff(X_contour) ** 2 + np.diff(Y_contour) ** 2) itr = 0 while itr < len(d_contour): From fa49d0e248706c690da7b2155fbe654f2c6c8540 Mon Sep 17 00:00:00 2001 From: Garud Snoep Date: Thu, 31 Oct 2024 20:53:26 +0100 Subject: [PATCH 7/9] localequilibrium: bugfix for H_inv being impossible when the least_squares Hessian is singular --- megpy/localequilibrium.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/megpy/localequilibrium.py b/megpy/localequilibrium.py index 624f058..8911cbb 100755 --- a/megpy/localequilibrium.py +++ b/megpy/localequilibrium.py @@ -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)) From f684497ed3add15c37399a063050e7b05dd13a10 Mon Sep 17 00:00:00 2001 From: Garud Snoep Date: Thu, 31 Oct 2024 20:54:06 +0100 Subject: [PATCH 8/9] equilibrium: removed interp_method = 'bounded_extrapolation' option from add_fluxsurfaces() --- megpy/equilibrium.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/megpy/equilibrium.py b/megpy/equilibrium.py index ca21702..ddc46e8 100755 --- a/megpy/equilibrium.py +++ b/megpy/equilibrium.py @@ -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. From a31f61c78c0316d351118add77ae89924d13c716 Mon Sep 17 00:00:00 2001 From: Garud Snoep Date: Wed, 20 Nov 2024 00:22:29 +0100 Subject: [PATCH 9/9] tracer: roll back contour_extrema changes for consistency --- megpy/tracer.py | 102 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 102 insertions(+) diff --git a/megpy/tracer.py b/megpy/tracer.py index ce5db83..446c348 100755 --- a/megpy/tracer.py +++ b/megpy/tracer.py @@ -618,6 +618,108 @@ def contour_center(c,tracer_diag='none'): def contour_extrema(c,tracer_diag='none'): """Find the (true) extrema in X and Y of a contour trace c given by c['X'], c['Y']. + Args: + `c` (dict): description of the contour, c={['X'],['Y'],['Y0'],['label'],...}, where: + - X,Y: the contour coordinates, + - Y0 (optional): is the average elevation, + - label is the normalised contour level label np.sqrt((level - center)/(threshold - center)). + + Returns: + (dict): the contour with the extrema information added + """ + # restack R_fs and Z_fs to get a continuous midplane outboard trace + X_out = np.hstack((c['X'][int(0.9*len(c['Y'])):],c['X'][:int(0.1*len(c['Y']))])) + Y_out = np.hstack((c['Y'][int(0.9*len(c['Y'])):],c['Y'][:int(0.1*len(c['Y']))])) + + X_in = c['X'][int(len(c['Y'])/2)-int(0.1*len(c['Y'])):int(len(c['Y'])/2)+int(0.1*len(c['Y']))] + Y_in = c['Y'][int(len(c['Y'])/2)-int(0.1*len(c['Y'])):int(len(c['Y'])/2)+int(0.1*len(c['Y']))] + + # find the approximate(!) extrema in Y of the contour + Y_max = np.max(c['Y']) + Y_min = np.min(c['Y']) + + # check if the midplane of the contour is provided + if 'Y0' not in c: + c['Y0'] = Y_min+((Y_max-Y_min)/2) + + # find the extrema in X of the contour at the midplane + c['X_out'] = float(interpolate.interp1d(Y_out,X_out,bounds_error=False)(c['Y0'])) + c['X_in'] = float(interpolate.interp1d(Y_in,X_in,bounds_error=False)(c['Y0'])) + + # in case level is out of bounds in these interpolations + if np.isnan(c['X_out']) or np.isinf(c['X_out']): + # restack X to get continuous trace on right side + X_ = np.hstack((c['X'][np.argmin(c['X']):],c['X'][:np.argmin(c['X'])])) + # take the derivative of X_ + dX_ = np.gradient(X_,edge_order=2) + # find X_out by interpolating the derivative of X to 0. + dX_out = dX_[np.argmax(dX_):np.argmin(dX_)] + X_out = X_[np.argmax(dX_):np.argmin(dX_)] + c['X_out'] = float(interpolate.interp1d(dX_out,X_out,bounds_error=False)(0.)) + if np.isnan(c['X_in']) or np.isinf(c['X_in']): + dX = np.gradient(c['X'],edge_order=2) + dX_in = dX[np.argmin(dX):np.argmax(dX)] + X_in = c['X'][np.argmin(dX):np.argmax(dX)] + c['X_in'] = float(interpolate.interp1d(dX_in,X_in,bounds_error=False)(0.)) + + # generate filter lists that take a representative slice of the max and min of the contour coordinates around the approximate Y_max and Y_min + alpha = (0.9+0.075*c['label']**2) # magic to ensure just enough points are + max_filter = [z > alpha*(Y_max-c['Y0']) for z in c['Y']-c['Y0']] + min_filter = [z < alpha*(Y_min-c['Y0']) for z in c['Y']-c['Y0']] + + # patch for the filter lists in case the filter criteria results in < 7 points (minimum of required for 5th order fit + 1) + i_Y_max = np.argmax(c['Y']) + i_Y_min = np.argmin(c['Y']) + + if np.array(max_filter).sum() < 7: + for i in range(i_Y_max-3,i_Y_max+4): + if c['Y'][i] >= (np.min(c['Y'])+np.max(c['Y']))/2: + max_filter[i] = True + + if np.array(min_filter).sum() < 7: + for i in range(i_Y_min-3,i_Y_min+4): + if c['Y'][i] <= (np.min(c['Y'])+np.max(c['Y']))/2: + min_filter[i] = True + + # fit the max and min slices of the contour, compute the gradient of these fits and interpolate to zero to find X_Ymax, Y_max, X_Ymin and Y_min + X_Ymax_fit = np.linspace(c['X'][max_filter][-1],c['X'][max_filter][0],5000) + try: + Y_max_fit = interpolate.UnivariateSpline(c['X'][max_filter][::-1],c['Y'][max_filter][::-1],k=5)(X_Ymax_fit) + except: + Y_max_fit = np.poly1d(np.polyfit(c['X'][max_filter][::-1],c['Y'][max_filter][::-1],5))(X_Ymax_fit) + Y_max_fit_grad = np.gradient(Y_max_fit,X_Ymax_fit) + + X_Ymax = interpolate.interp1d(Y_max_fit_grad,X_Ymax_fit,bounds_error=False)(0.) + Y_max = interpolate.interp1d(X_Ymax_fit,Y_max_fit,bounds_error=False)(X_Ymax) + + X_Ymin_fit = np.linspace(c['X'][min_filter][-1],c['X'][min_filter][0],5000) + try: + Y_min_fit = interpolate.UnivariateSpline(c['X'][min_filter],c['Y'][min_filter],k=5)(X_Ymin_fit) + except: + Y_min_fit = np.poly1d(np.polyfit(c['X'][min_filter][::-1],c['Y'][min_filter][::-1],5))(X_Ymin_fit) + Y_min_fit_grad = np.gradient(Y_min_fit,X_Ymin_fit) + + X_Ymin = interpolate.interp1d(Y_min_fit_grad,X_Ymin_fit,bounds_error=False)(0.) + Y_min = interpolate.interp1d(X_Ymin_fit,Y_min_fit,bounds_error=False)(X_Ymin) + + if tracer_diag =='fs': + # diagnostic plots + print(len(c['X'][max_filter]),len(c['X'][min_filter])) + + plt.plot(c['X'][max_filter],c['Y'][max_filter],'r.') + plt.plot(c['X'][min_filter],c['Y'][min_filter],'r.') + plt.plot(X_Ymax_fit,Y_max_fit,'g-') + plt.plot(X_Ymin_fit,Y_min_fit,'g-') + plt.axis('equal') + plt.show() + + c.update({'X_Ymax':float(X_Ymax),'Y_max':float(Y_max),'X_Ymin':float(X_Ymin),'Y_min':float(Y_min)}) + + return c + +def _contour_extrema(c,tracer_diag='none'): + """Find the (true) extrema in X and Y of a contour trace c given by c['X'], c['Y']. + Args: `c` (dict): description of the contour, c={['X'],['Y'],['Y0'],['label'],...}, where: - X,Y: the contour coordinates,