Files
Jungfraujoch/receiver/JFJochReceiverRugnux.cpp
T

485 lines
23 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "JFJochReceiverRugnux.h"
#include <cmath>
#include <filesystem>
#include <sys/resource.h>
#include <sys/syscall.h>
#include <unistd.h>
#include "spdlog/fmt/fmt.h"
#include "JFJochReceiverLite.h"
#include "../common/ParallelFor.h"
#include "../frame_serialize/CBORStream2Deserializer.h"
#include "../rugnux/RugnuxLive.h"
#include "../rugnux/ResultReport.h"
#include "../writer/RugnuxResultBlocks.h"
using namespace std::chrono_literals;
namespace {
int64_t SteadyNowNs() {
return std::chrono::duration_cast<std::chrono::nanoseconds>(
std::chrono::steady_clock::now().time_since_epoch()).count();
}
double SecondsSince(int64_t ns) {
return static_cast<double>(SteadyNowNs() - ns) / 1e9;
}
// A run reachable for Cut while it exists: registered in `slot` (under `m`, cancelled at once where the
// cut came first) and taken off again before it is destroyed, an exception included.
class Reachable {
std::mutex &m;
Rugnux *&slot;
public:
Reachable(std::mutex &m, Rugnux *&slot, Rugnux &run, const bool &cut, const std::atomic<bool> &cancelled)
: m(m), slot(slot) {
std::unique_lock ul(m);
slot = &run;
if (cut || cancelled)
run.Cancel();
}
~Reachable() {
std::unique_lock ul(m);
slot = nullptr;
}
};
// What rugnux reports as it goes: its per-image results into the receiver's plots and scan
// result, which a later pass over the same image overwrites - so the canonical pass, the last to
// read every image, is what stays - and its steps into the log.
class LiveObserver : public RugnuxObserver {
JFJochReceiverPlots &plots;
ScanResultGenerator &scan_result;
const AzimuthalIntegrationMapping &mapping;
RugnuxLiveStatus &status;
Logger &logger;
public:
LiveObserver(JFJochReceiverPlots &plots, ScanResultGenerator &scan_result,
const AzimuthalIntegrationMapping &mapping, RugnuxLiveStatus &status, Logger &logger)
: plots(plots), scan_result(scan_result), mapping(mapping), status(status), logger(logger) {}
void OnImageProcessed(const DataMessage &msg) override {
plots.Add(msg, AzimuthalIntegrationProfile(mapping));
scan_result.Add(msg);
}
void OnProgress(uint64_t done, uint64_t total) override { status.Progress(done, total); }
void OnStep(const std::string &title) override {
status.Step(title);
logger.Info("rugnux: {}", title);
}
void OnNote(const std::string &text) override {
status.Note(text);
logger.Info("rugnux: {}", text);
}
};
}
JFJochReceiverRugnux::JFJochReceiverRugnux(const DiffractionExperiment &in_experiment,
const PixelMask &in_pixel_mask,
ImagePuller &in_image_puller,
ImagePusher &in_image_pusher,
Logger &in_logger,
int64_t in_forward_nthreads,
int in_rugnux_nthreads,
const std::string &in_model_path,
double in_deadline_s,
RugnuxLiveStatus &in_live_status,
const SpotFindingSettings &in_spot_finding_settings,
JFJochReceiverCurrentStatus &in_current_status,
JFJochReceiverPlots &in_plots,
ImageBuffer &in_image_buffer,
ZMQPreviewSocket *in_zmq_preview_socket,
ZMQMetadataSocket *in_zmq_metadata_socket)
: JFJochReceiver(in_experiment, in_image_buffer, in_image_pusher, in_current_status, in_plots,
in_spot_finding_settings, in_logger, in_pixel_mask, in_zmq_preview_socket,
in_zmq_metadata_socket, nullptr),
image_puller(in_image_puller),
reader(in_image_buffer),
forward_nthreads(JFJochReceiverLite::NumberOfDataAnalysisThreads(in_forward_nthreads, in_experiment)),
rugnux_nthreads(in_rugnux_nthreads),
model_path(in_model_path),
deadline_s(in_deadline_s),
live_status(in_live_status),
measurement_started(1),
forward_started(forward_nthreads) {
logger.Info("Starting {} forward threads; rugnux runs on {} threads while the images arrive",
forward_nthreads, rugnux_nthreads);
for (int i = 0; i < forward_nthreads; i++)
forward_futures.emplace_back(std::async(std::launch::async, &JFJochReceiverRugnux::ForwardThread, this, i));
measurement = std::async(std::launch::async, &JFJochReceiverRugnux::MeasurementThread, this);
forward_started.wait();
}
JFJochReceiverRugnux::~JFJochReceiverRugnux() {
cancelled = true;
reader.Ended();
{
std::unique_lock ul(rugnux_m);
if (rugnux != nullptr)
rugnux->Cancel();
if (probe != nullptr)
probe->Cancel();
}
try {
if (measurement.valid())
measurement.get();
} catch (const std::exception &e) {
logger.Error("ReceiverRugnux teardown failed: {}", e.what());
}
}
void JFJochReceiverRugnux::MeasurementThread() {
std::future<void> rugnux_done, probe_done;
try {
// Waiting for the start message, as JFJochReceiverLite::MeasurementThread.
std::optional<std::chrono::steady_clock::time_point> give_up_at;
auto msg = image_puller.PollImage();
while (!cancelled && (!msg.has_value() || !msg->cbor || !msg->cbor->start_message)) {
if (detector_finished && !give_up_at)
give_up_at = std::chrono::steady_clock::now() + StartMessageGrace;
if (give_up_at && (std::chrono::steady_clock::now() > *give_up_at)) {
logger.Error("Detector finished, but no start message arrived - abandoning measurement");
cancelled = true;
break;
}
msg = image_puller.PollImage();
}
if (cancelled) {
current_status.SetProgress({});
current_status.SetStatus(GetStatus());
measurement_started.count_down();
return;
}
Configure(msg->cbor->start_message.value());
image_buffer.StartMeasurement(experiment);
start_time = std::chrono::system_clock::now();
SendStartMessage();
// Rugnux reads the images back from the buffer, so it needs one slot per image: a slot taken
// again would hand it a later image. Without that the collection is only forwarded.
const auto slots = image_buffer.GetStatus().total_slots;
live_status.Start(experiment.GetImageNum());
if (slots < experiment.GetImageNum()) {
logger.Error("Image buffer holds {} images, the collection has {}: no rugnux analysis for it",
slots, experiment.GetImageNum());
live_status.Finish(RugnuxLiveState::Skipped, "the image buffer cannot hold the whole collection");
} else {
// The start message as the writer gets it, so rugnux sees the dataset an HDF5 master states.
std::vector<uint8_t> start_cbor;
image_buffer.GetStartMessage(start_cbor);
reader.Open(CBORStream2Deserialize(start_cbor)->start_message.value());
char dir_template[] = "/tmp/jfjoch_rugnux_XXXXXX";
if (mkdtemp(dir_template) == nullptr)
throw JFJochException(JFJochExceptionCategory::FileWriteError, "Cannot create rugnux scratch directory");
scratch_dir = dir_template;
rugnux_done = std::async(std::launch::async, &JFJochReceiverRugnux::RugnuxThread, this);
probe_done = std::async(std::launch::async, &JFJochReceiverRugnux::ProbeThread, this);
}
measurement_started.count_down();
} catch (const JFJochException &e) {
logger.ErrorException(e);
Cancel(e);
measurement_started.count_down();
throw;
}
logger.Info("Receiving started");
for (auto &f: forward_futures)
f.get();
reader.Ended();
images_done = true;
live_status.LastImage();
const int64_t last_image = last_image_ns.load();
logger.Info("Receiving finished, {} images", images_collected.load());
if (rugnux_done.valid()) {
// The fixed cutoff: the analysis gets deadline_s after the last image, then gives what it has.
if (deadline_s > 0 && rugnux_done.wait_for(std::chrono::duration<double>(deadline_s))
== std::future_status::timeout)
Cut("deadline");
rugnux_done.wait();
logger.Info("rugnux live: done {:.2f} s after the last image", SecondsSince(last_image));
}
if (probe_done.valid())
probe_done.wait();
current_status.SetProgress(1.0);
current_status.SetStatus(GetStatus());
SendEndMessage();
const double end_sent_s = SecondsSince(last_image);
if (!image_buffer.CheckIfBufferReturned(SendFinalizeTimeout)) {
logger.Error("Send commands not finalized in {} seconds", SendFinalizeTimeout.count());
throw JFJochException(JFJochExceptionCategory::ZeroMQ, "Send commands not finalized in "
+ std::to_string(SendFinalizeTimeout.count()) + " seconds");
}
if (push_images_to_writer)
writer_error = image_pusher.Finalize();
if (last_image > 0 && rugnux_done.valid()) {
const double written_s = SecondsSince(last_image);
live_status.Latency(written_s);
logger.Info("rugnux live: last image -> all written {:.2f} s (end message sent at {:.2f} s)",
written_s, end_sent_s);
}
if (!scratch_dir.empty()) {
std::error_code ec;
std::filesystem::remove_all(scratch_dir, ec);
}
end_time = std::chrono::system_clock::now();
current_status.SetProgress({});
current_status.SetStatus(GetStatus());
}
void JFJochReceiverRugnux::RugnuxThread() {
// Below the forward threads: every thread rugnux starts inherits this, so an image arriving finds a
// core however busy rugnux keeps the machine.
// The process-wide ParallelFor pool is not started by this thread, so it is lowered separately.
setpriority(PRIO_PROCESS, static_cast<id_t>(syscall(SYS_gettid)), 10);
parallel_detail::WorkerPool::Instance().SetNice(10);
try {
const auto dataset = reader.GetDataset();
const RugnuxLiveSetup setup = MakeRugnuxLiveSetup(*dataset, rugnux_nthreads, scratch_dir + "/p", model_path);
LiveObserver observer(plots, scan_result, *az_int_mapping, live_status, logger);
ProcessResult result;
{
Rugnux process(reader, setup.experiment, *dataset->pixel_mask, setup.config);
Reachable reachable(rugnux_m, rugnux, process, cut, cancelled);
result = process.Run(&observer);
}
// The verdict and summary of the report a `rugnux` run writes, without writing it.
const ReportDocument doc = BuildReportDocument(scratch_dir + "/p", experiment.GetFilePrefix(),
setup.experiment, result);
std::vector<std::string> warnings;
for (const auto &w: doc.warnings)
if (!w.developer_only)
warnings.push_back(w.text);
live_status.Result(doc.verdict, doc.verdict_text, doc.summary_rows, warnings);
live_status.Finish(result.cancelled ? RugnuxLiveState::Cancelled : RugnuxLiveState::Done);
std::unique_lock ul(rugnux_m);
rugnux_result = std::move(result);
} catch (const std::exception &e) {
logger.Error("rugnux live: {}", e.what());
live_status.Finish(RugnuxLiveState::Failed, e.what());
}
if (!parallel_detail::WorkerPool::Instance().SetNice(0))
logger.Warning("rugnux live: the ParallelFor pool stays at nice 10 - raising it back needs CAP_SYS_NICE");
}
void JFJochReceiverRugnux::ProbeThread() {
// Rugnux waits for half the sweep before it decides on a lattice. In the meantime the machine is
// idle, so a small run of its own indexes the first 60 degrees and says at once whether there is
// a crystal - quick feedback for the person at the beamline, above all where there is none. It
// decides nothing for the run, which goes on as it would without it.
constexpr double PROBE_DEG = 60.0;
setpriority(PRIO_PROCESS, static_cast<id_t>(syscall(SYS_gettid)), 10);
try {
const auto dataset = reader.GetDataset();
const auto gonio = dataset->experiment.GetGoniometer();
if (!gonio || gonio->GetIncrement_deg() == 0.0f)
return;
const int64_t frames = std::lround(PROBE_DEG / std::abs(gonio->GetIncrement_deg()));
if (2 * frames > static_cast<int64_t>(reader.GetNumberOfImages()))
return; // a sweep this short is decided on soon enough by the run itself
RugnuxLiveSetup setup = MakeRugnuxLiveSetup(*dataset, rugnux_nthreads, "", "");
setup.config.end_image = static_cast<int>(frames);
setup.config.rotation_postrefine_geometry = false;
setup.config.run_scaling = false;
ProcessResult r;
{
Rugnux probe_run(reader, setup.experiment, *dataset->pixel_mask, setup.config);
Reachable reachable(rugnux_m, probe, probe_run, cut, cancelled);
r = probe_run.Run(nullptr);
}
if (r.consensus_cell && r.rotation_lattice_type) {
const auto &c = *r.consensus_cell;
live_status.Probe(fmt::format("lattice on the first {:.0f} deg: {}-centred {}, {:.1f} {:.1f} {:.1f} "
"{:.1f} {:.1f} {:.1f}, {:.0f}% of the images indexed", PROBE_DEG,
r.rotation_lattice_type->centering,
gemmi::crystal_system_str(r.rotation_lattice_type->crystal_system),
c.a, c.b, c.c, c.alpha, c.beta, c.gamma,
100.0 * r.indexing_rate.value_or(0.0f)));
} else
live_status.Probe(fmt::format("no lattice on the first {:.0f} deg", PROBE_DEG));
} catch (const NoLatticeFound &e) {
live_status.Probe(fmt::format("no lattice on the first {:.0f} deg: {}", PROBE_DEG, e.reason));
} catch (const std::exception &e) {
logger.Warning("rugnux live: the lattice probe failed: {}", e.what());
}
logger.Info("rugnux live: lattice probe - {}", live_status.Get().probe);
}
void JFJochReceiverRugnux::AddToEndMessage(EndMessage &message) {
std::unique_lock ul(rugnux_m);
if (!rugnux_result)
return;
const EndMessage &r = rugnux_result->end_message;
message.indexing_rate = r.indexing_rate;
message.bkg_estimate = r.bkg_estimate;
message.spindle_blind_fraction = r.spindle_blind_fraction;
message.ice_ring_score_mean = r.ice_ring_score_mean;
message.az_int_result = r.az_int_result;
message.rotation_lattice = r.rotation_lattice;
message.rotation_lattice_type = r.rotation_lattice_type;
message.rotation_extra_lattices = r.rotation_extra_lattices;
message.unit_cell = r.unit_cell;
message.space_group_name = r.space_group_name;
message.space_group_number = r.space_group_number;
message.reindex_matrix = r.reindex_matrix;
message.refined_beam_center_x = r.refined_beam_center_x;
message.refined_beam_center_y = r.refined_beam_center_y;
message.refined_poni_rot1 = r.refined_poni_rot1;
message.refined_poni_rot2 = r.refined_poni_rot2;
message.refined_poni_rot3 = r.refined_poni_rot3;
message.refined_rotation_axis = r.refined_rotation_axis;
// Merged data, map coefficients and placed model, for the writer to write beside the master. Only
// here, with rugnux_result set: every file in scratch_dir is then whole - rugnux writes synchronously
// and a cut never stops a write half-way - while a run that threw may have left one half-written.
FillRugnuxResultBlocks(scratch_dir + "/p", message);
}
void JFJochReceiverRugnux::Configure(const StartMessage &msg) {
// As JFJochReceiverLite::Configure.
if ((experiment.GetXPixelsNum() != msg.image_size_x)
|| (experiment.GetYPixelsNum() != msg.image_size_y))
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mismatch in detector size");
if (fabs(experiment.GetPixelSize_mm() - msg.pixel_size_x * 1e3) > 1e-7)
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mismatch in pixel size");
experiment.Detector().Description(msg.detector_description);
experiment.Detector().SerialNumber(msg.detector_serial_number);
experiment.Detector().SensorMaterial(msg.sensor_material);
experiment.Detector().SensorThickness_um(msg.sensor_thickness * 1e6);
experiment.Detector().SaturationLimit(SaturationLimitFromValue(msg.saturation_value));
experiment.Detector().CountRateCorrectionApplied(msg.countrate_correction_enabled);
experiment.Detector().CountRateCorrectionLookupTable(msg.countrate_correction_lookup_table);
experiment.Detector().FlatfieldApplied(msg.flatfield_enabled);
experiment.Detector().VirtualPixelInterpolationApplied(msg.virtual_pixel_interpolation_enabled);
experiment.ApplyPixelMask(false);
experiment.Detector().BitDepthImage(msg.bit_depth_image);
experiment.PixelSigned(msg.pixel_signed);
}
void JFJochReceiverRugnux::ForwardThread(uint32_t id) {
forward_started.count_down();
measurement_started.wait();
while (!cancelled && !end_message_received) {
try {
auto msg = image_puller.PollImage(std::chrono::milliseconds(5));
if (!msg.has_value() || !msg->cbor)
continue;
if (msg->cbor->end_message) {
end_message_received = true;
} else if (msg->cbor->calibration) {
if (push_images_to_writer)
image_pusher.SendCalibration(msg->cbor->calibration.value());
} else if (msg->cbor->data_message) {
last_image_ns = SteadyNowNs();
live_status.ImageReceived();
DataMessage data_msg = msg->cbor->data_message.value();
++images_collected;
const auto compressed_size_img = data_msg.image.GetCompressedSize();
const auto uncompressed_size_img = data_msg.image.GetUncompressedSize();
compressed_size += compressed_size_img;
uncompressed_size += uncompressed_size_img;
UpdateMaxImageReceived(data_msg.number);
data_msg.original_number = data_msg.number;
data_msg.user_data = experiment.GetImageAppendix();
data_msg.run_number = experiment.GetRunNumber();
data_msg.run_name = experiment.GetRunName();
data_msg.receiver_buf_available = image_buffer.GetAvailSlots();
data_msg.receiver_aq_dev_delay = image_puller.GetCurrentFifoUtilization();
if (compressed_size_img > 0)
data_msg.compression_ratio = static_cast<float>(uncompressed_size_img)
/ static_cast<float>(compressed_size_img);
auto loc = image_buffer.GetImageSlot();
if (loc == nullptr)
writer_queue_full = true;
else {
auto writer_buffer = (uint8_t *) loc->GetImage();
CBORStream2Serializer serializer(writer_buffer, experiment.GetImageBufferLocationSize());
serializer.SerializeImage(data_msg);
loc->SetImageNumber(data_msg.number);
loc->SetImageSize(serializer.GetBufferSize());
loc->SetIndexed(false);
loc->ReadyToSend();
if (zmq_preview_socket != nullptr)
zmq_preview_socket->SendImage(writer_buffer, serializer.GetBufferSize());
if (zmq_metadata_socket != nullptr)
zmq_metadata_socket->AddDataMessage(data_msg);
if (push_images_to_writer) {
if (image_pusher.SendImage(*loc))
++images_sent;
} else
loc->release();
UpdateMaxImageSent(data_msg.number);
}
}
UpdateMaxDelay(image_puller.GetCurrentFifoUtilization());
current_status.SetProgress(GetProgress());
current_status.SetStatus(GetStatus());
} catch (const JFJochException &e) {
logger.ErrorException(e);
Cancel(e);
}
}
}
void JFJochReceiverRugnux::Cancel(bool silent) {
detector_finished = true;
// The silent cancel is the detector reading idle at the end of every collection: rugnux goes on.
if (silent)
return;
// /cancel is the one way in for the beamline, which cannot tell the collection from its analysis: it
// ends the collection while the images still arrive, as JFJochReceiverLite does, and in every case
// cuts the analysis, so the end message goes out at once with what rugnux has established.
if (!images_done)
JFJochReceiver::Cancel(silent);
Cut("cancel");
}
void JFJochReceiverRugnux::Cut(const std::string &reason) {
std::unique_lock ul(rugnux_m);
if (cut || rugnux_result)
return;
cut = true;
reader.Ended(); // no image is waited for any more
if (rugnux != nullptr)
rugnux->Cancel();
if (probe != nullptr)
probe->Cancel();
live_status.Cut(reason);
logger.Info("rugnux live: analysis cut ({}) - the end message carries what is established", reason);
}
void JFJochReceiverRugnux::StopReceiver() {
if (measurement.valid()) {
measurement.get();
logger.Info("Receiver stopped");
}
}
float JFJochReceiverRugnux::GetEfficiency() const {
if (experiment.GetFrameNum() == 0)
return 0;
return static_cast<float>(images_collected) / static_cast<float>(experiment.GetFrameNum());
}
float JFJochReceiverRugnux::GetProgress() const {
if (experiment.GetFrameNum() == 0)
return 0.0;
return static_cast<float>(max_image_number_received) / static_cast<float>(experiment.GetFrameNum());
}