leonarski_fandClaude Opus 5.5
5028aac0ee
WritePlacedModel was handed the result of find_spacegroup_by_name unchecked, and PlacedModelStructure indexed every per-atom vector by atom_name's length. A malformed block from the wire could crash the writer. Now the per-atom vectors must all have atom_name's length and aniso must be empty or six per atom (PlacedModelStructure throws otherwise), and an unknown data space group is an error; WriteRugnuxResultFiles logs either and writes nothing for the model, as it already did for an MTZ table WriteMtz refuses. StreamWriter: "Received end message" is logged only with -v again, exactly as before this branch. The wall-clock times of the rugnux files are still logged where they are written, which only happens when the END carries them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
170 lines
7.9 KiB
C++
170 lines
7.9 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "WriteRugnuxFiles.h"
|
|
|
|
#include <chrono>
|
|
#include <stdexcept>
|
|
#include <filesystem>
|
|
|
|
#include <gemmi/mtz.hpp>
|
|
|
|
#include "WriteModel.h"
|
|
#include "../common/Logger.h"
|
|
|
|
void WriteMtz(const MtzTable &table, const std::string &filename) {
|
|
const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name(table.space_group_name);
|
|
if (sg == nullptr)
|
|
throw std::runtime_error(filename + ": unknown space group '" + table.space_group_name + "'");
|
|
const size_t rows = table.columns.empty() ? 0 : table.columns[0].values.size();
|
|
for (const auto &c : table.columns) {
|
|
if (c.values.size() != rows)
|
|
throw std::runtime_error(fmt::format("{}: column {} has {} rows, column {} {}", filename, c.label,
|
|
c.values.size(), table.columns[0].label, rows));
|
|
if (c.dataset_id < 0 || c.dataset_id >= static_cast<int32_t>(table.datasets.size()))
|
|
throw std::runtime_error(fmt::format("{}: column {} is in dataset {}, of {}", filename, c.label,
|
|
c.dataset_id, table.datasets.size()));
|
|
}
|
|
for (const auto &b : table.batches)
|
|
if (b.ints.size() != 29 || b.floats.size() != 156
|
|
|| b.ints[20] < 0 || b.ints[20] >= static_cast<int32_t>(table.datasets.size()))
|
|
throw std::runtime_error(fmt::format("{}: batch {} has no header of a dataset in the table", filename,
|
|
b.number));
|
|
|
|
gemmi::Mtz mtz;
|
|
mtz.title = table.title;
|
|
mtz.history = table.history;
|
|
mtz.spacegroup = sg;
|
|
mtz.set_cell_for_all(table.unit_cell);
|
|
mtz.sort_order = table.sort_order;
|
|
for (size_t i = 0; i < table.datasets.size(); i++) {
|
|
const MtzTable::Dataset &d = table.datasets[i];
|
|
mtz.datasets.push_back({static_cast<int>(i), d.project_name, d.crystal_name, d.dataset_name, mtz.cell,
|
|
d.wavelength_A});
|
|
}
|
|
for (const auto &c : table.columns)
|
|
mtz.add_column(c.label, c.type, c.dataset_id, -1, false);
|
|
const size_t w = table.columns.size();
|
|
mtz.nreflections = static_cast<int>(rows);
|
|
mtz.data.resize(rows * w);
|
|
for (size_t c = 0; c < w; c++)
|
|
for (size_t r = 0; r < rows; r++)
|
|
mtz.data[r * w + c] = table.columns[c].values[r];
|
|
for (const auto &b : table.batches) {
|
|
gemmi::Mtz::Batch &batch = mtz.batches.emplace_back();
|
|
batch.number = b.number;
|
|
batch.title = b.title;
|
|
batch.ints.assign(b.ints.begin(), b.ints.end());
|
|
batch.floats = b.floats;
|
|
batch.axes = b.axes;
|
|
}
|
|
mtz.write_to_file(filename);
|
|
}
|
|
|
|
gemmi::Structure PlacedModelStructure(const PlacedModelMessage &model) {
|
|
const size_t n = model.atom_name.size();
|
|
for (const size_t size : {model.model_number.size(), model.chain.size(), model.residue_name.size(),
|
|
model.seq_num.size(), model.icode.size(), model.het_flag.size(), model.element.size(),
|
|
model.altloc.size(), model.charge.size(), model.x.size(), model.y.size(),
|
|
model.z.size(), model.b_iso.size(), model.occupancy.size()})
|
|
if (size != n)
|
|
throw std::runtime_error(fmt::format("placed model: a per-atom field has {} entries for {} atoms",
|
|
size, n));
|
|
if (!model.aniso.empty() && model.aniso.size() != 6 * n)
|
|
throw std::runtime_error(fmt::format("placed model: {} anisotropic values for {} atoms",
|
|
model.aniso.size(), n));
|
|
gemmi::Structure st;
|
|
st.name = model.name;
|
|
st.cell = model.unit_cell;
|
|
st.spacegroup_hm = model.model_space_group_name;
|
|
const bool has_aniso = !model.aniso.empty();
|
|
gemmi::Model *m = nullptr;
|
|
gemmi::Chain *ch = nullptr;
|
|
gemmi::Residue *res = nullptr;
|
|
for (size_t i = 0; i < model.atom_name.size(); i++) {
|
|
if (m == nullptr || m->num != model.model_number[i]) {
|
|
m = &st.models.emplace_back(model.model_number[i]);
|
|
ch = nullptr;
|
|
}
|
|
if (ch == nullptr || ch->name != model.chain[i]) {
|
|
ch = &m->chains.emplace_back(model.chain[i]);
|
|
res = nullptr;
|
|
}
|
|
const gemmi::SeqId seqid(model.seq_num[i], static_cast<char>(model.icode[i]));
|
|
if (res == nullptr || res->name != model.residue_name[i] || res->seqid != seqid
|
|
|| res->het_flag != static_cast<char>(model.het_flag[i])) {
|
|
res = &ch->residues.emplace_back(gemmi::ResidueId{seqid, {}, model.residue_name[i]});
|
|
res->het_flag = static_cast<char>(model.het_flag[i]);
|
|
}
|
|
gemmi::Atom &a = res->atoms.emplace_back();
|
|
a.name = model.atom_name[i];
|
|
a.element = gemmi::Element(model.element[i]);
|
|
a.altloc = static_cast<char>(model.altloc[i]);
|
|
a.charge = static_cast<signed char>(model.charge[i]);
|
|
a.pos = gemmi::Position(model.x[i], model.y[i], model.z[i]);
|
|
a.b_iso = model.b_iso[i];
|
|
a.occ = model.occupancy[i];
|
|
if (has_aniso)
|
|
a.aniso = {model.aniso[6 * i], model.aniso[6 * i + 1], model.aniso[6 * i + 2],
|
|
model.aniso[6 * i + 3], model.aniso[6 * i + 4], model.aniso[6 * i + 5]};
|
|
}
|
|
st.setup_cell_images();
|
|
return st;
|
|
}
|
|
|
|
namespace {
|
|
|
|
double WallClockSeconds() {
|
|
return std::chrono::duration<double>(std::chrono::system_clock::now().time_since_epoch()).count();
|
|
}
|
|
|
|
// Rename tmp onto final, unless final is there already and may not be overwritten.
|
|
void RenameIntoPlace(const std::string &tmp, const std::string &final_name, bool overwrite, Logger &logger) {
|
|
if (!overwrite && std::filesystem::exists(final_name)) {
|
|
std::filesystem::remove(tmp);
|
|
logger.Error("{} already exists and overwrite is off - not written", final_name);
|
|
return;
|
|
}
|
|
std::filesystem::rename(tmp, final_name);
|
|
logger.Info("{} written, renamed into place at wall clock {:.3f} s", final_name, WallClockSeconds());
|
|
}
|
|
|
|
} // namespace
|
|
|
|
void WriteRugnuxResultFiles(const EndMessage &end, const std::string &file_prefix, bool overwrite,
|
|
Logger &logger) {
|
|
// A file that cannot be written costs only itself: the master is already in place.
|
|
try {
|
|
if (end.merged) {
|
|
WriteMtz(*end.merged, file_prefix + ".mtz.tmp");
|
|
RenameIntoPlace(file_prefix + ".mtz.tmp", file_prefix + ".mtz", overwrite, logger);
|
|
}
|
|
} catch (const std::exception &e) {
|
|
logger.Error("Cannot write {}.mtz: {}", file_prefix, e.what());
|
|
}
|
|
try {
|
|
if (end.map_coefficients) {
|
|
WriteMtz(*end.map_coefficients, file_prefix + "_maps.mtz.tmp");
|
|
RenameIntoPlace(file_prefix + "_maps.mtz.tmp", file_prefix + "_maps.mtz", overwrite, logger);
|
|
}
|
|
} catch (const std::exception &e) {
|
|
logger.Error("Cannot write {}_maps.mtz: {}", file_prefix, e.what());
|
|
}
|
|
try {
|
|
if (end.placed_model) {
|
|
const PlacedModelMessage &model = *end.placed_model;
|
|
const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name(model.space_group_name);
|
|
if (sg == nullptr)
|
|
throw std::runtime_error("unknown space group '" + model.space_group_name + "'");
|
|
// WritePlacedModel appends _model.cif and _model.pdb to the prefix it is given.
|
|
const std::string tmp_prefix = file_prefix + ".tmp";
|
|
WritePlacedModel(PlacedModelStructure(model), model.unit_cell, *sg, tmp_prefix, logger);
|
|
for (const std::string suffix : {"_model.cif", "_model.pdb"})
|
|
if (std::filesystem::exists(tmp_prefix + suffix))
|
|
RenameIntoPlace(tmp_prefix + suffix, file_prefix + suffix, overwrite, logger);
|
|
}
|
|
} catch (const std::exception &e) {
|
|
logger.Error("Cannot write {}_model.cif: {}", file_prefix, e.what());
|
|
}
|
|
}
|