def create_grid( F, R, Z, in_settings, critical, boundary=None, iter=None, fpsi=None, fast=None):#, # f(psi) = R*Bt current function #nrad_flexible, #single_rad_grid, #xpt_mindist, xpt_mul, strictbndry,debug): # if size(nrad_flexible) == 0 : # nrad_flexible = 0 if iter==None: iter = 0 if iter > 3: print("ERROR: Too many iterations") return #, {error:1} #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; # Check the settings # If a setting is missing, set a default value # inspect.getargspec(create_grid) # if N_PARAMS() LT 3 THEN BEGIN # PRINT, "ERROR: Need at least a 2D array of psi values, R and Z arrays" # RETURN, {error:1} #ENDIF ELSE IF N_PARAMS() LT 4 THEN BEGIN # ; Settings omitted. Set defaults # print "Settings not given -> using default values" # settings = {psi_inner:0.9, # psi_outer:1.1, # nrad:36, # npol:64, # rad_peaking:0.0, # pol_peaking:0.0, # parweight:0.0} # ENDIF ELSE BEGIN # print "Checking settings" settings = in_settings # So the input isn't changed # str_check_present, settings, 'psi_inner', 0.9 # str_check_present, settings, 'psi_outer', 1.1 # str_check_present, settings, 'nrad', 36 # str_check_present, settings, 'npol', 64 # str_check_present, settings, 'rad_peaking', 0.0 # str_check_present, settings, 'pol_peaking', 0.0 # str_check_present, settings, 'parweight', 0.0 #ENDELSE s = numpy.ndim(F) s1 = numpy.shape(F) if s != 2: print("ERROR: First argument must be 2D array of psi values") return #, {error:1} nx = s1[0] ny = s1[1] s = numpy.ndim(R) s1 = numpy.size(R) if s != 1 or s1 != nx : print("ERROR: Second argument must be 1D array of major radii") return # {error:1} s = numpy.ndim(Z) s1 = numpy.size(Z) if s != 1 or s1 != ny: print("ERROR: Second argument must be 1D array of heights") return # {error:1} # Get an even number of points for efficient FFTs if nx % 2 == 1: # odd number of points in R. Cut out last point R = R[0:(nx-1)] F = F[0:(nx-1), :] nx = nx - 1 if ny % 2 == 1: # Same for Z Z = Z[0:(ny-1)] F = F[:,0:(ny-1)] ny = ny - 1 if boundary != None: s = numpy.ndim(boundary) s1= numpy.shape(boundary) if s != 2 or s1[0] != 2: print("WARNING: boundary must be a 2D array: [2, n]. Ignoring") boundary = 0 else: # Calculate indices bndryi = numpy.zeros((2,1188)) bndryi[0,:] = numpy.interp(bndryi[0,:], R, numpy.arange(0.,nx)) bndryi[1,:] = numpy.interp(bndryi[1,:], Z, numpy.arange(0.,ny)) if bndryi == None : bndryi = numpy.zeros((2,4)) bndryi[0,:] = [1, nx-1, nx-1, 1] bndryi[1,:] = [1, 1, ny-1, ny-1] #;;;;;;;;;;;;;; Psi interpolation data ;;;;;;;;;;;;;; interp_data = Bunch(nx=nx, ny=ny, method=0, f= F) # Always include function if fast == 'fast': print("Using Fast settings") interp_data.method = 2 #;;;;;;;;;;;;;;; First plot ;;;;;;;;;;;;;;;; nlev = 100 minf = numpy.min(F) maxf = numpy.max(F) levels = numpy.arange(numpy.float(nlev))*(maxf-minf)/numpy.float(nlev-1) + minf Rr=numpy.tile(R,ny).reshape(ny,nx).T Zz=numpy.tile(Z,nx).reshape(nx,ny) contour( Rr, Zz, F, levels=levels) # arrange the plot on the screen #mngr = get_current_fig_manager() #geom = mngr.window.geometry() #x,y,dx,dy = geom.getRect() #mngr.window.setGeometry(100, 100, dx, dy) # if boundary != None : # plot(boundary[0,:],boundary[1,:],'r-') show(block=False) #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; n_opoint = critical.n_opoint n_xpoint = critical.n_xpoint primary_opt = critical.primary_opt inner_sep = critical.inner_sep opt_ri = critical.opt_ri opt_zi = critical.opt_zi opt_f = critical.opt_f xpt_ri = numpy.array(critical.xpt_ri).flatten() xpt_zi = numpy.array(critical.xpt_zi).flatten() xpt_f = numpy.array(critical.xpt_f).flatten() # Overplot the separatrices, O-points #oplot_critical, F, R, Z, critical # Psi normalisation factors faxis = opt_f[primary_opt] fnorm = xpt_f[inner_sep] - opt_f[primary_opt] # From normalised psi, get range of f f_inner = faxis + numpy.min(settings.psi_inner)*fnorm f_outer = faxis + numpy.max(settings.psi_outer)*fnorm # Check the number of x-points if critical.n_xpoint == 0 : #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; # Grid entirely in the core print("Generating grid entirely in the core") nrad = numpy.sum(settings.nrad) # Add up all points npol = numpy.sum(settings.npol) rad_peaking = settings.rad_peaking[0] # Just the first region pol_peaking = settings.pol_peaking[0] # work out where to put the surfaces in the core fvals = radial_grid(nrad, f_inner, f_outer, 1, 1, [xpt_f[inner_sep]], rad_peaking) fvals = fvals.flatten() # Create a starting surface sind = numpy.int(old_div(nrad, 2)) start_f = fvals[sind] # contour_lines( F, numpy.arange(nx).astype(float), numpy.arange(ny).astype(float), levels=[start_f]) cs=contour( Rr, Zz, F, levels=[start_f]) p = cs.collections[0].get_paths() # # You might get more than one contours for the same start_f. We need to keep the closed one vn=numpy.zeros(numpy.size(p)) v = p[0].vertices vn[0]=numpy.shape(v)[0] xx=[v[:,0]] yy=[v[:,1]] if numpy.shape(vn)[0] > 1: for i in range(1,numpy.shape(vn)[0]): v = p[i].vertices vn[i]=numpy.shape(v)[0] xx.append(v[:,0]) yy.append(v[:,1]) #xx = [xx,v[:,0]] #yy = [yy,v[:,1]] print("PRIMARY: ", primary_opt, opt_ri[primary_opt], opt_zi[primary_opt]) if numpy.shape(vn)[0] > 1 : # Find the surface closest to the o-point opt_r = numpy.interp(opt_ri[primary_opt], numpy.arange(len(R)), R) opt_z = numpy.interp(opt_zi[primary_opt], numpy.arange(len(Z)), Z) ind = closest_line(xx, yy, opt_r, opt_z) x=xx[ind] y=yy[ind] print("Contour: ", ind) else: ind = 0 x=xx[0] y=yy[0] # plot the start_f line zc = cs.collections[0] setp(zc, linewidth=4) clabel(cs, [start_f], # label the level inline=1, fmt='%9.6f', fontsize=14) draw() show(block=False) # ans=query_yes_no('Press enter to create grid') if ans != 1 : show() sys.exit() start_ri, start_zi=transform_xy(x,y,R,Z) ## Make sure that the line goes clockwise # m = numpy.argmax(numpy.interp(start_zi,numpy.arange(Z.size).astype(float), Z)) if (numpy.gradient(numpy.interp(start_ri, numpy.arange(R.size).astype(float), R)))[m] < 0.0: # R should be increasing at the top. Need to reverse start_ri = start_ri[::-1] start_zi = start_zi[::-1] print('points reversed') ## Last point should be the same as the first # # Smooth and refine the starting location np = numpy.size(start_ri) s = 3 ar=numpy.append(numpy.append(start_ri[(np-s-1):(np-1)], start_ri), start_ri[1:s+1]) start_ri = SMOOTH(ar, window_len=s)[s+1:(np+s+1)] az=numpy.append(numpy.append(start_zi[(np-s-1):(np-1)], start_zi), start_zi[1:s+1]) start_zi = SMOOTH(az, window_len=s)[s+1:(np+s+1)] for i in range (np) : ri1=0. zi1=0. out=follow_gradient( interp_data, R, Z, start_ri[i], start_zi[i], start_f, ri1, zi1 ) status=out.status ri1=out.rinext zi1=out.zinext start_ri[i] = ri1 start_zi[i] = zi1 a = grid_region(interp_data, R, Z, start_ri, start_zi, fvals, sind, npol, boundary=boundary, fpsi=fpsi, parweight=settings.parweight, oplot='oplot') plot( numpy.append(a.rxy[0,:], a.rxy[0,0]), numpy.append(a.zxy[0,:], a.zxy[0,0]), 'r') for i in range (1, nrad) : plot( numpy.append(a.rxy[i,:], a.rxy[i,0]), numpy.append(a.zxy[i,:], a.zxy[i,0]), 'r') for i in range (0, npol-1) : plot( a.rxy[:,i], a.zxy[:,i], 'r') draw() # Get other useful variables psixy = numpy.zeros((nrad, npol)) for i in range (0, npol) : psixy[:,i] = old_div((fvals - faxis),fnorm) # to get normalised psi # Calculate magnetic field components dpsidR = numpy.zeros((nrad, npol)) dpsidZ = numpy.zeros((nrad, npol)) interp_data.method = 2 for i in range (nrad) : for j in range (npol) : out = local_gradient( interp_data, a.rixy[i,j], a.zixy[i,j], status=0, dfdr=0., dfdz=0.) status=out.status dfdr=out.dfdr[0][0] dfdz=out.dfdz[0][0] # dfd* are derivatives wrt the indices. Need to multiply by dr/di etc dpsidR[i,j] = old_div(dfdr,numpy.interp(a.rixy[i,j], numpy.arange(R.size).astype(float),numpy.gradient(R))) dpsidZ[i,j] = old_div(dfdz,numpy.interp(a.zixy[i,j], numpy.arange(Z.size).astype(float),numpy.gradient(Z))) # Set topology to connect in the core yup_xsplit = [nrad] ydown_xsplit = [nrad] yup_xin = [0] yup_xout = [-1] ydown_xin = [0] ydown_xout = [-1] #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; # Create result structure result = Bunch(error=0, # Signal success psi_inner=settings.psi_inner, psi_outer=settings.psi_outer, # Unchanged psi range nrad=nrad, npol=npol, #Number of points in radial and poloidal direction Rixy=a.rixy, Zixy=a.zixy, # Indices into R and Z of each point Rxy=a.rxy, Zxy=a.zxy, # Location of each grid point psixy=psixy, # Normalised psi for each point dpsidR=dpsidR, dpsidZ=dpsidZ, # Psi derivatives (for Bpol) faxis=faxis, fnorm=fnorm, # Psi normalisation factors settings=settings, # Settings used to create grid critical=critical, # Critical points yup_xsplit=yup_xsplit, # X index where domain splits (number of points in xin) ydown_xsplit=ydown_xsplit, yup_xin=yup_xin, yup_xout=yup_xout, # Domain index to connect to ydown_xin=ydown_xin, ydown_xout=ydown_xout) return result
def grid_region ( interp_data, R, Z, ri, zi, # Starting line to grid. fvals, # Location of the surfaces sind, # Index in fvals of the starting line npar, # Number of points along the line slast=None, # Index in fvals of last successful point sfirst=None, oplot=None, boundary=None, ffirst=None, flast=None, fpsi=None, # f(psi) = R*Bt optional current function parweight=None, # Space equally in parallel (1) or poloidal (0) distance ydown_dist=None, yup_dist=None, ydown_space=None, yup_space=None ): nsurf = numpy.size(fvals) if sind >= 0 : # starting position is on one of the output surfaces f0 = fvals[sind] nin = sind else: # Starting position between surfaces n = old_div(numpy.size(ri),2) out=local_gradient (interp_data, ri[n], zi[n], status=0, f=f0) status=out.status f0=out.f[0][0] if fvals[0] < fvals[nsurf-1] : w = numpy.where(fvals < f0, nin) else: w = numpy.where(fvals >= f0, nin) nout = nsurf - nin - 1 sfirst = 0 # Innermost successful index slast = nsurf-1 # Last successful index ffirst = fvals[sfirst] flast = fvals[slast] print(" => Gridding range: ", numpy.min(fvals), numpy.max(fvals)) nr = interp_data.nx nz = interp_data.ny ind = poloidal_grid(interp_data, R, Z, ri, zi, npar, fpsi=fpsi, ydown_dist=ydown_dist, yup_dist=yup_dist, ydown_space=ydown_space, yup_space=yup_space, parweight=parweight) rii = numpy.interp(ind, numpy.arange(ri.size).astype(float), ri) zii = numpy.interp(ind, numpy.arange(zi.size).astype(float), zi) #rii = int_func(SMOOTH(deriv(rii), 3)) + rii[0] #zii = int_func(SMOOTH(deriv(zii), 3)) + zii[0] #STOP # Refine the location of the starting point for i in range (npar) : ri1=0. zi1=0. out=follow_gradient( interp_data, R, Z, rii[i], zii[i], f0, ri1, zi1 ) ri1=out.rinext zi1=out.zinext rii[i] = ri1 zii[i] = zi1 # From each starting point, follow gradient in both directions rixy = numpy.zeros((nsurf, npar)) zixy = numpy.zeros((nsurf, npar)) status=0 print('Starting') for i in range (npar) : # print 'i=', i if sind >= 0 : rixy[nin, i] = rii[i] zixy[nin, i] = zii[i] else: # fvals[nin] should be just outside the starting position ftarg = fvals[nin] rinext=0. zinext=0. out=follow_gradient( interp_data, R, Z, rii[i], zii[i], ftarg, rinext, zinext, status=0 ) status=out.status rinext=out.rinext zinext=out.zinext rixy[nin, i] = rinext zixy[nin, i] = zinext for j in range (nout) : # print nin+j+1 ftarg = fvals[nin+j+1] rinext=0. zinext=0. out=follow_gradient ( interp_data, R, Z, rixy[nin+j, i], zixy[nin+j, i], ftarg, rinext, zinext, status=status, boundary=boundary, fbndry=0.) status=out.status rinext=out.rinext zinext=out.zinext fbndry=out.fbndry #print "create_grid" #print ftarg, rinext, zinext if status == 1 : rixy[nin+j+1, i] = -1.0 if nin+j < slast : slast = nin+j # last good surface index fbndry = fvals[slast] if (fvals[1] - fvals[0])*(flast - fbndry) > 0 : flast = 0.95*fbndry + 0.05*f0 break elif status == 2 : # Hit a boundary rixy[nin+j+1, i] = rinext zixy[nin+j+1, i] = zinext if nin+j < slast : slast = nin+j # Set the last point if (fvals[1] - fvals[0])*(flast - fbndry) > 0 : flast = 0.95*fbndry + 0.05*f0 break else : rixy[nin+j+1, i] = rinext zixy[nin+j+1, i] = zinext for j in range (nin) : # print nin-j-1 ftarg = fvals[nin-j-1] rinext=0. zinext=0. out=follow_gradient ( interp_data, R, Z, rixy[nin-j, i], zixy[nin-j, i], ftarg, rinext, zinext, status=status, boundary=boundary, fbndry=fbndry ) status=out.status rinext=out.rinext zinext=out.zinext fbndry=out.fbndry #print ftarg, rinext, zinext if status == 1 : rixy[nin-j-1, i] = -1.0 if nin-j > sfirst : sfirst = nin-j fbndry = fvals[sfirst] if (fvals[1] - fvals[0])*(ffirst - fbndry) < 0 : ffirst = 0.95*fbndry + 0.05*f0 break rixy[nin-j-1, i] = rinext zixy[nin-j-1, i] = zinext if status == 2 : if nin-j > sfirst : sfirst = nin-j if (fvals[1] - fvals[0])*(ffirst - fbndry) < 0 : ffirst = 0.95*fbndry + 0.05*f0 break #print rixy[:,i], zixy[:,i] if oplot != None: numpy.interp(rixy[:, i], numpy.arange(R.size).astype(float), R) numpy.interp(zixy[:, i], numpy.arange(Z.size).astype(float), Z) figure (0) oplot_contour(rixy[:, i], zixy[:, i] , R, Z) return Bunch(rixy=rixy, zixy=zixy, rxy=numpy.interp(rixy, numpy.arange(R.size), R), zxy=numpy.interp(zixy, numpy.arange(Z.size), Z))