geometry.py

# File for handling the geometry and functions needed to sample points. 

# When using Miller geometry the parameterisation is:

# R = R0 + rho * A cos( alpha + delta * sin(alpha) ) + Delta( 1 - rho**2)
# Z = kappa * rho * A * sin(alpha)

# Where

# rho   = Normalised minor radius = a / a_sep
# alpha = angle

# R0    = Geometric axis radius [m]
# A     = Minor radius [m]
# delta = triangularity
# kappa = elongation
# Delta = Shafranov shift [m]

# => R = R(rho, alpha), Z = Z(rho,alpha)

import numpy as np
from math        import cos, sin, pi, atan2
from integration import rk4
from scipy.interpolate import interp1d, RectBivariateSpline
from matplotlib  import pyplot as plt

def get_R_Z( rho, alpha, params ):
    '''
    Calculates R and Z from rho and alpha.
    '''

    if( params['use_geqdsk'] ):
        
        # Interpolate bicubic splines 
        R = params['R_map'].__call__(rho, alpha, dx=0, dy=0, grid=False)
        Z = params['Z_map'].__call__(rho, alpha, dx=0, dy=0, grid=False)

    else:

        R0    = params['R0']
        A     = params['a_sep']
        delta = params['delta']
        kappa = params['kappa']
        shift = params['shift']  # Shafranov Shift = Delta.
        
        R = R0 + rho * A * cos( alpha + delta * sin( alpha ) ) + shift * (1.0 - rho**2.0)
        
        Z = kappa * rho * A * sin(alpha)
        
    return R, Z   

def derivatives_R_Z( rho, alpha, params ):
    '''
    Calculates the derivatives of R and Z wrt rho and alpha.
    '''

    if( params['use_geqdsk'] ):

        # Get derivatives from bicubic splines 
        dRdrho   = params['R_map'].__call__(rho, alpha, dx=1, dy=0, grid=False)
        dRdalpha = params['R_map'].__call__(rho, alpha, dx=0, dy=1, grid=False)

        dZdrho   = params['Z_map'].__call__(rho, alpha, dx=1, dy=0, grid=False)
        dZdalpha = params['Z_map'].__call__(rho, alpha, dx=0, dy=1, grid=False)

    else:
        
        A     = params['a_sep']
        delta = params['delta']
        kappa = params['kappa']
        shift = params['shift']  # Shafranov Shift = Delta.
        
        dRdrho = A * cos( alpha + delta * sin( alpha ) ) - 2.0 * shift * rho
        
        dRdalpha = -1.0 * rho * A * sin( alpha + delta * sin(alpha) ) * ( 1.0 + delta * cos(alpha) )
        
        dZdrho = kappa * A * sin(alpha)
        
        dZdalpha = kappa * rho * A * cos(alpha)
        
    return dRdrho, dRdalpha, dZdrho, dZdalpha

def derivatives_rho_alpha( rho, alpha, params ):
    '''
    Inverts the matrix of derivatives of R,Z wrt rho,alpha to get the 
    derivatives of rho,alpha wrt R,Z.
    '''

    dRdrho, dRdalpha, dZdrho, dZdalpha = derivatives_R_Z( rho, alpha, params )

    jacobian = np.array( [ [dRdrho, dZdrho], [dRdalpha, dZdalpha] ] )

    inverted_jacobian = np.linalg.inv(jacobian)

    drhodR   = inverted_jacobian[0,0]
    dalphadR = inverted_jacobian[0,1]
    drhodZ   = inverted_jacobian[1,0]
    dalphadZ = inverted_jacobian[1,1]

    if( params['DEBUG_JACOBIAN'] ):
        
        print('Checking Jacobian Inversion')
        print( np.matmul( jacobian, inverted_jacobian ) )
        print()

    return drhodR, dalphadR, drhodZ, dalphadZ

def mod_grad_rho( rho, alpha, params ):
    '''
    Constructs the modulus of the gradient of rho. Needed for the volume element in get_volume_element_1.
    '''

    drhodR, dalphadR, drhodZ, dalphadZ = derivatives_rho_alpha( rho, alpha, params )

    mod_drho = ( drhodR**2.0 + drhodZ**2.0 )**0.5
    return mod_drho

def get_volume_element_1(rho, alpha, drho, dalpha, params):
    '''
    Calculates volume element using grad(rho).
    dV = 2 * pi * R * dl * drho / |grad(rho)|.
    '''

    # Get R,Z
    R, Z = get_R_Z( rho, alpha, params )

    # Get |grad(rho)|
    mod_drho = mod_grad_rho( rho, alpha, params )

    # Need to get next point to calculate dl
    R_next, Z_next = get_R_Z( rho, alpha+dalpha, params )

    # Get line element length  
    dl = ( (R_next - R)**2.0 + (Z_next - Z)**2.0 )**0.5

    # Volume element - 2 pi R dl drho / |grad(rho)|
    dV = pi * (R + R_next) * dl * drho / mod_drho

    return dV

def cross2d(x, y):
    return x[0] * y[1] - x[1] * y[0]

def get_volume_element_2(rho, alpha, drho, dalpha, params):
    '''
    Calculates volume element using covariant basis vectors of alpha and rho.
    dV = 2 pi R | dr/drho x dr/dalpha | drho dalpha.
    '''

    # Get derivatives
    dRdrho, dRdalpha, dZdrho, dZdalpha = derivatives_R_Z( rho, alpha, params )
    
    drdrho   = np.array( [  dRdrho,   dZdrho] )
    drdalpha = np.array( [dRdalpha, dZdalpha] )

    # Get | dr/drho x dr/dalpha |
    area = abs( cross2d(drdrho, drdalpha) )

    # Get R,Z 
    R, Z = get_R_Z( rho, alpha, params )

    # Volume element
    dV = 2.0 * pi * R * area * drho * dalpha

    return dV

def get_flux_surfaces(n,params,nalpha=100,separatrix=False):
    '''
    Constructs n equidistant flux surfaces.
    '''

    # First generate some contours
    cont_Rs = np.zeros((int(n),nalpha))
    cont_Zs = np.zeros((int(n),nalpha))

    # Get n equidistant rho values
    rhos = np.linspace(0.0,1.0,n+1)[1:]

    if( not separatrix ):
        rhos[-1] = 0.99

    for ialpha in range(nalpha):

        alpha = 2.0 * pi * ( ialpha / (float(nalpha)-1.0) )

        for i in range(n):

            cont_Rs[i,ialpha], cont_Zs[i,ialpha] = get_R_Z( rhos[i], alpha, params )

    return cont_Rs, cont_Zs

# Routines for constructing coordinate maps with numerical equilibria -------------

def downsample( x, N ):
    '''
    Function for downsampling an array to between N and 2N points.
    '''

    ratio = int( len(x) / N )
    if ratio <= 1:
        return x

    xnew = []

    i = 0
    while( i < len(x) ):

        xnew.append( x[i] )
        i = i + ratio

    i = i - ratio
    if( i < len(x)-1 ):

        xnew.append( x[-1] )

    return xnew

def PsiContour(gfile,R,step,N,ndownsample):
    '''
    Function for finding a Psi contour in R,Z space from a
    given initial radius on the outboard midplane.

    R    = Initial Radial position.
    N    = Number of alpha points on contour.
    step = step size
    ndownsample = Ensure at most 2 ndownsample points on contour.
    '''

    # Store the trajectory
    Rs     = [R]
    Zs     = [gfile.z_axis]
    s      = [0.0]           # Path length - will be used to construct alpha coordinate
    thetas = [0.0]           # Normal poloidal angle - Note this is accumulating angle changes rather than storing poloidal angle.

    while( ( thetas[-1] - thetas[0] ) <= 2.0 * pi ):

        theta_old = atan2( Zs[-1] - gfile.z_axis, Rs[-1] - gfile.r_axis )

        # Step along
        Rnew, Znew = rk4(gfile,Rs[-1],Zs[-1],step)

        Rs.append(Rnew)
        Zs.append(Znew)
        s.append( s[-1]+step )

        theta_new = atan2(   Znew - gfile.z_axis,   Rnew - gfile.r_axis )

        # Need to handle branch cut
        if( Rnew < gfile.r_axis ):
            if( theta_old < 0.0 ):
                theta_old += 2.0 * pi
            if( theta_new < 0.0 ):
                theta_new += 2.0 * pi
            
        dtheta = abs( theta_new - theta_old )
        thetas.append( thetas[-1] + dtheta )

    # Downsample functions if needed
#    thetas = downsample( thetas, ndownsample )
#    s      = downsample( s,      ndownsample )
#    Rs     = downsample( Rs,     ndownsample )
#    Zs     = downsample( Zs,     ndownsample )

    # Need R,Z as functions of straight field line angle
    # Interpolate s(theta). Find s(2.0*pi)

    f = interp1d( thetas, s, kind='cubic', bounds_error = True )
    smax = f( 2.0 * pi )
    print('Length of contour: ' , smax )

    # Get Straight Field Line Angles
    sfla = [ xs * 2.0 * pi / smax for xs in s ]

    # Get R, Z as functions of straight field line angle
    fR = interp1d( sfla, Rs, kind='cubic', bounds_error=True )
    fZ = interp1d( sfla, Zs, kind='cubic', bounds_error=True )  

    # Get R,Z points at equal distances around contour
    alphas = np.linspace(0.0,2.0*pi,N)

    Rs = [ float( fR(x) ) for x in alphas ]
    Zs = [ float( fZ(x) ) for x in alphas ]

    return Rs, Zs

# Note this will not extend all the way down to psi = 0 but should be very close
def Rpsi(gfile,Nsteps):
    '''
    Calculate R(psin) along midplane.
    '''

    rmin = gfile.r_axis
    rmax = gfile.rleft + gfile.rdim

    Rs   = np.linspace(rmin, rmax, Nsteps)
    psis = []

    for R in Rs:

        Psi, psin, BR, BT, BZ  = gfile.get_fields(R,gfile.z_axis)
        psis.append(psin)

    psis = np.array(psis)

    # Ensure monotonicity at the beginning of the array
    while True:
        if psis[0] >= psis[1]:
            psis = psis[1:]
            Rs   = Rs[1:]
        else:
            break

    # Fix axis
    if( psis.size < Nsteps ):
        Rs   = np.insert(Rs,   0, gfile.r_axis, axis=0)
        psis = np.insert(psis, 0,          0.0, axis=0)

    # build R(psin)
    R_Psi = interp1d( psis, Rs, kind='linear', bounds_error=True )

    return R_Psi

def construct_charts(params):
    '''
    Construct the coordinate maps R,Z(rho, alpha) where rho = normalised poloidal flux
    and alpha = arc length normalised to 2pi.
    '''

    gfile = params['gfile']

    # Get R(psin)
    Rps = Rpsi( gfile, params['radial_steps'] )

    # Construct Psi grid for coordinate charts.
    # Throw away separatrix as this will not be tracked correctly.
    Psins   = np.linspace(0.0,1.0,params['npsi_map'])[:-1]

    if( params['sample_sqrt_psi'] ):

        # Generate N starting points equally spaced in Sqrt(Psi) ~ R.
        Psins     = [ x**2.0 for x in Psins ]

    Rs = [ float( Rps(x) ) for x in Psins[1:] ]  

    # Plot R(psin)
    plt.figure()
    plt.plot( Psins[1:], Rs )
    plt.xlabel('$\psi_{norm}$')
    plt.ylabel('R [m]')
    plt.title('Initial R( $\psi_{norm}$ )')
    plt.savefig(params['output_dir']+'/R_vs_psi_grid.pdf')
    plt.close()

    alphas = np.linspace( 0.0, 2.0*pi, params['nalpha_map'] )

    # First index is rho, Second index is alpha
    # Set rho = 0 to (r_axis,z_axis) for all alpha
    Rmap = [ [gfile.r_axis for x in range(params['nalpha_map']) ] ]
    Zmap = [ [gfile.z_axis for x in range(params['nalpha_map']) ] ]

    # Get contours
    print( 'Making Contours' )

    for iR, R in enumerate(Rs):
        
        print( iR+1 , ' / ' , len(Rs) )

        if( Psins[iR+1] < params['rho_low'] ):
            Rm, Zm = PsiContour(gfile,R, params['step_low'],params['nalpha_map'], params['nalpha_downsample'])
        else:
            Rm, Zm = PsiContour(gfile,R,params['step_high'],params['nalpha_map'], params['nalpha_downsample'])

        Rmap.append( Rm )
        Zmap.append( Zm )

    # Construct Bicubic splines
    params['R_map'] = RectBivariateSpline(Psins, alphas, np.array(Rmap), kx=3, ky=3, s=0)
    params['Z_map'] = RectBivariateSpline(Psins, alphas, np.array(Zmap), kx=3, ky=3, s=0)

    # Plot coordinate charts
    plot_charts( Psins, alphas, np.array(Rmap), np.array(Zmap), params )

# Diagnostic check ---------------------------------------------------------------

def check_flux_surface(rho,params):
    '''
    Checks for a single flux surface that grad_rho / mod(grad_rho) is a unit vector
    pointing orthogonally to the flux surface.
    '''

    nalphas = 10000

    cont_Rs = np.zeros(nalphas)
    cont_Zs = np.zeros(nalphas)

    drhodRs = np.zeros(nalphas)
    drhodZs = np.zeros(nalphas)

    dRdalphas = np.zeros(nalphas)
    dZdalphas = np.zeros(nalphas)

    modgradrhos = np.zeros(nalphas)

    # Construct a contour
    for ialpha in range(nalphas):
        
        alpha = 2.0 * pi * ( ialpha / (nalphas - 1.0) )

        cont_Rs[ialpha], cont_Zs[ialpha] = get_R_Z( rho, alpha, params )

        dRdrho, dRdalpha, dZdrho, dZdalpha = derivatives_R_Z( rho, alpha, params )

        dRdalphas[ialpha] = dRdalpha
        dZdalphas[ialpha] = dZdalpha

        drhodR, dalphadR, drhodZ, dalphadZ = derivatives_rho_alpha( rho, alpha, params )

        drhodRs[ialpha] = drhodR
        drhodZs[ialpha] = drhodZ
        
        modgradrhos[ialpha] = mod_grad_rho( rho, alpha, params )

    # Loop round and check orthonormality

    max_dot_product_numeric  = 0.0
    max_dot_product_analytic = 0.0

    for ialpha in range(nalphas-1):

        # Arc element
        dl = np.array( [ cont_Rs[ialpha+1] - cont_Rs[ialpha], cont_Zs[ialpha+1] - cont_Zs[ialpha] ] )
        dl = dl / ( dl[0]**2.0 + dl[1]**2.0 )**0.5

        # Alpha contravariant basis vector
        dRZdalpha = np.array( [ dRdalphas[ialpha], dZdalphas[ialpha] ] )

        # Normalised Rho Gradient Vector
        drho = np.array( [ drhodRs[ialpha], drhodZs[ialpha] ] )

        # Calculate dot product using calculated arc length
        numeric_dot = ( dl[0]*drho[0] + dl[1]*drho[1] ) / modgradrhos[ialpha]

        # Calculate analytic dot product
        analytic_dot = ( dRZdalpha[0]*drho[0] + dRZdalpha[1]*drho[1] )

        if( abs(numeric_dot) > max_dot_product_numeric ):
            max_dot_product_numeric = abs(numeric_dot)

        if( abs(analytic_dot) > max_dot_product_analytic ):
            max_dot_product_analytic = abs(analytic_dot)

    print( 'Maximum numeric dot product = ',  max_dot_product_numeric )
    print( 'Maximum analytic dot product = ', max_dot_product_analytic )
    print()
    input("Press Enter to continue...")
    print()

def plot_charts( Psin, Alpha, R, Z, params ):
    '''
    Plots coordinate charts generated from numerical equilibrium.
    '''

    fig, (ax1,ax2) = plt.subplots(ncols=2)

    CS = ax1.contourf(Psin, Alpha, R.transpose(), 20, cmap=plt.cm.OrRd)
    cbar = fig.colorbar(CS,ax=ax1)

    ax1.set_title('R [m]')
    ax1.set_xlabel('Rho')
    ax1.set_ylabel('Alpha')

    CS = ax2.contourf(Psin, Alpha, Z.transpose(), 20, cmap=plt.cm.OrRd)
    cbar = fig.colorbar(CS,ax=ax2)

    ax2.set_title('Z [m]')
    ax2.set_xlabel('Rho')
    ax2.set_ylabel('Alpha')

    plt.savefig(params['output_dir']+'/charts.pdf')

    if( params['show-plots'] ):
        plt.show()

    plt.close()