mirror of
https://github.com/paulscherrerinstitute/sf_daq_broker.git
synced 2026-04-30 14:32:23 +02:00
Data Backend account
79 lines
2.8 KiB
Python
79 lines
2.8 KiB
Python
import numpy
|
|
import h5py
|
|
import numpy
|
|
import zmq
|
|
class DetectorRawWriter(object):
|
|
def __init__(self, output_file, n_images, n_modules, metadata):
|
|
self.output_file = output_file
|
|
self.n_modules = n_modules
|
|
self.metadata = metadata
|
|
# TODO: Put real name here, please.
|
|
self.detector_name = "detector"
|
|
|
|
self.pulse_id_cache = numpy.zeros(shape=n_images, dtype="uint64")
|
|
self.frame_index_cache = numpy.zeros(shape=n_images, dtype="uint64")
|
|
self.daq_rec_cache = numpy.zeros(shape=n_images, dtype="uint32")
|
|
self.is_good_frame_cache = numpy.zeros(shape=n_images, dtype="uint8")
|
|
|
|
self.file = h5py.File(self.output_file, "w")
|
|
|
|
self._create_metadata_datasets(metadata)
|
|
|
|
self.current_write_index = 0
|
|
|
|
self._image_dataset = self.file.create_dataset(
|
|
name="/data/" + self.detector_name + "/data",
|
|
dtype="uint16",
|
|
shape=(n_images, self.n_modules * 512, 1024),
|
|
chunks=(1, self.n_modules * 512, 1024)
|
|
)
|
|
|
|
def _create_metadata_datasets(self, metadata):
|
|
|
|
for key, value in metadata.items():
|
|
self.file.create_dataset(key, data=numpy.string_(value))
|
|
|
|
def write(self, pulse_id, meta_buffer, data_buffer):
|
|
self._image_dataset.id.write_direct_chunk((self.current_write_index, 0, 0),
|
|
data_buffer[:])
|
|
|
|
self.pulse_id_cache[self.current_write_index] = meta_buffer["pulse_id"]
|
|
self.frame_index_cache[self.current_write_index] = meta_buffer["frame_index"]
|
|
self.daq_rec_cache[self.current_write_index] = meta_buffer["daq_rec"]
|
|
self.is_good_frame_cache[self.current_write_index] = meta_buffer["is_good_image"]
|
|
|
|
self.current_write_index += 1
|
|
|
|
def _write_metadata(self):
|
|
self.file["/data/" + self.detector_name + "/pulse_id"] = self.pulse_id_cache
|
|
self.file["/data/" + self.detector_name + "/frame_index"] = self.frame_index_cache
|
|
self.file["/data/" + self.detector_name + "/daq_rec"] = self.daq_rec_cache
|
|
self.file["/data/" + self.detector_name + "/is_good_frame"] = self.is_good_frame_cache
|
|
|
|
def close(self):
|
|
self._write_metadata()
|
|
self.file.close()
|
|
|
|
|
|
output_file = "/sf/alvra/data/p17502/raw/test_python_direct.h5"
|
|
n_images = 10000
|
|
n_modules = 4
|
|
metadata = {}
|
|
|
|
meta_buffer = {
|
|
"pulse_id": 12345,
|
|
"daq_rec": 12346,
|
|
"frame_index": 12345,
|
|
"is_good_image": True
|
|
}
|
|
|
|
source = numpy.zeros(dtype="uint16", shape=[512 * n_modules, 1024]) + 5
|
|
data_buffer = source.tobytes()
|
|
|
|
detector_writer = DetectorRawWriter(output_file=output_file, n_images=n_images, n_modules=n_modules, metadata=metadata)
|
|
|
|
for i in range(n_images):
|
|
detector_writer.write(i, meta_buffer, data_buffer)
|
|
|
|
detector_writer.close()
|