diff --git a/extra_foam/gui/ctrl_widgets/pump_probe_ctrl_widget.py b/extra_foam/gui/ctrl_widgets/pump_probe_ctrl_widget.py index 1689d8ff..d54001bb 100644 --- a/extra_foam/gui/ctrl_widgets/pump_probe_ctrl_widget.py +++ b/extra_foam/gui/ctrl_widgets/pump_probe_ctrl_widget.py @@ -65,6 +65,8 @@ def __init__(self, *args, **kwargs): self._abs_difference_cb.setChecked(True) self._reset_btn = QPushButton("Reset") + self.save_btn = QPushButton("Save to file") + self.save_btn.setFixedWidth(150) self.initUI() self.initConnections() @@ -88,6 +90,7 @@ def initUI(self): layout.addWidget(self._off_pulse_le, 1, 3) layout.addWidget(self._abs_difference_cb, 0, 4, AR) + layout.addWidget(self.save_btn, 2, 4, AR) self.setLayout(layout) diff --git a/extra_foam/gui/windows/pump_probe_w.py b/extra_foam/gui/windows/pump_probe_w.py index d42bca3e..f697d0bf 100644 --- a/extra_foam/gui/windows/pump_probe_w.py +++ b/extra_foam/gui/windows/pump_probe_w.py @@ -7,8 +7,13 @@ Copyright (C) European X-Ray Free-Electron Laser Facility GmbH. All rights reserved. """ +from enum import Enum +import os +import os.path as osp + +import numpy as np from PyQt5.QtCore import Qt -from PyQt5.QtWidgets import QFrame, QVBoxLayout, QSplitter +from PyQt5.QtWidgets import QFileDialog, QFrame, QVBoxLayout, QSplitter from .base_window import _AbstractPlotWindow from ..ctrl_widgets import PumpProbeCtrlWidget @@ -91,6 +96,11 @@ def refresh(self): self._plot.setData(x, y) +class SaveFile(Enum): + NPZ = "NumPy Binary File (*.npz)" + TXT = "Text File (*.txt)" + + class PumpProbeWindow(_AbstractPlotWindow): """PumpProbeWindow class.""" _title = "Pump-probe" @@ -138,8 +148,61 @@ def initUI(self): def initConnections(self): """Override.""" - pass + self._ctrl_widget.save_btn.clicked.connect(self.saveToFile) def closeEvent(self, QCloseEvent): self._ctrl_widget.resetAnalysisType() super().closeEvent(QCloseEvent) + + def saveToFile(self): + # Get the current data + if not len(self._queue): + return + data = self._queue[0] + + # Open file dialog + filepath, filter = QFileDialog.getSaveFileName( + caption="Save image", + directory=osp.expanduser("~"), + filter=f"{SaveFile.NPZ.value};;{SaveFile.TXT.value}") + filter = SaveFile(filter) + + # Validate filepath + if not filepath: + return + suffix = f".{filter.name.lower()}" + if not filepath.lower().endswith(suffix): + filepath += suffix + + # Copy reference file from tmp folder to desired destination + os.makedirs(osp.dirname(filepath), exist_ok=True) + + # Save the data + pp = data.pp + train_id = data.tid + position = pp.x + intensity_on = pp.y_on + intensity_off = pp.y_off + intensity_sub = pp.y + + if position is None: + return + + if filter is SaveFile.NPZ: + np.savez(filepath, + trainId=train_id, + position=position, + intensity_on=intensity_on, + intensity_off=intensity_off, + intensity_subtracted=intensity_sub) + elif filter is SaveFile.TXT: + # write out conversion to tabular ASCII + with open(filepath, 'w') as f: + f.write(f" Train-ID: {train_id}\n\n") + f.write(' q I_on ' + ' I_off I_diff\n\n') + for i, q in enumerate(position): + f.write(f"{q:12f}" + f"{intensity_on[i]:16f}" + f"{intensity_off[i]:16f}" + f"{intensity_sub[i]:16f}\n")