def test_run_get_dataframe(mock_fxe_run): run = RunDirectory(mock_fxe_run) df = run.get_dataframe(fields=[("*_XGM/*", "*.i[xy]Pos*")]) assert len(df.columns) == 4 assert "SA1_XTD2_XGM/DOOCS/MAIN/beamPosition.ixPos" in df.columns df2 = run.get_dataframe(fields=[("*_XGM/*", "*.i[xy]Pos*")], timestamps=True) assert len(df2.columns) == 8 assert "SA1_XTD2_XGM/DOOCS/MAIN/beamPosition.ixPos" in df2.columns assert "SA1_XTD2_XGM/DOOCS/MAIN/beamPosition.ixPos.timestamp" in df2.columns
from time import monotonic from karabo_data import RunDirectory print("Opening raw run...") start = monotonic() run = RunDirectory('/gpfs/exfel/exp/SA1/201830/p900025/raw/r0150/') delta = monotonic() - start print(len(run.files), "files") print(delta, "seconds") print() print("Retrieving data frame for XGM ixPos & iyPos...") start = monotonic() df = run.get_dataframe(fields=[("*_XGM/*", "*.i[xy]Pos"), ("*_XGM/*", "*.photonFlux")]) delta = monotonic() - start print(delta, "seconds") print(df.head())
class XasAnalyzer(abc.ABC): """Abstract class for Xray Absoprtion Spectroscopy analysis.""" def __init__(self, run_folder): """Initialization. :param str run_folder: full path of the run folder. """ self._run = RunDirectory(run_folder) self._sources = { 'MONO': 'SA3_XTD10_MONO/MDL/PHOTON_ENERGY', 'XGM': 'SCS_BLU_XGM/XGM/DOOCS', 'XGM_OUTPUT': 'SCS_BLU_XGM/XGM/DOOCS:output', 'SA3_XGM': 'SA3_XTD10_XGM/XGM/DOOCS', 'SA3_XGM_OUTPUT': 'SA3_XTD10_XGM/XGM/DOOCS:output' } # get the DataFrame for XGM control data self._xgm_df = self._run.get_dataframe(fields=[(self._sources['XGM'], '*value')]) self._xgm_df.rename(columns=lambda x: x.split('/')[-1], inplace=True) self._sa3_xgm_df = self._run.get_dataframe( fields=[(self._sources['SA3_XGM'], '*value')]) self._sa3_xgm_df.rename(columns=lambda x: x.split('/')[-1], inplace=True) # get the DataFrame for SoftMono control data self._mono_df = self._run.get_dataframe(fields=[(self._sources['MONO'], '*value')]) self._mono_df.rename(columns=lambda x: x.split('/')[-1], inplace=True) self._photon_energies = None # photon energies for each pulse self._I0 = None self._I1 = OrderedDict() self._data = None # pulse-resolved data in DataFrame def info(self): """Print out information of the run(s).""" first_train = self._run.train_ids[0] last_train = self._run.train_ids[-1] train_count = len(self._run.train_ids) span_sec = (last_train - first_train) / 10 span_txt = str(datetime.timedelta(seconds=span_sec)) photon_energies = self._mono_df['actualEnergy'] print('# of trains: ', train_count) print('Duration: ', span_txt) print('First train ID: ', first_train) print('Last train ID: ', last_train) print('Min photon energy: ', round(photon_energies.min(), 4), 'eV') print('Max photon energy: ', round(photon_energies.max(), 4), 'eV') print('MCP channels:') for ch, value in self._channels.items(): print(' - {}: {}'.format(ch, value['raw'])) def _check_sources(self): """Check all the required sources are in the data.""" sources = self._run.all_sources for src in self._sources.values(): if src not in sources: raise ValueError("Source not found: {}!".format(src)) def plot_xgm_run(self, *, figsize=(8, 5.6)): """Plot the train resolved data from XGM. :param tuple figsize: figure size. """ import matplotlib.pyplot as plt plt.rcParams['font.size'] = 12 fig, (ax1, ax2) = plt.subplots(2, 1, figsize=figsize) ax1_tw = ax1.twinx() ln1 = ax1.plot(self._xgm_df['pulseEnergy.photonFlux'], label=r"Pulse energy ($\mu$J)") # "nummberOfBrunches" is indeed the name in the Karabo Device # implementation. For more details, please check # https://git.xfel.eu/gitlab/karaboDevices/xgmDoocs number_of_bunches = self._xgm_df['pulseEnergy.nummberOfBrunches'] ln2 = ax1_tw.plot(number_of_bunches, label="Number of pulses", c='g') lns = ln1 + ln2 labels = [l.get_label() for l in lns] ax1.legend(lns, labels) ax1.set_ylabel(r"Pulse energy ($\mu$J)") ax1_tw.set_ylabel("Number of pulses") if number_of_bunches.max() - number_of_bunches.min() < 5: mean_n_bunches = int(number_of_bunches.mean()) ax1_tw.set_ylim((mean_n_bunches - 4.5, mean_n_bunches + 4.5)) ax2.plot(1000 * self._xgm_df['beamPosition.ixPos'], label="x") ax2.plot(1000 * self._xgm_df['beamPosition.iyPos'], label="y") ax2.set_xlabel("Train ID") ax2.set_ylabel(r"Beam position ($\mu$m)") ax2.legend() fig.tight_layout() return fig, (ax1, ax1_tw, ax2) def plot_xgm_train(self, *, index=0, train_id=None, figsize=(8, 5.6)): """Plot xgm measurement in a given train. :param int index: train index. Ignored if train_id is given. :param int train_id: train ID. :param tuple figsize: figure size. """ import matplotlib.pyplot as plt plt.rcParams['font.size'] = 12 key = "data.intensityTD" filtered = self._run.select("*XGM/*", key) if train_id is None: tid, data = filtered.train_from_index(index) else: tid, data = filtered.train_from_id(train_id) fig, (ax1, ax2) = plt.subplots(2, 1, figsize=figsize) ax1.plot(data[self._sources['SA3_XGM_OUTPUT']][key], marker='.') ax2.plot(data[self._sources['XGM_OUTPUT']][key], marker='.') for ax in (ax1, ax2): ax.set_ylabel(r"Pulse energy ($\mu$J)") ax.set_xlim((-0.5, 100.5)) ax1.set_title("SA3 XGM") ax2.set_title("SCS XGM") ax2.set_xlabel("Pulse ID") fig.suptitle("Train ID: {}".format(tid)) fig.tight_layout(rect=[0, 0.03, 1, 0.95]) return fig, (ax1, ax2) @abc.abstractmethod def process(self, *args, **kwargs): """Process the run data. :return: the current instance. """ pass def select(self, keys, lower=-np.inf, upper=np.inf): """Select data within the given boundaries. It modifies the internal data inplace. :param str/list/tuple/numpy.ndarray: key(s) for applying the filter. :param float lower: lower boundary (included). :param float upper: higher boundary (included). :return: the current instance. """ n0 = len(self._data) if isinstance(keys, (list, tuple, np.ndarray)): # TODO: remove this for loop for key in keys: self._data.query("{} <= {} <= {}".format(lower, key, upper), inplace=True) else: self._data.query("{} <= {} <= {}".format(lower, keys, upper), inplace=True) print("{} out of {} data are selected!".format(len(self._data), n0)) return self @property @abc.abstractmethod def data(self): """Get the pulse-resolved data in pandas.DataFrame.""" pass @abc.abstractmethod def compute_total_absorption(self): """Compute absorption for all data.""" pass @abc.abstractmethod def compute_spectrum(self, n_bins=20, point_wise=False): """Compute spectrum. :param int n_bins: number of energy bins. :param bool point_wise: if True, calculate the absorption point wise and then average. Otherwise, average over I0 and I1 first and then calculate the absorption. Default = False """ pass @abc.abstractmethod def plot_correlation(self, *args, **kwargs): """Generate correlation plots.""" pass @abc.abstractmethod def plot_spectrum(self, *args, **kwargs): """Generate spectrum plots."""