Skip to content
Snippets Groups Projects
Commit 0b5a3917 authored by Vadim Gubaidulin's avatar Vadim Gubaidulin
Browse files

Added class from postmortem postprocessing.

parent 7082f04b
No related branches found
No related tags found
No related merge requests found
......@@ -13,6 +13,184 @@ Q_X = 18.16
Q_Y = 10.22
OMEGA_REV = 2*pi*c/CIRCUMFERENCE
r_p = e**2/(m_p*c**2)
import h5py as hp
import numpy as np
import pandas as pd
import os
import nafflib
import argparse
import matplotlib.pyplot as plt
from scipy.signal import periodogram, lombscargle, hilbert
from FITX import fit_risetime
class FBTPostmortemViewer():
def __init__(self, filename, h_rf):
self.filename = filename
self.filling_pattern_calibration = 100
self.data_type = np.float32
self.h_rf = h_rf
def read_data(self):
_, extension = os.path.splitext(self.filename)
# self.h_rf = h_rf
# print(filename, extension)
if extension == ".nxs":
with hp.File(filename, "r") as file:
print(file['post-mortem/rcm/ans/rf/bbfdataviewer.1/'].keys())
data = np.array(file['post-mortem/rcm/ans/rf/bbfdataviewer.1/clean_data'][:])
print('Shape of nexus raw data', data.shape)
len_data = len(data)
len_common = int(len_data // self.h_rf)
resize_data = int(len_common * self.h_rf)
data_bunch_x = data[:resize_data].reshape((len_common, self.h_rf)).astype(np.float32)
data_bunch_y = data_bunch_x
# data_bunch_x = data_bunch_x.T
# data_bunch_y = data_bunch_y.T
# data_bunch -= np.mean(data_bunch, axis=0, dtype=np.int64)
elif extension == ".txt":
data_pd = pd.read_csv(self.filename, sep="\t", header=None)
data_bunch_x = data_pd.to_numpy()
data_bunch_y = data_bunch_x
# data_bunch -= np.mean(data_bunch, axis=0, dtype=np.int64)
elif extension == ".hdf5":
with hp.File(self.filename) as f:
# print(f['V'].shape)
data_bunch_x = np.array(f['H'], dtype=np.int16)[:, :]
# data_bunch_x -= np.mean(data_bunch_x, axis=1, dtype=np.int64, keepdims=True)
data_bunch_y = np.array(f['V'], dtype=np.int16)[:, :]
# data_bunch_y -= np.mean(data_bunch_y, axis=1, dtype=np.int64, keepdims=True)
data_bunch_x = data_bunch_x.T
data_bunch_y = data_bunch_y.T
else:
raise ValueError()
# return data_bunch_x, data_bunch_y
self.data_bunch_x = data_bunch_x.astype(np.float32)
self.data_bunch_y = data_bunch_y.astype(np.float32)
print(f'Shape of the data is' ,data_bunch_x.shape)
return 0
def process_data(self):
print(self.data_bunch_y.shape)
self.filling_pattern = np.mean(self.data_bunch_y, axis=0, dtype=np.float32)
self.data_bunch_x -= np.mean(self.data_bunch_x, axis=0, dtype=np.float32, keepdims=True)
self.data_bunch_y -= np.mean(self.data_bunch_y, axis=0, dtype=np.float32, keepdims=True)
self.spectrum_x = periodogram(self.data_bunch_x.flatten(), fs=self.h_rf)
self.spectrum_y = periodogram(self.data_bunch_y.flatten(), fs=self.h_rf)
self.spectrum_x1 = periodogram(self.data_bunch_x.flatten()[0::4], fs=self.h_rf//4)
self.spectrum_x2 = periodogram(self.data_bunch_x.flatten()[1::4], fs=self.h_rf//4)
self.spectrum_x3 = periodogram(self.data_bunch_x.flatten()[2::4], fs=self.h_rf//4)
self.spectrum_x4 = periodogram(self.data_bunch_x.flatten()[3::4], fs=self.h_rf//4)
self.spectrum_y1 = periodogram(self.data_bunch_y.flatten()[0::4], fs=self.h_rf//4)
self.spectrum_y2 = periodogram(self.data_bunch_y.flatten()[1::4], fs=self.h_rf//4)
self.spectrum_y3 = periodogram(self.data_bunch_y.flatten()[2::4], fs=self.h_rf//4)
self.spectrum_y4 = periodogram(self.data_bunch_y.flatten()[3::4], fs=self.h_rf//4)
# self.tunes_x = [nafflib.get_tune(self.data_bunch_x[:, i]) for i in range(self.h_rf)]
# self.tunes_y = [nafflib.get_tune(self.data_bunch_y[:, i]) for i in range(self.h_rf)]
self.fft_x = np.asarray([np.fft.rfft(self.data_bunch_x[:, i]) for i in range(self.h_rf)]).T
self.fft_y = np.asarray([np.fft.rfft(self.data_bunch_y[:, i]) for i in range(self.h_rf)]).T
self.fft_x = np.fft.rfft(self.data_bunch_x, axis=0).T
self.fft_y = np.fft.rfft(self.data_bunch_x, axis=1).T
self.fftfreq = np.fft.rfftfreq(self.data_bunch_x.shape[0])
# self.bunch_phase = np.zeros((self.h_rf,))
# self.dominant_freqs = np.zeros((self.h_rf,))
# magnitudes = np.abs(self.fft_x)
# for i in range(self.h_rf):
# idx = np.argmax(magnitudes[1:]) + 1
# self.dominant_freqs[i] = self.fftfreq[idx]
# self.bunch_phase[i] = np.angle(fft_result[idx])
def save_processed_data(self):
pass
def plot_spectrum(self, ax=None, plane='y'):
plane = list(plane)
if ax == None:
_, ax = plt.subplots(1, 1)
if 'x' in plane:
ax.plot(self.spectrum_x[0], self.spectrum_x[1], label='horizontal')
if 'y' in plane:
ax.plot(self.spectrum_y[0], self.spectrum_y[1], label='vertical')
ax.set_xlim(0, )
ax.set_ylim(0, )
ax.legend()
def plot_partial_spectrum(self, ax=None, plane='y'):
plane = list(plane)
if ax == None:
_, ax = plt.subplots(1, 1)
if 'x' in plane:
ax.plot(self.spectrum_x1[0], self.spectrum_x1[1], label='horizontal, ADC1')
ax.plot(self.spectrum_x2[0], self.spectrum_x2[1], label='horizontal, ADC2')
ax.plot(self.spectrum_x3[0], self.spectrum_x3[1], label='horizontal, ADC3')
ax.plot(self.spectrum_x4[0], self.spectrum_x4[1], label='horizontal, ADC4')
if 'y' in plane:
ax.plot(self.spectrum_y1[0], self.spectrum_y1[1], label='vertical, ADC1')
ax.plot(self.spectrum_y2[0], self.spectrum_y2[1], label='vertical, ADC2')
ax.plot(self.spectrum_y3[0], self.spectrum_y3[1], label='vertical, ADC3')
ax.plot(self.spectrum_y4[0], self.spectrum_y4[1], label='vertical, ADC4')
ax.legend()
def plot_offset(self, ax=None, plane='y'):
plane = list(plane)
if ax == None:
_, ax = plt.subplots(1, 1)
if 'y' in plane:
ax.plot(self.data_bunch_y.flatten())
def plot_bunch_offset(self, bunch_numbers, axs=None, plane='y'):
plane = list(plane)
if axs == None:
fig, axs = plt.subplots(len(bunch_numbers), 1)
for i in range(len(bunch_numbers)):
if 'y' in plane:
axs[i].plot(self.data_bunch_y[:, bunch_numbers[i]], label=f'vertical, bunch {bunch_numbers[i]}')
if 'x' in plane:
axs[i].plot(self.data_bunch_x[:, bunch_numbers[i]], label=f'horizontal, bunch {bunch_numbers[i]}')
fig.legend()
def plot_fitexp(self, ax=None, plane='y'):
if ax == None:
_, ax = plt.subplots(1, 1)
if plane == 'y':
ax.plot(self.data_bunch_y.flatten())
def plot_filling_pattern(self, ax=None):
if ax == None:
_, ax = plt.subplots(1, 1)
ax.plot(np.abs(self.filling_pattern) / self.filling_pattern_calibration)
ax.set_xlabel('Bunch number')
ax.set_ylabel('Bunch current (?mA?)')
def plot_tunes(self, ax=None):
if ax == None:
_, ax = plt.subplots(1, 1)
# ax.plot(self.fftfreq, self.fft_x)
ax.plot(self.fftfreq, np.abs(self.fft_y))
def plot_mean_fft(self, ax=None):
if ax == None:
_, ax = plt.subplots(1, 1)
ax.plot(self.fftfreq, np.mean(np.abs(self.fft_x), axis=1))
def plot_tune_vs_bunch_number(self, ax=None):
if ax == None:
_, ax = plt.subplots(1, 1)
ax.plot(self.tunes_x, label='horizontal', alpha=0.5)
ax.plot(self.tunes_y, label='vertical', alpha=0.5)
ax.set_xlabel('Bunch number')
ax.set_ylabel('Fractional tune')
ax.legend()
def fit_growth_rate(self):
pass
class GrowthDampDataAnalyzer():
def __init__(self, folder):
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment