plotting.py

from math         import exp, pi, floor
from geometry     import get_R_Z, get_flux_surfaces
from matplotlib   import pyplot as plt
from profiles     import profiles, evaluate_all_profiles
from reactivities import bosch_hale_reactivity, sadler_van_belle_reactivity
from velocity     import sample_isotropic
from mpl_toolkits import mplot3d

import numpy as np

def plot_plasma_profiles( params, save=None, show=False ):
    '''
    Plots plasma density, temperature and reactivity against normalised minor radius.
    '''

    rho_array, n_array, T_array, R_array = evaluate_all_profiles( params )
    
    fig, (ax1,ax2,ax3) = plt.subplots(3, 1,sharex=True)
    
    ax1.plot(rho_array, n_array, color='k', linewidth=1.2)
    ax1.set_ylabel(r'$N_{i}$ $[m^{-3}]$')
    
    ax2.plot(rho_array, T_array, color='k', linewidth=1.2)
    ax2.set_ylabel(r'$T_{i}$ $[keV]$')
    
    ax3.plot(rho_array, R_array, color='k', linewidth=1.2)
    ax3.set_xlabel(r'$\rho$')
    ax3.set_ylabel(r'$R_{vol}$ $[m^{-3}s^{-1}]$')

    if( save is not None ):
        plt.savefig(save)

    if( show ):
        plt.show()

    plt.close()

def plot_plasma_profiles_2D( params, save=None, show=False ):
    '''
    Plots plasma density and temperature on a poloidal cross section.
    '''

    # Define grid resolution
    R_resolution = 50
    Z_resolution = 50

    # Loop over (rho,alpha) with fine spacing and bin function values to average
    nres_rho   = 1000
    nres_alpha = 1000 

    # First generate some contours
    ncontours = 5
    cont_Rs, cont_Zs = get_flux_surfaces(ncontours,params)

    # Get Bounds
    Rmin = np.amin( cont_Rs ) * 0.90
    Rmax = np.amax( cont_Rs ) * 1.05
    Zmin = np.amin( cont_Zs ) * 1.10
    Zmax = np.amax( cont_Zs ) * 1.10 

    dR = (Rmax-Rmin) / float(R_resolution)
    dZ = (Zmax-Zmin) / float(Z_resolution)

    # Get axes
    Raxis = np.linspace( Rmin + 0.5*dR, Rmax - 0.5*dR, R_resolution )
    Zaxis = np.linspace( Zmin + 0.5*dZ, Zmax - 0.5*dZ, Z_resolution )

    # Build profile arrays
    ni = np.zeros((R_resolution,Z_resolution))
    Ti = np.zeros((R_resolution,Z_resolution))

    # Count points in each cell for averaging
    count = np.zeros((R_resolution,Z_resolution))

    for i in range(nres_rho):

        rho = 1.0 * i / (float(nres_rho) - 1.0)

        for j in range(nres_alpha):

            alpha = 2.0 * pi * j / (float(nres_alpha) - 1.0)
            
            # Get R, Z values
            R, Z = get_R_Z( rho, alpha, params )

            # Get profile values
            ni_, Ti_ = profiles( rho, params )

            # Bin values
            iR = floor( (R - Rmin) / dR )
            iZ = floor( (Z - Zmin) / dZ )

            ni[iR,iZ] += ni_
            Ti[iR,iZ] += Ti_

            count[iR,iZ] += 1.0

    # Calculate average values
    ni = ni / count
    Ti = Ti / count

    # Plot density
    fig, ax = plt.subplots()
    fig.set_figwidth(3.5)
    CS = ax.contourf(Raxis, Zaxis, ni.transpose(), 20, cmap=plt.cm.OrRd)
    cbar = fig.colorbar(CS)
    ax.set_aspect('equal')

    # Plot contours
    for i in range(ncontours-1):
        plt.plot( cont_Rs[i,:], cont_Zs[i,:], '--', c='black')
    plt.plot( cont_Rs[-1,:], cont_Zs[-1,:], c='black')

    plt.xlabel('R [m]')
    plt.ylabel('Z [m]')

    if( save ):
        plt.savefig(params['output_dir']+'/ni_2D.pdf')

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

    plt.close()

    # Plot temperature
    fig, ax = plt.subplots()
    fig.set_figwidth(3.5)
    CS = ax.contourf(Raxis, Zaxis, Ti.transpose(), 20, cmap=plt.cm.OrRd)
    cbar = fig.colorbar(CS)
    ax.set_aspect('equal')

    # Plot contours
    for i in range(ncontours-1):
        plt.plot( cont_Rs[i,:], cont_Zs[i,:], '--', c='black')
    plt.plot( cont_Rs[-1,:], cont_Zs[-1,:], c='black')

    plt.xlabel('R [m]')
    plt.ylabel('Z [m]')

    if( save ):
        plt.savefig(params['output_dir']+'/Ti_2D.pdf')

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

    plt.close()

def plot_samples( R_array, Z_array, params, save=False ):
    '''
    Plots sampled points on equilibrium.
    '''

    fig, ax = plt.subplots()
    fig.set_figwidth(3.5)

    # First generate some contours
    ncontours = 5
    cont_Rs, cont_Zs = get_flux_surfaces(ncontours,params,nalpha=100)

    # Get Bounds
    Rmin = min( np.amin( cont_Rs ), np.amin( R_array ) ) * 0.90
    Rmax = max( np.amax( cont_Rs ), np.amax( R_array ) ) * 1.05
    Zmin = min( np.amin( cont_Zs ), np.amin( Z_array ) ) * 1.10
    Zmax = max( np.amax( cont_Zs ), np.amax( Z_array ) ) * 1.10 

    print()
    print('Generating plot of samples...')
    print()

    # Define grid resolution
    R_resolution = 50
    Z_resolution = 50

    dR = (Rmax-Rmin) / float(R_resolution)
    dZ = (Zmax-Zmin) / float(Z_resolution)

    # Get axes
    Raxis = np.linspace( Rmin + 0.5*dR, Rmax - 0.5*dR, R_resolution )
    Zaxis = np.linspace( Zmin + 0.5*dZ, Zmax - 0.5*dZ, Z_resolution )

    # Build array
    rates = np.zeros((R_resolution,Z_resolution))

    for i in range(len(R_array)):

        R = R_array[i]
        Z = Z_array[i]

        iR = floor( (R - Rmin) / dR )
        iZ = floor( (Z - Zmin) / dZ )

        rates[iR,iZ] += 1

    # Normalise to 1 
    rates = rates / np.max(rates)

    # Plot density
    CS = ax.contourf(Raxis, Zaxis, rates.transpose(), 20, cmap=plt.cm.OrRd)
    ax.set_aspect('equal')
    cbar = fig.colorbar(CS)

    # Plot contours
    for i in range(ncontours-1):
        plt.plot( cont_Rs[i,:], cont_Zs[i,:], '--', c='black')
    plt.plot( cont_Rs[-1,:], cont_Zs[-1,:], c='black')

    plt.xlabel('R [m]')
    plt.ylabel('Z [m]')

    if( save ):
        plt.savefig(params['output_dir']+'/samples.pdf')

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

    plt.close()

def plot_cdf(grid, cdf, xlabel, title, save=None, show=False):
    '''
    Just plots a single CDF.
    '''
    
    plt.figure()
    plt.plot( grid, cdf )
    plt.xlabel(xlabel)
    plt.ylabel('CDF')
    plt.title(title)

    if( show ):
        plt.show()

    if( save is not None ):
        plt.savefig(save)
    
    plt.close()

def check_cdf(rho_grid, original_cdf, sampled_cdf, params, save=True):
    '''
    Checks that the sampled CDF against rho matches the original CDF calculated from the reactivity.
    '''
    
    plt.figure()
    plt.plot(rho_grid,  sampled_cdf, 'black', label= 'Sampled CDF')
    plt.plot(rho_grid, original_cdf,  'blue', label='Original CDF')
    plt.xlabel('Rho')
    plt.ylabel('CDF')
    plt.legend()

    if( save ):
        plt.savefig(params['output_dir']+'/CDF_check.pdf')

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

    plt.close()


def plot_reactivity(params,save=False):
    '''
    Plots the reactivity being used on linear and log scales.
    '''

    Ts  = np.linspace(0,100,101)

    if( params['bosch-hale'] ):
        Rs = np.array( [ bosch_hale_reactivity( x )       for x in Ts ])
        label = 'Bosch-Hale Reactivity'
    else:
        Rs = np.array( [ sadler_van_belle_reactivity( x ) for x in Ts ])
        label = 'Sadler-Van-Belle Reactivity'

    # Linear plot

    plt.figure()
    plt.plot(Ts, Rs, color='black', linewidth=1.2)
    plt.xlabel(r'$T_{i} [keV]$')
    plt.ylabel(r'$<\sigma v>$ $[m^{3}s^{-1}]$')
    plt.title(label)

    if( save ):
        plt.savefig( params['output_dir']+'/Reactivity.pdf' )

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

    plt.close()

    # Log plot

    plt.figure()
    plt.plot(Ts, Rs, color='black', linewidth=1.2)
    plt.xlabel(r'$T_{i} [keV]$')
    plt.ylabel(r'$<\sigma v>$ $[m^{3}s^{-1}]$')
    plt.title(label)
    plt.yscale('log')

    if( save ):
        plt.savefig( params['output_dir']+'/Reactivity_log.pdf' )

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

    plt.close()

def compare_reactivities(params,save=False):
    '''
    Compares the Bosch-Hale and Sadler-Van-Belle reactivity fits.
    Also plots data included in the Bosch-Hale paper to confirm correctness. 
    '''

    Ts  = np.linspace(0,100,101)
    BH  = np.array( [ bosch_hale_reactivity( x )       for x in Ts ])
    SVB = np.array( [ sadler_van_belle_reactivity( x ) for x in Ts ])

    # Paper quoted values
    BH_x = [0.2,0.3,0.4,0.5,0.6,0.7,0.8,1.0,1.3,1.5,1.8,2.0,2.5,3.0,4.0,5.0,6.0,8.0,10.0,12.0,15.0,20.0,30.0,40.0,50.0]
    BH_y = [1.254e-26,7.292e-25,9.344e-24,5.697e-23,2.253e-22,6.740e-22,1.662e-21,6.857e-21,2.546e-20,6.923e-20,1.539e-19,2.977e-19,8.425e-19,
            1.867e-18,5.974e-18,1.366e-17,2.554e-17,6.222e-17,1.136e-16,1.747e-16,2.740e-16,4.330e-16,6.681e-16,7.998e-16,8.649e-16]

    BH_y = [ x/1.0e6 for x in BH_y]

    # Linear plot

    plt.figure()
    plt.plot(Ts,  BH,   color='black', linewidth=1.2, label='Bosch-Hale')
    plt.plot(BH_x,BH_y, color='red',   linewidth=1.2, label='Bosch-Hale (paper)')
    plt.plot(Ts, SVB,   color='blue',  linewidth=1.2, label='Sadler-Van-Belle')


    plt.xlabel(r'$T_{i} [keV]$')
    plt.ylabel(r'$<\sigma v>$ $[m^{3}s^{-1}]$')
    plt.legend(loc="lower right")

    if( save ):
        plt.savefig( params['output_dir']+'/Reactivity_comparison.pdf' )

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

    plt.close()

    # Log plot

    plt.figure()
    plt.plot(Ts,  BH, color='black', linewidth=1.2, label='Bosch-Hale')
    plt.plot(BH_x,BH_y, color='red',   linewidth=1.2, label='Bosch-Hale (paper)')
    plt.plot(Ts, SVB, color='blue',  linewidth=1.2, label='Sadler-Van-Belle')

    plt.xlabel(r'$T_{i} [keV]$')
    plt.ylabel(r'$<\sigma v>$ $[m^{3}s^{-1}]$')
    plt.legend(loc="lower right")
    plt.yscale('log')

    if( save ):
        plt.savefig( params['output_dir']+'/Reactivity_comparison_log.pdf' )

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

    plt.close()

def test_isotropic_sampling():
    '''
    Test isotropic sampling by sampling and plotting some points.
    '''

    fig = plt.figure()
    ax = plt.axes(projection='3d')

    x = []
    y = []
    z = []

    for i in range(1000):

        v = sample_isotropic()
        x.append(v[0])
        y.append(v[1])
        z.append(v[2])

    ax.scatter(x, y, z, marker='o')

    ax.set_xlabel('Vx')
    ax.set_ylabel('Vy')
    ax.set_zlabel('Vz')

    plt.show()
    plt.close()