THz-TDS Analysis

thz.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.fftpack import fft, fftfreq
from scipy import signal
import glob
from scipy import constants
import pandas as pd
import glob
import matplotlib
matplotlib.style.use('seaborn-poster')

def windowing(data, variance):
    '''
    @params
    data : numpy array | the data to be windowed
    variance : integer | the variance of the gaussian winoww
    @return
    windowed : numpy array | array of the windowed data of the same length
    '''
    index = np.argmax(np.abs(data))
    l = 2*index
    window = signal.gaussian(l, variance)
    windowed = np.zeros_like(data)
    windowed[:l] += window
    windowed *= data
    return windowed

def paramaters(f, trans_phi, trans_magnitude, d=1.002/1000):
    '''
    @params
    f : array | the frequency in Hz
    trans_phi : array | phase of the transmission
    trans_magnitude : array | magnitude of the transmission
    d : float | thickness of the sample
    @return arrays
    n |  real part refractive index
    kappa | imag. part refractive index
    epsilon_r | real dielectric function
    epsilon_i | imag. dielectric function
    cond_r | real conductivity
    cond_i | imag. conductivity
    '''
    n = trans_phi * 3e8 / (2 * np.pi * f * d) + 1
    kappa = -3e8 / (2 * np.pi * f * d) * np.log(trans_magnitude * (n+1)**2 / (4*n))
    epsilon_r = n**2 - kappa**2
    epsilon_i = 2*n*kappa
    cond = (epsilon_r+1j*epsilon_i-1)*2*np.pi*f*constants.epsilon_0/1j
    cond_r, cond_i = np.real(cond), np.imag(cond)
    return n, kappa, epsilon_r, epsilon_i, cond_r, cond_i

def main():
    temps = [5, 6, 7, 10, 25, 50, 75, 100, 125, 150, 175, 200, 225, 250, 275, 300]

    epsilon_all_r = []
    cond_all_r = []
    epsilon_all_i = []
    index = -1 
    for temp in temps[:index]:
        loc_empty = str(temp) + 'K_empty.txt'
        loc_sample = str(temp) + 'K_MgO.txt'
        loc_glob = str(temp) + 'K'

        d_E = []
        d_S = []
        for loc in glob.glob('Part_II/'+str(temp)+'K/*_'+loc_empty):
            data_empty = pd.read_csv(loc, skiprows=16, header=None, usecols=[0,1], delimiter='\t')
            d_E.append(data_empty[1])    
        empty_v = np.mean(np.asarray(d_E), axis=0)

        for loc in glob.glob('Part_II/'+str(temp)+'K/*_'+loc_sample):
            data_sample = pd.read_csv(loc, skiprows=16, header=None, usecols=[0,1], delimiter='\t')
            d_S.append(data_sample[1])
        sample_v = np.mean(np.asarray(d_S), axis=0)

        time = data_sample[0]
        time *= 1e-12
        empty_v_win = windowing(empty_v, 150)
        sample_v_win = windowing(sample_v, 150)

        n = len(time)
        t = time[1]
        empty_v_fft = fft(empty_v_win)
        sample_v_fft = fft(sample_v_win)
        
        f = np.linspace(0.0, 1/(2.0*t), n//2)
        # plt.plot(f, 2.0/n * np.abs(sample_v_fft[:n//2]))
        # plt.plot(f, 2.0/n * np.abs(empty_v_fft[:n//2]))
        # plt.show()    
        trans =  sample_v_fft[:n//2] / empty_v_fft[:n//2]
        trans_phi = np.angle(trans)
        
        p = -trans_phi

        threshold = np.pi
        for i in range(1,len(p)-1):
            if p[i+1]-p[i] > threshold:
                for j in range(i+1,len(p)):
                    p[j] = p[j]-2*np.pi
            elif p[i+1]-p[i] < -threshold:
                for j in range(i+1,len(p)):
                    p[j] = p[j]+2*np.pi
        
        # plt.plot(f[:n//2], abs(trans), label=str(temp)+'K')
        # plt.legend(loc='upper right')
        # plt.xlabel('frequency (THz)')
        # plt.ylabel('Transmission amplitude')
        # plt.show()
            
        # plt.plot(f[:n//2], p, label=str(temp)+'K')
        # plt.legend(loc='upper right')
        # plt.xlabel('frequency (THz)')
        # plt.ylabel('Transmission phase')
        # plt.show()
        N, kappa, epsilon_r, epsilon_i, cond_r, cond_i = parameters(f[100:n//2], p[100:], np.abs(trans)[100:])
        epsilon_all_r.append(epsilon_r)
        epsilon_all_i.append(epsilon_i)
        cond_all_r.append(cond_r)

    # np.savetxt('data', np.asarray(f[100:int(n/2)], epsilon_all_r[i, cond_all_r[i]).T)
    for i, temp in enumerate(temps[:index]):
        plt.plot(f[100:int(n/2)], epsilon_all_r[i], label=str(temp)+'K')    
    plt.xlim(0.2*1e12, 5*1e12)
    plt.legend()
    plt.xlabel('Frequency (Hz)')
    plt.ylabel('$\epsilon_{1}(\omega)$')
    plt.show()
    for i, temp in enumerate(temps[:index]):
        plt.plot(f[100:int(n/2)], epsilon_all_i[i], label=str(temp)+'K')    
    plt.xlim(0.2*1e12, 5*1e12)
    plt.legend()
    plt.xlabel('Frequency (Hz)')
    plt.ylabel('$\epsilon_{2}(\omega)$')
    plt.show()
    for i, temp in enumerate(temps[:index]):
        plt.plot(f[100:int(n/2)], cond_all_r[i], label=str(temp)+'K')    
    # plt.legend()
    plt.xlim(0.2*1e12, 5*1e12)
    plt.legend()
    plt.xlabel('Frequency (Hz)')
    plt.ylabel('$\sigma_{1}(\omega)$')
    plt.show()

if __name__=='__main__':
    main()