# =============================================================================
# A couple of functions to compute a cross-correlation between the signals from
# pins of the double rake probe and then plot the result. Modified version:
# this code computes a fit for 3 closest to the plasma pairs of pins. The only
# thing that should be done is to write the number of the shot and set left and
# right points in time that indicates interval when plasma is quasi-stable.
# =============================================================================
import numpy as np
from matplotlib import pyplot as plt
from urllib.request import urlopen
plt.ion()
# =============================================================================
# A function for data collection from the pins of the double rake probe
# =============================================================================
def collect(shot, pin):
print('Data from shot %i, pin %s' % (shot, pin))
url = f'http://golem.fjfi.cvut.cz/utils/data/{shot}/{pin}'
ourl = urlopen(url)
data_all = np.loadtxt(ourl)
data = np.array([row[1] for row in data_all])
ourl.close()
return data
# =============================================================================
# The function will compute cross-correlation between the signals from two pins
# of the double rake probe. 'left' and 'right' are points on x-axis (time)
# that should be taken from 'U_loop' diagnostics and are points indicating a
# plato. Data from that interval gives better results.
# =============================================================================
def cross_corr(shot, pin1, pin2, left, right):
dt = 1e-6 # time step computed from the frequency
data1 = collect(shot, pin1)
data2 = collect(shot, pin2)
time_axis = np.arange(data1.size)*dt
mask_ft = (left < time_axis) & (time_axis < right)
data1 = data1[mask_ft]
data2 = data2[mask_ft]
norm = np.sqrt(np.sum(np.correlate(data1, data1)) * np.sum(np.correlate(data2, data2)))
corr = np.correlate(data1, data2, 'same') / norm
N = len(corr)
time_lag = (np.arange(N) - N//2) * dt # time lag between two pins
return time_lag, corr
# =============================================================================
# A function for fitting the cross-correlation for more precise time lag.
# Fitting is done by a polynomial function.
# =============================================================================
def fit(time_lag, corr):
center = np.argmax(corr)
time = time_lag[center - 5:center + 5]
corr = corr[center - 5:center + 5]
coefs = np.polyfit(time, corr, 9)
ffit = np.poly1d(coefs)
time_new = np.linspace(time[0], time[-1], 100)
ffit_new = ffit(time_new)
return time_new, ffit_new
# =============================================================================
# A function for plotting a close-up of the cross-correlation and the fit
# around the peak
# =============================================================================
def plot_fit(lag, corr, lag_new, ffit):
# plt.figure(); plt.clf()
if np.array_equal(corr, corr1):
plt.title('Cross-correlation and fit of pins %s and %s, shot %i' % (pin16, pin12, shot))
elif np.array_equal(corr, corr2):
plt.title('Cross-correlation and fit of pins %s and %s, shot %i' % (pin15, pin11, shot))
else:
plt.title('Cross-correlation and fit of pins %s and %s, shot %i' % (pin14, pin10, shot))
plt.xlabel('Time lag, [us]')
plt.ylabel('Cross-correlation coeficient, [-]')
c = np.argmax(corr)
plt.plot(lag[c-5:c+5]*1e6, corr[c-5:c+5], 'o', label = 'corr')
plt.plot(lag_new*1e6, ffit, label = 'fit')
# plt.axvline(lag[np.argmax(corr)]*1e6, color = 'k')
plt.axvline(lag_new[np.argmax(ffit)]*1e6, color='g', label = 'fit lag')
plt.legend(loc='upper left')
plt.show()
# =============================================================================
# Example: plot and results
# =============================================================================
shot = 30371
left, right = 8e-3, 10e-3
pin16, pin12, pin15, pin11, pin14, pin10 = 'drppin16', 'drppin12', 'drppin15', \
'drppin11', 'drppin14', 'drppin10'
#shot, pin1, pin2 = 21757, 'a1', 'b1'
lag1, corr1 = cross_corr(shot, pin16, pin12, left, right)
lag_new1, ffit1 = fit(lag1, corr1)
lag2, corr2 = cross_corr(shot, pin15, pin11, left, right)
lag_new2, ffit2 = fit(lag2, corr2)
lag3, corr3 = cross_corr(shot, pin14, pin10, left, right)
lag_new3, ffit3 = fit(lag3, corr3)
fig = plt.figure(1)
plot_fit(lag1, corr1, lag_new1, ffit1)
fig.set_size_inches(9.5, 6., forward=True)
fig.savefig('Fit, shot %i, pins %s and %s' % (shot, pin16, pin12))
fig = plt.figure(2)
plot_fit(lag2, corr2, lag_new2, ffit2)
fig.set_size_inches(9.5, 6., forward=True)
fig.savefig('Fit, shot %i, pins %s and %s' % (shot, pin15, pin11))
fig = plt.figure(3)
plot_fit(lag3, corr3, lag_new3, ffit3)
fig.set_size_inches(9.5, 6., forward=True)
fig.savefig('Fit, shot %i, pins %s and %s' % (shot, pin14, pin10))
#print("Time shift between the pins:", round(lag_new1[np.argmax(ffit)]*1e6, 3), "[us]")
print("Pins %s and %s, poloidal velocity:" % (pin16[-2:], pin12[-2:]), round(2.5e-3
/ lag_new1[np.argmax(ffit1)], 3), "[m/s]")
print("Pins %s and %s, poloidal velocity:" % (pin15[-2:], pin11[-2:]), round(2.5e-3
/ lag_new2[np.argmax(ffit2)], 3), "[m/s]")
print("Pins %s and %s, poloidal velocity:" % (pin14[-2:], pin10[-2:]), round(2.5e-3
/ lag_new3[np.argmax(ffit3)], 3), "[m/s]")
