def plot_ave_freq(self, Res, N, title, color='blue', y_lim=[0, 1.7], label='whole', save_path=False): if save_path: with open(os.path.join(save_path, 'AveFreq-%s.txt' % title), 'w') as f: f.write('Frequency (kHz), Amplitude\n') for j in range(Res.shape[0]): f.write('{}, {}\n'.format(self.grid[j] / 1000, Res[j] / max(Res))) fig = plt.figure(figsize=(6, 4.1), num='Average Frequency--%s' % title) ax = fig.add_subplot() ax.plot(self.grid / 1000, Res / max(Res), lw=1, color=color, label=label) plot_norm(ax, x_lim=[0, 800], y_lim=y_lim, xlabel='Frequency (kHz)', ylabel='Normalized Amplitude', title='Average Frequency')
def plot_wave_frequency(self, TRAI_select, pop): fig = plt.figure(figsize=(6.5, 10), num='Waveform & Frequency--pop%s' % pop) for idx, j in enumerate(TRAI_select): i = self.data_tra[j - 1] valid_time, valid_data = self.waveform.cal_wave(i, valid=False) half_frq, normalization_half = self.cal_frequency(j - 1, valid=False) ax = fig.add_subplot(5, 2, 1 + idx * 2) ax.plot(valid_time, valid_data) ax.axhline(abs(i[2]), 0, valid_data.shape[0], linewidth=1, color="black") ax.axhline(-abs(i[2]), 0, valid_data.shape[0], linewidth=1, color="black") plot_norm(ax, 'Time (μs)', 'Amplitude (μV)', legend=False, grid=True) ax = fig.add_subplot(5, 2, 2 + idx * 2) ax.plot(half_frq, normalization_half) plot_norm(ax, 'Freq (Hz)', '|Y(freq)|', x_lim=[0, pow(10, 6)], legend=False)
def plot_ave_freq(self, Res, N, title): fig = plt.figure(figsize=(6, 4.1), num='Average Frequency--%s' % title) ax = fig.add_subplot() ax.plot(self.grid, Res / N) plot_norm(ax, xlabel='Freq (Hz)', ylabel='|Y(freq)|', title='Average Frequency', legend=False)
def cal_ML(self, fig, tmp, xlabel, ylabel, COLOR='black', ECOLOR=None): """ Calculate the maximum likelihood function distribution :param tmp: Energy/Amplitude/Duration in order of magnitude of original data :param xlabel: 'Amplitude (μV)', 'Duration (μs)', 'Energy (aJ)' :param ylabel: 'ML (A)', 'ML (D)', 'ML (E)' :param COLOR: Color when drawing with original data, population I and population II respectively :param ECOLOR: Line color of error bar, corresponding parameter COLOR :return: """ if not ECOLOR: ECOLOR = [0.7, 0.7, 0.7] N = len(tmp) fig.subplots_adjust(left=0.131, bottom=0.179, right=0.975, top=0.944) fig.text(0.96, 0.2, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }, horizontalalignment="right") ax = fig.add_subplot() ax.set_xscale("log", nonposx='clip') ML_y, Error_bar = [], [] for j in range(N): valid_x = sorted(tmp)[j:] E0 = valid_x[0] Sum = np.sum(np.log(valid_x / E0)) N_prime = N - j alpha = 1 + N_prime / Sum error_bar = (alpha - 1) / pow(N_prime, 0.5) ML_y.append(alpha) Error_bar.append(error_bar) ax.errorbar(sorted(tmp), ML_y, yerr=Error_bar, fmt='o', ecolor=ECOLOR, color=COLOR, elinewidth=1, capsize=2, ms=3) plot_norm(ax, xlabel, ylabel, y_lim=[1.25, 3], legend=False) fig.canvas.draw() fig.canvas.flush_events() with open( '/'.join([self.output, self.status]) + '_ML(%s).txt' % xlabel[0], 'w') as f: f.write('{}, {}, Error bar\n'.format(xlabel, ylabel)) for i, j, k in zip(tmp, ML_y, Error_bar): f.write('{}, {}, {}\n'.format(i, j, k))
def plot_wave_TRAI(self, k, valid=False, color='blue'): # Waveform with specific TRAI i = self.data_tra[k - 1] if i[0 if self.device == 'stream' else -1] != k: return str( 'Error: TRAI %d in data_tra is inconsistent with %d by input!' % (i[0 if self.device == 'stream' else -1], k)) time, sig = self.cal_wave(i, valid=valid) for tmp_tail, s in enumerate(sig[::-1]): if s != 0: tail = -tmp_tail if tmp_tail > 0 else None break time, sig = time[:tail], sig[:tail] fig = plt.figure(figsize=(6, 4.1), num='Waveform--TRAI %d (%s)' % (k, valid)) fig.text(0.95, 0.17, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }, horizontalalignment="right") ax = fig.add_subplot(1, 1, 1) ax.plot(time, sig, lw=1, color=color) if self.device == 'vallen': plt.axhline(abs(i[2]), 0, sig.shape[0], linewidth=1, color="black") plt.axhline(-abs(i[2]), 0, sig.shape[0], linewidth=1, color="black") elif self.device == 'pac': plt.axhline(abs(self.thr), 0, sig.shape[0], linewidth=1, color="black") plt.axhline(-abs(self.thr), 0, sig.shape[0], linewidth=1, color="black") plot_norm(ax, 'Time (μs)', 'Amplitude (μV)', title='TRAI:%d' % k, legend=False, grid=True)
def plot_freq_TRAI(self, k, valid=False): # Frequency with specific TRAI half_frq, normalization_half = self.cal_frequency(k - 1, valid=valid) fig = plt.figure(figsize=(6, 4.1), num='Frequency--TRAI:%d (%s)' % (k, valid)) ax = plt.subplot() ax.plot(half_frq, normalization_half) plot_norm(ax, 'Freq (Hz)', '|Y(freq)|', x_lim=[0, pow(10, 6)], title='TRAI:%d' % k, legend=False)
def plot_wave_frequency(self, TRAI, valid=False, t_lim=100, n=3, wtpacket=False, wtpacket_eng=False): i = self.data_tra[TRAI - 1] valid_time, valid_data = self.waveform.cal_wave(i, valid=valid) half_frq, normalization_half, time = self.cal_frequency(TRAI - 1, valid=valid) if time[-1] < t_lim: return fig = plt.figure(figsize=(9.2, 3), num='Waveform & Frequency--TRAI %d' % TRAI) ax = fig.add_subplot(1, 2, 1) ax.plot(valid_time, valid_data, lw=1) if self.device != 'stream': ax.axhline(abs(i[2]), 0, valid_data.shape[0], linewidth=1, color="black") ax.axhline(-abs(i[2]), 0, valid_data.shape[0], linewidth=1, color="black") plot_norm(ax, 'Time (μs)', 'Amplitude (μV)', legend=False, grid=True) ax = fig.add_subplot(1, 2, 2) ax.plot(half_frq / 1000, normalization_half, lw=1) plot_norm(ax, 'Freq (kHz)', '|Y(freq)|', x_lim=[0, pow(10, 3)], legend=False) if wtpacket: fig = plt.figure(figsize=(15, 7), num='WaveletPacket--TRAI %d' % TRAI) map, wp, energy = self.cal_wtpacket(valid_data, 'db8', n) for i in range(2, n + 2): level_num = pow(2, i - 1) # ['aaa', 'aad', 'add', 'ada', 'dda', 'ddd', 'dad', 'daa'] re = [node.path for node in wp.get_level(i - 1, 'freq')] for j in range(1, level_num + 1): ax = fig.add_subplot(n, level_num, level_num * (i - 2) + j) ax.plot(map[re[j - 1]]) plot_norm(ax, '', '', legend=False) if wtpacket_eng: fig = plt.figure(figsize=(4.6, 3), num='WaveletPacket Energy--TRAI %d' % TRAI) ax = fig.add_subplot() values = [i / sum(energy) for i in energy] index = np.arange(pow(2, n)) ax.bar(index, values, 0.45, color="#87CEFA") plot_norm(ax, 'Clusters', 'Reviews (%)', legend=False)
def plot_ave_freq(self, Res, N, title, color='blue', y_lim=[0, 1.7], label='whole'): fig = plt.figure(figsize=(6, 4.1), num='Average Frequency--%s' % title) ax = fig.add_subplot() ax.plot(self.grid / 1000, Res / N, lw=1, color=color, label=label) plot_norm(ax, x_lim=[0, 800], y_lim=y_lim, xlabel='Freq (kHz)', ylabel='|Y(freq)|', title='Average Frequency')
def plot_XXX_Freq(self, freq, feature, ylabel, marker='o', markersize=10, color='blue'): fig = plt.figure(figsize=[6, 3.9]) ax = fig.add_subplot() ax.set_yscale("log", nonposy='clip') plt.scatter(freq, feature, marker=marker, s=markersize, color=color) plot_norm(ax, 'Peak Frequency (kHz)', ylabel, x_lim=[0, 800], legend=False)
def plot_freqDomain(self, ALL_TRAI, t_lim=100, lw=1): fig = plt.figure(figsize=[6, 3.9], num='Frequency domain') z = 0 ax = fig.add_subplot(111, projection='3d') for i in ALL_TRAI: half_frq, normalization_half, t = self.cal_frequency(i - 1) if t[-1] < t_lim: continue valid_idx = np.where((half_frq / 1000) < 1000)[0] # ax.plot(half_frq[valid_idx] / 1000, [z] * valid_idx.shape[0], normalization_half[valid_idx], lw=lw) ax.scatter(half_frq[valid_idx] / 1000, [z] * valid_idx.shape[0], normalization_half[valid_idx]) z += 1 plot_norm(ax, 'Freq (kHz)', 'Points', '|Y(freq)|', x_lim=[0, 1000], y_lim=[0, z], legend=False)
def plot_2cls_freq(self, TRAI_1, TRAI_2, same): fig = plt.figure(figsize=(6.5, 10), num='Frequency with same %s' % same) for idx, k in enumerate(TRAI_1): half_frq, normalization_half = self.cal_frequency(k - 1) ax = fig.add_subplot(5, 2, 1 + idx * 2) ax.plot(half_frq, normalization_half) plot_norm(ax, 'Freq (Hz)', '|Y(freq)|', x_lim=[0, pow(10, 6)], legend=False) half_frq, normalization_half = self.cal_frequency(TRAI_2[idx] - 1) ax2 = fig.add_subplot(5, 2, 2 + idx * 2) ax2.plot(half_frq, normalization_half) plot_norm(ax2, 'Freq (Hz)', '|Y(freq)|', x_lim=[0, pow(10, 6)], legend=False)
def plot_correlation(self, fig, tmp_1, tmp_2, xlabel, ylabel, COLOR='black'): fig.subplots_adjust(left=0.115, bottom=0.17, right=0.975, top=0.95) fig.text(0.96, 0.2, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }, horizontalalignment="right") ax = fig.add_subplot() if self.device == 'PAC-self': ax.scatter(tmp_1, tmp_2, marker='o', s=15, color=COLOR, edgecolors=COLOR) else: ax.loglog(tmp_1, tmp_2, '.', Marker='.', markersize=8, color=COLOR) plot_norm(ax, xlabel, ylabel, legend=False) fig.canvas.draw() fig.canvas.flush_events() with open( '/'.join([self.output, self.status]) + '_%s-%s.txt' % (ylabel, xlabel), 'w') as f: f.write('{}, {}\n'.format(xlabel, ylabel)) for i, j in zip(tmp_1, tmp_2): f.write('{}, {}\n'.format(i, j))
def cal_BathLaw(self, fig, tmp, xlabel, ylabel, INTERVAL_NUM=None, COLOR='black', bin_method='log'): if INTERVAL_NUM is None: INTERVAL_NUM = 8 fig.subplots_adjust(left=0.145, bottom=0.19, right=0.975, top=0.962) fig.text(0.12, 0.2, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }) ax = fig.add_subplot() tmp_max = int(max(tmp)) if bin_method == 'linear': x = np.array([]) inter = [pow(10, i) for i in range(0, len(str(tmp_max)))] for i in inter: x = np.append( x, np.linspace(i, i * 10, INTERVAL_NUM, endpoint=False)) elif bin_method == 'log': x, x_eny = np.array([]), np.array([]) inter = self.__cal_log_interval(tmp) for i in inter: if i < 0: continue logspace = np.logspace(i, i + 1, INTERVAL_NUM, endpoint=False) x = np.append(x, logspace) tmp_xx = [(logspace[i + 1] + logspace[i]) / 2 for i in range(len(logspace) - 1)] tmp_xx.append((10 * logspace[0] + logspace[-1]) / 2) x_eny = np.append(x_eny, np.array(tmp_xx)) y = [] for k in range(x.shape[0]): N, Naft = self.__cal_N_Naft( tmp, [x[k], x[k + 1] ]) if k != x.shape[0] - 1 else self.__cal_N_Naft( tmp, [x[k], float('inf')]) if Naft != 0 and N != 0: y.append(np.log10(N / Naft)) else: y.append(float('inf')) y = np.array(y) if bin_method == 'linear': x, y = x[y != float('inf')], y[y != float('inf')] x_eny = np.zeros(x.shape[0]) for idx in range(len(x) - 1): x_eny[idx] = (x[idx] + x[idx + 1]) / 2 x_eny[-1] = x[-1] + pow(10, len(str(int( x[-1])))) * (0.9 / INTERVAL_NUM) / 2 elif bin_method == 'log': x_eny, y = x_eny[y != float('inf')], y[y != float('inf')] ax.semilogx(x_eny, y, color=COLOR, marker='o', markersize=8, mec=COLOR, mfc='none') ax.axhline(1.2, ls='-.', linewidth=1, color="black") plot_norm(ax, xlabel, ylabel, y_lim=[-1, 4], legend=False) fig.canvas.draw() fig.canvas.flush_events() with open( '/'.join([self.output, self.status]) + '_BathLaw(%s).txt' % xlabel[0], 'w') as f: f.write('{}, {}\n'.format(xlabel, ylabel)) for i, j in zip(x_eny, y): f.write('{}, {}\n'.format(i, j))
def plot_feature_time(self, fig, tmp, ylabel, x_max=29000, color_tmp='black', color_stretcher='r', width=55, stretcher_data=None, smooth=False): if stretcher_data: fig.subplots_adjust(left=0.12, bottom=0.157, right=0.877, top=0.853) fig.text(0.86, 0.18, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }, horizontalalignment="right") else: fig.subplots_adjust(left=0.115, bottom=0.17, right=0.975, top=0.95) fig.text(0.96, 0.19, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }, horizontalalignment="right") ax = fig.add_subplot() if stretcher_data: time, displace, load, strain, stress = load_stress(stretcher_data) if smooth: stress = smooth_curve(time, stress) strain_max = strain[-1] * x_max / self.Time[-1] ax.bar(self.Time, tmp, color=color_tmp, width=width, log=False if self.device == 'PAC-self' else True) plot_norm(ax, 'Time (s)', ylabel, x_lim=[0, x_max], y_lim=[0, 80 if self.device == 'PAC-self' else 15000], legend=False) ax2 = ax.twinx() # ax2.plot(time, stress, color_stretcher, lw=3) plot_norm(ax2, 'Time (s)', 'Stress (MPa)', x_lim=[0, x_max], y_lim=[0, 700], legend=False, font_color=color_stretcher) ax3 = ax2.twiny() for key in ['right', 'top']: ax3.spines[key].set_color(color_stretcher) ax3.plot(strain, stress, color_stretcher, lw=3) plot_norm(ax3, 'Strain (%)', x_lim=[0, strain_max], y_lim=[0, 700], legend=False, font_color=color_stretcher) else: ax.bar(self.Time, tmp, color=color_tmp, width=width, log=False if self.device == 'PAC-self' else True) plot_norm(ax, 'Time (s)', ylabel, x_lim=[0, x_max], y_lim=[0, 80 if self.device == 'PAC-self' else 15000], legend=False) fig.canvas.draw() fig.canvas.flush_events() with open( '/'.join([self.output, self.status]) + '_%s-Time.txt' % ylabel.split()[0], 'w') as f: f.write('TRAI, Time (s), {}\n'.format(ylabel)) for i, j, k in zip(self.TRAI, self.Time, tmp): f.write('{:.0f}, {}, {}\n'.format(i, j, k))
def cal_contour(self, fig, tmp_1, tmp_2, xlabel, ylabel, x_lim, y_lim, size_x=40, size_y=40, method='linear_bin', padding=False, colorbar=False, clabel=False): tmp_1, tmp_2 = 20 * np.log10(tmp_1), 20 * np.log10(tmp_2) if method == 'Log bin': sum_x, sum_y = x_lim[1] - x_lim[0], y_lim[1] - y_lim[0] arry_x = np.logspace(np.log10(sum_x + 10), 1, size_x) / ( sum(np.logspace(np.log10(sum_x + 10), 1, size_x)) / sum_x) arry_y = np.logspace(np.log10(sum_y + 10), 1, size_y) / ( sum(np.logspace(np.log10(sum_y + 10), 1, size_y)) / sum_y) x, y = [], [] for tmp, res, arry in zip([x_lim[0], y_lim[0]], [x, y], [arry_x, arry_y]): for i in arry: res.append(tmp) tmp += i x, y = np.array(x), np.array(y) elif method == 'Linear bin': x, y = np.linspace(x_lim[0], x_lim[1], size_x), np.linspace(y_lim[0], y_lim[1], size_y) X, Y = np.meshgrid(x, y) height = np.zeros([X.shape[0], Y.shape[1]]) linestyles = ['solid'] * 8 + ['--'] * 4 levels = [1, 2, 3, 6, 12, 24, 48, 96, 192, 384, 768, 1536] colors = [[1, 0, 1], [0, 0, 1], [0, 1, 0], [1, 0, 0], [0.5, 0.5, 0.5], [1, 0.3, 0], [0, 0, 0], [0, 1, 1], [1, 0, 1], [0, 0, 1], [0, 1, 0], [1, 0, 0]] for i in range(X.shape[1] - 1): valid_x = np.where((tmp_1 < X[0, i + 1]) & (tmp_1 >= X[0, i]))[0] for j in range(Y.shape[0] - 1): valid_y = np.where((tmp_2 < Y[j + 1, 0]) & (tmp_2 >= Y[j, 0]))[0] height[j, i] = np.intersect1d(valid_x, valid_y).shape[0] fig.subplots_adjust(left=0.115, bottom=0.17, right=0.975, top=0.95) if colorbar: fig.text(0.78, 0.2, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }, horizontalalignment="right") else: fig.text(0.96, 0.2, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }, horizontalalignment="right") ax = fig.add_subplot() if padding: ct = ax.contourf(X, Y, height, levels, colors=colors, extend='max') if colorbar: cbar = fig.colorbar(ct) else: ct = ax.contour(X, Y, height, levels, colors=colors, linewidths=1, linestyles=linestyles) if colorbar: cbar = fig.colorbar(ct) if clabel: ax.clabel(ct, inline=True, colors='k', fmt='%.1f') plot_norm(ax, xlabel, ylabel, legend=False) fig.canvas.draw() fig.canvas.flush_events()
def cal_CCDF(self, fig, tmp, xlabel, ylabel, LIM=None, COLOR='black', FIT=False): """ Calculate Complementary Cumulative Distribution Function :param tmp: Energy/Amplitude/Duration in order of magnitude of original data :param xlabel: 'Amplitude (μV)', 'Duration (μs)', 'Energy (aJ)' :param ylabel: 'CCDF (A)', 'CCDF (D)', 'CCDF (E)' :param LIM: Use in function fitting, support specific values or indexes, value: [0, float('inf')], [100, 900], ... index: [0, None], [11, -2], ... :param FIT: Whether to fit parameters, support True or False :param COLOR: Color when drawing with original data, population I and population II respectively :return: """ if LIM is None: LIM = [0, float('inf')] N = len(tmp) fig.subplots_adjust(left=0.133, bottom=0.179, right=0.975, top=0.962) fig.text(0.15, 0.2, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }) ax = fig.add_subplot() xx, yy = [], [] for i in range(N - 1): xx.append(np.mean([tmp[i], tmp[i + 1]])) yy.append((N - i + 1) / N) if FIT: xx, yy = np.array(xx), np.array(yy) fit_lim = np.where((xx > LIM[0]) & (xx < LIM[1]))[0] fit = np.polyfit(np.log10(xx[fit_lim[0]:fit_lim[-1]]), np.log10(yy[fit_lim[0]:fit_lim[-1]]), 1) alpha, b = fit[0], fit[1] fit_x = np.linspace(xx[fit_lim[0]], xx[fit_lim[-1]], 100) fit_y = self.convert(fit_x, alpha, b) ax.plot(fit_x, fit_y, '-.', lw=1, color=COLOR) ax.loglog(xx, yy, color=COLOR, label='slope-{:.2f}'.format(abs(alpha))) plot_norm(ax, xlabel, ylabel, legend_loc='upper right') else: ax.loglog(xx, yy, color=COLOR) plot_norm(ax, xlabel, ylabel, legend=False) fig.canvas.draw() fig.canvas.flush_events() with open( '/'.join([self.output, self.status]) + '_CCDF(%s).txt' % xlabel[0], 'w') as f: f.write('{}, {}\n'.format(xlabel, ylabel)) for i, j in zip(xx, yy): f.write('{}, {}\n'.format(i, j))
def stream(file, t_str, t_end, staLen=5, overlap=1, staWin='hamming', IZCRT=0.7, ITU=550, alpha=1.7, t_backNoise=1e4): # ====================================================== 数据读取 ====================================================== with open(file, 'r') as f: for _ in range(4): f.readline() fs = int(f.readline().strip().split()[-1]) * 1e-3 sig_initial = np.array( list(map(lambda x: float(x.strip()) * 1e4, f.readlines()[4:-1]))) t_initial = np.array([i / fs for i in range(sig_initial.shape[0])]) # ====================================================== 计算结果 ====================================================== t = t_initial[int(t_str // t_initial[1]):int(t_end // t_initial[1]) + 1] - t_initial[int(t_str // t_initial[1])] sig = sig_initial[int(t_str // t_initial[1]):int(t_end // t_initial[1]) + 1] width = int(fs * staLen) stride = int(width) - overlap t_stE, stE = shortTermEny(sig, width, stride, fs, staWin) t_zcR, zcR = zerosCrossingRate(sig, width, stride, fs, staWin) stE_dev = cal_deriv(t_stE, stE) start, end = find_wave(stE, stE_dev, zcR, t_stE, IZCRT=IZCRT, ITU=ITU, alpha=alpha, t_backNoise=t_backNoise) # ====================================================== 图形展示 ====================================================== x = [t, t_stE, t_stE, t_zcR] y = [sig, stE, stE_dev, zcR] color = ['black', 'green', 'gray', 'purple'] ylabel = [ r'$Amplitude$ $(μV)$', r'$STEnergy$ $(μV^2 \cdot μs)$', r'$S\dot{T}E$ $(μV^2)$', r'$ST\widehat{Z}CR$ $(\%)$' ] fig, axes = plt.subplots(4, 1, sharex=True, figsize=(10, 8)) for idx, ax in enumerate(axes): ax.plot(x[idx], y[idx], lw=0.5, color=color[idx]) if idx == 0: for s, e in tqdm(zip(start, end)): ax.plot(t[int(t_stE[s] // t[1]) + 1:int(t_stE[e] // t[1]) + 2], sig[int(t_stE[s] // t[1]) + 1:int(t_stE[e] // t[1]) + 2], lw=0.5, color='red') ax.grid() plot_norm(ax, r'$Time$ $(μs)$' if idx == 3 else '', ylabel[idx], legend=False) plt.subplots_adjust(wspace=0, hspace=0)
def cal_PDF(self, fig, tmp, xlabel, ylabel, LIM=None, INTERVAL_NUM=None, COLOR='black', FIT=False, bin_method='log'): """ Calculate Probability Density Distribution Function :param tmp: Energy/Amplitude/Duration in order of magnitude of original data :param xlabel: 'Amplitude (μV)', 'Duration (μs)', 'Energy (aJ)' :param ylabel: 'PDF (A)', 'PDF (D)', 'PDF (E)' :param LIM: Use in function fitting, support specific values or indexes, value: [0, float('inf')], [100, 900], ... index: [0, None], [11, -2], ... :param INTERVAL_NUM: Number of bins divided in each order of magnitude :param COLOR: Color when drawing with original data, population I and population II respectively :param FIT: Whether to fit parameters, support True or False :param bin_method: Method to divide the bin, Support linear partition and logarithmic partition :return: """ if INTERVAL_NUM is None: INTERVAL_NUM = 6 if LIM is None: LIM = [0, None] fig.subplots_adjust(left=0.133, bottom=0.179, right=0.975, top=0.962) fig.text(0.15, 0.2, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }) ax = fig.add_subplot() if bin_method == 'linear': inter, mid = self.__cal_linear_interval(tmp, INTERVAL_NUM) xx, yy = self.__cal_linear(tmp, inter, mid, INTERVAL_NUM) elif bin_method == 'log': inter = self.__cal_log_interval(tmp) xx, yy = self.__cal_log(tmp, inter, INTERVAL_NUM) if FIT: fit = np.polyfit(np.log10(xx[LIM[0]:LIM[1]]), np.log10(yy[LIM[0]:LIM[1]]), 1) alpha, b = fit[0], fit[1] fit_x = np.linspace(xx[LIM[0]], xx[-1], 100) fit_y = self.convert(fit_x, alpha, b) ax.plot(fit_x, fit_y, '-.', lw=1, color=COLOR) ax.loglog(xx, yy, '.', marker='.', markersize=8, color=COLOR, label='slope-{:.2f}'.format(abs(alpha))) plot_norm(ax, xlabel, ylabel, legend_loc='upper right', legend=True) else: ax.loglog(xx, yy, '.', marker='.', markersize=8, color=COLOR) plot_norm(ax, xlabel, ylabel, legend=False) fig.canvas.draw() fig.canvas.flush_events() with open('/'.join([self.output, self.status]) + '_%s.txt' % ylabel, 'w') as f: f.write('{}, {}\n'.format(xlabel, ylabel)) for i, j in zip(xx, yy): f.write('{}, {}\n'.format(i, j))
def plot_2cls_wave(self, TRAI_select_1, TRAI_select_2, same, value, valid=False): fig = plt.figure(figsize=(9.2, 3), num='Waveforms with same %s--%d μV' % (same, value)) fig.text(0.48, 0.24, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }, horizontalalignment="right") fig.text(0.975, 0.24, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }, horizontalalignment="right") i = self.data_tra[TRAI_select_1 - 1] if i[-1] != TRAI_select_1: print( 'Error: TRAI %d in data_tra is inconsistent with %d by input!' % (i[-1], TRAI_select_1)) return valid_time, valid_data = self.cal_wave(i, valid=valid) ax = fig.add_subplot(1, 2, 1) ax.plot(valid_time, valid_data, lw=0.5, color=self.color_1) ax.axhline(abs(i[2]), 0, valid_data.shape[0], linewidth=1, color="black") ax.axhline(-abs(i[2]), 0, valid_data.shape[0], linewidth=1, color="black") plot_norm(ax, xlabel='Time (μs)', ylabel='Amplitude (μV)', legend=False, grid=True) ax2 = fig.add_subplot(1, 2, 2) i = self.data_tra[TRAI_select_2 - 1] if i[-1] != TRAI_select_2: print( 'Error: TRAI %d in data_tra is inconsistent with %d by input!' % (i[-1], TRAI_select_2)) return valid_time, valid_data = self.cal_wave(i, valid=valid) ax2.plot(valid_time, valid_data, lw=0.5, color=self.color_2) ax2.axhline(abs(i[2]), 0, valid_data.shape[0], linewidth=1, color="black") ax2.axhline(-abs(i[2]), 0, valid_data.shape[0], linewidth=1, color="black") plot_norm(ax2, xlabel='Time (μs)', ylabel='Amplitude (μV)', legend=False, grid=True)
def cal_WaitingTime(self, fig, time, xlabel, ylabel, INTERVAL_NUM=None, COLOR='black', FIT=False, LIM=None, bin_method='log'): if INTERVAL_NUM is None: INTERVAL_NUM = 8 if LIM is None: LIM = [0, None] fig.subplots_adjust(left=0.139, bottom=0.189, right=0.975, top=0.962) fig.text(0.16, 0.22, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }) ax = fig.add_subplot() res = time[1:] - time[:-1] res = res[np.where(res != 0)[0]] if bin_method == 'linear': inter, mid = self.__cal_negtive_interval(res, 0.9 / INTERVAL_NUM) xx, yy = self.__cal_linear(sorted(np.array(res)), inter, mid, INTERVAL_NUM) elif bin_method == 'log': inter = self.__cal_log_interval(res) xx, yy = self.__cal_log(sorted(np.array(res)), inter, INTERVAL_NUM) if FIT: xx, yy = np.array(xx), np.array(yy) fit = np.polyfit(np.log10(xx[LIM[0]:LIM[1]]), np.log10(yy[LIM[0]:LIM[1]]), 1) alpha, b = fit[0], fit[1] fit_x = np.linspace(xx[0], xx[-1], 100) fit_y = self.convert(fit_x, alpha, b) ax.plot(fit_x, fit_y, '-.', lw=1, color=COLOR) ax.loglog(xx, yy, '.', markersize=8, marker='o', mec=COLOR, mfc='none', color=COLOR, label='slope-{:.2f}'.format(abs(alpha))) plot_norm(ax, xlabel, ylabel, legend_loc='upper right') else: ax.loglog(xx, yy, '.', markersize=8, marker='o', mec=COLOR, mfc='none', color=COLOR) plot_norm(ax, xlabel, ylabel, legend=False) fig.canvas.draw() fig.canvas.flush_events() with open('/'.join([self.output, self.status]) + '_WaitingTime.txt', 'w') as f: f.write('{}, {}\n'.format(xlabel, ylabel)) for i, j in zip(xx, yy): f.write('{}, {}\n'.format(i, j))
horizontalalignment="right") ax = plt.subplot() ax.loglog(Amp[idx_1], Eny[idx_1], '.', Marker='.', markersize=3, color='pink', label='Pop 1') ax.loglog(Amp[idx_2], Eny[idx_2], '.', Marker='.', markersize=3, color='r', label='Pop 2') ax.loglog(Amp[idx_3], Eny[idx_3], '.', Marker='.', markersize=3, color='b', label='Pop 3') ax.loglog(fit_x, fit_y1, '.', Marker='.', markersize=0.5, color='black') ax.loglog(fit_x, fit_y2, '.', Marker='.', markersize=0.5, color='black') plot_norm(ax, xlabelz[0], xlabelz[2], legend=True) for k, idx in enumerate([idx_1, idx_2, idx_3], 1): with open('%s-pop%d.txt' % (fold, k), 'w') as f: f.write('SetID, TRAI, Time, Chan, Thr, Amp, RiseT, Dur, Eny, RMS, Counts\n') for i in data_pri[np.where((data_pri[:, 2] == 3) & (data_pri[:, 1] < 5600))[0]][idx]: f.write('{:.0f}, {:.0f}, {:.8f}, {:.0f}, {:.7f}, {:.7f}, {:.2f}, {:.2f}, {:.7f}, {:.3f}, {:.0f}\n'.format( i[0], i[-1], i[1], i[2], i[3], i[4], i[5], i[6], i[7], i[8], i[9])) ''' ''' fig = plt.figure(figsize=[6, 3.9]) ax = plt.subplot() for k in TRAI[cls_KKM[0]]: tmp = data_tra[k - 1] sig = np.multiply(array.array('h', bytes(tmp[-2])), tmp[-3] * 1000) sig = (sig / max(sig)) time = np.linspace(0, pow(tmp[-5], -1) * (tmp[-4] - 1) * pow(10, 6), tmp[-4]) if time[-1] < 100: continue hx = fftpack.hilbert(sig) ax.semilogy(time, np.sqrt(sig**2 + hx**2), '.', Marker='.', color=color_1) for k in TRAI[cls_KKM[1]]: tmp = data_tra[k - 1] sig = np.multiply(array.array('h', bytes(tmp[-2])), tmp[-3] * 1000) sig = (sig / max(sig))
def plot_wave_TRAI(self, fig, k, data_pri, show_features=False, valid=False, cwt=False): # Waveform with specific TRAI try: if self.device == 'VALLEN': i = self.data_tra[k - 1] else: i = self.data_tra[k - self.data_tra[0][-1]] except IndexError: return str('Error: TRAI %d can not be found in data!' % k) if i[-1] != k: return str( 'Error: TRAI %d in data_tra is inconsistent with %d by input!' % (i[-1], k)) time, sig = self.cal_wave(i, valid=valid) for tmp_tail, s in enumerate(sig[::-1]): if s != 0: tail = -tmp_tail if tmp_tail > 0 else None break time, sig = time[:tail], sig[:tail] if cwt: fig.subplots_adjust(left=0.076, bottom=0.205, right=0.984, top=0.927, hspace=0.2, wspace=0.26) fig.text(0.47, 0.25, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }, horizontalalignment="right") ax = fig.add_subplot(1, 2, 2) ax.cla() Twxo, Wxo, ssq_freqs, *_ = ssq_cwt(sig, wavelet='morlet', scales='log-piecewise', fs=i[3], t=time) ax.contourf(time, ssq_freqs * 1000, pow(abs(Twxo), 0.5), cmap='cubehelix_r') plot_norm(ax, 'Time (μs)', 'Frequency (kHz)', y_lim=[min(ssq_freqs * 1000), 1000], legend=False) ax = fig.add_subplot(1, 2, 1) ax.cla() ax.plot(time, sig, lw=1) else: fig.subplots_adjust(left=0.115, bottom=0.17, right=0.975, top=0.95) fig.text(0.96, 0.2, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }, horizontalalignment="right") ax = fig.add_subplot() ax.cla() ax.plot(time, sig, lw=1) if self.device == 'vallen': if show_features: try: string = data_pri[np.where(data_pri[:, -1] == i[-1])][0] except IndexError: return str('Error: TRAI %d can not be found in data!' % k) print("=" * 23 + " Waveform information " + "=" * 23) for info, value, r in zip([ 'SetID', 'Time', 'Chan', 'Thr', 'Amp', 'RiseT', 'Dur', 'Eny', 'RMS', 'Counts', 'TRAI' ], [j for j in string], [0, 8, 0, 8, 8, 2, 2, 8, 8, 0, 0]): if r == 0: print('%s: %d' % (info, int(value))) else: print('%s: %s' % (info, round(value, r))) ax.axhline(abs(i[2]), 0, sig.shape[0], linewidth=1, color="black") ax.axhline(-abs(i[2]), 0, sig.shape[0], linewidth=1, color="black") elif self.device == 'pac': if show_features: # time, channel_num, sample_interval, points_num, dataset, hit_num # ID, Time(s), Chan, Thr(μV), Thr(dB), Amp(μV), Amp(dB), RiseT(s), Dur(s), Eny(aJ), RMS(μV), Frequency(Hz), Counts string = data_pri[np.where(data_pri[:, 0] == i[-1])][0] print("=" * 23 + " Waveform information " + "=" * 23) for info, value, r in zip([ 'Hit number', 'Time', 'Chan', 'Thr', 'Amp', 'RiseT', 'Dur', 'Eny', 'RMS', 'Counts' ], [ j for j in string[np.array( [0, 1, 2, 3, 5, 7, 8, 9, 10, 12])] ], [0, 7, 0, 8, 8, 7, 7, 8, 8, 0]): if r == 0: print('%s: %d' % (info, int(value))) else: print('%s: %s' % (info, round(value, r))) ax.axhline(abs(self.thr_μV), 0, sig.shape[0], linewidth=1, color="black") ax.axhline(-abs(self.thr_μV), 0, sig.shape[0], linewidth=1, color="black") plot_norm(ax, 'Time (μs)', 'Amplitude (μV)', legend=False, grid=True) # ================================================= 画图重绘与刷新 ================================================ fig.canvas.draw() fig.canvas.flush_events() with open( '/'.join([self.output, self.status]) + '-%d' % i[-1] + '.txt', 'w') as f: f.write('Time, Signal\n') for k in range(sig.shape[0]): f.write("{}, {}\n".format(time[k], sig[k]))
def plot_filtering(self, TRAI, N, CutoffFreq, btype, valid=False, originWave=False, filteredWave=True): """ Signal Filtering :param TRAI: :param N: Order of filter :param CutoffFreq: Cutoff frequency, Wn = 2 * cutoff frequency / sampling frequency, len(Wn) = 2 if btype in ['bandpass', 'bandstop'] else 1 :param btype: Filter Types, {'lowpass', 'highpass', 'bandpass', 'bandstop'} :param originWave: Whether to display the original waveform :param filteredWave: Whether to display the filtered waveform :param valid: Whether to truncate the waveform according to the threshold :return: """ tmp = self.data_tra[int(TRAI - 1)] if TRAI != tmp[-1]: print('Error: TRAI is incorrect!') time, sig = self.cal_wave(tmp, valid=valid) b, a = butter(N, list(map(lambda x: 2 * x * 1e3 / tmp[3], CutoffFreq)), btype) sig_filter = filtfilt(b, a, sig) if originWave: fig = plt.figure(figsize=(9.2, 3)) ax = fig.add_subplot(1, 2, 1) ax.plot(time, sig, lw=1, color='blue') plot_norm(ax, 'Time (μs)', 'Amplitude (μV)', title='TRAI: %d' % TRAI, legend=False, grid=True) ax = fig.add_subplot(1, 2, 2) Twxo, Wxo, ssq_freqs, *_ = ssq_cwt(sig, wavelet='morlet', scales='log-piecewise', fs=tmp[3], t=time) plt.contourf(time, ssq_freqs * 1000, abs(Twxo), cmap='jet') plot_norm(ax, r'Time (μs)', r'Frequency (kHz)', y_lim=[min(ssq_freqs * 1000), 1000], legend=False) if filteredWave: if btype in ['lowpass', 'highpass']: label = 'Frequency %s %d kHz' % ('<' if btype == 'lowpass' else '>', CutoffFreq) elif btype == 'bandpass': label = '%d kHz < Frequency < %d kHz' % (CutoffFreq[0], CutoffFreq[1]) else: label = 'Frequency < %d kHz or > %d kHz' % (CutoffFreq[0], CutoffFreq[1]) fig = plt.figure(figsize=(9.2, 3)) ax = fig.add_subplot(1, 2, 1) ax.plot(time, sig_filter, lw=1, color='gray', label=label) plot_norm(ax, 'Time (μs)', 'Amplitude (μV)', title='TRAI: %d (%s)' % (TRAI, btype), grid=True, frameon=False, legend_loc='upper right') ax = fig.add_subplot(1, 2, 2) Twxo, Wxo, ssq_freqs, *_ = ssq_cwt(sig_filter, wavelet='morlet', scales='log-piecewise', fs=tmp[3], t=time) plt.contourf(time, ssq_freqs * 1000, abs(Twxo), cmap='jet') plot_norm(ax, r'Time (μs)', r'Frequency (kHz)', y_lim=[min(ssq_freqs * 1000), 1000], legend=False) return sig_filter
def plot_envelope(self, TRAI, COLOR, features_path, valid=False, method='hl', xlog=False): fig = plt.figure(figsize=[6, 3.9]) fig.text(0.95, 0.17, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }, horizontalalignment="right") ax = plt.subplot() for idx, [trai, color] in enumerate(zip(TRAI, COLOR)): XX, YY = [], [] for k in tqdm(trai): tmp = self.data_tra[k - 1] time, sig = self.cal_wave(tmp, valid=valid) if time[-1] < 50: continue if method == 'se': sig = (sig / max(sig)) X_pos_frontier, X_neg_frontier = se.get_frontiers(sig, 0) XX.extend(np.linspace(0, time[-1], len(X_pos_frontier) - 2)) YY.extend(sig[X_pos_frontier[2:]]**2) if not xlog: ax.semilogy(np.linspace(0, time[-1], len(X_pos_frontier) - 2), sig[X_pos_frontier[2:]]**2, '.', Marker='.', color=color) else: ax.loglog(np.linspace(0, time[-1], len(X_pos_frontier) - 2), sig[X_pos_frontier[2:]]**2, '.', Marker='.', color=color) else: sig = sig**2 / max(sig**2) high_idx, low_idx = hl_envelopes_idx(sig, dmin=60, dmax=60) XX.extend(time[low_idx]) YY.extend(sig[low_idx]) if not xlog: ax.semilogy(time[low_idx], sig[low_idx], '.', Marker='.', color=color) else: ax.loglog(time[low_idx], sig[low_idx], '.', Marker='.', color=color) with open( '%s_Decay Function_Pop %d.txt' % (features_path[:-4], idx + 1), 'w') as f: f.write('Time (μs), Normalized A$^2$\n') for j in range(len(XX)): f.write('{}, {}\n'.format(XX[j], YY[j])) plot_norm(ax, 'Time (μs)', 'Normalized A$^2$', legend=False)
def plot_stream(self, k, staLen=3, overlap=1, staWin='hamming', IZCRT=0.3, ITU=150, alpha=1, t_backNoise=0, plot=True, classify=False, valid=False, t_str=0, t_end=float('inf')): """ Waveform stream segmentation program :param k: :param staLen: The duration of the window function, in μs The width of the window function = sampling frequency (MHz) * duration, so this value reflects the duration of the window function on the microsecond scale. :param overlap: The overlap coefficient of the window function, stride = the width of the window function - the overlap coefficient :param staWin: The name of window function,{'hamming', 'hanning', 'blackman', 'bartlett'} :param IZCRT: Identification Zero Crossing Threshold :param ITU: Identification Threshold Upper :param alpha: weighted factor for the standard deviation of the STE signal, where this a factor can be estimated along with the previous calibration for the thresholds, if a heavy background noise is expected the value for this weighting value must be incremented. :param t_backNoise: Used to evaluate background noise time :param plot: Whether to draw a figure to show :param classify: Whether to display the segmentation results in the figure :param valid: Whether to pre-cut the waveform stream according to the threshold :param t_str: Waveform stream data segmentation start time :param t_end: Waveform stream data segmentation cut-off time :return: """ tmp = self.data_tra[int(k - 1)] if tmp[0 if self.device == 'stream' else -1] != k: return str( 'Error: TRAI %d in data_tra is inconsistent with %d by input!' % (tmp[0 if self.device == 'stream' else -1], k)) time, sig = self.cal_wave(tmp, valid=valid) if t_str: range_idx = np.where((time >= t_str) & (time < t_end))[0] time = time[range_idx] - time[range_idx[0]] sig = sig[range_idx] width = int(tmp[3] * pow(10, -6) * staLen) stride = int(width) - overlap t_stE, stE = shortTermEny(sig, width, stride, tmp[3] * pow(10, -6), staWin) t_zcR, zcR = zerosCrossingRate(sig, width, stride, tmp[3] * pow(10, -6), staWin) stE_dev = cal_deriv(t_stE, stE) start, end = find_wave(stE, stE_dev, zcR, t_stE, IZCRT=IZCRT, ITU=ITU, alpha=alpha, t_backNoise=t_backNoise) if plot: x = [time, t_stE, t_stE, t_zcR] y = [sig, stE, stE_dev, zcR] color = ['black', 'green', 'gray', 'purple'] ylabel = [ r'$Amplitude$ $(μV)$', r'$STEnergy$ $(μV^2 \cdot μs)$', r'$S\dot{T}E$ $(μV^2)$', r'$ST\widehat{Z}CR$ $(\%)$' ] fig, axes = plt.subplots(4, 1, sharex=True, figsize=(10, 8)) for idx, ax in enumerate(axes): ax.plot(x[idx], y[idx], color=color[idx]) if classify: if idx == 0: for s, e in zip(start, end): ax.plot(time[np.where((time >= t_stE[s]) & (time <= t_stE[e]))[0]], sig[np.where((time >= t_stE[s]) & (time <= t_stE[e]))[0]], lw=1, color='red') ax.grid() plot_norm(ax, r'$Time$ $(μs)$', ylabel[idx], legend=False) return start, end, time, sig, t_stE
def cal_OmoriLaw(self, fig, tmp, xlabel, ylabel, INTERVAL_NUM=None, FIT=False, bin_method='log'): if INTERVAL_NUM is None: INTERVAL_NUM = 8 eny_lim = [[0.01, 0.1], [0.1, 1], [1, 10], [10, 1000], [1000, 10000]] tmp = self.__cal_OmiroLaw_helper(tmp, eny_lim) fig.subplots_adjust(left=0.115, bottom=0.17, right=0.975, top=0.95) fig.text(0.16, 0.21, self.status, fontdict={ 'family': 'Arial', 'fontweight': 'bold', 'fontsize': 12 }) ax = fig.add_subplot() for i, [marker, color, label] in enumerate( zip(['>', 'o', 'p', 'h', 'H'], [[1, 0, 1], [0, 0, 1], [ 0, 1, 0 ], [1, 0, 0], [0.5, 0.5, 0.5]], [ '$10^{-2}aJ<E_{MS}<10^{-1}aJ$', '$10^{-1}aJ<E_{MS}<10^{0}aJ$', '$10^{0}aJ<E_{MS}<10^{1}aJ$', '$10^{1}aJ<E_{MS}<10^{3}aJ$', '$10^{3}aJ<E_{MS}<10^{4}aJ$' ])): if len(tmp[i]): tmp[i] = np.array(tmp[i]) tmp[i] = tmp[i][np.where(tmp[i] != 0)[0]] if bin_method == 'linear': inter, mid = self.__cal_negtive_interval( tmp[i], 0.9 / INTERVAL_NUM) xx, yy = self.__cal_linear(sorted(np.array(tmp[i])), inter, mid, INTERVAL_NUM) elif bin_method == 'log': inter = self.__cal_log_interval(tmp[i]) xx, yy = self.__cal_log(sorted(np.array(tmp[i])), inter, INTERVAL_NUM) if FIT: xx, yy = np.array(xx), np.array(yy) fit = np.polyfit(np.log10(xx), np.log10(yy), 1) alpha, b = fit[0], fit[1] fit_x = np.linspace(xx[0], xx[-1], 100) fit_y = self.convert(fit_x, alpha, b) ax.plot(fit_x, fit_y, '-.', lw=1, color=color) ax.loglog(xx, yy, markersize=8, marker=marker, mec=color, mfc='none', color=color, label='{}--{:.2f}'.format(label, abs(alpha))) else: ax.loglog(xx, yy, markersize=8, marker=marker, mec=color, mfc='none', color=color, label=label) plot_norm(ax, xlabel, ylabel, legend_loc='upper right') fig.canvas.draw() fig.canvas.flush_events() with open( '/'.join([self.output, self.status]) + '_OmoriLaw_(%s).txt' % label[1:-1].replace('<', ' ').replace('>', ' '), 'w') as f: f.write('t-t_{MS} (s), r_{AS}(t-t_{MS})(s^{-1})\n') for i, j in zip(xx, yy): f.write('{}, {}\n'.format(i, j))
def linear_matching(tmp1, tmp2, xlabel, ylabel, slope, intercept): idx_1, idx_2 = [], [] formula = lambda x, a, b: pow(x, a) * pow(10, b) fit_x = cal_fitx(tmp1, 1000) if len(slope) == 1: fit_y = [formula(i, slope[0], intercept[0]) for i in fit_x] label = cal_label(tmp1, tmp2, formula, slope, intercept) idx_1 = np.where(np.array(label) == 0)[0] idx_2 = np.where(np.array(label) == 1)[0] fig = plt.figure(figsize=[6, 3.9]) ax = plt.subplot() ax.loglog(tmp1[idx_1], tmp2[idx_1], '.', Marker='.', markersize=8, color='red', label='Pop 1') ax.loglog(tmp1[idx_2], tmp2[idx_2], '.', Marker='.', markersize=8, color='blue', label='Pop 2') ax.loglog(fit_x, fit_y, '.', Marker='.', markersize=0.5, color='black') plot_norm(ax, xlabel, ylabel, legend=True) elif len(slope) == 2: fit_y1 = [formula(i, slope[0], intercept[0]) for i in fit_x] fit_y2 = [formula(i, slope[1], intercept[1]) for i in fit_x] label = cal_label(tmp1, tmp2, formula, slope, intercept) idx_1 = np.where(np.array(label) == 0)[0] idx_2 = np.where(np.array(label) == 1)[0] idx_3 = np.where(np.array(label) == 2)[0] fig = plt.figure(figsize=[6, 3.9]) ax = plt.subplot() ax.loglog(tmp1[idx_1], tmp2[idx_1], '.', Marker='.', markersize=8, color='black', label='Pop 1') ax.loglog(tmp1[idx_2], tmp2[idx_2], '.', Marker='.', markersize=8, color='r', label='Pop 2') ax.loglog(tmp1[idx_3], tmp2[idx_3], '.', Marker='.', markersize=8, color='b', label='Pop 3') ax.loglog(fit_x, fit_y1, '.', Marker='.', markersize=0.5, color='black') ax.loglog(fit_x, fit_y2, '.', Marker='.', markersize=0.5, color='black') plot_norm(ax, xlabel, ylabel, legend=True) else: print("Current function don't support fit more than two lines.") return idx_1, idx_2