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))
