# =============================================================================
# A couple of functions to compute a cross-correlation between the signals from
# pins of double rake probe and then plot the result
# =============================================================================
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():
# plt.figure(); plt.clf()
plt.title('Cross-correlation and fit of pins %s and %s, shot %i' % (pin1, pin2, 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, pin1, pin2 = 30375, 'drppin15', 'drppin11'
#shot, pin1, pin2 = 21757, 'a1', 'b1'
lag, corr = cross_corr(shot, pin1, pin2, 10e-3, 12e-3)
lag_new, ffit = fit(lag, corr)
fig = plt.figure()
plot_fit()
fig.set_size_inches(9.5, 6., forward=True)
fig.savefig('Fit, pins %s and %s, shot %i' % (pin1, pin2, shot))
print("Time shift between the pins:", round(lag_new[np.argmax(ffit)]*1e6, 3), "[us]")
print("Poloidal velocity:", round(2.5e-3 / lag_new[np.argmax(ffit)], 3), "[m/s]")
