88 lines
2.7 KiB
C++
88 lines
2.7 KiB
C++
// Copyright (2019-2023) Paul Scherrer Institute
|
|
|
|
#include "HDF5Writer.h"
|
|
#include "HDF5NXmx.h"
|
|
#include <nlohmann/json.hpp>
|
|
|
|
HDF5Writer::HDF5Writer(const StartMessage &request)
|
|
: images_per_file(request.images_per_file),
|
|
file_prefix(request.file_prefix),
|
|
max_spot_count(request.max_spot_count),
|
|
az_int_bin_to_q(request.az_int_bin_to_q),
|
|
user_data(request.user_data) {}
|
|
|
|
void HDF5Writer::Write(const DataMessage& message) {
|
|
std::lock_guard<std::mutex> lock(hdf5_mutex);
|
|
if (message.image.size == 0)
|
|
return;
|
|
|
|
if (message.number < 0)
|
|
throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds, "No support for negative images");
|
|
|
|
size_t file_number = 0;
|
|
size_t image_number = message.number;
|
|
if (images_per_file > 0) {
|
|
file_number = message.number / images_per_file;
|
|
image_number = message.number % images_per_file;
|
|
}
|
|
if (files.size() <= file_number)
|
|
files.resize(file_number + 1);
|
|
|
|
if (!files[file_number])
|
|
files[file_number] = std::make_unique<HDF5DataFile>(HDF5Metadata::DataFileName(file_prefix, file_number),
|
|
az_int_bin_to_q,
|
|
file_number * images_per_file,
|
|
max_spot_count);
|
|
// Ignore zero size images
|
|
if (message.image.size > 0)
|
|
files[file_number]->Write(message, image_number);
|
|
|
|
if (files[file_number]->GetNumImages() == images_per_file)
|
|
AddStats(files[file_number]->Close());
|
|
}
|
|
|
|
std::vector<HDF5DataFileStatistics> HDF5Writer::Finalize() {
|
|
std::lock_guard<std::mutex> lock(hdf5_mutex);
|
|
for (auto &f: files) {
|
|
if (f)
|
|
AddStats(f->Close());
|
|
}
|
|
return stats;
|
|
}
|
|
|
|
void HDF5Writer::AddStats(const std::optional<HDF5DataFileStatistics>& s) {
|
|
if (!s)
|
|
return;
|
|
|
|
stats.push_back(*s);
|
|
if (socket) {
|
|
nlohmann::json j;
|
|
j["filename"] = s->filename;
|
|
j["nimages"] = s->total_images;
|
|
if (!user_data.empty()) {
|
|
nlohmann::json j_userdata;
|
|
|
|
// if user_data is valid json, interpret it as such, otherwise embed as string
|
|
try {
|
|
j_userdata = nlohmann::json::parse(user_data);
|
|
} catch (...) {
|
|
j_userdata = user_data;
|
|
}
|
|
j["user_data"] = j_userdata;
|
|
}
|
|
socket->Send(j.dump());
|
|
}
|
|
}
|
|
|
|
void HDF5Writer::SetupSocket(const std::string &addr) {
|
|
socket = std::make_unique<ZMQSocket>(ZMQSocketType::Pub);
|
|
socket->Bind(addr);
|
|
}
|
|
|
|
std::optional<std::string> HDF5Writer::GetZMQAddr() {
|
|
if (socket) {
|
|
return socket->GetEndpointName();
|
|
} else
|
|
return {};
|
|
}
|