spec_analysis.py

"""
NOT UPDATED WITH LATEST CHANGES

This program was designed to run with python 3 in a Spyder environnement, with usual packages local 
libraries found in the folder utils of the environment. 
Only the imput parameter section should be modified.

They are: - [shot] int
              the number of the shot
          - [path] str
              the name of the folder in which all the code architecture is
          - [spectrum_type] str
              the type of Fourier transform applied to data. 
              Can be frequency or spatial.
          - [I_TF] float
              the value of the current in the toroidal coils, only relevant when [data_origin] is simulation
          - [data_origin] str
              the way data were collected. 
              Can be experiment or simulation.
          - [av_size] int
              the length of the average segment
              Recommended values: 20 experiment and simulation
          - [lb] int
              the low born for the domain of the model to fit spectral index
          - [hb] int
              the high born for the domain of the model to fit spectral index
              Recommended values lb, hb = 60, 310 (simulation) and lb, hb = 200, 600 (experiment)
              This may change from one spectrum to another

Take the data from the .txt files generated by extract_all_data.py or the DDAQ file.
Analyse spatial or temporal frequencies of the collected signals, only ion saturation current measurement are 
fast enough for that type of study though. Evaluate the spectral index for different magnetic fields and try to 
find a trend/relation between those quantities.
"""

#Find the path to local libraries
import sys
sys.path.append('E:\\north_diagnostics\\diagnostics')
sys.path.append('E:\\north_diagnostics\\utils')
sys.path.append('E:\\north_diagnostics')

#Import all useful libraries
import numpy as np
import matplotlib.pyplot as plt
import os
#import scipy.signal as scs
import diagnostics.probe
import utils.dau

#Input parameters
shot = 10164
path = 'E:/north_diagnostics'
spectrum_type = 'frequency' 
I_TF = np.array([800, 900, 1000, 1100, 1200])
data_origin = 'experiment'
av_size = 100
lb = 250 #50 250
hb = 550 #200 550






"""
Main program: from now, all shall remain untouched.
"""
#Extract and cleaning data from .txt file
if data_origin == 'experiment':
      #Activated probes for that type of data
      studied_probes = np.array([1,2,3,4,5,6,7,8,9,10,11,13,14,15,16,17,18,19,20,21,22,23,24,25,
                               26,27,28,29,30,31,32,33,38,39,40,41,42,43,44,45,46,47,48,49])
      all_data = utils.dau.read_probe_data(shot, path, 'ion_saturation_current', 1/75, 
                                         studied_probes-1, 0, 1, False)
      machine_data = utils.dau.read_machine_data(shot, path)
      print(f"Data shot {shot} loaded with success")
    
      #Reshape the data array in a more practicable way
      I_TF = np.zeros(13)
      window_length = 10000
      data = np.zeros((len(I_TF), window_length, 51))
      start_times = np.linspace(200E-3, 800E-3, len(I_TF))
      for i in range(len(I_TF)):
          ind_start, ind_end = int(start_times[i]*1E6), int(start_times[i]*1E6+window_length)
          data[i, :, :] = all_data[ind_start:ind_end, :]
          I_TF[i] = machine_data[int(start_times[i]*100), 2]
    
      #Create a folder to store data if it doesn't exist
      if not os.path.exists(f"{path}/Figures/{shot}_{spectrum_type}_Btorsweep/"):
          os.makedirs(f"{path}/Figures/{shot}_{spectrum_type}_Btorsweep/")
          print("A new directory for storing data was created")
    
      #The time step between two recorded points
      time_step = 1E-6  
        
elif data_origin == 'simulation': 
      #Activated probes for that type of data
      studied_probes = np.array([1,2,3,4,5,6,7,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,
                               26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50])
    
      all_data = np.genfromtxt(f"{path}/Data/probeIsat_Btorsweep_fhz.txt", skip_header=5)
      print("Data probe loaded with success")
      l, h = np.shape(all_data)
      
      #Reshape the data array in a more practicable way
      n_sweep = int((I_TF[-1]-800)//100)
      data = np.zeros((n_sweep+1, l//(n_sweep+1)-300, 51))
      for i in range(n_sweep+1):
          data[i, :, :] = all_data[i*l//(n_sweep+1)+100:(i+1)*l//(n_sweep+1)-200, 1:]
      
      #Create a folder to store data if it doesn't exist
      if not os.path.exists(f"{path}/Figures/simu_{spectrum_type}_Btorsweep/"):
        os.makedirs(f"{path}/Figures/simu_{spectrum_type}_Btorsweep/")
        print("A new directory for storing data was created")
    
      #The time step between two recorded points
      time_step = 1E-5

#Send an error message if their is a mispelling in the data_origin variable    
else:
    print('WARNING: the data origin is not recognized')
    sys.exit()

#Machine parameter
mu_0 = 4*np.pi*1E-7 # in H/m
N_TF = 8
N_winding = 12
R0 = 250E-3 # in m

#Convert current into the magnetic field
B_0_tor = mu_0*(N_TF*N_winding)*I_TF/(2*np.pi*R0)


#Plot the data of one probe and its Fourier transform
fig = plt.figure(figsize=(10,10))

#Collect the spectral indices data
spec_ind = np.zeros((len(studied_probes), len(I_TF)))

for i in range(len(studied_probes)):
    print(f"Probe {studied_probes[i]}:")
    
    for j in range(len(I_TF)):
        print(f"   Coil current: {int(I_TF[j])} A")
        
        #Calculation and smoothing of the Fourier transform
        freq, spectrum = utils.dau.frequency_fft(data[j, :, studied_probes[i]], time_step) 
        spectrum_filt = utils.dau.rolling_average(spectrum, av_size)
        #freq, spectrum = scs.periodogram(data[j, :, studied_probes[i]], nfft=1000, fs=1/time_step, return_onesided=True, scaling='density')
        #spectrum_filt = spectrum
      
        #Plot the signal
        plt.subplot(1,2,1)
         
        if data_origin == 'experiment':
            plt.plot(data[j, :, 0]*1E3, data[j, :, studied_probes[i]]*1E3, 
                     label ='Temporal signal for B0 = {B_0_tor[j]:.2f} T') 
        else:
            plt.plot(data[j, :, 0], data[j, :, studied_probes[i]]*1E3, 
                     label ='Temporal signal for B0 = {B_0_tor[j]:.2f} T')
        plt.xlabel("time (ms)") 
        plt.ylabel("Ion saturation current (mA)")

        #Plot the Fourier transform
        plt.subplot(1,2,2)
        plt.plot(freq, spectrum_filt, label ='Temporal spectrum for B0 = {B_0_tor[j]:.2f} T') 
        model = np.polyfit(np.log(freq[lb:hb]),np.log(spectrum_filt[lb:hb]),1)
        spectrum_model = np.exp(model[1])*(freq[lb:hb])**model[0]
        spec_ind[i, j] = - model[0]
        plt.plot(freq[lb:hb], spectrum_model, 
                 label=f'model for B0 = {B_0_tor[j]:.2f} T: spectral index {model[0]:.2f}')
        #plt.ylim((1E-20,1E-11))
        
    #Plot the limits of the modelling domain
    plt.axvline(x=freq[lb], color='red', linestyle='--')
    plt.axvline(x=freq[hb], color='red', linestyle='--')
    
    #Add some legend
    plt.xlabel("Frequencies (Hz)")  
    plt.ylabel("Amplitude (UA)") 
    plt.xscale("log")
    plt.yscale("log")
    plt.suptitle(f"Signal and spectrum of probe {studied_probes[i]}")
      
    #Save the data
    if data_origin == 'experiment':
        plt.savefig(f"{path}/Figures/{shot}_{spectrum_type}_Btorsweep/probe{studied_probes[i]}.png")
    elif data_origin == 'simulation':
        plt.savefig(f"{path}/Figures/simu_{spectrum_type}_Btorsweep/probe{studied_probes[i]}.png")
      
    #Prepare to plot a new figure
    plt.clf() 

#Get the probe positions
R = []
for i in studied_probes:
  probe = diagnostics.Probe(path = f"{path}/Data", shot = shot, number = i, caching = True)
  x_p, y_p = probe.position['r'], probe.position['z']
  R.append(x_p*1E-3)  
R = np.array(R)

     
for i in range(len(studied_probes)): 
    print(f"Probe {studied_probes[i]}")
    #Plot the spectral index vs toroidal magnetic field
    B = B_0_tor*R0/R[i]
    plt.scatter(B, spec_ind[i, :], label='Data')
    model = np.polyfit(B, spec_ind[i, :], 1)
    spec_ind_model = model[1]+B*model[0]
    #plt.plot(B, spec_ind_model, label=f'model: slope {model[0]:.2f}; y-intercept {model[1]:.2f}')
    
    #Add some legend
    plt.xlabel("Magnetic field (T)")  
    plt.ylabel("Spectral indices") 
    plt.suptitle(f"Spectral index scaling for probe {studied_probes[i]}")
    plt.legend()
    
    #Save the data
    if data_origin == 'experiment':
      plt.savefig(f"{path}/Figures/{shot}_{spectrum_type}_Btorsweep/spec_ind_probe{studied_probes[i]}.png")
    elif data_origin == 'simulation':
      plt.savefig(f"{path}/Figures/simu_{spectrum_type}_Btorsweep/spec_ind_probe{studied_probes[i]}.png")
    plt.clf()
  
plt.close()