# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Boeuf results

# +
# Scientific libraries
import numpy as np
import scipy


# Graphic libraries

import matplotlib.pyplot as plt
# %matplotlib widget

plt.style.use("presentation")
plt.rcParams["figure.figsize"] = (4, 3)

# Creating alias for magic commands

# LPPview Classes
from LPPview import *
from LPPview.Classes.LPPic_temporal import History
from LPPview.Classes.LPPic_fields import field
from LPPview.Classes.LPPic_fields import field as Field
from LPPview.Classes.LPPic_newwalls import newwalls as Newwalls
from LPPview.Classes.LPPic_temporal import History

from plasmapy.physics import Debye_length
from plasmapy.physics import plasma_frequency
from astropy import units as u

from scipy.ndimage import gaussian_filter1d as smooth
from scipy.ndimage import gaussian_filter

from tqdm import tqdm_notebook as tqdm

# +
path_ref = "/DATA/tavant/266_Boeuf_166Thomas/"
path_L2 = "/DATA/tavant/158_Beauf_fakeR/"
path_L4 = "/DATA/tavant/163_Beauf_fakeR2/"
paths = [path_ref, path_L4, path_L2]
names = ["no $L_R$", "$L_R$=4cm", "$L_R$=2cm"]
colors = ["k","b", "r"]
fields = [Field(path) for path in paths ]
histories = [History(path)for path in paths ]
walls = [Newwalls(path) for path in paths ]


# -

# # Wave energy density 

def return_wave_energy(field : fields, i=-1):
    """Return the wave energy:
    $ W = 3 eps_0 | d E|^2 $
    
    By estimating dE = std(E) \sqrt{2} """
    tab = field.return_fromkey(i, "Ej(1)")
    
    std_E = tab.std(axis=0)
    
    return 3*field.eps0*2*std_E**2


# +
fig, ax = plt.subplots(1, 1, figsize=(3.5, 3))

vmax = 0
for f, n in zip(fields, names):
    value = return_wave_energy(f, -1)

    plt.plot(value, label=n)

ax.legend()    
ax.set_xlabel("Axial position $z$ [cm]")
ax.set_ylabel("Wave energy $W$ [SI]")

#plt.savefig("Boeuf_Ex_snapshot.png", dpi=400)
# -

# # Wave amplitude evolution

def return_W_2D_zt(field, nt_max=None):
    """return the spatio_temporal evolution of the wave energy"""
    if nt_max is None:
        nt_max = field._nT
        
    w_2D = np.zeros( shape=( nt_max, field._ymax+1))
    
    for i in tqdm(range(nt_max)):
        value = return_wave_energy(field, i)
        w_2D[i] = value
        
    return w_2D


import ipywidgets as widgets
widgets.IntSlider()

# +
nt_max = 9999
for f in fields:
    nt_max = min(nt_max, f._nT)
    
w2D_0, w2D_1, w2D_2 = [return_W_2D_zt(f, nt_max) for f in fields]


# +
def W_2D_evolution(smooth_gaussian_2D=False):

    fig, axarr = plt.subplots(1, 3, figsize=(8, 3))

    icut = 10

    nt_max = 900
    for f in fields:
        nt_max = min( nt_max, f._nT)

    datas = [w2D_0, w2D_1, w2D_2]

    if smooth_gaussian_2D :
        datas = [ gaussian_filter(d, sigma=2, order=0) for d in datas]

    vmax = max( [ np.ma.masked_invalid(np.abs(d)).max() for d in datas])
    print(vmax)

    vmax *= 0.5
    for f, ax, n, data in zip(fields, axarr, names, datas):
        f.definecoords()


        im =ax.imshow(data/data.max(), extent=(0,f._Ly*100, 0, nt_max*f._dT*f._Na*1e6), vmin=0, vmax=1, cmap="plasma")
        ax.set_title(n, fontsize=12)


        ax.set_xlabel("Axial position $z$ [cm]")


    axarr[0].set_ylabel("Time [$\mu$s]")
    cb = plt.colorbar(im, ax=axarr[2], fraction=0.05, aspect=30, ticks=[0,1])
    cb.ax.set_ylabel(" Wave energy $W$ ")
    #cb.ax.set_yticks([0,1])
    cb.ax.set_yticklabels(["min", "max"])
    fig.suptitle("Axial-temporal evolution of $W$", fontsize=14)

#plt.savefig("Boeuf_Ex_snapshot.png", dpi=400)
W_2D_evolution()
plt.savefig("Boeuf_Fake_R_W.png", dpi=400)


W_2D_evolution(True)
plt.savefig("Boeuf_Fake_R_W_smoothed.png", dpi=400)


# -

axial_velocity = ( 2-   1.5 ) / (7.94-7.62) * u.cm/u.microsecond
axial_velocity.to(u.m/u.second)


# +
def W_2D_evolution(smooth_gaussian_2D=False):

    fig, ax = plt.subplots(1, 1, figsize=(8, 3))

    icut = 10

    nt_max = 900
    for f in fields:
        nt_max = min( nt_max, f._nT)

    datas = [w2D_0, w2D_1, w2D_2]

    if smooth_gaussian_2D :
        datas = [ gaussian_filter(d, sigma=2, order=0) for d in datas]

    vmax = max( [ np.ma.masked_invalid(np.abs(d)).max() for d in datas])
    print(vmax)

    vmax *= 0.5
    for f, n, data in zip(fields, names, datas):
        f.definecoords()

        ax.plot(f.tab_y, data[ 300:400, :].mean(axis=0), label=n)
        #im =ax.imshow(data/data.max(), extent=(0,f._Ly*100, 0, nt_max*f._dT*f._Na*1e6), vmin=0, vmax=1, cmap="plasma")
    ax.legend( fontsize=12)


    ax.set_xlabel("Axial position $z$ [cm]")


    ax.set_ylabel(" Wave energy $W$ ")
    fig.suptitle("Axial evolution of $W$", fontsize=14)

#plt.savefig("Boeuf_Ex_snapshot.png", dpi=400)
W_2D_evolution()


W_2D_evolution(True)


# -

# ## Temporal evolution

# +
def W_evolution(smooth_gaussian_2D=False, z_cm_list=[0.4, 0.7, 1.5], scaled=False):

    fig, axarr = plt.subplots(1, 3, figsize=(8, 3))

    icut = 10

    nt_max = 900
    for f in fields:
        nt_max = min( nt_max, f._nT)
    tab_time = np.arange(0, nt_max)*fields[0]._dT*fields[0]._Na*1e6
    
    datas = [w2D_0, w2D_1, w2D_2]

    if smooth_gaussian_2D :
        datas = [ gaussian_filter(d, sigma=2, order=0) for d in datas]

    z_cm_list = [0.4, 0.8, 1.5]
    z_index_list = [int( z*1e-2 / fields[0]._dX) for z in z_cm_list]
    print(z_index_list)
    
    dat_max = 0
    for f, ax, n, data in zip(fields, axarr, names, datas):
        f.definecoords()

        for ax, z in zip(axarr, z_index_list):
            ax.plot(tab_time, data[:, z], label=n)
            dat_max = max(dat_max, data[:, z].max())
            
    for ax in axarr:
        ax.set_xlabel("Time [$\mu$s]")
        ax.set_xlim(0, 10)
        
        if scaled:
            ax.set_ylim(top=dat_max*1.1)
     
    for ax, z in zip(axarr, z_cm_list):
        ax.annotate(f"z={z:2.1f} cm", (6.5, 0.035))
        
    axarr[0].set_ylabel("Wave energy W [SI]")
    
    axarr[1].legend(loc='upper center', bbox_to_anchor=(0.5, 1.2),
                    ncol=3,
                   # fancybox=True,
                   # shadow=True,
                   )
    plt.subplots_adjust(wspace=0.3, bottom=0.2, top=0.85)    
    
#    fig.suptitle("Axial-temporal evolution of $W$", fontsize=14)
W_evolution(scaled=True)
plt.savefig("Boeuf_Fake_R_W_temporal_smoothed.png", dpi=400)

W_evolution(smooth_gaussian_2D=True, scaled=True)

plt.savefig("Boeuf_Fake_R_W_temporal_smoothed.png", dpi=400)

# -



# # Wave group velocity

# +

def return_vg(field:fields, i=-1):
    """The group velocity is the axial ion velocity"""
    
    Ji = field.return_fromkey(-1, "Ji(2)")
    ni = field.return_fromkey(-1, "Numi")
    
    vi = Ji/(ni * field.qe)
    
    return vi.mean(axis=0)

fig, ax = plt.subplots(1, 1, figsize=(3.5, 3))

vmax = 0
for f, n in zip(fields, names):
    value = return_vg(f, -1)

    plt.plot(f.tab_y, value*1e-3, label=n)

ax.set_xlim(0, 2.5)    
ax.set_xlabel("Axial position $z$ [cm]")
ax.set_ylabel("Group velocity $v_g$ [km/s]")

plt.legend()

plt.savefig("Boeuf_v_g.png", dpi=400)

ax.axhline(axial_velocity.to(u.m/u.second).value*1e-3, linestyle="--", c="k", label="Axial velocity of $W$")
plt.legend()

plt.savefig("Boeuf_v_g_with_observed.png", dpi=400)

# -

# # Growth rate from PIC

# +
def return_growth_rate_axial(field, i=-1, smooth_W=False, smooth_sigma=2):
    """Compute the growth rate \gamma by Eq. 20 of Lafleur 2018:
                     d_t W + \div( v_g W)
            \gamma = --------------------
                          2 W
    
    - d_t W is computed by upwind difference,
     - \div(v_g W) is computed by center difference
    """
    
    if i == -1 or i == field._nT-1:
        i1 = i - 1
        i2 = i
    else:
        i1 = i
        i2 = i+1
        
    
    W1 = return_wave_energy(field, i1)
    
    try:
        W2 = return_wave_energy(field, i2)
    except IndexError:
        print(i2, field._nT)
        raise RuntimeError
    
    if smooth_W :
        W1 = smooth(W1, sigma=smooth_sigma)
        W2 = smooth(W2, sigma=smooth_sigma)
    
    vg = return_vg(field, i1)
    
    dt_W = ( W2 - W1)/ (field._dT*field._Na)
    
    div_vgW = np.gradient(vg*W1, field._dX)
    
    return ( dt_W + div_vgW)/(2 * W1), dt_W
    
    
# -

# ## Axial profile at t=10 $\mu$s

# +
fig, ax = plt.subplots(1, 1, figsize=(3.5, 3))

icut = 10

nt0 = 400
for f, n in zip(fields, names):
    f.definecoords()
    
    value, dt_W = return_growth_rate_axial(f, nt0, smooth_W=True, smooth_sigma=5)
    Nt_av = 100
    print(f"averaged over {Nt_av*f._dT*f._Na*1e6:f} $\mu$s")  

    for i in range(1,Nt_av):
        a, b= return_growth_rate_axial(f, nt0-i, smooth_W=True, smooth_sigma=5)
        
        value += a
        dt_W += b
        
    value /= Nt_av
    dt_W /= Nt_av
    
    
    if True:
        value = smooth(value, sigma=1)
    
    plt.plot(f.tab_y[icut: -1-icut], value[icut: -1-icut]*1e-6, label=n)


ax.set_xlabel("Axial position $z$ [cm]")
ax.set_ylabel("Growth rate $\gamma$ [$\mu$s$^{-1}$]")
plt.legend()
#plt.savefig("Boeuf_Ex_snapshot.png", dpi=400)
# -

# ## Temporal evolution

# +

def temporal_evolution_2D(field, smooth_W=True, smooth_sigma=5, nt_max = None):
    """return the temporal evolution as a 2D table Z-t"""
    
    if nt_max is None:
        nt_max = field._nT
        
    growthrate_2D = np.zeros((nt_max, field._ymax+1))
    dtW_2D = np.zeros((nt_max, field._ymax+1))
    
    for i in tqdm(range(nt_max)):
        gamma, dt_W = return_growth_rate_axial(field, i, smooth_W=smooth_W, smooth_sigma=smooth_sigma)
        growthrate_2D[i] = gamma
        dtW_2D[i] = dt_W
        
    return growthrate_2D, dtW_2D
    
    
def temporal_evolution(field, z_cm=[0.7]):
    """return the temporal evolution at the positions choosen"""
    pass


# +
nt_max = 900
for f in fields:
    nt_max = min( nt_max, f._nT)
        
growthrate_2D_0, dtW_2D_0 =  temporal_evolution_2D(fields[0], smooth_W=True, smooth_sigma=2, nt_max=nt_max)
growthrate_2D_1, dtW_2D_1 =  temporal_evolution_2D(fields[1], smooth_W=True, smooth_sigma=2, nt_max=nt_max)
growthrate_2D_2, dtW_2D_2 =  temporal_evolution_2D(fields[2], smooth_W=True, smooth_sigma=2, nt_max=nt_max)


# -

# # 2D Axial-temporal evolutions

# ## $\partial_t W$

# +
def dtW_2D_evolution(smooth_gaussian_2D=False):

    fig, axarr = plt.subplots(1, 3, figsize=(8, 3))

    icut = 10

    nt_max = 900
    for f in fields:
        nt_max = min( nt_max, f._nT)

    datas = [dtW_2D_0, dtW_2D_1, dtW_2D_2]

    if smooth_gaussian_2D :
        datas = [ gaussian_filter(d, sigma=2, order=0) for d in datas]

    vmax = min( [ np.ma.masked_invalid(np.abs(d)).max() for d in datas])
    print(vmax)

    vmax *= 1
    for f, ax, n, data in zip(fields, axarr, names, datas):
        f.definecoords()


        im =ax.imshow(data, extent=(0,f._Ly*100, 0, nt_max*f._dT*f._Na*1e6), vmin=-vmax, vmax=vmax, cmap="seismic")
        ax.set_title(n, fontsize=12)


        ax.set_xlabel("Axial position $z$ [cm]")


    axarr[0].set_ylabel("Time [$\mu$s]")
    cb = plt.colorbar(im, ax=axarr[2], fraction=0.05, aspect=30)
    cb.ax.set_ylabel(" $\partial_t W$  [SI]")
    fig.suptitle("Axial-temporal evolution of $\partial_t W$", fontsize=14)

#plt.savefig("Boeuf_Ex_snapshot.png", dpi=400)
dtW_2D_evolution()
plt.savefig("Boeuf_Fake_R_dtW.png", dpi=400)


dtW_2D_evolution(True)
plt.savefig("Boeuf_Fake_R_dtW_smoothed.png", dpi=400)



# +
fig, axarr = plt.subplots(1, 3, figsize=(8, 3))
smooth_gaussian_2D = True

datas = datas = [dtW_2D_0, dtW_2D_1, dtW_2D_2]

if smooth_gaussian_2D :
    datas = [ gaussian_filter(d, sigma=5, order=0) for d in datas]
    
tab_time = np.arange(0, nt_max)*fields[0]._dT*fields[0]._Na*1e6

z_cm_list = [0.4, 0.7, 1.5]
z_index_list = [int( z*1e-2 / fields[0]._dX) for z in z_cm_list]

for f, ax, n, data in zip(fields, axarr, names, datas):
    f.definecoords()

    for ax, z in zip(axarr, z_index_list):
        ax.plot(tab_time, 1e-3*data[:, z], label=n)

for ax, z in zip(axarr, z_cm_list):
    ax.set_title(f"Axial position $z={z}$ [cm]", fontsize=11)
        
for ax in axarr:
    ax.legend(labelspacing=0.02)
    ax.set_xlabel("Time $t$ [$\\mu$s]")
    
axarr[0].set_ylabel("$\partial_t W$ [10$^{-3}$ SI]")
plt.savefig("Boeuf_dtW_toporal.pdf")


# -

# ## Growth rate $\gamma$

# +

def growth_rate_2D_evolution(smooth_gaussian_2D=False, datas=[growthrate_2D_0, growthrate_2D_1, growthrate_2D_2]):
    fig, axarr = plt.subplots(1, 3, figsize=(8, 3))

    icut = 10

    nt_max = 900
    for f in fields:
        nt_max = min( nt_max, f._nT)
        
   

    if smooth_gaussian_2D :
        datas = [ gaussian_filter(d, sigma=2, order=0) for d in datas]
    vmax = min( [ np.ma.masked_invalid(d).max() for d in datas ])
    print(vmax)

    vmax *= 1
    for f, ax, n, data in zip(fields, axarr, names, datas):
        f.definecoords()


        im =ax.imshow(data*1e-6, extent=(0,f._Ly*100, 0, nt_max*f._dT*f._Na*1e6), vmin=-vmax*1e-6, vmax=vmax*1e-6, cmap="seismic")
        ax.set_title(n, fontsize=12)


        ax.set_xlabel("Axial position $z$ [cm]")


    axarr[0].set_ylabel("Time [$\mu$s]")
    cb = plt.colorbar(im, ax=axarr[2], fraction=0.05, aspect=30)
    cb.ax.set_ylabel("Growth rate $\gamma$  [$\mu$s$^{-1}$]")
    fig.suptitle("Axial-temporal evolution of $\gamma$", fontsize=14)


growth_rate_2D_evolution()
plt.savefig("Boeuf_Fake_R_gamma.png", dpi=400)

growth_rate_2D_evolution(True)
plt.savefig("Boeuf_Fake_R_gamma_smoothed.png", dpi=400)


# +
def gamma_evolution(smooth_gaussian_2D=False, z_cm_list=[0.4, 0.7, 1.5], scaled=False, datas=[growthrate_2D_0, growthrate_2D_1, growthrate_2D_2]):

    fig, axarr = plt.subplots(1, 3, figsize=(8, 3))

    icut = 10

    nt_max = 900
    for f in fields:
        nt_max = min( nt_max, f._nT)
    tab_time = np.arange(0, nt_max)*fields[0]._dT*fields[0]._Na*1e6
    
    
    
    if smooth_gaussian_2D :
        datas = [ gaussian_filter(d, sigma=2, order=0) for d in datas]
        
    datas = [ 1e-6*d for d in datas]

    z_cm_list = [0.4, 0.8, 1.5]
    z_index_list = [int( z*1e-2 / fields[0]._dX) for z in z_cm_list]
    print(z_index_list)
    
    dat_max = 0
    dat_min = 0
    for f, ax, n, data in zip(fields, axarr, names, datas):
        f.definecoords()

        for ax, z in zip(axarr, z_index_list):
            ax.plot(tab_time, data[:, z], label=n)
            dat_max = max(dat_max, np.ma.masked_invalid(data[:, z]).max())
            dat_min = min(dat_max, np.ma.masked_invalid(data[:, z]).min())
            
    for ax in axarr:
        ax.set_xlabel("Time [$\mu$s]")
        ax.set_xlim(0, 10)
        
        if scaled:
            ax.set_ylim(bottom=1.1*dat_min, top=dat_max*1.1)
     
    for ax, z in zip(axarr, z_cm_list):
        ax.annotate(f"z={z:2.1f} cm", (5, 5))
        
    axarr[0].set_ylabel("Growth rate $\gamma$ [$\mu$s$^{-1}$]")
    
    axarr[1].legend(loc='upper center', bbox_to_anchor=(0.5, 1.2),
                    ncol=3,
                   # fancybox=True,
                   # shadow=True,
                   )
    plt.subplots_adjust(wspace=0.3, bottom=0.2, top=0.85)    
    
#    fig.suptitle("Axial-temporal evolution of $W$", fontsize=14)
gamma_evolution(scaled=True)
plt.savefig("Boeuf_Fake_R_gamma_temporal.pdf")

gamma_evolution(smooth_gaussian_2D=True, scaled=True)
plt.savefig("Boeuf_Fake_R_gamma_temporal_smoothed.pdf")


# -

# ## Axial profile of $\gamma$

# +
def gamma_axial_evolution(smooth_gaussian_2D=False, scaled=False, tmin_mus=9, plot_zpos=False, datas =  [growthrate_2D_0, growthrate_2D_1, growthrate_2D_2]):

    fig, ax = plt.subplots(1, 1, figsize=(3.5, 3))

    icut = 10

    nt_max = 900
    for f in fields:
        nt_max = min( nt_max, f._nT)
        
    tmin = int( tmin_mus *nt_max/10)

    
    
    if smooth_gaussian_2D :
        datas = [ gaussian_filter(d, sigma=2, order=0) for d in datas]
        
    datas = [ 1e-6*d for d in datas]
    to_pots = [ d[tmin:, :].mean(axis=0) for d in datas]

    z_cm_list = [0.4, 0.8, 1.5]
    z_index_list = [int( z*1e-2 / fields[0]._dX) for z in z_cm_list]

    
    dat_max = 0
    dat_min = 0
    
    for f, n, data in zip(fields, names, to_pots):
        f.definecoords()
        ax.plot(f.tab_y, data, label=n)
        
        dat_max = max(dat_max, np.ma.masked_invalid(data).max())
        dat_min = min(dat_max, np.ma.masked_invalid(data).min())
            
    ax.set_xlabel("Axial position $z$ [cm]")
    ax.set_xlim(0, 2.5)
        
    if scaled:
        ax.set_ylim(bottom=-0.5*dat_max, top=dat_max*1.1)
     
    if plot_zpos:
        for z in z_cm_list:
            ax.axvline(z, c='k', linestyle="--", alpha=0.5)
        
    ax.set_ylabel("Growth rate $\gamma$ [$\mu$s$^{-1}$]")
    
    #ax.legend(loc='upper center', bbox_to_anchor=(0.5, 1.2),
    #               ncol=3,
    #               # fancybox=True,
    #               # shadow=True,
    #               )
    ax.legend()
    #plt.subplots_adjust(wspace=0.3, bottom=0.2, top=0.85)    
    
#    fig.suptitle("Axial-temporal evolution of $W$", fontsize=14)
gamma_axial_evolution(scaled=True, tmin_mus=3)
plt.savefig("Boeuf_Fake_R_gamma_axial.pdf")

gamma_axial_evolution(smooth_gaussian_2D=True, scaled=True, tmin_mus=3)
plt.savefig("Boeuf_Fake_R_gamma_axial_smoothed.pdf")


# -

# # Integration of the growth rate
#
# In the stationnary case, we have
# $$ 2 \gamma W = \nabla \cdot ( \vec{v_g} W) $$
#
# Lets see if we can integrate it, and obtaine the same mean profile.
#
# Unfortunatly, we need a _seed_ to start.

def integrate_gamma(field, gamma_vect, z_cm_start=0, W_start=1):
    
        
    v_g = return_vg(field)
    dz = field._dX
    z_ind_start = int(z_cm_start*1e-2 / dz)
    
    w_vector = np.zeros(field._ymax+1)
    w_vector[z_ind_start] = W_start
    
    for z in range(z_ind_start+1, field._ymax+1):
        w_vector[z] = (2*dz*gamma_vect[z-1] + v_g[z-1])*w_vector[z-1]/v_g[z]
        
    return w_vector


# +
tmin_mus = 7
tmin_ind = int( tmin_mus*1e-6/(fields[0]._dT*fields[0]._Na) )
gamma_vect = growthrate_2D_0[tmin_ind:, :].mean(axis=0)
z_cm_start=0.05

w_vector = integrate_gamma(fields[0], gamma_vect, z_cm_start=z_cm_start, W_start=1)

plt.figure()
plt.plot(fields[0].tab_y, w_vector, label="By integration")

PIC_W = w2D_0[tmin_ind:, :].mean(axis=0)
plt.plot(fields[0].tab_y, PIC_W/PIC_W[int(z_cm_start*1e-2 / fields[0]._dX)], label="PIC measurement")

plt.ylabel("Wave energy $W$ normalized")
plt.ylabel("Axial position $z$ [cm]")
plt.legend()
plt.title("Case no $L_R$, \n integration between $t=7$ and $10\mu$s")
# -

# # Ion accoutic wave growth rate
# Because the saturation may not be responssible for the observations, lets compare $\gamma_{\rm PIC}$ with the theory:
# $$\gamma_{IAW} = \sqrt{ \frac{\pi m_e}{8 m_i} } \frac{ \vec{k}\cdot\vec{u_e}}{ 1 + k^2 \lambda_{De}^2)^{3/2}} $$
#
# Plausible hypotheses on $k$:
# 1. _local parameters_ : The maximum growth rate correspond to the wavevector $k \lambda_{De} = 1/\sqrt{2}$
# 2. _Convected wave_ : the wave rises in the maximum growing place (z ~ 0.5cm) then is transported.

# +
from plasmapy.physics import Debye_length
from plasmapy.physics import plasma_frequency
from astropy import units as u



# +
def get_ue(run: Field):
    """return the mean azimuthal electron velocity"""
    Je_x = np.array(run.meanfield("Je(1)", mean_axis="x", imin=280 ))
    ne = np.array(run.meanfield("Nume", mean_axis="x", imin=280))
    v_e =  Je_x/run.qe/ne
    
    return v_e

def get_lde_wpi(f):
    
    ne_vect = f.meanfield("Nume", "x")
    Te_vect = f.meanfield("Eke(1)", "x")
    wpi_vect = plasma_frequency(ne_vect/u.m**3, particle="Xe+")
    lDe_vect = Debye_length(Te_vect*u.eV, ne_vect/u.m**3,)

    return wpi_vect, lDe_vect


# +
field = fields[0]

u_e = get_ue(field)
wpi_vect, lDe_vect = get_lde_wpi(field)

factor = np.sqrt(field.me * np.pi/(8*field.mi))
factor


# -

# ## First hypothese: local wave

# +

def compute_gamma_IAW(field, case=1):
    
    u_e = -get_ue(field)
    wpi_vect, lDe_vect = get_lde_wpi(field)

    factor = np.sqrt(field.me * np.pi/(8*field.mi))

    if case==1:
        k_vect = 1/(np.sqrt(2) * lDe_vect.value)
    elif case==2:
        k_vect = np.ones_like(lDe_vect.value) *  1/(np.sqrt(2) * lDe_vect[int(0.5*1e-2/field._dX)].value)
        

    gamma_theo = factor * ( k_vect * u_e) / ( 1 + k_vect**2 * lDe_vect.value**2)**(3/2)
    
    return gamma_theo
# -

tmin_ind

# +
gamma_theo_IAW = compute_gamma_IAW(field)
gamma_PIC = growthrate_2D_0[tmin_ind:, :].mean(axis=0)

plt.figure()

plt.plot(field.tab_y,1e-6* gamma_theo_IAW, label = "Theo IAW")

plt.plot(field.tab_y,1e-6* gamma_PIC, label = "PIC measurment")
plt.legend()
plt.ylabel("Growth rate $\gamma$ [$\mu$s$^{-1}$]")

plt.xlabel("Azimuthal position $z$ [cm]")
plt.ylim(bottom=-5)
plt.xlim(0, 2.5)

# +
gamma_theo_IAW = compute_gamma_IAW(field, case=2)
gamma_PIC = growthrate_2D_0[tmin_ind:, :].mean(axis=0)

plt.figure()

plt.plot(field.tab_y, gamma_theo_IAW, label = "Theo IAW")

plt.plot(field.tab_y, gamma_PIC, label = "PIC measurment")
plt.legend()
plt.ylabel("Growth rate $\gamma$")

plt.xlabel("Azimuthal position $z$ [cm]")
plt.ylim(bottom=-0.5e7)
plt.xlim(0, 2.5)
# -



# # All the same, but not smoothed
# let see what it can change

# +
nt_max = 900
for f in fields:
    nt_max = min( nt_max, f._nT)
        
growthrate_2D_0_raw, dtW_2D_0_raw =  temporal_evolution_2D(fields[0], smooth_W=False, smooth_sigma=2, nt_max=nt_max)
growthrate_2D_1_raw, dtW_2D_1_raw =  temporal_evolution_2D(fields[1], smooth_W=False, smooth_sigma=2, nt_max=nt_max)
growthrate_2D_2_raw, dtW_2D_2_raw =  temporal_evolution_2D(fields[2], smooth_W=False, smooth_sigma=2, nt_max=nt_max)

# +
datas_raw=[growthrate_2D_0_raw, growthrate_2D_1_raw, growthrate_2D_2_raw ]

growth_rate_2D_evolution(smooth_gaussian_2D=False, datas=datas_raw)
growth_rate_2D_evolution(smooth_gaussian_2D=True, datas=datas_raw)

# +

gamma_axial_evolution(scaled=True, tmin_mus=3, datas=datas_raw)

gamma_axial_evolution(smooth_gaussian_2D=True, scaled=True, tmin_mus=3, datas=datas_raw)
# -

gamma_evolution(smooth_gaussian_2D=True, z_cm_list=[0.4, 0.8, 1.6], scaled=True, datas=datas_raw)



StarcoderdataJupyterScriptsDedupFiltered