Ejemplo n.º 1
0
def blob_velocity(n, **kwargs):
    import numpy as np

    from boututils import calculus as Calc
    # Calculate blob velocity in normalized time and normalized grid spacing
    #
    # Input: Blob density as a 3D vector in the form  n[t,x,z] where t is time and x,z are the perpendicular spatial coordinates
    #
    # Keywords:
    #
    #           type='peak' -> Calculate velocity of the peak density
    #           type='COM' -> Calculate centre of mass velocity
    #           Index=True -> return indices used to create velocity
    #
    # Default: Peak velocity with no index returning

    size = n.shape

    try:
        v_type = kwargs['type']
    except:
        v_type = 'peak'  #Default to peak velocity calculation
    try:
        return_index = kwargs['Index']
    except:
        return_index = False  #Default to no index returning

    if v_type == 'peak':
        x = np.zeros(size[0])
        z = np.zeros(size[0])
        for i in np.arange(size[0]):
            nmax, nmin = np.amax((n[i, :, :])), np.amin((n[i, :, :]))
            xpos, zpos = np.where(n[i, :, :] == nmax)
            x[i] = xpos[0]
            z[i] = zpos[0]

    if v_type == 'COM':
        x = np.zeros(size[0])
        z = np.zeros(size[0])
        for i in np.arange(size[0]):
            data = n[i, :, :] - n[0, 0, 0]  #use corner cell rather than nmin
            ntot = np.sum(data[:, :])

            z[i] = old_div(
                np.sum(np.sum(data[:, :], axis=0) * (np.arange(size[2]))),
                ntot)
            x[i] = old_div(
                np.sum(np.sum(data[:, :], axis=1) * (np.arange(size[1]))),
                ntot)

    vx = Calc.deriv(x)
    vz = Calc.deriv(z)

    if return_index:
        return vx, vz, x, z
    else:
        return vx, vz
Ejemplo n.º 2
0
def blob_velocity(n,**kwargs):
	import numpy as np

	from boututils import calculus as Calc
	# Calculate blob velocity in normalized time and normalized grid spacing
	# 
	# Input: Blob density as a 3D vector in the form  n[t,x,z] where t is time and x,z are the perpendicular spatial coordinates
	# 
	# Keywords: 
	#	    
	# 	    type='peak' -> Calculate velocity of the peak density
	#           type='COM' -> Calculate centre of mass velocity
	# 	    Index=True -> return indices used to create velocity
	#
	# Default: Peak velocity with no index returning

	size = n.shape
	
	try:
		v_type = kwargs['type']
	except:
		v_type = 'peak'		#Default to peak velocity calculation
	try: 		
		return_index = kwargs['Index']
	except:
		return_index = False	#Default to no index returning

	
	if v_type == 'peak':
		x = np.zeros(size[0])
		z = np.zeros(size[0])
		for i in np.arange(size[0]):
			nmax,nmin = np.amax((n[i,:,:])),np.amin((n[i,:,:]))
			xpos,zpos = np.where(n[i,:,:]==nmax)		
			x[i] = xpos[0]
			z[i] = zpos[0]	
		
	if v_type == 'COM':
		x = np.zeros(size[0])
		z = np.zeros(size[0])
		for i in np.arange(size[0]):            	
			data = n[i,:,:] - n[0,0,0]   #use corner cell rather than nmin
			ntot = np.sum(data[:,:])
		
			z[i] = old_div(np.sum(np.sum(data[:,:],axis=0)*(np.arange(size[2]))),ntot)
			x[i] = old_div(np.sum(np.sum(data[:,:],axis=1)*(np.arange(size[1]))),ntot)
				
	vx = Calc.deriv(x)
	vz = Calc.deriv(z)

	if return_index:
		return vx,vz,x,z
	else:
		return vx,vz
def calc_curvilinear_curvature(fname, field, grid):
    from scipy.signal import savgol_filter

    f = DataFile(str(fname), write=True)
    B = f.read("B")
    dBydz = np.zeros(np.shape(B))
    dBydx = np.zeros(np.shape(B))
    dBxdz = np.zeros(np.shape(B))
    dBzdx = np.zeros(np.shape(B))
    dx = grid.metric()["dx"]
    dz = grid.metric()["dz"]
    g_11 = grid.metric()["g_xx"]
    g_22 = grid.metric()["g_yy"]
    g_33 = grid.metric()["g_zz"]
    g_12 = 0.0
    g_13 = grid.metric()["g_xz"]
    g_23 = 0.0
    J = np.sqrt(g_11 * (g_22 * g_33 - g_23 * g_23) + g_12 *
                (g_13 * g_23 - g_12 * g_33) + g_13 *
                (g_12 * g_23 - g_22 * g_23))
    Bx_smooth = np.zeros(B.shape)
    By_smooth = np.zeros(B.shape)
    Bz_smooth = np.zeros(B.shape)

    for y in np.arange(0, B.shape[1]):
        pol, _ = grid.getPoloidalGrid(y)
        R = pol.R
        Z = pol.Z
        for x in np.arange(0, B.shape[0]):
            Bx_smooth[x, y, :] = savgol_filter(
                field.Bxfunc(R[x, :], y, Z[x, :]),
                np.int(np.ceil(B.shape[-1] / 21) // 2 * 2 + 1), 5)
            By_smooth[x, y, :] = savgol_filter(
                field.Byfunc(R[x, :], y, Z[x, :]),
                np.int(np.ceil(B.shape[-1] / 21) // 2 * 2 + 1), 5)

            dBydz[x, y, :] = calc.deriv(By_smooth[x, y, :]) / dz[x, y, :]
            dBxdz[x, y, :] = calc.deriv(Bx_smooth[x, y, :]) / dz[x, y, :]
        for z in np.arange(0, B.shape[-1]):
            Bz_smooth[:, y, z] = savgol_filter(
                field.Bzfunc(R[:, z], y, Z[:, z]),
                np.int(np.ceil(B.shape[0] / 7) // 2 * 2 + 1), 5)
            dBzdx[:, y, z] = calc.deriv(Bz_smooth[:, y, z]) / dx[:, y, z]
            dBydx[:, y, z] = calc.deriv(By_smooth[:, y, z]) / dx[:, y, z]

    bxcvx = (-1 / J) * (dBydz / B**2.)
    bxcvy = (1 / J) * ((dBxdz - dBzdx) / B**2.)
    bxcvz = (1 / J) * (dBydx / B**2.)

    f.write('bxcvz', bxcvz)
    f.write('bxcvx', bxcvx)
    f.write('bxcvy', bxcvy)
    f.close()
Ejemplo n.º 4
0
def pdiff_xy(nr, nz, r, z, f):
    # Get field components
    dfdR = numpy.zeros((nr, nz))
    dfdZ = numpy.zeros((nr, nz))

    for i in range(nz):
        dfdR[:, i] = deriv(f[:, i], r)

    for i in range(nr):
        dfdZ[i, :] = deriv(f[i, :], z)

    return Bunch(r=dfdR, z=dfdZ, phi=0.0)
def field_line_curvature(xin, yin, zin, configuration=1):

    tracer = Client(
        'http://esb.ipp-hgw.mpg.de:8280/services/FieldLineProxy?wsdl')

    pos = tracer.types.Points3D()

    config = tracer.types.MagneticConfig()
    config.configIds = configuration

    lineTask = tracer.types.LineTracing()
    lineTask.numSteps = 512

    task = tracer.types.Task()
    task.step = 2 * np.pi * (np.sqrt(xin**2 + yin**2)) / 512
    task.lines = lineTask

    pos.x1 = xin
    pos.x2 = yin
    pos.x3 = zin

    res_forward = tracer.service.trace(pos, config, task, None, None)

    x = np.asarray(res_forward.lines[0].vertices.x1)
    y = np.asarray(res_forward.lines[0].vertices.x2)
    z = np.asarray(res_forward.lines[0].vertices.x3)

    # config.inverseField=1

    # res_backward = tracer.service.trace(pos, config, task, None, None)

    # x = np.insert(x,0,res_backward.lines[0].vertices.x1[:])
    # y = np.insert(y,0,res_backward.lines[0].vertices.x2[:])
    # z = np.insert(z,0,res_backward.lines[0].vertices.x3[:])

    r = np.sqrt(x**2 + y**2)
    phi = np.arctan2(y, x)
    dphi = phi[1] - phi[0]

    drdphi = calc.deriv(r) / dphi
    d2rdphi2 = calc.deriv(drdphi) / dphi
    k = (r**2 + 2 * drdphi**2 - r * d2rdphi2) / ((r**2 + drdphi**2)**(1.5))

    # dzdphi = calc.deriv(z)/dphi
    # d2zdphi2 = calc.deriv(dzdphi)/dphi

    # Ktot = np.sqrt(d2zdphi2**2 + (d2rdphi2*dzdphi - d2zdphi2*drdphi)**2 - (d2rdphi2)**2) / (drdphi**2+dzdphi**2)**1.5

    # kz = Ktot - k

    return r, phi, x, y, z, k
Ejemplo n.º 6
0
def pdiff_xy ( nr, nz, r, z, f ):
    # Get field components
    dfdR = numpy.zeros((nr, nz))
    dfdZ = numpy.zeros((nr, nz))
  
    for i in range (nz) :
        dfdR[:,i] = deriv(f[:,i], r)
     
  
    for i in range (nr) :
        dfdZ[i,:] = deriv(f[i,:], z)
     
  
    return Bunch(r=dfdR, z=dfdZ, phi=0.0)
Ejemplo n.º 7
0
def DDX(psi, var):
    s = numpy.shape(var)
    nx = s[0]
    ny = s[1]

    dv = numpy.zeros((nx, ny))

    for i in range(ny):

        dv[:, i] = deriv(psi[:, i], var[:, i])

    return dv
Ejemplo n.º 8
0
def DDX(psi, var):
    s = numpy.shape(var)
    nx = s[0]
    ny = s[1]

    dv = numpy.zeros((nx, ny))

    for i in range(ny):

        dv[:, i] = deriv(psi[:, i], var[:, i])

    return dv
Ejemplo n.º 9
0
def pdiff(nr, nz, r, z, f):
    print("Calculating DCT...")
    dctf = dct2dslow(f)
    print("Finished DCT")

    drdi = deriv(r)
    dzdi = deriv(z)

    # Get field components
    dfdR = numpy.zeros((nr, nz))
    dfdZ = numpy.zeros((nr, nz))
    for i in range(nr):
        for j in range(nz):
            g = local_gradient(dctf, i, j, status=None)
            status = g.status
            dfdr = g.dfdr[0][0]
            dfdz = g.dfdz[0][0]

            # dfd* are derivatives wrt the indices. Need to divide by dr/di etc
            dfdR[i, j] = old_div(dfdr, drdi[i])
            dfdZ[i, j] = old_div(dfdz, dzdi[j])

    return Bunch(r=dfdR, z=dfdZ, phi=0.0)
Ejemplo n.º 10
0
def pdiff ( nr, nz, r, z, f):
    print("Calculating DCT...")
    dctf= dct2dslow( f )
    print("Finished DCT")

    drdi = deriv(r)
    dzdi = deriv(z)

    # Get field components
    dfdR = numpy.zeros((nr, nz))
    dfdZ = numpy.zeros((nr, nz))
    for i in range (nr) :
        for j in range (nz) :
            g = local_gradient(dctf, i, j, status=None)
            status=g.status
            dfdr=g.dfdr[0][0]
            dfdz=g.dfdz[0][0]

            # dfd* are derivatives wrt the indices. Need to divide by dr/di etc
            dfdR[i,j] = old_div(dfdr,drdi[i])
            dfdZ[i,j] = old_div(dfdz,dzdi[j])
    
    
    return Bunch(r=dfdR, z=dfdZ, phi=0.0)
Ejemplo n.º 11
0
def new_hfunc(h, psi, a, b, h0, fixpos):
    # global psi, fixpos, h0, a, b

    if fixpos == 0:
        h2 = numpy.append(h0, h)
    elif fixpos == numpy.size(psi) - 1:
        h2 = numpy.append(h, h0)
    else:
        h2 = numpy.append(numpy.append(h[0:(fixpos)], h0), h[fixpos::])

    f = a * h2 + b * deriv(psi, h2)

    if fixpos == 0:
        f = f[1::]
    elif fixpos == numpy.size(psi) - 1:
        f = f[0:(numpy.size(f) - 1)]
    else:
        f = numpy.append(f[0:(fixpos)], f[(fixpos + 1)::])

    return f
Ejemplo n.º 12
0
def new_hfunc ( h, psi, a, b, h0, fixpos ):
   # global psi, fixpos, h0, a, b


    if fixpos == 0 :
        h2 = numpy.append(h0, h)
    elif fixpos == numpy.size(psi)-1 :
        h2 = numpy.append(h, h0)
    else:
        h2 = numpy.append(numpy.append(h[0:(fixpos)], h0), h[fixpos::])

    f = a*h2 + b*deriv( psi, h2)

    if fixpos == 0 :
        f = f[1::]
    elif fixpos == numpy.size(psi)-1 :
        f = f[0:(numpy.size(f)-1)]
    else:
        f = numpy.append(f[0:(fixpos)], f[(fixpos+1)::])


    return f
Ejemplo n.º 13
0
def Bt_func(Bt, psi, a, b):
    #global  psi, a, b

    return deriv(psi, Bt) + a * Bt + old_div(b, Bt)
Ejemplo n.º 14
0
def surface_average(var, g, area=None):
    """Average a variable over a surface

    Parameters
    ----------
    var : array_like
        3-D or 4D variable to integrate (either [x,y,z] or [t,x,y,z])
    g : dict
        A dictionary of various grid quantities
    area : bool
        Average by flux-surface area = (B/Bp)*dl * R*dz

    Returns
    -------
    float
        Surface average of variable

    """

    s = np.ndim(var)

    if s == 4:
        nx = np.shape(var)[1]
        ny = np.shape(var)[2]
        nt = np.shape(var)[0]

        result = np.zeros((nx, nt))
        for t in range(nt):

            result[:, t] = surface_average(var[t, :, :, :], g, area=area)

        return result
    elif s != 3:
        print("ERROR: surface_average var must be 3 or 4D")
        return 0

# 3D [x,y,z]
    nx = np.shape(var)[0]
    ny = np.shape(var)[1]
    #    nz = np.shape(var)[2]

    # Use bunch to create grid structure
    grid = bunchify(g)

    # Calculate poloidal angle from grid
    theta = np.zeros((nx, ny))

    #status = gen_surface(mesh=grid) ; Start generator
    xi = -1
    yi = np.arange(0, ny, dtype=int)
    last = 0
    while True:
        #yi = gen_surface(last=last, xi=xi, period=periodic)
        xi = xi + 1
        if xi == nx - 1:
            last = 1

        dtheta = 2. * np.pi / np.float(ny)
        r = grid.Rxy[xi, yi]
        z = grid.Zxy[xi, yi]
        n = np.size(r)

        dl = old_div(np.sqrt(deriv(r)**2 + deriv(z)**2), dtheta)
        if area:
            dA = (old_div(grid.Bxy[xi, yi], grid.Bpxy[xi, yi])) * r * dl
            A = int_func(np.arange(n), dA)
            theta[xi, yi] = 2. * np.pi * A / A[n - 1]
        else:
            nu = dl * (grid.Btxy[xi, yi]) / ((grid.Bpxy[xi, yi]) * r)
            theta[xi, yi] = int_func(np.arange(n) * dtheta, nu)
            theta[xi, yi] = 2. * np.pi * theta[xi, yi] / theta[xi, yi[n - 1]]

        if last == 1: break

    vy = np.zeros(ny)
    result = np.zeros(nx)
    for x in range(nx):
        for y in range(ny):
            vy[y] = np.mean(var[x, y, :])

        result[x] = old_div(idl_tabulate(theta[x, :], vy), (2. * np.pi))

    return result
Ejemplo n.º 15
0
def volume_integral(var, g, xr=False):
    """Integrate a variable over a volume

    Parameters
    ----------
    var : array_like
        Variable to integrate
    g : dict
        A dictionary of various grid quantities
    xr : (int, int), optional
        Range of x indices (default: all of x)

    Returns
    -------
    float
        Volumne integral of variable

    """

    s = np.ndim(var)

    grid = bunchify(g)

    if s == 4:
        # 4D [t,x,y,z] - integrate for each t
        nx = np.shape(var)[1]
        ny = np.shape(var)[2]
        nt = np.shape(var)[0]

        result = np.zeros(nt)
        for t in range(nt):
            result[t] = volume_integral(var[t, :, :, :], g, xr=xr)
        return result

    elif s == 3:
        # 3D [x,y,z] - average in Z
        nx = np.shape(var)[0]
        ny = np.shape(var)[1]
        #       nz = np.shape(var)[2]

        zi = np.zeros((nx, ny))
        for x in range(nx):
            for y in range(ny):
                zi[x, y] = np.mean(var[x, y, :])

        return volume_integral(zi, g, xr=xr)

    elif s != 2:
        print("ERROR: volume_integral var must be 2, 3 or 4D")

    # 2D [x,y]
    nx = np.shape(var)[0]
    ny = np.shape(var)[1]

    if xr == False: xr = [0, nx - 1]

    result = 0.0

    #status = gen_surface(mesh=grid) ; Start generator
    xi = -1
    yi = np.arange(0, ny, dtype=int)
    last = 0
    #  iy = np.zeros(nx)
    while True:
        #yi = gen_surface(last=last, xi=xi, period=periodic)
        xi = xi + 1
        if xi == nx - 1: last = 1

        if (xi >= np.min(xr)) & (xi <= np.max(xr)):
            dtheta = 2. * np.pi / np.float(ny)
            r = grid.Rxy[xi, yi]
            z = grid.Zxy[xi, yi]
            n = np.size(r)
            dl = old_div(np.sqrt(deriv(r)**2 + deriv(z)**2), dtheta)

            # Area of flux-surface
            dA = (grid.Bxy[xi, yi] / grid.Bpxy[xi, yi] * dl) * (r * 2. * np.pi)
            # Volume
            if xi == nx - 1:
                dpsi = (grid.psixy[xi, yi] - grid.psixy[xi - 1, yi])
            else:
                dpsi = (grid.psixy[xi + 1, yi] - grid.psixy[xi, yi])

            dV = dA * dpsi / (r *
                              (grid.Bpxy[xi, yi]))  # May need factor of 2pi
            dV = np.abs(dV)

            result = result + np.sum(var[xi, yi] * dV)

        if last == 1: break

    return result
Ejemplo n.º 16
0
def calc_com_velocity(path=".",
                      fname="rot_ell.curv.68.16.128.Ic_02.nc",
                      tmax=-1,
                      track_peak=False):
    """
   
    input:

    return:

    """

    n = collect("Ne", path=path, tind=[0, tmax], info=False)
    t_array = collect("t_array", path=path, tind=[0, tmax], info=False)
    wci = collect("Omega_ci", path=path, tind=[0, tmax], info=False)
    dt = (t_array[1] - t_array[0]) / wci

    nt = n.shape[0]
    nx = n.shape[1]
    ny = n.shape[2]
    nz = n.shape[3]

    if fname is not None:
        fdata = DataFile(fname)

        R = fdata.read("R")
        Z = fdata.read("Z")

    else:
        R = np.zeros((nx, ny, nz))
        Z = np.zeros((nx, ny, nz))
        rhos = collect('rho_s0', path=path, tind=[0, tmax])
        Rxy = collect("R0", path=path, info=False) * rhos
        dx = (collect('dx', path=path, tind=[0, tmax], info=False) * rhos *
              rhos / (Rxy))[0, 0]
        dz = (collect('dz', path=path, tind=[0, tmax], info=False) * Rxy)
        for i in np.arange(0, nx):
            for j in np.arange(0, ny):
                R[i, j, :] = dx * i
                for k in np.arange(0, nz):
                    Z[i, j, k] = dz * k

    max_ind = np.zeros((nt, ny))
    fwhd = np.zeros((nt, nx, ny, nz))
    xval = np.zeros((nt, ny), dtype='int')
    zval = np.zeros((nt, ny), dtype='int')
    xpeakval = np.zeros((nt, ny))
    zpeakval = np.zeros((nt, ny))
    Rpos = np.zeros((nt, ny))
    Zpos = np.zeros((nt, ny))
    pos = np.zeros((nt, ny))
    vr = np.zeros((nt, ny))
    vz = np.zeros((nt, ny))
    vtot = np.zeros((nt, ny))
    pos_fit = np.zeros((nt, ny))
    v_fit = np.zeros((nt, ny))
    Zposfit = np.zeros((nt, ny))
    RZposfit = np.zeros((nt, ny))

    for y in np.arange(0, ny):
        for t in np.arange(0, nt):

            data = n[t, :, y, :]
            nmax, nmin = np.amax((data[:, :])), np.amin((data[:, :]))
            data[data < (nmin + 0.368 * (nmax - nmin))] = 0
            fwhd[t, :, y, :] = data
            ntot = np.sum(data[:, :])
            zval_float = np.sum(np.sum(data[:, :], axis=0) *
                                (np.arange(nz))) / ntot
            xval_float = np.sum(np.sum(data[:, :], axis=1) *
                                (np.arange(nx))) / ntot

            xval[t, y] = int(np.round(xval_float))
            zval[t, y] = int(np.round(zval_float))

            xpos, zpos = np.where(data[:, :] == nmax)
            xpeakval[t, y] = xpos[0]
            zpeakval[t, y] = zpos[0]

            if track_peak:
                Rpos[t, y] = R[int(xpeakval[t, y]), y, int(zpeakval[t, y])]
                Zpos[t, y] = Z[int(xpeakval[t, y]), y, int(zpeakval[t, y])]
            else:
                Rpos[t, y] = R[xval[t, y], y, zval[t, y]]
                Zpos[t, y] = Z[xval[t, y], y, zval[t, y]]

        pos[:, y] = np.sqrt((Rpos[:, y] - Rpos[0, y])**2)
        z1 = np.polyfit(t_array[:], pos[:, y], 5)
        f = np.poly1d(z1)
        pos_fit[:, y] = f(t_array[:])

        t_cross = np.where(pos_fit[:, y] > pos[:, y])[0]
        t_cross = 0  # t_cross[0]

        pos_fit[:t_cross, y] = pos[:t_cross, y]

        z1 = np.polyfit(t_array[:], pos[:, y], 5)
        f = np.poly1d(z1)
        pos_fit[:, y] = f(t_array[:])

        v_fit[:, y] = calc.deriv(pos_fit[:, y]) / dt

        posunique, pos_index = np.unique(pos[:, y], return_index=True)
        pos_index = np.sort(pos_index)
        XX = np.vstack(
            (t_array[:]**5, t_array[:]**4, t_array[:]**3, t_array[:]**2,
             t_array[:], pos[pos_index[0], y] * np.ones_like(t_array[:]))).T

        pos_fit_no_offset = np.linalg.lstsq(XX[pos_index, :-2], pos[pos_index,
                                                                    y])[0]
        pos_fit[:, y] = np.dot(pos_fit_no_offset, XX[:, :-2].T)
        v_fit[:, y] = calc.deriv(pos_fit[:, y]) / dt

    return v_fit[:, 0], pos_fit[:, 0], pos[:, 0], Rpos[:, 0], Zpos[:,
                                                                   0], t_cross
Ejemplo n.º 17
0
def surface_average(var, g, area=None):
    """Average a variable over a surface

    Parameters
    ----------
    var : array_like
        3-D or 4D variable to integrate (either [x,y,z] or [t,x,y,z])
    g : dict
        A dictionary of various grid quantities
    area : bool
        Average by flux-surface area = (B/Bp)*dl * R*dz

    Returns
    -------
    float
        Surface average of variable

    """

    s = np.ndim(var)



    if s == 4 :
        nx = np.shape(var)[1]
        ny = np.shape(var)[2]
        nt = np.shape(var)[0]

        result = np.zeros((nx,nt))
        for t in range (nt):

            result[:,t] = surface_average(var[t,:,:,:], g, area=area)

        return result
    elif s != 3 :
        print("ERROR: surface_average var must be 3 or 4D")
        return 0


  # 3D [x,y,z]
    nx = np.shape(var)[0]
    ny = np.shape(var)[1]
#    nz = np.shape(var)[2]

# Use bunch to create grid structure
    grid=bunchify(g)


  # Calculate poloidal angle from grid
    theta = np.zeros((nx,ny))

  #status = gen_surface(mesh=grid) ; Start generator
    xi = -1
    yi = np.arange(0,ny,dtype=int)
    last = 0
    while True:
    #yi = gen_surface(last=last, xi=xi, period=periodic)
        xi = xi + 1
        if xi == nx-1 :
            last = 1

        dtheta = 2.*np.pi / np.float(ny)
        r = grid.Rxy[xi,yi]
        z = grid.Zxy[xi,yi]
        n = np.size(r)

        dl = old_div(np.sqrt( deriv(r)**2 + deriv(z)**2 ), dtheta)
        if area:
            dA = (old_div(grid.Bxy[xi,yi],grid.Bpxy[xi,yi]))*r*dl
            A = int_func(np.arange(n),dA)
            theta[xi,yi] = 2.*np.pi*A/A[n-1]
        else:
            nu = dl * (grid.Btxy[xi,yi]) / ((grid.Bpxy[xi,yi]) * r )
            theta[xi,yi] = int_func(np.arange(n)*dtheta,nu)
            theta[xi,yi] = 2.*np.pi*theta[xi,yi] / theta[xi,yi[n-1]]

        if last==1 : break

    vy = np.zeros(ny)
    result = np.zeros(nx)
    for x in range(nx) :
        for y in range(ny) :
            vy[y] = np.mean(var[x,y,:])

        result[x] = old_div(idl_tabulate(theta[x,:], vy), (2.*np.pi))

    return result
Ejemplo n.º 18
0
def Bt_func ( Bt , psi, a, b):
    #global  psi, a, b

    return deriv( psi, Bt ) + a*Bt + old_div(b, Bt)
Ejemplo n.º 19
0
gfile='./cbm18_dens8.grid_nx68ny64.nc'


g = file_import(gfile)


Dphi0 = collect("Dphi0", path=path0)
phi0 = collect("phi0", path=path1) # needs diamagnetic effects
#
psixy=g.get('psixy')
PSI_AXIS=g.get('psi_axis')
PSI_BNDRY=g.get('psi_bndry')
#
psix=old_div((psixy[:,32]-PSI_AXIS),(PSI_BNDRY-PSI_AXIS))
Epsi=-deriv(phi0[:,32],psix)
#
#
fig=figure()
plot(psix,-Dphi0[:,32], 'r', linewidth=5)
plot(psix,Epsi,'k',linewidth=5)
annotate('w/o flow', xy=(.3, .7),  xycoords='axes fraction',horizontalalignment='center', verticalalignment='center', size=30)
annotate('w/ flow', xy=(.7, .4),  xycoords='axes fraction',horizontalalignment='center', verticalalignment='center', color='r', size=30)
xlabel('Radial $\psi$',fontsize=25)
ylabel('$\Omega(\psi)/\omega_A$',fontsize=25)
ylim([-.05,0])
xlim([0.4,1.2])
fig.set_tight_layout(True)
show(block=False)

p_f0 = collect("P", path=path0)
Ejemplo n.º 20
0
def GPI_curvature(x,y,z):
    phi0 = 4.787
    
    from osa import Client
    from boututils import calculus as calc
    
    nx = x.shape[0]
    nz = x.shape[-1]
    tracer = Client('http://esb.ipp-hgw.mpg.de:8280/services/FieldLineProxy?wsdl')

    config = tracer.types.MagneticConfig()
    config.configIds = [1]


    #### Expand in phi to take derivatives
    phi_list = np.linspace(.99*phi0, 1.01*phi0, 9)
    r = x*np.cos(phi0) + y*np.sin(phi0)

    dz = z[0,1]-z[0,0]
    dx = x[1,0]-x[1,1]
    dphi = phi_list[1]-phi_list[0]

    r_extended = np.zeros((nx,9,nz))
    z_extended = np.zeros((nx,9,nz))
    phi_extended = np.zeros((nx,9,nz))
    
    
    for j in np.arange(0,phi_list.shape[0]):
        r_extended[:,j,:] = x*np.cos(phi_list[j]) + y*np.sin(phi_list[j])
        z_extended[:,j,:] = z
        phi_extended[:,j,:] = np.ones((nx,nz))*phi_list[j]
        
    points = tracer.types.Points3D()
    points.x1 = (r_extended*np.cos(phi_extended)).flatten()
    points.x2 = (r_extended*np.sin(phi_extended)).flatten()
    points.x3 = z_extended.flatten()

    res = tracer.service.magneticField(points, config)
    
    ## Reshape to 3d array
    Bx = np.ndarray.reshape(np.asarray(res.field.x1),(nx,9,nz))
    By = np.ndarray.reshape(np.asarray(res.field.x2),(nx,9,nz))
    
    bphi = -Bx*np.sin(phi_extended) + By*np.cos(phi_extended)
    Bz = np.ndarray.reshape(np.asarray(res.field.x3),(nx,9,nz))
    b2 = (Bx**2 + By**2 + Bz**2)
    Bx /= b2
    By /= b2
    Bz /= b2
    bphi /= b2

    dbphidz = np.zeros((nx,nz))
    dbrdz = np.zeros((nx,nz))
    dbzdphi = np.zeros((nx,nz))
    dbzdr = np.zeros((nx,nz))
    dbphidr = np.zeros((nx,nz))
    dbrdphi = np.zeros((nx,nz))
    curlbr = np.zeros((nx,nz))
    curlbphi = np.zeros((nx,nz))
    curlbz = np.zeros((nx,nz))
    
    for i in np.arange(0,nx):
        dbphidz[i,:] = calc.deriv(bphi[i,4,:])/dz
        dbrdz[i,:] = calc.deriv(Bx[i,4,:])/dz

    for k in np.arange(0,nz):
        dbzdr[:,k] = calc.deriv(Bz[:,4,k])/dx
        dbphidr[:,k] =  calc.deriv(bphi[:,4,k])/dx
        for i in np.arange(0,nx):
            dbzdphi[i,k] = calc.deriv(Bz[i,:,k])[4]/dphi
            dbrdphi[i,k] = calc.deriv(Bx[i,:,k])[4]/dphi
            curlbr[i,k]   = (dbzdphi[i,k] - dbphidz[i,k])  
            curlbphi[i,k] = (dbrdz[i,k] - dbzdr[i,k])
            curlbz[i,k]   = (dbphidr[i,k] - dbrdphi[i,k])


    return curlbr, curlbphi, curlbz
Ejemplo n.º 21
0
def calc_curvilinear_curvature(fname, field, grid, maps):
    from scipy.signal import savgol_filter

    f = DataFile(str(fname), write=True)
    B = f.read("B")

    dx = grid.metric()["dx"]
    dz = grid.metric()["dz"]
    g_11 = grid.metric()["g_xx"]
    g_22 = grid.metric()["g_yy"]
    g_33 = grid.metric()["g_zz"]
    g_12 = 0.0
    g_13 = grid.metric()["g_xz"]
    g_23 = 0.0

    GR = np.zeros(B.shape)
    GZ = np.zeros(B.shape)
    Gphi = np.zeros(B.shape)
    dRdz = np.zeros(B.shape)
    dZdz = np.zeros(B.shape)
    dRdx = np.zeros(B.shape)
    dZdx = np.zeros(B.shape)

    for y in np.arange(0, B.shape[1]):
        pol, _ = grid.getPoloidalGrid(y)
        R = pol.R
        Z = pol.Z
        # G = \vec{B}/B, here in cylindrical coordinates
        GR[:, y, :] = field.Bxfunc(R, y, Z) / ((B[:, y, :])**2)
        GZ[:, y, :] = field.Bzfunc(R, y, Z) / ((B[:, y, :])**2)
        Gphi[:, y, :] = field.Byfunc(R, y, Z) / ((B[:, y, :])**2)
        for x in np.arange(0, B.shape[0]):
            dRdz[x, y, :] = calc.deriv(R[x, :]) / dz[x, y, :]
            dZdz[x, y, :] = calc.deriv(Z[x, :]) / dz[x, y, :]
        for z in np.arange(0, B.shape[-1]):
            dRdx[:, y, z] = calc.deriv(R[:, z]) / dx[:, y, z]
            dZdx[:, y, z] = calc.deriv(Z[:, z]) / dx[:, y, z]

    R = f.read("R")
    Z = f.read("Z")
    dy = f.read("dy")

    ## calculate Jacobian and contravariant terms in curvilinear coordinates
    J = R * (dZdz * dRdx - dZdx * dRdz)
    Gx = (GR * dZdz - GZ * dRdz) * (R / J)
    Gz = (GZ * dRdx - GR * dZdx) * (R / J)

    G_x = Gx * g_11 + Gphi * g_12 + Gz * g_13
    G_y = Gx * g_12 + Gphi * g_22 + Gz * g_23
    G_z = Gx * g_13 + Gphi * g_23 + Gz * g_33

    dG_zdy = np.zeros(B.shape)
    dG_ydz = np.zeros(B.shape)
    dG_xdz = np.zeros(B.shape)
    dG_zdx = np.zeros(B.shape)
    dG_ydx = np.zeros(B.shape)
    dG_xdy = np.zeros(B.shape)
    for y in np.arange(0, B.shape[1]):
        for x in np.arange(0, B.shape[0]):
            dG_ydz[x, y, :] = calc.deriv(G_y[x, y, :]) / dz[x, y, :]
            dG_xdz[x, y, :] = calc.deriv(G_x[x, y, :]) / dz[x, y, :]
        for z in np.arange(0, B.shape[-1]):
            dG_ydx[:, y, z] = calc.deriv(G_y[:, y, z]) / dx[:, y, z]
            dG_zdx[:, y, z] = calc.deriv(G_z[:, y, z]) / dx[:, y, z]

    #this should really use the maps...
    for x in np.arange(0, B.shape[0]):
        for z in np.arange(0, B.shape[-1]):
            dG_zdy[x, :, z] = calc.deriv(G_z[x, :, z]) / dy[x, :, z]
            dG_xdy[x, :, z] = calc.deriv(G_x[x, :, z]) / dy[x, :, z]

    bxcvx = (dG_zdy - dG_ydz) / J
    bxcvy = (dG_xdz - dG_zdx) / J
    bxcvz = (dG_ydx - dG_xdy) / J
    bxcv = g_11 * (bxcvx**2) + g_22 * (bxcvy**2) + g_33 * (bxcvz**2) + 2 * (
        bxcvz * bxcvx * g_13)
    f.write('bxcvx', bxcvx)
    f.write('bxcvy', bxcvy)
    f.write('bxcvz', bxcvz)
    f.write('J', J)
    f.close()
def island_fluxtube_grid(nx, ny, nz, turns=1, configuration=1):
    ### Copy GPI Code:
    # Generates a nx by nz rectangular grid at the location of the
    # outboard midplane island in standard configuration
    ####
    bottom_left = [6.1, 0, -0.5]
    bottom_right = [6.2, 0, -0.5]
    top_left = [6.1, 0, 0.5]
    top_right = [6.2, 0, 0.5]
    x_array = np.zeros((nx, ny, nz))
    z_array = np.zeros((nx, ny, nz))
    y_array = np.zeros((nx, ny, nz))
    for y in np.arange(0, ny):
        left_edge_x = np.linspace(bottom_left[0], top_left[0], nz)
        left_edge_z = np.linspace(bottom_left[2], top_left[2], nz)
        left_edge_y = np.linspace(bottom_left[1], top_left[1], nz)

        right_edge_x = np.linspace(bottom_right[0], top_right[0], nz)
        right_edge_z = np.linspace(bottom_right[2], top_right[2], nz)
        right_edge_y = np.linspace(bottom_right[1], top_right[1], nz)

        for z in np.arange(0, nz):
            line_x = np.linspace(left_edge_x[z], right_edge_x[z], nx)
            line_z = np.linspace(left_edge_z[z], right_edge_z[z], nx)
            line_y = np.linspace(left_edge_y[z], right_edge_y[z], nx)
            z_array[:, y, z] = line_z
            x_array[:, y, z] = line_x
            y_array[:, y, z] = line_y

    ### Have grid, must calculate curvature ####

    tracer = Client(
        'http://esb.ipp-hgw.mpg.de:8280/services/FieldLineProxy?wsdl')

    pos = tracer.types.Points3D()

    config = tracer.types.MagneticConfig()
    config.configIds = configuration

    lineTask = tracer.types.LineTracing()
    lineTask.numSteps = turns * ny

    task = tracer.types.Task()
    task.step = turns * 2 * np.pi * np.mean(x_array) / lineTask.numSteps
    task.lines = lineTask

    pos.x1 = x_array.flatten()
    pos.x2 = y_array.flatten()
    pos.x3 = z_array.flatten()

    res = tracer.service.trace(pos, config, task, None, None)

    x = np.asarray(res.lines[0].vertices.x1)
    y = np.asarray(res.lines[0].vertices.x2)
    z = np.asarray(res.lines[0].vertices.x3)

    r = np.sqrt(x**2 + y**2)
    phi = np.arctan2(y, x)
    dphi = phi[1] - phi[0]

    drdphi = calc.deriv(r) / dphi
    d2rdphi2 = calc.deriv(drdphi) / dphi
    k = (r**2 + 2 * drdphi**2 - r * d2rdphi2) / ((r**2 + drdphi**2)**(1.5))

    k = np.ndarray.reshape(k, (nx, ny, nz))
    # dzdphi = calc.deriv(z)/dphi
    # d2zdphi2 = calc.deriv(dzdphi)/dphi

    # Ktot = np.sqrt(d2zdphi2**2 + (d2rdphi2*dzdphi - d2zdphi2*drdphi)**2 - (d2rdphi2)**2) / (drdphi**2+dzdphi**2)**1.5

    # kz = Ktot - k

    return r, phi, x, y, z, k
Ejemplo n.º 23
0
def process_grid(rz_grid,
                 mesh,
                 output=None,
                 poorquality=None,
                 gui=None,
                 parent=None,
                 reverse_bt=None,
                 curv=None,
                 smoothpressure=None,
                 smoothhthe=None,
                 smoothcurv=None,
                 settings=None):

    if settings == None:
        # Create an empty structure
        settings = Bunch(dummy=0)

        # Check settings
        settings.calcp = -1
        settings.calcbt = -1
        settings.calchthe = -1
        settings.calcjpar = -1

# ;CATCH, err
# ;IF err NE 0 THEN BEGIN
# ;  PRINT, "PROCESS_GRID failed"
#;  PRINT, "   Error message: "+!ERROR_STATE.MSG
# ;  CATCH, /cancel
# ;  RETURN
# ;ENDIF

    MU = 4.e-7 * numpy.pi

    poorquality = 0

    if output == None: output = "bout.grd.nc"

    # Size of the mesh
    nx = numpy.int(numpy.sum(mesh.nrad))
    ny = numpy.int(numpy.sum(mesh.npol))

    # Find the midplane
    ymid = 0
    status = gen_surface(mesh=mesh)  # Start generator

    while True:
        period, yi, xi, last = gen_surface(period=None, last=None, xi=None)
        if period:
            rm = numpy.max(mesh.Rxy[xi, yi])
            ymidindx = numpy.argmax(mesh.Rxy[xi, yi])
            ymid = yi[ymidindx]
            break

        if last == 1: break

    Rxy = numpy.asarray(mesh.Rxy)
    Zxy = numpy.asarray(mesh.Zxy)
    psixy = mesh.psixy * mesh.fnorm + mesh.faxis  # Non-normalised psi

    pressure = numpy.zeros((nx, ny))

    # Use splines to interpolate pressure profile
    status = gen_surface(mesh=mesh)  # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=last, xi=xi)
        if period:
            # Pressure only given on core surfaces
            # pressure[xi,yi] = SPLINE(rz_grid.npsigrid, rz_grid.pres, mesh.psixy[xi,yi[0]], /double)
            sol = interpolate.UnivariateSpline(rz_grid.npsigrid,
                                               rz_grid.pres,
                                               s=1)
            pressure[xi, yi] = sol(mesh.psixy[xi, yi[0]])

        else:

            pressure[xi, yi] = rz_grid.pres[numpy.size(rz_grid.pres) - 1]

        if last == 1: break

    # Add a minimum amount
    if numpy.min(pressure) < 1.0e-2 * numpy.max(pressure):
        print("****Minimum pressure is very small:", numpy.min(pressure))
        print("****Setting minimum pressure to 1% of maximum")
        pressure = pressure + 1e-2 * numpy.max(pressure)

    if smoothpressure != None:
        p0 = pressure[:, ymid]  # Keep initial pressure for comparison
        while True:
            #!P.multi=[0,0,2,0,0]
            fig = figure()
            plot(p0,
                 xtitle="X index",
                 ytitle="pressure at y=" + numpy.strip(numpy.str(ymid), 2) +
                 " dashed=original",
                 color=1,
                 lines=1)
            plot(pressure[:, ymid], color=1)
            plot(deriv(p0),
                 xtitle="X index",
                 ytitle="DERIV(pressure)",
                 color=1,
                 lines=1)
            plot(deriv(pressure[:, ymid]), color=1)
            sm = query_yes_no(
                "Smooth pressure profile?")  #, gui=gui, dialog_parent=parent)
            if sm:
                # Smooth the pressure profile

                p2 = pressure
                for i in range(6):
                    status = gen_surface(mesh=mesh)  # Start generator
                    while True:
                        # Get the next domain
                        period, yi, xi, last = gen_surface(period=period,
                                                           last=last,
                                                           xi=xi)

                        if (xi > 0) and (xi < (nx - 1)):
                            for j in range(numpy.size(yi)):
                                p2[xi,
                                   yi[j]] = (0.5 * pressure[xi, yi[j]] + 0.25 *
                                             (pressure[xi - 1, yi[j]] +
                                              pressure[xi + 1, yi[j]]))

                        # Make sure it's still constant on flux surfaces
                        p2[xi, yi] = numpy.mean(p2[xi, yi])
                        if last != None: break
                    pressure = p2

            if sm == 0: break

    if numpy.min(pressure) < 0.0:
        print("")
        print("============= WARNING ==============")
        print("Poor quality equilibrium: Pressure is negative")
        print("")
        poorquality = 1

    dpdpsi = DDX(psixy, pressure)

    #;IF MAX(dpdpsi)*mesh.fnorm GT 0.0 THEN BEGIN
    #;  PRINT, ""
    #;  PRINT, "============= WARNING =============="
    #;  PRINT, "Poor quality equilibrium: Pressure is increasing radially"
    #;  PRINT, ""
    #;  poorquality = 1
    #;ENDIF

    # Grid spacing
    dx = numpy.zeros((nx, ny))
    for y in range(ny):
        dx[0:(nx - 1), y] = psixy[1::, y] - psixy[0:(nx - 1), y]
        dx[nx - 1, y] = dx[nx - 2, y]

    # Sign
    bpsign = 1.
    xcoord = psixy
    if numpy.min(dx) < 0.:
        bpsign = -1.
        dx = -dx  # dx always positive
        xcoord = -xcoord

    dtheta = 2. * numpy.pi / numpy.float(ny)
    dy = numpy.zeros((nx, ny)) + dtheta

    # B field components
    # Following signs mean that psi increasing outwards from
    # core to edge results in Bp clockwise in the poloidal plane
    # i.e. in the positive Grad Theta direction.

    Brxy = old_div(mesh.dpsidZ, Rxy)
    Bzxy = old_div(-mesh.dpsidR, Rxy)
    Bpxy = numpy.sqrt(Brxy**2 + Bzxy**2)

    # Determine direction (dot B with grad y vector)

    dot = (Brxy[0, ymid] * (Rxy[0, ymid + 1] - Rxy[0, ymid - 1]) +
           Bzxy[0, ymid] * (Zxy[0, ymid + 1] - Zxy[0, ymid - 1]))

    if dot < 0.:
        print(
            "**** Poloidal field is in opposite direction to Grad Theta -> Bp negative"
        )
        Bpxy = -Bpxy
        if bpsign > 0: sys.exit()  # Should be negative
        bpsign = -1.0
    else:
        if bpsign < 0: sys.exit()  # Should be positive
        bpsign = 1.

# Get toroidal field from poloidal current function fpol
    Btxy = numpy.zeros((nx, ny))
    fprime = numpy.zeros((nx, ny))
    fp = deriv(rz_grid.npsigrid * (rz_grid.sibdry - rz_grid.simagx),
               rz_grid.fpol)

    status = gen_surface(mesh=mesh)  # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=period, xi=xi)

        if period:
            # In the core
            #fpol = numpy.interp(rz_grid.fpol, rz_grid.npsigrid, mesh.psixy[xi,yi], /spline)

            sol = interpolate.UnivariateSpline(rz_grid.npsigrid,
                                               rz_grid.fpol,
                                               s=1)
            #  fpol = SPLINE(rz_grid.npsigrid, rz_grid.fpol, mesh.psixy[xi,yi[0]], 'double')
            fpol = sol(mesh.psixy[xi, yi[0]])

            sol = interpolate.UnivariateSpline(rz_grid.npsigrid, fp, s=1)
            # fprime[xi,yi] = SPLINE(rz_grid.npsigrid, fp, mesh.psixy[xi,yi[0]], 'double')
            fprime[xi, yi] = sol(mesh.psixy[xi, yi[0]])

        else:
            # Outside core. Could be PF or SOL
            fpol = rz_grid.fpol[numpy.size(rz_grid.fpol) - 1]
            fprime[xi, yi] = 0.

        Btxy[xi, yi] = old_div(fpol, Rxy[xi, yi])

        if last == 1: break

    # Total B field
    Bxy = numpy.sqrt(Btxy**2 + Bpxy**2)

    #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    # Go through the domains to get a starting estimate
    # of hthe
    hthe = numpy.zeros((nx, ny))

    #   Pick a midplane index
    status = gen_surface(mesh=mesh)  # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=last, xi=xi)

        if period:
            # In the core
            rmax = numpy.argmax(Rxy[xi, yi])
            ymidplane = yi[rmax]
            break

        if last == 1: break

    status = gen_surface(mesh=mesh)  # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=last, xi=xi)

        n = numpy.size(yi)

        # Get distance along this line

        if period:

            # Periodic, so can use FFT
            #drdi = REAL_PART(fft_deriv(Rxy[xi, yi]))
            #dzdi = REAL_PART(fft_deriv(Zxy[xi, yi]))
            line = numpy.append(Rxy[xi, yi[n - 1::]], Rxy[xi, yi])
            line = numpy.append(line, Rxy[xi, yi[0:1]])

            drdi = deriv(line)[1:n + 1]

            line = numpy.append(Zxy[xi, yi[n - 1::]], Zxy[xi, yi])
            line = numpy.append(line, Zxy[xi, yi[0:1]])

            dzdi = deriv(line)[1:n + 1]
        else:
            # Non-periodic
            drdi = numpy.gradient(Rxy[xi, yi])
            dzdi = numpy.gradient(Zxy[xi, yi])

        dldi = numpy.sqrt(drdi**2 + dzdi**2)

        if 0:

            # Need to smooth to get sensible results
            if period:
                n = numpy.size(dldi)
                line = numpy.append(dldi[(n - 2)::], dldi)  # once
                line = numpy.append(line, dldi[0:2])
                dldi = SMOOTH(line, 5)[4:(n + 4)]

                line = numpy.append(dldi[(n - 2)::], dldi)  #twice
                line = numpy.append(line, dldi[0:2])
                dldi = SMOOTH(line, 5)[4:(n + 4)]

                line = numpy.append(dldi[(n - 2)::], dldi)  # three
                line = numpy.append(line, dldi[0:2])
                dldi = SMOOTH(line, 5)[4:(n + 4)]

            else:
                line = dldi
                dldi = SMOOTH(line, 5)[2:n + 2]
                line = dldi
                dldi = SMOOTH(line, 5)[2:n + 2]
                line = dldi
                dldi = SMOOTH(dldi, 5)[2:n + 2]

        hthe[xi, yi] = old_div(dldi, dtheta)  # First estimate of hthe

        # Get outboard midplane
        if period and xi == 0:
            m = numpy.argmax(Rxy[0, yi])
            ymidplane = yi[m]

        if last == 1: break

    print("Midplane index ", ymidplane)

    fb0 = force_balance(psixy, Rxy, Bpxy, Btxy, hthe, pressure)
    print("Force imbalance: ", numpy.mean(numpy.abs(fb0)),
          numpy.max(numpy.abs(fb0)))

    #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    # Correct pressure using hthe

    print("Calculating pressure profile from force balance")

    try:

        # Calculate force balance
        dpdx = old_div((-Bpxy * DDX(xcoord, Bpxy * hthe) -
                        Btxy * hthe * DDX(xcoord, Btxy) -
                        (Btxy * Btxy * hthe / Rxy) * DDX(xcoord, Rxy)),
                       (MU * hthe))

        # Surface average
        dpdx2 = surface_average(dpdx, mesh)

        pres = numpy.zeros((nx, ny))
        # Integrate to get pressure
        for i in range(ny):
            pres[:, i] = int_func(psixy[:, i], dpdx2[:, i])
            pres[:, i] = pres[:, i] - pres[nx - 1, i]

        status = gen_surface(mesh=mesh)  # Start generator
        while True:
            # Get the next domain
            period, yi, xi, last = gen_surface(period=None, last=None, xi=None)

            ma = numpy.max(pres[xi, yi])

            for i in range(numpy.size(yi)):
                pres[:, yi[i]] = pres[:, yi[i]] - pres[xi, yi[i]] + ma

            if last == 1: break

        pres = pres - numpy.min(pres)

        # Some sort of smoothing here?

        fb0 = force_balance(psixy, Rxy, Bpxy, Btxy, hthe, pres)
        print("Force imbalance: ", numpy.mean(numpy.abs(fb0)),
              numpy.max(numpy.abs(fb0)))

        #!P.MULTI=[0,0,2,0,0]
        fig = figure(figsize=(7, 11))
        subplots_adjust(left=.07,
                        bottom=.07,
                        right=0.95,
                        top=0.95,
                        wspace=.3,
                        hspace=.25)

        SURFACE(pressure, fig, xtitle="X", ytitle="Y", var='Pa', sub=[2, 1, 1])
        title("Input pressure")
        SURFACE(pres, fig, xtitle="X", ytitle="Y", var='Pa', sub=[2, 1, 2])
        title("New pressure")
        #  arrange the plot on the screen
        #      mngr = get_current_fig_manager()
        #      geom = mngr.window.geometry()
        #      x,y,dx,dy = geom.getRect()
        #      mngr.window.setGeometry(0, 0, dx, dy)
        #
        show(block=False)

        calcp = settings.calcp

        if calcp == -1:
            calcp = query_yes_no(
                "Keep new pressure?")  #, gui=gui, dialog_parent=parent)
        else:
            time.sleep(2)
        if calcp == 1:
            pressure = pres
            dpdpsi = dpdx2

    except Exception:
        print("WARNING: Pressure profile calculation failed: "
              )  #, !ERROR_STATE.MSG
        pass

    #CATCH, /cancel

#;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
# Correct f = RBt using force balance

    calcbt = settings.calcbt
    if calcbt == -1:
        calcbt = query_yes_no("Correct f=RBt using force balance?"
                              )  #, gui=gui, dialog_parent=parent)
    if calcbt == 1:

        new_Btxy = newton_Bt(psixy, Rxy, Btxy, Bpxy, pres, hthe, mesh)

        fb0 = force_balance(psixy, Rxy, Bpxy, new_Btxy, hthe, pressure)
        print("force imbalance: ", numpy.mean(numpy.abs(fb0)),
              numpy.max(numpy.abs(fb0)))

        fig = figure(figsize=(7, 11))
        subplots_adjust(left=.07,
                        bottom=.07,
                        right=0.95,
                        top=0.95,
                        wspace=.3,
                        hspace=.25)

        subplot(211)
        SURFACE(Btxy, fig, xtitle="X", ytitle="Y", var='T', sub=[2, 1, 1])
        title("Input Bt")
        subplot(212)
        SURFACE(new_Btxy, fig, xtitle="X", ytitle="Y", var='T', sub=[2, 1, 2])
        title("New Bt")
        #  arrange the plot on the screen
        #mngr = get_current_fig_manager()
        #geom = mngr.window.geometry()
        #x,y,dx,dy = geom.getRect()
        #mngr.window.setGeometry(600, 0, dx, dy)

        show(block=False)

        calcbt = settings.calcbt
        if calcbt == -1:
            calcbt = query_yes_no(
                "Keep new Bt?")  #, gui=gui, dialog_parent=parent)
        if calcbt == 1:
            Btxy = new_Btxy
            Bxy = numpy.sqrt(Btxy**2 + Bpxy**2)

#;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
# CALCULATE HTHE
# Modify hthe to fit force balance using initial guess
# Does not depend on signs
#;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

    calchthe = settings.calchthe
    if calchthe == -1:
        calchthe = query_yes_no("Adjust hthe using force balance?"
                                )  #, gui=gui, dialog_parent=parent)
    if calchthe == 1:
        # This doesn't behave well close to the x-points
        fixhthe = numpy.int(old_div(nx, 2))
        nh = correct_hthe(Rxy,
                          psixy,
                          Btxy,
                          Bpxy,
                          hthe,
                          pressure,
                          fixhthe=fixhthe)

        fb0 = force_balance(psixy, Rxy, Bpxy, Btxy, nh, pressure)
        print("Force imbalance: ", numpy.mean(numpy.abs(fb0)),
              numpy.max(numpy.abs(fb0)))

        print("numpy.maximum difference in hthe: ",
              numpy.max(numpy.abs(hthe - nh)))
        print("numpy.maximum percentage difference: ",
              100. * numpy.max(numpy.abs(old_div((hthe - nh), hthe))))

        #!P.multi=[0,0,1,0,0]
        fig = figure(figsize=(7, 4))
        title("Poloidal arc length at midplane. line is initial estimate")
        plot(hthe[:, 0], '-')
        plot(nh[:, 0], 'r-+')
        #  arrange the plot on the screen
        #mngr = get_current_fig_manager()
        #geom = mngr.window.geometry()
        #x,y,dx,dy = geom.getRect()
        #mngr.window.setGeometry(0, 1150, dx, dy)

        show(block=False)

        if query_yes_no(
                "Keep new hthe?") == 1:  #, gui=gui, dialog_parent=parent) :
            hthe = nh

    if smoothhthe != None:
        # Smooth hthe to prevent large jumps in X or Y. This
        # should be done by creating a better mesh in the first place

        # Need to smooth in Y and X otherwise smoothing in X
        # produces discontinuities in Y
        hold = hthe

        if 1:
            # Nonlinear smoothing. Tries to smooth only regions with large
            # changes in gradient

            hthe = 0.  # smooth_nl(hthe, mesh)

        else:
            # Just use smooth in both directions

            for i in range(ny):
                hthe[:, i] = SMOOTH(SMOOTH(hthe[:, i], 10), 10)

        status = gen_surface(mesh=mesh)  # Start generator
        while True:
            # Get the next domain
            period, yi, xi, last = gen_surface(period=None, last=None, xi=None)

            n = numpy.size(yi)

            if period:
                hthe[xi, yi] = (SMOOTH([
                    hthe[xi, yi[(n - 4):(n - 1)]], hthe[xi, yi], hthe[xi,
                                                                      yi[0:3]]
                ], 4))[4:(n + 3)]
            else:
                hthe[xi, yi] = SMOOTH(hthe[xi, yi], 4)

            if last == 1: break

    # Calculate field-line pitch
    pitch = hthe * Btxy / (Bpxy * Rxy)

    # derivative with psi
    dqdpsi = DDX(psixy, pitch)

    qinty, qloop = int_y(pitch,
                         mesh,
                         loop=0,
                         nosmooth='nosmooth',
                         simple='simple')
    qinty = qinty * dtheta
    qloop = qloop * dtheta

    sinty = int_y(dqdpsi, mesh, nosmooth='nosmooth', simple='simple') * dtheta

    # NOTE: This is only valid in the core
    pol_angle = numpy.zeros((nx, ny))
    for i in range(nx):
        pol_angle[i, :] = 2.0 * numpy.pi * qinty[i, :] / qloop[i]

    #;;;;;;;;;;;;;;;;;;; THETA_ZERO ;;;;;;;;;;;;;;;;;;;;;;
    # re-set zshift to be zero at the outboard midplane

    print("MIDPLANE INDEX = ", ymidplane)

    status = gen_surface(mesh=mesh)  # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=None, last=None, xi=None)

        w = numpy.size(numpy.where(yi == ymidplane))
        if w > 0:
            # Crosses the midplane
            qinty[xi, yi] = qinty[xi, yi] - qinty[xi, ymidplane]
            sinty[xi, yi] = sinty[xi, yi] - sinty[xi, ymidplane]
        else:
            # Doesn't include a point at the midplane
            qinty[xi, yi] = qinty[xi, yi] - qinty[xi, yi[0]]
            sinty[xi, yi] = sinty[xi, yi] - sinty[xi, yi[0]]

        if last == 1: break

    print("")
    print("==== Calculating curvature ====")

    #;;;;;;;;;;;;;;;;;;; CURVATURE ;;;;;;;;;;;;;;;;;;;;;;;
    # Calculating b x kappa

    if curv == None:

        print("*** Calculating curvature in toroidal coordinates")

        thetaxy = numpy.zeros((nx, ny))
        status = gen_surface(mesh=mesh)  # Start generator
        while True:
            # Get the next domain
            period, yi, xi, last = gen_surface(period=None, last=None, xi=None)
            thetaxy[xi,
                    yi] = numpy.arange(numpy.size(yi)).astype(float) * dtheta
            if last == 1: break

        bxcv = curvature(nx,
                         ny,
                         Rxy,
                         Zxy,
                         Brxy,
                         Bzxy,
                         Btxy,
                         psixy,
                         thetaxy,
                         hthe,
                         mesh=mesh)

        bxcvx = bpsign * bxcv.psi
        bxcvy = bxcv.theta
        bxcvz = bpsign * (bxcv.phi - sinty * bxcv.psi - pitch * bxcv.theta)

        # x borders
        bxcvx[0, :] = bxcvx[1, :]
        bxcvx[nx - 1, :] = bxcvx[nx - 2, :]

        bxcvy[0, :] = bxcvy[1, :]
        bxcvy[nx - 1, :] = bxcvy[nx - 2, :]

        bxcvz[0, :] = bxcvz[1, :]
        bxcvz[nx - 1, :] = bxcvz[nx - 2, :]

    elif curv == 1:
        # Calculate on R-Z mesh and then interpolate onto grid
        # ( cylindrical coordinates)

        print("*** Calculating curvature in cylindrical coordinates")

        bxcv = rz_curvature(rz_grid)

        # DCT methods cause spurious oscillations
        # Linear interpolation seems to be more robust
        bxcv_psi = numpy.interp(bxcv.psi, mesh.Rixy, mesh.Zixy)
        bxcv_theta = old_div(numpy.interp(bxcv.theta, mesh.Rixy, mesh.Zixy),
                             hthe)
        bxcv_phi = numpy.interp(bxcv.phi, mesh.Rixy, mesh.Zixy)

        # If Bp is reversed, then Grad x = - Grad psi
        bxcvx = bpsign * bxcv_psi
        bxcvy = bxcv_theta
        bxcvz = bpsign * (bxcv_phi - sinty * bxcv_psi - pitch * bxcv_theta)
    elif curv == 2:
        # Curvature from Curl(b/B)

        bxcvx = bpsign * (Bpxy * Btxy * Rxy * DDY(old_div(1., Bxy), mesh) /
                          hthe)
        bxcvy = -bpsign * Bxy * Bpxy * DDX(xcoord,
                                           Btxy * Rxy / Bxy ^ 2) / (2. * hthe)
        bxcvz = Bpxy ^ 3 * DDX(xcoord, old_div(
            hthe, Bpxy)) / (2. * hthe * Bxy) - Btxy * Rxy * DDX(
                xcoord, old_div(Btxy, Rxy)) / (2. * Bxy) - sinty * bxcvx

    else:
        # calculate in flux coordinates.

        print("*** Calculating curvature in flux coordinates")

        dpb = numpy.zeros((nx, ny))  # quantity used for y and z components

        for i in range(ny):
            dpb[:, i] = MU * dpdpsi / Bxy[:, i]

        dpb = dpb + DDX(xcoord, Bxy)

        bxcvx = bpsign * (Bpxy * Btxy * Rxy * DDY(old_div(1., Bxy), mesh) /
                          hthe)
        bxcvy = bpsign * (Bpxy * Btxy * Rxy * dpb / (hthe * Bxy ^ 2))
        bxcvz = -dpb - sinty * bxcvx

    if smoothcurv:
        # Smooth curvature to prevent large jumps

        # Nonlinear smoothing. Tries to smooth only regions with large
        # changes in gradient

        bz = bxcvz + sinty * bxcvx

        print("Smoothing bxcvx...")
        bxcvx = 0.  #smooth_nl(bxcvx, mesh)
        print("Smoothing bxcvy...")
        bxcvy = 0.  #smooth_nl(bxcvy, mesh)
        print("Smoothing bxcvz...")
        bz = 0.  #smooth_nl(bz, mesh)

        bxcvz = bz - sinty * bxcvx

#;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
# CALCULATE PARALLEL CURRENT
#
# Three ways to calculate Jpar0:
# 1. From fprime and pprime
# 2. From Curl(B) in field-aligned coords
# 3. From the curvature
#
# Provides a way to check if Btor should be reversed
#
#;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

    print("")
    print("==== Calculating parallel current ====")

    jpar0 = -Bxy * fprime / MU - Rxy * Btxy * dpdpsi / Bxy

    # Set to zero in PF and SOL
    status = gen_surface(mesh=mesh)  # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=None, last=None, xi=None)

        if period == None: jpar0[xi, yi] = 0.0
        if last == 1: break

# Curl(B) expression for Jpar0 (very noisy usually)
    j0 = (bpsign * ((Bpxy * Btxy * Rxy / (Bxy * hthe)) *
                    (DDX(xcoord, Bxy**2 * hthe / Bpxy) - bpsign * Btxy * Rxy *
                     DDX(xcoord, Btxy * hthe /
                         (Rxy * Bpxy))) - Bxy * DDX(xcoord, Btxy * Rxy)) / MU)

    # Create a temporary mesh structure to send to adjust_jpar
    tmp_mesh = Bunch(mesh,
                     bxcvx=bxcvx,
                     bxcvy=bxcvy,
                     bxcvz=bxcvz,
                     Bpxy=Bpxy,
                     Btxy=Btxy,
                     Bxy=Bxy,
                     dx=dx,
                     dy=dy,
                     hthe=hthe,
                     jpar0=jpar0,
                     pressure=pressure)
    tmp_mesh.psixy = psixy

    jpar = adjust_jpar(tmp_mesh, noplot='noplot')

    #!P.multi=[0,2,2,0,0]

    fig = figure(figsize=(15, 11))
    subplots_adjust(left=.07,
                    bottom=.07,
                    right=0.95,
                    top=0.95,
                    wspace=.3,
                    hspace=.25)

    subplot(221)
    SURFACE(jpar0, fig, xtitle="X", ytitle="Y", var='A', sub=[2, 2, 1])
    title("Jpar from F' and P'")

    subplot(222)
    SURFACE(jpar, fig, xtitle="X", ytitle="Y", var='A', sub=[2, 2, 2])
    title("Jpar from curvature")

    subplot(223)
    plot(jpar0[0, :], '-', jpar[0, :], '+')
    ylim([
        numpy.min([jpar0[0, :], jpar[0, :]]),
        numpy.max([jpar0[0, :], jpar[0, :]])
    ])
    title("jpar at x=0. Solid from f' and p'")

    subplot(224)
    plot(jpar0[:, ymidplane], '-', jpar[:, ymidplane], '+')
    ylim([
        numpy.min([jpar0[:, ymidplane], jpar[:, ymidplane]]),
        numpy.max([jpar0[:, ymidplane], jpar[:, ymidplane]])
    ])

    title("Jpar at y=" + numpy.str(ymidplane) + " Solid from f' and p'")

    #  arrange the plot on the screen
    #mngr = get_current_fig_manager()
    #geom = mngr.window.geometry()
    #x,y,dx,dy = geom.getRect()
    #mngr.window.setGeometry(1350, 0, dx, dy)

    show(block=False)

    # !P.multi=0

    calcjpar = settings.calcjpar
    if calcjpar == -1:
        calcjpar = query_yes_no(
            "Use Jpar from curvature?")  #, gui=gui, dialog_parent=parent)
    if calcjpar == True:
        jpar0 = jpar

    if 0:

        # Try smoothing jpar0 in psi, preserving zero points and maxima
        jps = jpar0
        for y in range(ny):
            j = jpar0[:, y]
            js = j
            ma = numpy.max(numpy.abs(j))
            ip = numpy.argmax(numpy.abs(j))
            if (ma < 1.e-4) or (ip == 0):
                jps[:, y] = j

            level = 1.
            #i0 = MAX(WHERE(ABS(j[0:ip]) LT level))
            i1 = numpy.min(numpy.where(numpy.abs(j[ip::]) < level))

            #IF i0 LE 0 THEN i0 = 1
            i0 = 1

            if i1 == -1:
                i1 = nx - 2
            else:
                i1 = i1 + ip

            if (ip <= i0) or (ip >= i1):

                # Now preserve starting and end points, and peak value
                div = numpy.int(old_div(
                    (i1 - i0),
                    10)) + 1  # reduce number of points by this factor

                inds = [i0]  # first point
                for i in [i0 + div, ip - div, div]:
                    inds = [inds, i]
                inds = [inds, ip]  # Put in the peak point

                # Calculate spline interpolation of inner part

                js[0:ip] = spline_mono(inds,
                                       j[inds],
                                       numpy.arange(ip + 1),
                                       yp0=(j[i0] - j[i0 - 1]),
                                       ypn_1=0.0)

                inds = [ip]  # peak point
                for i in [ip + div, i1 - div, div]:
                    inds = [inds, i]

                inds = [inds, i1]  # Last point
                js[ip:i1] = spline_mono(inds,
                                        j[inds],
                                        ip + numpy.arange(i1 - ip + 1),
                                        yp0=0.0,
                                        ypn_1=(j[i1 + 1] - j[i1]))

                jps[:, y] = js

#;;;;;;;;;;;;;;;;;;; TOPOLOGY ;;;;;;;;;;;;;;;;;;;;;;;
# Calculate indices for backwards-compatibility

    nr = numpy.size(mesh.nrad)
    np = numpy.size(mesh.npol)
    if (nr == 2) and (np == 3):
        print("Single null equilibrium")

        ixseps1 = mesh.nrad[0]
        ixseps2 = nx

        jyseps1_1 = mesh.npol[0] - 1
        jyseps1_2 = mesh.npol[0] + numpy.int(old_div(mesh.npol[1], 2))
        ny_inner = jyseps1_2
        jyseps2_1 = jyseps1_2
        jyseps2_2 = ny - mesh.npol[2] - 1

    elif (nr == 3) and (np == 6):
        print("Double null equilibrium")

        ixseps1 = mesh.nrad[0]
        ixseps2 = ixseps1 + mesh.nrad[1]

        jyseps1_1 = mesh.npol[0] - 1
        jyseps2_1 = jyseps1_1 + mesh.npol[1]

        ny_inner = jyseps2_1 + mesh.npol[2] + 1

        jyseps1_2 = ny_inner + mesh.npol[3] - 1
        jyseps2_2 = jyseps1_2 + mesh.npol[4]

    elif (nr == 1) and (np == 1):

        print("Single domain")

        ixseps1 = nx
        ixseps2 = nx

        jyseps1_1 = -1
        jyseps1_2 = numpy.int(old_div(ny, 2))
        jyseps2_1 = numpy.int(old_div(ny, 2))
        ny_inner = numpy.int(old_div(ny, 2))
        jyseps2_2 = ny - 1

    else:
        print("***************************************")
        print("* WARNING: Equilibrium not recognised *")
        print("*                                     *")
        print("*  Check mesh carefully!              *")
        print("*                                     *")
        print("*  Contact Ben Dudson                 *")
        print("*      [email protected]     *")
        print("***************************************")
        ixseps1 = -1
        ixseps2 = -1

        jyseps1_1 = -1
        jyseps1_2 = numpy.int(old_div(ny, 2))
        jyseps2_1 = numpy.int(old_div(ny, 2))
        ny_inner = numpy.int(old_div(ny, 2))
        jyseps2_2 = ny - 1

    print("Generating plasma profiles:")

    print("  1. Flat temperature profile")
    print("  2. Flat density profile")
    print("  3. Te proportional to density")
    while True:
        opt = input("Profile option:")
        if eval(opt) >= 1 and eval(opt) <= 3: break

    if eval(opt) == 1:
        # flat temperature profile

        print("Setting flat temperature profile")
        while True:
            Te_x = eval(input("Temperature (eV):"))

            # get density
            Ni = old_div(pressure, (2. * Te_x * 1.602e-19 * 1.0e20))

            print("numpy.maximum density (10^20 m^-3):", numpy.max(Ni))

            done = query_yes_no("Is this ok?")
            if done == 1: break

        Te = numpy.zeros((nx, ny)) + Te_x
        Ti = Te
        Ni_x = numpy.max(Ni)
        Ti_x = Te_x
    elif eval(opt) == 2:
        print("Setting flat density profile")

        while True:
            Ni_x = eval(input("Density [10^20 m^-3]:"))

            # get temperature
            Te = old_div(pressure, (2. * Ni_x * 1.602e-19 * 1.0e20))

            print("numpy.maximum temperature (eV):", numpy.max(Te))
            if query_yes_no("Is this ok?") == 1: break

        Ti = Te
        Ni = numpy.zeros((nx, ny)) + Ni_x
        Te_x = numpy.max(Te)
        Ti_x = Te_x
    else:
        print("Setting te proportional to density")

        while True:
            Te_x = eval(input("Maximum temperature [eV]:"))

            Ni_x = old_div(numpy.max(pressure),
                           (2. * Te_x * 1.602e-19 * 1.0e20))

            print("Maximum density [10^20 m^-3]:", Ni_x)

            Te = Te_x * pressure / numpy.max(pressure)
            Ni = Ni_x * pressure / numpy.max(pressure)
            if query_yes_no("Is this ok?") == 1: break
        Ti = Te
        Ti_x = Te_x

    rmag = numpy.max(numpy.abs(Rxy))
    print("Setting rmag = ", rmag)

    bmag = numpy.max(numpy.abs(Bxy))
    print("Setting bmag = ", bmag)

    #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    # save to file
    # open a new netCDF file for writing.
    handle = file_open(output)

    print("Writing grid to file " + output)

    # Size of the grid

    s = file_write(handle, "nx", nx)
    s = file_write(handle, "ny", ny)

    # Topology for original scheme
    s = file_write(handle, "ixseps1", ixseps1)
    s = file_write(handle, "ixseps2", ixseps2)
    s = file_write(handle, "jyseps1_1", jyseps1_1)
    s = file_write(handle, "jyseps1_2", jyseps1_2)
    s = file_write(handle, "jyseps2_1", jyseps2_1)
    s = file_write(handle, "jyseps2_2", jyseps2_2)
    s = file_write(handle, "ny_inner", ny_inner)

    # Grid spacing

    s = file_write(handle, "dx", dx)
    s = file_write(handle, "dy", dy)

    s = file_write(handle, "ShiftAngle", qloop)
    s = file_write(handle, "zShift", qinty)
    s = file_write(handle, "pol_angle", pol_angle)
    s = file_write(handle, "ShiftTorsion", dqdpsi)

    s = file_write(handle, "Rxy", Rxy)
    s = file_write(handle, "Zxy", Zxy)
    s = file_write(handle, "Bpxy", Bpxy)
    s = file_write(handle, "Btxy", Btxy)
    s = file_write(handle, "Bxy", Bxy)
    s = file_write(handle, "hthe", hthe)
    s = file_write(handle, "sinty", sinty)
    s = file_write(handle, "psixy", psixy)

    # Topology for general configurations
    s = file_write(handle, "yup_xsplit", mesh.yup_xsplit)
    s = file_write(handle, "ydown_xsplit", mesh.ydown_xsplit)
    s = file_write(handle, "yup_xin", mesh.yup_xin)
    s = file_write(handle, "yup_xout", mesh.yup_xout)
    s = file_write(handle, "ydown_xin", mesh.ydown_xin)
    s = file_write(handle, "ydown_xout", mesh.ydown_xout)
    s = file_write(handle, "nrad", mesh.nrad)
    s = file_write(handle, "npol", mesh.npol)

    # plasma profiles

    s = file_write(handle, "pressure", pressure)
    s = file_write(handle, "Jpar0", jpar0)
    s = file_write(handle, "Ni0", Ni)
    s = file_write(handle, "Te0", Te)
    s = file_write(handle, "Ti0", Ti)

    s = file_write(handle, "Ni_x", Ni_x)
    s = file_write(handle, "Te_x", Te_x)
    s = file_write(handle, "Ti_x", Ti_x)
    s = file_write(handle, "bmag", bmag)
    s = file_write(handle, "rmag", rmag)

    # Curvature
    s = file_write(handle, "bxcvx", bxcvx)
    s = file_write(handle, "bxcvy", bxcvy)
    s = file_write(handle, "bxcvz", bxcvz)

    # Psi range
    s = file_write(handle, "psi_axis", mesh.faxis)
    psi_bndry = mesh.faxis + mesh.fnorm
    s = file_write(handle, "psi_bndry", psi_bndry)

    file_close, handle
    print("DONE")
Ejemplo n.º 24
0
def volume_integral(var, g, xr=False):
    """Integrate a variable over a volume

    Parameters
    ----------
    var : array_like
        Variable to integrate
    g : dict
        A dictionary of various grid quantities
    xr : (int, int), optional
        Range of x indices (default: all of x)

    Returns
    -------
    float
        Volumne integral of variable

    """

    s = np.ndim(var)

    grid=bunchify(g)


    if s == 4 :
        # 4D [t,x,y,z] - integrate for each t
        nx = np.shape(var)[1]
        ny = np.shape(var)[2]
        nt = np.shape(var)[0]

        result = np.zeros(nt)
        for t in range(nt) :
            result[t] = volume_integral(var[t,:,:,:],g,xr=xr)
        return result

    elif s == 3 :
        # 3D [x,y,z] - average in Z
        nx = np.shape(var)[0]
        ny = np.shape(var)[1]
 #       nz = np.shape(var)[2]

        zi = np.zeros((nx, ny))
        for x in range(nx):
            for y in range(ny):
                zi[x,y] = np.mean(var[x,y,:])

        return volume_integral(zi, g, xr=xr)


    elif s != 2 :
        print("ERROR: volume_integral var must be 2, 3 or 4D")


    # 2D [x,y]
    nx = np.shape(var)[0]
    ny = np.shape(var)[1]

    if xr == False : xr=[0,nx-1]

    result = 0.0

    #status = gen_surface(mesh=grid) ; Start generator
    xi = -1
    yi = np.arange(0,ny,dtype=int)
    last = 0
  #  iy = np.zeros(nx)
    while True:
        #yi = gen_surface(last=last, xi=xi, period=periodic)
        xi = xi + 1
        if xi == nx-1 : last = 1

        if (xi >= np.min(xr)) & (xi <= np.max(xr)) :
            dtheta = 2.*np.pi / np.float(ny)
            r = grid.Rxy[xi,yi]
            z = grid.Zxy[xi,yi]
            n = np.size(r)
            dl = old_div(np.sqrt( deriv(r)**2 + deriv(z)**2 ), dtheta)

      # Area of flux-surface
            dA = (grid.Bxy[xi,yi]/grid.Bpxy[xi,yi]*dl) * (r*2.*np.pi)
      # Volume
            if xi == nx-1 :
                dpsi = (grid.psixy[xi,yi] - grid.psixy[xi-1,yi])
            else:
                dpsi = (grid.psixy[xi+1,yi] - grid.psixy[xi,yi])

            dV = dA * dpsi / (r*(grid.Bpxy[xi,yi])) # May need factor of 2pi
            dV = np.abs(dV)

            result = result + np.sum(var[xi,yi] * dV)

        if last==1 : break

    return result
Ejemplo n.º 25
0
period = 15

gfile = './cbm18_dens8.grid_nx68ny64.nc'

g = file_import(gfile)

Dphi0 = collect("Dphi0", path=path0)
phi0 = collect("phi0", path=path1)  # needs diamagnetic effects
#
psixy = g.get('psixy')
PSI_AXIS = g.get('psi_axis')
PSI_BNDRY = g.get('psi_bndry')
#
psix = old_div((psixy[:, 32] - PSI_AXIS), (PSI_BNDRY - PSI_AXIS))
Epsi = -deriv(phi0[:, 32], psix)
#
#
fig = figure()
plot(psix, -Dphi0[:, 32], 'r', linewidth=5)
plot(psix, Epsi, 'k', linewidth=5)
annotate('w/o flow',
         xy=(.3, .7),
         xycoords='axes fraction',
         horizontalalignment='center',
         verticalalignment='center',
         size=30)
annotate('w/ flow',
         xy=(.7, .4),
         xycoords='axes fraction',
         horizontalalignment='center',
Ejemplo n.º 26
0
def curvature(nx,
              ny,
              Rxy,
              Zxy,
              BRxy,
              BZxy,
              BPHIxy,
              PSIxy,
              THETAxy,
              hthexy,
              CURLB=None,
              JXB=None,
              CURVEC=None,
              BXCURVEC=None,
              BXCV=None,
              DEBUG=None,
              mesh=None):
    #;
    #; Calculate the magnetic field curvature and other related quantities
    #;--------------------------------------------------------------------

    print('Calculating curvature-related quantities...')

    #;;-vector quantities are stored as 2D arrays of structures {r,phi,z}
    vec = Bunch(r=0., phi=0., z=0.)
    curlb = numpy.tile(vec, (nx, ny))
    jxb = numpy.tile(vec, (nx, ny))
    curvec = numpy.tile(vec, (nx, ny))
    bxcurvec = numpy.tile(vec, (nx, ny))

    bxcv = Bunch()
    bxcv.psi = numpy.zeros((nx, ny))
    bxcv.theta = numpy.zeros((nx, ny))
    bxcv.phi = numpy.zeros((nx, ny))

    status = gen_surface(mesh=mesh)  # Start generator

    while True:
        period, yi, xi, last = gen_surface(period=None, last=None, xi=None)
        nys = numpy.size(yi)
        x = xi

        # Get vector along the surface
        if period == 1:
            dr = fft_deriv(Rxy[x, yi])
            dz = fft_deriv(Zxy[x, yi])
        else:
            dr = deriv(Rxy[x, yi])
            dz = deriv(Zxy[x, yi])

        dl = numpy.sqrt(dr**2 + dz**2)

        dr = old_div(dr, dl)
        dz = old_div(dz, dl)

        for j in range(nys):
            y = yi[j]

            if period:
                yp = yi[(j + 1) % nys]
                ym = yi[(j - 1 + nys) % nys]
            else:
                yp = yi[numpy.min([j + 1, nys - 1])]
                ym = yi[numpy.max([j - 1, 0])]

            grad_Br = pdiff_rz(Rxy, Zxy, BRxy, x, y, yp, ym)
            grad_Bz = pdiff_rz(Rxy, Zxy, BZxy, x, y, yp, ym)
            grad_Bphi = pdiff_rz(Rxy, Zxy, BPHIxy, x, y, yp, ym)

            grad_Psi = pdiff_rz(Rxy, Zxy, PSIxy, x, y, yp, ym)

            #grad_Theta = pdiff_rz(Rxy, Zxy, THETAxy, x, y, yp, ym)
            grad_Theta = Bunch(r=old_div(dr[j], hthexy[x, y]),
                               z=old_div(dz[j], hthexy[x, y]),
                               phi=0.0)

            grad_Phi = Bunch(r=0.0, z=0.0, phi=old_div(
                1., Rxy[x, y]))  #-gradient of the toroidal angle

            vecR = Bunch(r=Rxy[x, y], z=Zxy[x, y])
            vecB = Bunch(r=BRxy[x, y], z=BZxy[x, y], phi=BPHIxy[x, y])

            curlb[x, y] = curlcyl(vecR, vecB, grad_Br, grad_Bphi, grad_Bz)

            jxb[x, y] = xprod(curlb[x, y], vecB)

            #-magnitude of B at 5 locations in cell
            bstrength = numpy.sqrt(BRxy**2 + BZxy**2 + BPHIxy**2)

            #-unit B vector at cell center
            vecB_unit = Bunch(r=old_div(BRxy[x, y], bstrength[x, y]),
                              z=old_div(BZxy[x, y], bstrength[x, y]),
                              phi=old_div(BPHIxy[x, y], bstrength[x, y]))

            #-components of gradient of unit B vector at 5 locations in cell
            grad_Br_unit = pdiff_rz(Rxy, Zxy, old_div(BRxy, bstrength), x, y,
                                    yp, ym)

            grad_Bz_unit = pdiff_rz(Rxy, Zxy, old_div(BZxy, bstrength), x, y,
                                    yp, ym)

            grad_Bphi_unit = pdiff_rz(Rxy, Zxy, old_div(BPHIxy, bstrength), x,
                                      y, yp, ym)

            #-curl of unit B vector at cell center
            curlb_unit = curlcyl(vecR, vecB_unit, grad_Br_unit, grad_Bphi_unit,
                                 grad_Bz_unit)

            #-curvature vector at cell center
            curvec[x, y] = xprod(vecB_unit, curlb_unit, minus='MINUS')

            #-unit b cross curvature vector at cell center
            bxcurvec[x, y] = xprod(vecB_unit, curvec[x, y])

            #-calculate bxcurvec dotted with grad_psi, grad_theta, and grad_phi
            bxcv.psi[x, y] = dotprod(bxcurvec[x, y], grad_Psi)
            bxcv.theta[x, y] = numpy.real(dotprod(bxcurvec[x, y], grad_Theta))
            bxcv.phi[x, y] = dotprod(bxcurvec[x, y], grad_Phi)

        if last == 1: break

#   if DEBUG : sys.exit()

    print('...done')

    return bxcv
Ejemplo n.º 27
0
def curvature( nx, ny, Rxy, Zxy, BRxy, BZxy, BPHIxy, PSIxy, THETAxy, hthexy,
               CURLB=None, JXB=None, CURVEC=None, BXCURVEC=None, BXCV=None,
               DEBUG=None, mesh=None):
#;
#; Calculate the magnetic field curvature and other related quantities
#;--------------------------------------------------------------------

    print('Calculating curvature-related quantities...')

#;;-vector quantities are stored as 2D arrays of structures {r,phi,z}
    vec=Bunch( r=0.,phi=0.,z=0.)
    curlb=numpy.tile(vec,(nx,ny))
    jxb=numpy.tile(vec,(nx,ny))
    curvec=numpy.tile(vec,(nx,ny))
    bxcurvec=numpy.tile(vec,(nx,ny))

    bxcv=Bunch()
    bxcv.psi=numpy.zeros((nx,ny))
    bxcv.theta=numpy.zeros((nx,ny))
    bxcv.phi=numpy.zeros((nx,ny))


    status = gen_surface(mesh=mesh) # Start generator


    while True:
        period, yi, xi, last = gen_surface(period=None, last=None, xi=None)
        nys = numpy.size(yi)
        x=xi


     # Get vector along the surface
        if period ==1 :
            dr = fft_deriv(Rxy[x,yi])
            dz = fft_deriv(Zxy[x,yi])
        else:
            dr = deriv(Rxy[x,yi])
            dz = deriv(Zxy[x,yi])

        dl = numpy.sqrt(dr**2 + dz**2)

        dr = old_div(dr, dl)
        dz = old_div(dz, dl)



        for j in range (nys) :
            y = yi[j]

            if period :
                yp = yi[ (j+1)     % nys ]
                ym = yi[ (j-1+nys) % nys ]
            else:
                yp = yi[ numpy.min([j+1 , nys-1]) ]
                ym = yi[ numpy.max([j-1 , 0]) ]


            grad_Br   = pdiff_rz(Rxy, Zxy, BRxy, x, y, yp, ym)
            grad_Bz   = pdiff_rz(Rxy, Zxy, BZxy, x, y, yp, ym)
            grad_Bphi = pdiff_rz(Rxy, Zxy, BPHIxy, x, y, yp, ym)



            grad_Psi  = pdiff_rz(Rxy, Zxy, PSIxy, x, y, yp, ym)


            #grad_Theta = pdiff_rz(Rxy, Zxy, THETAxy, x, y, yp, ym)
            grad_Theta = Bunch( r=old_div(dr[j],hthexy[x,y]), z=old_div(dz[j],hthexy[x,y]), phi=0.0 )

            grad_Phi=Bunch( r=0.0,z=0.0,phi=old_div(1.,Rxy[x,y]) ) #-gradient of the toroidal angle

            vecR=Bunch(  r=Rxy[x,y],z=Zxy[x,y] )
            vecB=Bunch( r=BRxy[x,y],z=BZxy[x,y],phi=BPHIxy[x,y] )


            curlb[x,y]=curlcyl(vecR, vecB, grad_Br, grad_Bphi, grad_Bz)


            jxb[x,y]=xprod(curlb[x,y], vecB)


            #-magnitude of B at 5 locations in cell
            bstrength = numpy.sqrt(BRxy**2 + BZxy**2 + BPHIxy**2)

            #-unit B vector at cell center
            vecB_unit=Bunch( r=old_div(BRxy[x,y],bstrength[x,y]),
                  z=old_div(BZxy[x,y],bstrength[x,y]),
                  phi=old_div(BPHIxy[x,y],bstrength[x,y]) )

            #-components of gradient of unit B vector at 5 locations in cell
            grad_Br_unit = pdiff_rz(Rxy, Zxy, old_div(BRxy,bstrength), x, y, yp, ym)

            grad_Bz_unit = pdiff_rz(Rxy, Zxy, old_div(BZxy,bstrength), x, y, yp, ym)

            grad_Bphi_unit = pdiff_rz(Rxy, Zxy, old_div(BPHIxy,bstrength), x, y, yp, ym)

            #-curl of unit B vector at cell center
            curlb_unit=curlcyl(vecR, vecB_unit, grad_Br_unit, grad_Bphi_unit, grad_Bz_unit)

            #-curvature vector at cell center
            curvec[x,y]=xprod(vecB_unit,curlb_unit,minus='MINUS')

            #-unit b cross curvature vector at cell center
            bxcurvec[x,y]=xprod(vecB_unit,curvec[x,y])

            #-calculate bxcurvec dotted with grad_psi, grad_theta, and grad_phi
            bxcv.psi[x,y]=dotprod(bxcurvec[x,y],grad_Psi)
            bxcv.theta[x,y]=numpy.real(dotprod(bxcurvec[x,y],grad_Theta))
            bxcv.phi[x,y]=dotprod(bxcurvec[x,y],grad_Phi)


        if last==1 : break

#   if DEBUG : sys.exit()

    print('...done')

    return bxcv
Ejemplo n.º 28
0
def follow_gradient(interp_data,
                    R,
                    Z,
                    ri0,
                    zi0,
                    ftarget,
                    ri,
                    zi,
                    status=0,
                    boundary=None,
                    fbndry=None,
                    ibndry=None):

    global rd_com, idata, lastgoodf, lastgoodpos, Rpos, Zpos, ood, tol, Ri, Zi, dR, dZ

    tol = 0.1

    Rpos = R
    Zpos = Z

    Ri = numpy.arange(Rpos.size).astype(float)
    Zi = numpy.arange(Zpos.size).astype(float)
    dR = deriv(Rpos)
    dZ = deriv(Zpos)

    ibndry = -1

    idata = interp_data

    if boundary != None:
        bndry = boundary
        ri0c = ri0
        zi0c = zi0
    else:
        bndry = 0

    ood = 0

    if ftarget == None: print(ftarget)

    # Get starting f
    out = local_gradient(interp_data,
                         ri0,
                         zi0,
                         status=status,
                         f=0.,
                         dfdr=None,
                         dfdz=None)
    status = out.status
    f0 = out.f
    if status == 1:
        ri = ri0
        zi = zi0
        status = 1
        return Bunch(status=status, ri=ri, zi=zi)

    fmax = ftarget  # Target (with maybe boundary in the way)

    # Call LSODE to follow gradient
    rzold = [ri0, zi0]
    rcount = 0

    solver = lsode(radial_differential, f0, rzold)
    rznew = solver.integrate(ftarget)
    nsteps = solver.steps

    lstat = 0
    #    print 'nsteps=',nsteps
    #print rzold, rznew
    #
    #sys.exit()
    #
    #   #         if nsteps > 100 : lstat = -1
    #            if lstat == -1 :
    #                print "  -> Excessive work "+str(f0)+" to "+str(ftarget)+" Trying to continue..."
    #                lstat = 2 # continue
    #                rcount = rcount + 1
    #                if rcount > 3 :
    #                    print "   -> Too many repeats. Giving Up."
    #
    #                    ri = lastgoodpos[0]
    #                    zi = lastgoodpos[1]
    #                    fmax = lastgoodf
    #
    #                    return Bunch(status=status,ri=ri,zi=zi)
    #
    #            # Get f at this new location
    #                out=local_gradient( interp_data, rznew[0], rznew[1], status=status, f=f0, dfdr=None, dfdz=None)
    #                status=out.status
    #                f0=out.f
    #
    #                if status == 1 :
    #                    ri = ri0
    #                    zi = zi0
    #                    status = 1
    #                    return Bunch(status=status, rinext=ri, zinext=zi)
    #
    #                rzold = rznew
    #
    #
    #            else :
    #                print "I break"
    #                break
    #
    #        print 'am I here?'
    #
    #        if status==0:
    #            print 'I break from try?'
    #            break
    #
    #        if lstat < 0 :
    #            print "Error in LSODE routine when following psi gradient."
    #            print "LSODE status: ", lstat
    #           #STOP
    #
    #
    #    except Exception as theError:
    #        print theError

    ri = rznew[0]
    zi = rznew[1]

    #   else:
    #   # An error occurred in LSODE.
    #   # lastgoodf contains the last known good f value
    #   #PRINT, "CAUGHT ERROR "+!ERROR_STATE.MSG
    #   #CATCH, /cancel
    #       ri = lastgoodpos[0]
    #       zi = lastgoodpos[1]
    #       fmax = lastgoodf
    #       if ood :
    #       # Gone Out Of Domain
    #           status = 2
    #           fbndry = lastgoodf
    #       #PRINT, "Out of domain at f = ", fbndry
    #       # Repeat to verify that this does work
    #           rzold = [ri0, zi0]
    #           try :
    #               fbndry = lastgoodf - 0.1*(ftarget - f0)
    #               if theError != 0 :
    #                   print "   Error again at ", fbndry
    #
    #
    #               solver=lsode(radial_differential, f0, rzold)
    #               rznew=solver.integrate(fbndry - f0)
    #           except Exception as theError:
    #               print theError
    #
    #           return Bunch(status=status, rinext=ri, zinext=zi)
    #
    #   # Otherwise just crossed a boundary
    #
    #   #CATCH, /cancel

    #if boundary != None:
    ## Check if the line crossed a boundary
    ##PRINT, "Checking boundary ", boundary[*,1:2], [ri0, ri], [zi0, zi]
    #    cpos, ncross, inds2 = line_crossings([ri0, ri], [zi0, zi], 0,
    #                      boundary[0,:], boundary[1,:], 1, ncross=0, inds2=0)
    #    if (ncross % 2) == 1 : # Odd number of boundary crossings
    #        if numpy.sqrt( (ri - cpos[0,0])**2 + (zi - cpos[1,0])**2 ) > 0.1 :
    #    #PRINT, "FINDING BOUNDARY", SQRT( (ri - cpos[0,0])^2 + (zi - cpos[1,0])^2 )
    #    # Use divide-and-conquer to find crossing point
    #
    #            tol = 1e-4 # Make the boundary crossing stricter
    #
    #            ibndry = inds2[0] # Index in boundary where hit
    #
    #            fcur = f0 # Current known good position
    #            rzold = [ri0,zi0]
    #            rzcur = rzold
    #            while True:
    #                fbndry = (fcur + fmax) / 2
    #     # Try to go half-way to fmax
    #                #CATCH, theError
    #                if theError != 0 :
    #        # Crossed boundary. Change fmax
    #                    #CATCH, /cancel
    #                    fmax = fbndry
    #                    ibndry = inds2[0] # refined boundary index
    #                else:
    #                    solver=lsode(radial_differential, f0, rzold)
    #                    rznew=solver.integrate(fbndry - f0)
    #
    #       # Didn't cross. Make this the new current location
    #                    fcur = fbndry
    #                    rzcur = rznew
    #
    #                    #CATCH, /cancel
    #
    #                if numpy.abs(fmax - fcur) < 0.01*numpy.abs(ftarget - f0):
    #                    break
    #            ri = rzcur[0]
    #            zi = rzcur[1]
    #            fbndry = fcur
    #
    #    #PRINT, "Hit boundary", ri, zi, " f =", fbndry
    #            status = 2
    #            return Bunch(status=status, rinext=ri, zinext=zi, fbndry=fbndry, ibndry=ibndry)

    #print "follow_gradient"
    #print ri, zi
    return Bunch(status=0, rinext=ri, zinext=zi, fbndry=fbndry, ibndry=ibndry)
Ejemplo n.º 29
0
def follow_gradient(interp_data, R, Z, ri0, zi0, ftarget, ri, zi, status=0, boundary=None, fbndry=None, ibndry=None):

    global rd_com, idata, lastgoodf, lastgoodpos, Rpos, Zpos, ood, tol, Ri, Zi, dR, dZ

    tol = 0.1

    Rpos = R
    Zpos = Z

    Ri = numpy.arange(Rpos.size).astype(float)
    Zi = numpy.arange(Zpos.size).astype(float)
    dR = deriv(Rpos)
    dZ = deriv(Zpos)

    ibndry = -1

    idata = interp_data

    if boundary != None:
        bndry = boundary
        ri0c = ri0
        zi0c = zi0
    else:
        bndry = 0

    ood = 0

    if ftarget == None:
        print(ftarget)

    # Get starting f
    out = local_gradient(interp_data, ri0, zi0, status=status, f=0.0, dfdr=None, dfdz=None)
    status = out.status
    f0 = out.f
    if status == 1:
        ri = ri0
        zi = zi0
        status = 1
        return Bunch(status=status, ri=ri, zi=zi)

    fmax = ftarget  # Target (with maybe boundary in the way)

    # Call LSODE to follow gradient
    rzold = [ri0, zi0]
    rcount = 0

    solver = lsode(radial_differential, f0, rzold)
    rznew = solver.integrate(ftarget)
    nsteps = solver.steps

    lstat = 0
    #    print 'nsteps=',nsteps
    # print rzold, rznew
    #
    # sys.exit()
    #
    #   #         if nsteps > 100 : lstat = -1
    #            if lstat == -1 :
    #                print "  -> Excessive work "+str(f0)+" to "+str(ftarget)+" Trying to continue..."
    #                lstat = 2 # continue
    #                rcount = rcount + 1
    #                if rcount > 3 :
    #                    print "   -> Too many repeats. Giving Up."
    #
    #                    ri = lastgoodpos[0]
    #                    zi = lastgoodpos[1]
    #                    fmax = lastgoodf
    #
    #                    return Bunch(status=status,ri=ri,zi=zi)
    #
    #            # Get f at this new location
    #                out=local_gradient( interp_data, rznew[0], rznew[1], status=status, f=f0, dfdr=None, dfdz=None)
    #                status=out.status
    #                f0=out.f
    #
    #                if status == 1 :
    #                    ri = ri0
    #                    zi = zi0
    #                    status = 1
    #                    return Bunch(status=status, rinext=ri, zinext=zi)
    #
    #                rzold = rznew
    #
    #
    #            else :
    #                print "I break"
    #                break
    #
    #        print 'am I here?'
    #
    #        if status==0:
    #            print 'I break from try?'
    #            break
    #
    #        if lstat < 0 :
    #            print "Error in LSODE routine when following psi gradient."
    #            print "LSODE status: ", lstat
    #           #STOP
    #
    #
    #    except Exception as theError:
    #        print theError

    ri = rznew[0]
    zi = rznew[1]

    #   else:
    #   # An error occurred in LSODE.
    #   # lastgoodf contains the last known good f value
    #   #PRINT, "CAUGHT ERROR "+!ERROR_STATE.MSG
    #   #CATCH, /cancel
    #       ri = lastgoodpos[0]
    #       zi = lastgoodpos[1]
    #       fmax = lastgoodf
    #       if ood :
    #       # Gone Out Of Domain
    #           status = 2
    #           fbndry = lastgoodf
    #       #PRINT, "Out of domain at f = ", fbndry
    #       # Repeat to verify that this does work
    #           rzold = [ri0, zi0]
    #           try :
    #               fbndry = lastgoodf - 0.1*(ftarget - f0)
    #               if theError != 0 :
    #                   print "   Error again at ", fbndry
    #
    #
    #               solver=lsode(radial_differential, f0, rzold)
    #               rznew=solver.integrate(fbndry - f0)
    #           except Exception as theError:
    #               print theError
    #
    #           return Bunch(status=status, rinext=ri, zinext=zi)
    #
    #   # Otherwise just crossed a boundary
    #
    #   #CATCH, /cancel

    # if boundary != None:
    ## Check if the line crossed a boundary
    ##PRINT, "Checking boundary ", boundary[*,1:2], [ri0, ri], [zi0, zi]
    #    cpos, ncross, inds2 = line_crossings([ri0, ri], [zi0, zi], 0,
    #                      boundary[0,:], boundary[1,:], 1, ncross=0, inds2=0)
    #    if (ncross % 2) == 1 : # Odd number of boundary crossings
    #        if numpy.sqrt( (ri - cpos[0,0])**2 + (zi - cpos[1,0])**2 ) > 0.1 :
    #    #PRINT, "FINDING BOUNDARY", SQRT( (ri - cpos[0,0])^2 + (zi - cpos[1,0])^2 )
    #    # Use divide-and-conquer to find crossing point
    #
    #            tol = 1e-4 # Make the boundary crossing stricter
    #
    #            ibndry = inds2[0] # Index in boundary where hit
    #
    #            fcur = f0 # Current known good position
    #            rzold = [ri0,zi0]
    #            rzcur = rzold
    #            while True:
    #                fbndry = (fcur + fmax) / 2
    #     # Try to go half-way to fmax
    #                #CATCH, theError
    #                if theError != 0 :
    #        # Crossed boundary. Change fmax
    #                    #CATCH, /cancel
    #                    fmax = fbndry
    #                    ibndry = inds2[0] # refined boundary index
    #                else:
    #                    solver=lsode(radial_differential, f0, rzold)
    #                    rznew=solver.integrate(fbndry - f0)
    #
    #       # Didn't cross. Make this the new current location
    #                    fcur = fbndry
    #                    rzcur = rznew
    #
    #                    #CATCH, /cancel
    #
    #                if numpy.abs(fmax - fcur) < 0.01*numpy.abs(ftarget - f0):
    #                    break
    #            ri = rzcur[0]
    #            zi = rzcur[1]
    #            fbndry = fcur
    #
    #    #PRINT, "Hit boundary", ri, zi, " f =", fbndry
    #            status = 2
    #            return Bunch(status=status, rinext=ri, zinext=zi, fbndry=fbndry, ibndry=ibndry)

    # print "follow_gradient"
    # print ri, zi
    return Bunch(status=0, rinext=ri, zinext=zi, fbndry=fbndry, ibndry=ibndry)
Ejemplo n.º 30
0
    ny = dims[2]
    nz = dims[3]

    ##### Calculate geometric and physical quantities
    lZeta  = 1e2*zmax*2*np.pi*grid['R0']    # toroidal range [cm]
    lbNorm = lZeta*(grid['Bpxy'][0, ny // 2] / grid['Bxy'][0, ny // 2])  # binormal coord range [cm]
    zPerp  = lbNorm*np.arange(nz)/(nz-1)    # binormal coordinate [cm]

    cLight = 3e10                        # speed of light [cm/s]
    vTe    = 4.2e7*np.sqrt(grid['Te_x']) # electron thermal speed [cm/s]
    kperp  = 2*np.pi/lbNorm              # binormal wavenumber, [cm-1]
    wce    = 1.76e7*1e4*grid['bmag']     # electron cyclotron frequency, [rad/s]

    # Ni scale length [cm]
    Ln = np.mean(np.abs(grid['Ni0'][:, ny // 2] /
                        deriv(grid['Rxy'][:, ny // 2]*1e2, grid['Ni0'][:, ny // 2])))

    vPe   = (vTe)**2 / (wce*Ln) # electron diamagnetic drift speed [cm/s]
    wstar = vPe*kperp

    logLam = 24. - np.log(np.sqrt(grid['Ni_x']*1e14 / grid['Te_x']))
    nuei   = zeff*2.91e-6*(grid['Ni_x']*1e14)*logLam/(grid['Te_x'])**1.5
    #wci=9.58e3*(1./d.AA)*1e4*du.Bmag

    lpar = np.sum(((grid['Bxy']/grid['Bpxy']))*grid['dlthe']) / grid['nx'] # [m], average over flux surfaces
    kpar = 2*np.pi/(1e2*lpar) # cm-1
    spar = (kpar / kperp)**2 * wci * wce / (0.51 * nuei) # [1/s]
    sparn = spar / wstar

    wpe   = 5.64e4*np.sqrt(1e14*grid['Ni_x']) # electron plasma frequency, [rad/s]
    mu    = (cLight*kperp/wpe)**2
Ejemplo n.º 31
0
def process_grid( rz_grid, mesh, output=None, poorquality=None,
                  gui=None, parent=None, reverse_bt=None,
                  curv=None, smoothpressure=None,
                  smoothhthe=None, smoothcurv=None,
                  settings=None):

    if settings==None :
        # Create an empty structure
        settings = Bunch(dummy=0)

    # Check settings
        settings.calcp= -1
        settings.calcbt= -1
        settings.calchthe= -1
        settings.calcjpar= -1

   # ;CATCH, err
   # ;IF err NE 0 THEN BEGIN
   # ;  PRINT, "PROCESS_GRID failed"
        #;  PRINT, "   Error message: "+!ERROR_STATE.MSG
   # ;  CATCH, /cancel
   # ;  RETURN
   # ;ENDIF

    MU = 4.e-7*numpy.pi

    poorquality = 0

    if output==None : output="bout.grd.nc"

    # Size of the mesh
    nx = numpy.int(numpy.sum(mesh.nrad))
    ny = numpy.int(numpy.sum(mesh.npol))

    # Find the midplane
    ymid = 0
    status = gen_surface(mesh=mesh) # Start generator

    while True:
        period, yi, xi, last = gen_surface(period=None, last=None, xi=None)
        if period :
            rm = numpy.max(mesh.Rxy[xi,yi])
            ymidindx = numpy.argmax(mesh.Rxy[xi,yi])
            ymid = yi[ymidindx]
            break

        if last==1: break


    Rxy = numpy.asarray(mesh.Rxy)
    Zxy = numpy.asarray(mesh.Zxy)
    psixy = mesh.psixy*mesh.fnorm + mesh.faxis # Non-normalised psi

    pressure = numpy.zeros((nx, ny))



    # Use splines to interpolate pressure profile
    status = gen_surface(mesh=mesh) # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=last, xi=xi)
        if period :
            # Pressure only given on core surfaces
           # pressure[xi,yi] = SPLINE(rz_grid.npsigrid, rz_grid.pres, mesh.psixy[xi,yi[0]], /double)
            sol=interpolate.UnivariateSpline(rz_grid.npsigrid, rz_grid.pres,s=1)
            pressure[xi,yi] =sol(mesh.psixy[xi,yi[0]])

        else:

            pressure[xi,yi] = rz_grid.pres[numpy.size(rz_grid.pres)-1]

        if last==1 : break


    # Add a minimum amount
    if numpy.min(pressure) < 1.0e-2*numpy.max(pressure) :
        print("****Minimum pressure is very small:", numpy.min(pressure))
        print("****Setting minimum pressure to 1% of maximum")
        pressure = pressure + 1e-2*numpy.max(pressure)


    if smoothpressure != None :
        p0 = pressure[:,ymid] # Keep initial pressure for comparison
        while True :
            #!P.multi=[0,0,2,0,0]
            fig=figure()
            plot( p0, xtitle="X index", ytitle="pressure at y="+numpy.strip(numpy.str(ymid),2)+" dashed=original", color=1, lines=1)
            plot( pressure[:,ymid], color=1)
            plot( deriv(p0), xtitle="X index", ytitle="DERIV(pressure)", color=1, lines=1)
            plot( deriv(pressure[:,ymid]), color=1 )
            sm = query_yes_no("Smooth pressure profile?")#, gui=gui, dialog_parent=parent)
            if sm :
                # Smooth the pressure profile

                p2 = pressure
                for i in range (6) :
                    status = gen_surface(mesh=mesh) # Start generator
                    while True :
                        # Get the next domain
                        period, yi, xi, last = gen_surface(period=period, last=last, xi=xi)

                        if (xi > 0) and (xi < (nx-1)) :
                            for j in range (numpy.size(yi)) :
                                p2[xi,yi[j]] = ( 0.5*pressure[xi,yi[j]] +
                                                0.25*(pressure[xi-1,yi[j]] + pressure[xi+1,yi[j]])
                                                )



                        # Make sure it's still constant on flux surfaces
                        p2[xi,yi] = numpy.mean(p2[xi,yi])
                        if last != None : break
                    pressure = p2


            if sm == 0 : break


    if numpy.min(pressure) < 0.0 :
        print("")
        print("============= WARNING ==============")
        print("Poor quality equilibrium: Pressure is negative")
        print("")
        poorquality = 1


    dpdpsi = DDX(psixy, pressure)


    #;IF MAX(dpdpsi)*mesh.fnorm GT 0.0 THEN BEGIN
    #;  PRINT, ""
    #;  PRINT, "============= WARNING =============="
    #;  PRINT, "Poor quality equilibrium: Pressure is increasing radially"
    #;  PRINT, ""
    #;  poorquality = 1
    #;ENDIF

    # Grid spacing
    dx = numpy.zeros((nx, ny))
    for y in range (ny) :
        dx[0:(nx-1),y] = psixy[1::,y] - psixy[0:(nx-1),y]
        dx[nx-1,y] = dx[nx-2,y]


    # Sign
    bpsign = 1.
    xcoord = psixy
    if numpy.min(dx) < 0. :
        bpsign = -1.
        dx = -dx # dx always positive
        xcoord = -xcoord


    dtheta = 2.*numpy.pi / numpy.float(ny)
    dy = numpy.zeros((nx, ny)) + dtheta


    # B field components
    # Following signs mean that psi increasing outwards from
    # core to edge results in Bp clockwise in the poloidal plane
    # i.e. in the positive Grad Theta direction.

    Brxy = old_div(mesh.dpsidZ, Rxy)
    Bzxy = old_div(-mesh.dpsidR, Rxy)
    Bpxy = numpy.sqrt(Brxy**2 + Bzxy**2)


    # Determine direction (dot B with grad y vector)

    dot = ( Brxy[0,ymid]*(Rxy[0,ymid+1] - Rxy[0,ymid-1]) +
            Bzxy[0,ymid]*(Zxy[0,ymid+1] - Zxy[0,ymid-1])
            )

    if dot < 0. :
        print("**** Poloidal field is in opposite direction to Grad Theta -> Bp negative")
        Bpxy = -Bpxy
        if bpsign > 0 : sys.exit() # Should be negative
        bpsign = -1.0
    else:
        if bpsign < 0 : sys.exit() # Should be positive
        bpsign = 1.


  # Get toroidal field from poloidal current function fpol
    Btxy = numpy.zeros((nx, ny))
    fprime = numpy.zeros((nx, ny))
    fp = deriv(rz_grid.npsigrid*(rz_grid.sibdry - rz_grid.simagx), rz_grid.fpol)


    status = gen_surface(mesh=mesh) # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=period, xi=xi)

        if period :
            # In the core
            #fpol = numpy.interp(rz_grid.fpol, rz_grid.npsigrid, mesh.psixy[xi,yi], /spline)

            sol=interpolate.UnivariateSpline(rz_grid.npsigrid, rz_grid.fpol,s=1)
         #  fpol = SPLINE(rz_grid.npsigrid, rz_grid.fpol, mesh.psixy[xi,yi[0]], 'double')
            fpol = sol(mesh.psixy[xi,yi[0]])

            sol=interpolate.UnivariateSpline(rz_grid.npsigrid, fp ,s=1)
           # fprime[xi,yi] = SPLINE(rz_grid.npsigrid, fp, mesh.psixy[xi,yi[0]], 'double')
            fprime[xi,yi] = sol(mesh.psixy[xi,yi[0]])

        else:
            # Outside core. Could be PF or SOL
            fpol = rz_grid.fpol[numpy.size(rz_grid.fpol)-1]
            fprime[xi,yi] = 0.

        Btxy[xi,yi] = old_div(fpol, Rxy[xi,yi])

        if last ==1 : break

    # Total B field
    Bxy = numpy.sqrt(Btxy**2 + Bpxy**2)


  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  # Go through the domains to get a starting estimate
  # of hthe
    hthe = numpy.zeros((nx, ny))

  #   Pick a midplane index
    status = gen_surface(mesh=mesh) # Start generator
    while True:
    # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=last, xi=xi)

        if period :
      # In the core
            rmax = numpy.argmax(Rxy[xi,yi])
            ymidplane = yi[rmax]
            break

        if last == 1: break

    status = gen_surface(mesh=mesh) # Start generator
    while True:
    # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=last, xi=xi)

        n = numpy.size(yi)

        # Get distance along this line

        if period :

            # Periodic, so can use FFT
            #drdi = REAL_PART(fft_deriv(Rxy[xi, yi]))
            #dzdi = REAL_PART(fft_deriv(Zxy[xi, yi]))
            line=numpy.append(Rxy[xi,yi[n-1::]], Rxy[xi,yi])
            line=numpy.append(line,Rxy[xi,yi[0:1]])

            drdi = deriv(line)[1:n+1]

            line=numpy.append(Zxy[xi,yi[n-1::]], Zxy[xi,yi])
            line=numpy.append(line,Zxy[xi,yi[0:1]])

            dzdi = deriv(line)[1:n+1]
        else:
        # Non-periodic
            drdi = numpy.gradient(Rxy[xi, yi])
            dzdi = numpy.gradient(Zxy[xi, yi])


        dldi = numpy.sqrt(drdi**2 + dzdi**2)


        if 0 :

        # Need to smooth to get sensible results
            if period :
                n = numpy.size(dldi)
                line=numpy.append(dldi[(n-2)::], dldi) # once
                line=numpy.append(line,dldi[0:2])
                dldi = SMOOTH(line, 5)[4:(n+4)]

                line=numpy.append(dldi[(n-2)::], dldi) #twice
                line=numpy.append(line,dldi[0:2])
                dldi = SMOOTH(line, 5)[4:(n+4)]

                line=numpy.append(dldi[(n-2)::], dldi) # three
                line=numpy.append(line,dldi[0:2])
                dldi = SMOOTH(line, 5)[4:(n+4)]

            else:
                line = dldi
                dldi = SMOOTH(line, 5)[2:n+2]
                line = dldi
                dldi = SMOOTH(line, 5)[2:n+2]
                line = dldi
                dldi = SMOOTH(dldi, 5)[2:n+2]


        hthe[xi, yi] = old_div(dldi, dtheta) # First estimate of hthe

        # Get outboard midplane
        if period and xi == 0 :
            m = numpy.argmax(Rxy[0,yi])
            ymidplane = yi[m]

        if last == 1 : break

    print("Midplane index ", ymidplane)

    fb0 = force_balance(psixy, Rxy, Bpxy, Btxy, hthe, pressure)
    print("Force imbalance: ", numpy.mean(numpy.abs(fb0)), numpy.max(numpy.abs(fb0)))




  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  # Correct pressure using hthe

    print("Calculating pressure profile from force balance")

    try:

    # Calculate force balance
        dpdx = old_div(( -Bpxy*DDX(xcoord, Bpxy * hthe) - Btxy*hthe*DDX(xcoord, Btxy) - (Btxy*Btxy*hthe/Rxy)*DDX(xcoord, Rxy) ), (MU*hthe))

        # Surface average
        dpdx2 = surface_average(dpdx, mesh)

        pres = numpy.zeros((nx, ny))
        # Integrate to get pressure
        for i in range (ny) :
            pres[:,i] = int_func(psixy[:,i], dpdx2[:,i])
            pres[:,i] = pres[:,i] - pres[nx-1,i]



        status = gen_surface(mesh=mesh) # Start generator
        while True:
      # Get the next domain
            period, yi, xi, last = gen_surface(period=None, last=None, xi=None)

            ma = numpy.max(pres[xi,yi])

            for i in range (numpy.size(yi)) :
                pres[:,yi[i]] = pres[:,yi[i]] - pres[xi,yi[i]] + ma

            if last == 1 : break


        pres = pres - numpy.min(pres)

    # Some sort of smoothing here?


        fb0 = force_balance(psixy, Rxy, Bpxy, Btxy, hthe, pres)
        print("Force imbalance: ", numpy.mean(numpy.abs(fb0)), numpy.max(numpy.abs(fb0)))


       #!P.MULTI=[0,0,2,0,0]
        fig=figure(figsize=(7, 11))
        subplots_adjust(left=.07, bottom=.07, right=0.95, top=0.95,
                wspace=.3, hspace=.25)

        SURFACE( pressure, fig, xtitle="X", ytitle="Y", var='Pa', sub=[2,1,1])
        title("Input pressure")
        SURFACE( pres, fig, xtitle="X", ytitle="Y", var='Pa', sub=[2,1,2])
        title("New pressure")
  #  arrange the plot on the screen
  #      mngr = get_current_fig_manager()
  #      geom = mngr.window.geometry()
  #      x,y,dx,dy = geom.getRect()
  #      mngr.window.setGeometry(0, 0, dx, dy)
  #
        show(block=False)


        calcp = settings.calcp

        if calcp == -1 :
            calcp = query_yes_no("Keep new pressure?")#, gui=gui, dialog_parent=parent)
        else: time.sleep( 2 )
        if calcp == 1 :
            pressure = pres
            dpdpsi = dpdx2


    except Exception:
        print("WARNING: Pressure profile calculation failed: ")#, !ERROR_STATE.MSG
        pass

    #CATCH, /cancel

  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  # Correct f = RBt using force balance

    calcbt = settings.calcbt
    if calcbt == -1 : calcbt = query_yes_no("Correct f=RBt using force balance?")#, gui=gui, dialog_parent=parent)
    if calcbt == 1 :

        new_Btxy = newton_Bt(psixy, Rxy, Btxy, Bpxy, pres, hthe, mesh)

        fb0 = force_balance(psixy, Rxy, Bpxy, new_Btxy, hthe, pressure)
        print("force imbalance: ", numpy.mean(numpy.abs(fb0)), numpy.max(numpy.abs(fb0)))


        fig=figure(figsize=(7, 11))
        subplots_adjust(left=.07, bottom=.07, right=0.95, top=0.95,
                wspace=.3, hspace=.25)

        subplot(211)
        SURFACE( Btxy, fig, xtitle="X", ytitle="Y", var='T', sub=[2,1,1])
        title("Input Bt")
        subplot(212)
        SURFACE( new_Btxy, fig, xtitle="X", ytitle="Y", var='T', sub=[2,1,2])
        title("New Bt")
          #  arrange the plot on the screen
        #mngr = get_current_fig_manager()
        #geom = mngr.window.geometry()
        #x,y,dx,dy = geom.getRect()
        #mngr.window.setGeometry(600, 0, dx, dy)


        show(block=False)

        calcbt = settings.calcbt
        if calcbt == -1 : calcbt = query_yes_no("Keep new Bt?")#, gui=gui, dialog_parent=parent)
        if calcbt == 1 :
            Btxy = new_Btxy
            Bxy = numpy.sqrt(Btxy**2 + Bpxy**2)

  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  # CALCULATE HTHE
  # Modify hthe to fit force balance using initial guess
  # Does not depend on signs
  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

    calchthe = settings.calchthe
    if calchthe == -1 : calchthe = query_yes_no("Adjust hthe using force balance?")#, gui=gui, dialog_parent=parent)
    if calchthe == 1 :
        # This doesn't behave well close to the x-points
        fixhthe = numpy.int(old_div(nx, 2))
        nh = correct_hthe(Rxy, psixy, Btxy, Bpxy, hthe, pressure, fixhthe=fixhthe)

        fb0 = force_balance(psixy, Rxy, Bpxy, Btxy, nh, pressure)
        print("Force imbalance: ", numpy.mean(numpy.abs(fb0)), numpy.max(numpy.abs(fb0)))

        print("numpy.maximum difference in hthe: ", numpy.max(numpy.abs(hthe - nh)))
        print("numpy.maximum percentage difference: ", 100.*numpy.max(numpy.abs(old_div((hthe - nh),hthe))))

       #!P.multi=[0,0,1,0,0]
        fig=figure(figsize=(7, 4))
        title("Poloidal arc length at midplane. line is initial estimate")
        plot( hthe[:,0], '-' )
        plot( nh[:,0], 'r-+' )
                  #  arrange the plot on the screen
        #mngr = get_current_fig_manager()
        #geom = mngr.window.geometry()
        #x,y,dx,dy = geom.getRect()
        #mngr.window.setGeometry(0, 1150, dx, dy)

        show(block=False)

        if query_yes_no("Keep new hthe?") == 1:#, gui=gui, dialog_parent=parent) :
            hthe = nh



    if smoothhthe != None :
    # Smooth hthe to prevent large jumps in X or Y. This
    # should be done by creating a better mesh in the first place

    # Need to smooth in Y and X otherwise smoothing in X
    # produces discontinuities in Y
        hold = hthe

        if 1 :
      # Nonlinear smoothing. Tries to smooth only regions with large
      # changes in gradient

            hthe =0.# smooth_nl(hthe, mesh)

        else:
      # Just use smooth in both directions

            for i in range (ny) :
                hthe[:,i] = SMOOTH(SMOOTH(hthe[:,i],10),10)


        status = gen_surface(mesh=mesh) # Start generator
        while True:
        # Get the next domain
            period, yi, xi, last = gen_surface(period=None, last=None, xi=None)

            n = numpy.size(yi)

            if period :
                hthe[xi,yi] = (SMOOTH([hthe[xi,yi[(n-4):(n-1)]], hthe[xi,yi], hthe[xi,yi[0:3]]], 4))[4:(n+3)]
            else:
                hthe[xi,yi] = SMOOTH(hthe[xi,yi], 4)

            if last == 1: break



    # Calculate field-line pitch
    pitch = hthe * Btxy / (Bpxy * Rxy)


    # derivative with psi
    dqdpsi = DDX(psixy, pitch)


    qinty, qloop = int_y(pitch, mesh, loop=0, nosmooth='nosmooth', simple='simple')
    qinty = qinty * dtheta
    qloop = qloop * dtheta


    sinty = int_y(dqdpsi, mesh, nosmooth='nosmooth', simple='simple') * dtheta



    # NOTE: This is only valid in the core
    pol_angle = numpy.zeros((nx,ny))
    for i in range (nx) :  pol_angle[i, :] = 2.0*numpy.pi * qinty[i,:] / qloop[i]


  #;;;;;;;;;;;;;;;;;;; THETA_ZERO ;;;;;;;;;;;;;;;;;;;;;;
  # re-set zshift to be zero at the outboard midplane

    print("MIDPLANE INDEX = ", ymidplane)

    status = gen_surface(mesh=mesh) # Start generator
    while True:
    # Get the next domain
        period, yi, xi, last = gen_surface(period=None, last=None, xi=None)

        w = numpy.size(numpy.where(yi == ymidplane))
        if w > 0 :
      # Crosses the midplane
            qinty[xi, yi] = qinty[xi, yi] - qinty[xi, ymidplane]
            sinty[xi, yi] = sinty[xi, yi] - sinty[xi, ymidplane]
        else:
      # Doesn't include a point at the midplane
            qinty[xi, yi] = qinty[xi, yi] - qinty[xi,yi[0]]
            sinty[xi, yi] = sinty[xi, yi] - sinty[xi,yi[0]]

        if last ==1 : break

    print("")
    print("==== Calculating curvature ====")

  #;;;;;;;;;;;;;;;;;;; CURVATURE ;;;;;;;;;;;;;;;;;;;;;;;
  # Calculating b x kappa

    if curv == None :

        print("*** Calculating curvature in toroidal coordinates")

        thetaxy = numpy.zeros((nx, ny))
        status = gen_surface(mesh=mesh) # Start generator
        while True:
            # Get the next domain
            period, yi, xi, last = gen_surface(period=None, last=None, xi=None)
            thetaxy[xi,yi] = numpy.arange(numpy.size(yi)).astype(float)*dtheta
            if last ==1 : break


        bxcv = curvature( nx, ny, Rxy,Zxy, Brxy, Bzxy, Btxy,
                    psixy, thetaxy, hthe,
                     mesh=mesh)

        bxcvx = bpsign*bxcv.psi
        bxcvy= bxcv.theta
        bxcvz = bpsign*(bxcv.phi - sinty*bxcv.psi - pitch*bxcv.theta)


        # x borders
        bxcvx[0,:] = bxcvx[1,:]
        bxcvx[nx-1,:] = bxcvx[nx-2,:]

        bxcvy[0,:] = bxcvy[1,:]
        bxcvy[nx-1,:] = bxcvy[nx-2,:]

        bxcvz[0,:] = bxcvz[1,:]
        bxcvz[nx-1,:] = bxcvz[nx-2,:]

    elif curv == 1 :
        # Calculate on R-Z mesh and then interpolate onto grid
        # ( cylindrical coordinates)

        print("*** Calculating curvature in cylindrical coordinates")

        bxcv = rz_curvature(rz_grid)

        # DCT methods cause spurious oscillations
        # Linear interpolation seems to be more robust
        bxcv_psi = numpy.interp(bxcv.psi, mesh.Rixy, mesh.Zixy)
        bxcv_theta = old_div(numpy.interp(bxcv.theta, mesh.Rixy, mesh.Zixy), hthe)
        bxcv_phi = numpy.interp(bxcv.phi, mesh.Rixy, mesh.Zixy)

        # If Bp is reversed, then Grad x = - Grad psi
        bxcvx = bpsign*bxcv_psi
        bxcvy = bxcv_theta
        bxcvz = bpsign*(bxcv_phi - sinty*bxcv_psi - pitch*bxcv_theta)
    elif curv == 2 :
        # Curvature from Curl(b/B)

        bxcvx = bpsign*(Bpxy * Btxy*Rxy * DDY(old_div(1., Bxy), mesh) / hthe)
        bxcvy = -bpsign*Bxy*Bpxy * DDX(xcoord, Btxy*Rxy/Bxy^2) / (2.*hthe)
        bxcvz = Bpxy^3 * DDX(xcoord, old_div(hthe,Bpxy)) / (2.*hthe*Bxy) - Btxy*Rxy*DDX(xcoord, old_div(Btxy,Rxy)) / (2.*Bxy) - sinty*bxcvx

    else:
        # calculate in flux coordinates.

        print("*** Calculating curvature in flux coordinates")

        dpb = numpy.zeros((nx, ny))      # quantity used for y and z components

        for i in range (ny) :
            dpb[:,i] = MU*dpdpsi/Bxy[:,i]

        dpb = dpb + DDX(xcoord, Bxy)

        bxcvx = bpsign*(Bpxy * Btxy*Rxy * DDY(old_div(1., Bxy), mesh) / hthe)
        bxcvy = bpsign*(Bpxy*Btxy*Rxy*dpb / (hthe*Bxy^2))
        bxcvz = -dpb - sinty*bxcvx



    if smoothcurv:
        # Smooth curvature to prevent large jumps

        # Nonlinear smoothing. Tries to smooth only regions with large
        # changes in gradient

        bz = bxcvz + sinty * bxcvx

        print("Smoothing bxcvx...")
        bxcvx = 0.#smooth_nl(bxcvx, mesh)
        print("Smoothing bxcvy...")
        bxcvy = 0.#smooth_nl(bxcvy, mesh)
        print("Smoothing bxcvz...")
        bz = 0.#smooth_nl(bz, mesh)

        bxcvz = bz - sinty * bxcvx


  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  # CALCULATE PARALLEL CURRENT
  #
  # Three ways to calculate Jpar0:
  # 1. From fprime and pprime
  # 2. From Curl(B) in field-aligned coords
  # 3. From the curvature
  #
  # Provides a way to check if Btor should be reversed
  #
  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

    print("")
    print("==== Calculating parallel current ====")

    jpar0 = - Bxy * fprime / MU - Rxy*Btxy * dpdpsi / Bxy


    # Set to zero in PF and SOL
    status = gen_surface(mesh=mesh) # Start generator
    while True:
    # Get the next domain
        period, yi, xi, last = gen_surface(period=None, last=None, xi=None)

        if period == None : jpar0[xi,yi] = 0.0
        if last == 1 : break

  # Curl(B) expression for Jpar0 (very noisy usually)
    j0 = ( bpsign*((Bpxy*Btxy*Rxy/(Bxy*hthe))*( DDX(xcoord, Bxy**2*hthe/Bpxy) - bpsign*Btxy*Rxy*DDX(xcoord,Btxy*hthe/(Rxy*Bpxy)) )
        - Bxy*DDX(xcoord, Btxy*Rxy)) / MU )



  # Create a temporary mesh structure to send to adjust_jpar
    tmp_mesh = Bunch(mesh,
                           bxcvx=bxcvx, bxcvy=bxcvy,  bxcvz=bxcvz,
                            Bpxy=Bpxy,  Btxy=Btxy,  Bxy=Bxy,
                            dx=dx,  dy=dy,
                            hthe=hthe,  jpar0=jpar0,  pressure=pressure)
    tmp_mesh.psixy = psixy

    jpar = adjust_jpar( tmp_mesh, noplot='noplot')


   #!P.multi=[0,2,2,0,0]

    fig=figure(figsize=(15, 11))
    subplots_adjust(left=.07, bottom=.07, right=0.95, top=0.95,
                wspace=.3, hspace=.25)

    subplot(221)
    SURFACE( jpar0, fig, xtitle="X", ytitle="Y", var='A', sub=[2,2,1])
    title("Jpar from F' and P'")

    subplot(222)
    SURFACE( jpar, fig, xtitle="X", ytitle="Y", var='A', sub=[2,2,2])
    title("Jpar from curvature")

    subplot(223)
    plot( jpar0[0,:],'-', jpar[0,:] ,'+' )
    ylim([numpy.min([jpar0[0,:],jpar[0,:]]), numpy.max([jpar0[0,:],jpar[0,:]])])
    title("jpar at x=0. Solid from f' and p'")

    subplot(224)
    plot(jpar0[:,ymidplane],'-' , jpar[:,ymidplane] , '+' )
    ylim([numpy.min([jpar0[:,ymidplane],jpar[:,ymidplane]]),numpy.max([jpar0[:,ymidplane],jpar[:,ymidplane]])])

    title("Jpar at y="+numpy.str(ymidplane)+" Solid from f' and p'")

        #  arrange the plot on the screen
    #mngr = get_current_fig_manager()
    #geom = mngr.window.geometry()
    #x,y,dx,dy = geom.getRect()
    #mngr.window.setGeometry(1350, 0, dx, dy)


    show(block=False)

 # !P.multi=0

    calcjpar = settings.calcjpar
    if calcjpar == -1 : calcjpar = query_yes_no("Use Jpar from curvature?")#, gui=gui, dialog_parent=parent)
    if calcjpar == True :
        jpar0 = jpar


    if 0 :

    # Try smoothing jpar0 in psi, preserving zero points and maxima
        jps = jpar0
        for y in range ( ny ):
            j = jpar0[:,y]
            js = j
            ma = numpy.max(numpy.abs(j))
            ip = numpy.argmax(numpy.abs(j))
            if (ma < 1.e-4) or (ip == 0) :
                jps[:,y] = j

            level = 1.
            #i0 = MAX(WHERE(ABS(j[0:ip]) LT level))
            i1 = numpy.min(numpy.where(numpy.abs(j[ip::]) < level))

            #IF i0 LE 0 THEN i0 = 1
            i0 = 1

            if i1 == -1 :
                i1 = nx-2
            else:
                i1 = i1 + ip

            if (ip <= i0) or (ip >= i1) :

      # Now preserve starting and end points, and peak value
                div = numpy.int(old_div((i1-i0),10))+1 # reduce number of points by this factor

                inds = [i0] # first point
                for i in [i0+div, ip-div, div] : inds = [inds, i]
                inds = [inds, ip] # Put in the peak point

        # Calculate spline interpolation of inner part

                js[0:ip] = spline_mono(inds, j[inds], numpy.arange(ip+1),
                             yp0=(j[i0] - j[i0-1]), ypn_1=0.0)

                inds = [ip] # peak point
                for i in [ip+div, i1-div, div] :
                    inds = [inds, i]


                inds = [inds, i1]  # Last point
                js[ip:i1] = spline_mono(inds, j[inds], ip+numpy.arange(i1-ip+1),
                              yp0=0.0, ypn_1=(j[i1+1]-j[i1]))

                jps[:,y] = js



  #;;;;;;;;;;;;;;;;;;; TOPOLOGY ;;;;;;;;;;;;;;;;;;;;;;;
  # Calculate indices for backwards-compatibility

    nr = numpy.size(mesh.nrad)
    np = numpy.size(mesh.npol)
    if (nr == 2) and (np == 3) :
        print("Single null equilibrium")

        ixseps1 = mesh.nrad[0]
        ixseps2 = nx

        jyseps1_1 = mesh.npol[0]-1
        jyseps1_2 = mesh.npol[0] + numpy.int(old_div(mesh.npol[1],2))
        ny_inner = jyseps1_2
        jyseps2_1 = jyseps1_2
        jyseps2_2 = ny - mesh.npol[2]-1

    elif (nr == 3) and (np == 6) :
        print("Double null equilibrium")

        ixseps1 = mesh.nrad[0]
        ixseps2 = ixseps1 + mesh.nrad[1]

        jyseps1_1 = mesh.npol[0]-1
        jyseps2_1 = jyseps1_1 + mesh.npol[1]

        ny_inner = jyseps2_1 + mesh.npol[2] + 1

        jyseps1_2 = ny_inner + mesh.npol[3] - 1
        jyseps2_2 = jyseps1_2 + mesh.npol[4]

    elif (nr == 1) and (np == 1) :

        print("Single domain")

        ixseps1 = nx
        ixseps2 = nx

        jyseps1_1 = -1
        jyseps1_2 = numpy.int(old_div(ny,2))
        jyseps2_1 = numpy.int(old_div(ny,2))
        ny_inner = numpy.int(old_div(ny,2))
        jyseps2_2 = ny - 1

    else:
        print("***************************************")
        print("* WARNING: Equilibrium not recognised *")
        print("*                                     *")
        print("*  Check mesh carefully!              *")
        print("*                                     *")
        print("*  Contact Ben Dudson                 *")
        print("*      [email protected]     *")
        print("***************************************")
        ixseps1 = -1
        ixseps2 = -1

        jyseps1_1 = -1
        jyseps1_2 = numpy.int(old_div(ny,2))
        jyseps2_1 = numpy.int(old_div(ny,2))
        ny_inner = numpy.int(old_div(ny,2))
        jyseps2_2 = ny - 1


    print("Generating plasma profiles:")

    print("  1. Flat temperature profile")
    print("  2. Flat density profile")
    print("  3. Te proportional to density")
    while True:
        opt = input("Profile option:")
        if eval(opt) >= 1 and eval(opt) <= 3 : break


    if eval(opt) == 1 :
        # flat temperature profile

        print("Setting flat temperature profile")
        while True:
            Te_x = eval(input("Temperature (eV):"))


        # get density
            Ni = old_div(pressure, (2.* Te_x* 1.602e-19*1.0e20))

            print("numpy.maximum density (10^20 m^-3):", numpy.max(Ni))

            done = query_yes_no("Is this ok?")
            if done == 1 : break

        Te = numpy.zeros((nx, ny))+Te_x
        Ti = Te
        Ni_x = numpy.max(Ni)
        Ti_x = Te_x
    elif eval(opt) == 2 :
        print("Setting flat density profile")

        while True:
            Ni_x = eval(input("Density [10^20 m^-3]:"))

            # get temperature
            Te = old_div(pressure, (2.* Ni_x * 1.602e-19*1.0e20))

            print("numpy.maximum temperature (eV):", numpy.max(Te))
            if query_yes_no("Is this ok?") == 1 : break

        Ti = Te
        Ni = numpy.zeros((nx, ny)) + Ni_x
        Te_x = numpy.max(Te)
        Ti_x = Te_x
    else:
        print("Setting te proportional to density")

        while True:
            Te_x = eval(input("Maximum temperature [eV]:"))


            Ni_x = old_div(numpy.max(pressure), (2.*Te_x * 1.602e-19*1.0e20))

            print("Maximum density [10^20 m^-3]:", Ni_x)

            Te = Te_x * pressure / numpy.max(pressure)
            Ni = Ni_x * pressure / numpy.max(pressure)
            if query_yes_no("Is this ok?") == 1 : break
        Ti = Te
        Ti_x =  Te_x


    rmag = numpy.max(numpy.abs(Rxy))
    print("Setting rmag = ", rmag)

    bmag = numpy.max(numpy.abs(Bxy))
    print("Setting bmag = ", bmag)

    #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    # save to file
    # open a new netCDF file for writing.
    handle = file_open(output)

    print("Writing grid to file "+output)

    # Size of the grid

    s = file_write(handle, "nx", nx)
    s = file_write(handle, "ny", ny)


    # Topology for original scheme
    s = file_write(handle, "ixseps1", ixseps1)
    s = file_write(handle, "ixseps2", ixseps2)
    s = file_write(handle, "jyseps1_1", jyseps1_1)
    s = file_write(handle, "jyseps1_2", jyseps1_2)
    s = file_write(handle, "jyseps2_1", jyseps2_1)
    s = file_write(handle, "jyseps2_2", jyseps2_2)
    s = file_write(handle, "ny_inner", ny_inner);

    # Grid spacing

    s = file_write(handle, "dx", dx)
    s = file_write(handle, "dy", dy)

    s = file_write(handle, "ShiftAngle", qloop)
    s = file_write(handle, "zShift", qinty)
    s = file_write(handle, "pol_angle", pol_angle)
    s = file_write(handle, "ShiftTorsion", dqdpsi)

    s = file_write(handle, "Rxy",  Rxy)
    s = file_write(handle, "Zxy",  Zxy)
    s = file_write(handle, "Bpxy", Bpxy)
    s = file_write(handle, "Btxy", Btxy)
    s = file_write(handle, "Bxy",  Bxy)
    s = file_write(handle, "hthe", hthe)
    s = file_write(handle, "sinty", sinty)
    s = file_write(handle, "psixy", psixy)

    # Topology for general configurations
    s = file_write(handle, "yup_xsplit", mesh.yup_xsplit)
    s = file_write(handle, "ydown_xsplit", mesh.ydown_xsplit)
    s = file_write(handle, "yup_xin", mesh.yup_xin)
    s = file_write(handle, "yup_xout", mesh.yup_xout)
    s = file_write(handle, "ydown_xin", mesh.ydown_xin)
    s = file_write(handle, "ydown_xout", mesh.ydown_xout)
    s = file_write(handle, "nrad", mesh.nrad)
    s = file_write(handle, "npol", mesh.npol)

    # plasma profiles

    s = file_write(handle, "pressure", pressure)
    s = file_write(handle, "Jpar0", jpar0)
    s = file_write(handle, "Ni0", Ni)
    s = file_write(handle, "Te0", Te)
    s = file_write(handle, "Ti0", Ti)


    s = file_write(handle, "Ni_x", Ni_x)
    s = file_write(handle, "Te_x", Te_x)
    s = file_write(handle, "Ti_x", Ti_x)
    s = file_write(handle, "bmag", bmag)
    s = file_write(handle, "rmag", rmag)

    # Curvature
    s = file_write(handle, "bxcvx", bxcvx)
    s = file_write(handle, "bxcvy", bxcvy)
    s = file_write(handle, "bxcvz", bxcvz)

    # Psi range
    s = file_write(handle, "psi_axis", mesh.faxis)
    psi_bndry = mesh.faxis + mesh.fnorm
    s = file_write(handle, "psi_bndry", psi_bndry)

    file_close, handle
    print("DONE")