leonarski_fandClaude Opus 5.5
6efc42360a
gemmi reads any text it does not recognise as PDB and can hand back a structure with an empty model in it, so 'no models' let garbage through. Count the atom sites instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
100 lines
3.8 KiB
C++
100 lines
3.8 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "RugnuxModelFile.h"
|
|
|
|
#include <fstream>
|
|
#include <iterator>
|
|
#include <stdexcept>
|
|
|
|
#include <zlib.h>
|
|
#include <nlohmann/json.hpp>
|
|
#include <gemmi/mmread.hpp>
|
|
#include <gemmi/calculate.hpp> // count_atom_sites
|
|
|
|
namespace {
|
|
|
|
// The bytes of a gzip stream inflated, every member of it.
|
|
std::string Gunzip(const std::string &bytes) {
|
|
z_stream s{};
|
|
if (inflateInit2(&s, 15 + 32) != Z_OK)
|
|
throw std::runtime_error("cannot start zlib");
|
|
std::string out;
|
|
char buf[1 << 16];
|
|
s.next_in = reinterpret_cast<Bytef *>(const_cast<char *>(bytes.data()));
|
|
s.avail_in = static_cast<uInt>(bytes.size());
|
|
int ret = Z_OK;
|
|
while (ret != Z_STREAM_END || s.avail_in > 0) {
|
|
if (ret == Z_STREAM_END)
|
|
inflateReset(&s);
|
|
s.next_out = reinterpret_cast<Bytef *>(buf);
|
|
s.avail_out = sizeof buf;
|
|
ret = inflate(&s, Z_NO_FLUSH);
|
|
if (ret != Z_OK && ret != Z_STREAM_END) {
|
|
inflateEnd(&s);
|
|
throw std::runtime_error("the gzip stream is damaged");
|
|
}
|
|
out.append(buf, sizeof buf - s.avail_out);
|
|
if (ret == Z_OK && s.avail_in == 0 && s.avail_out != 0) {
|
|
inflateEnd(&s);
|
|
throw std::runtime_error("the gzip stream is cut short");
|
|
}
|
|
}
|
|
inflateEnd(&s);
|
|
return out;
|
|
}
|
|
|
|
std::string ReadFile(const std::string &path) {
|
|
std::ifstream f(path, std::ios::binary);
|
|
if (!f)
|
|
throw std::runtime_error("cannot open " + path);
|
|
return {std::istreambuf_iterator<char>(f), std::istreambuf_iterator<char>()};
|
|
}
|
|
|
|
} // namespace
|
|
|
|
RugnuxModel ParseRugnuxModel(const std::string &bytes, const std::string &name, const std::string &pae_json) {
|
|
RugnuxModel m;
|
|
m.name = name;
|
|
std::string text = bytes.size() >= 2 && bytes[0] == '\x1f' && bytes[1] == '\x8b' ? Gunzip(bytes) : bytes;
|
|
gemmi::cif::Document doc;
|
|
m.structure = gemmi::read_structure_from_memory(text.data(), text.size(), name, gemmi::CoorFormat::Detect, &doc);
|
|
if (gemmi::count_atom_sites(m.structure) == 0)
|
|
throw std::runtime_error(name + " holds no atoms");
|
|
m.modelcif = m.structure.input_format == gemmi::CoorFormat::Mmcif && !doc.blocks.empty()
|
|
&& doc.blocks[0].has_mmcif_category("_ma_qa_metric");
|
|
if (!pae_json.empty()) {
|
|
try {
|
|
const nlohmann::json j = nlohmann::json::parse(pae_json);
|
|
const nlohmann::json &pae = j.is_array() ? j.at(0).at("predicted_aligned_error")
|
|
: j.at("predicted_aligned_error");
|
|
std::vector<std::vector<double>> rows;
|
|
for (const auto &row : pae)
|
|
rows.push_back(row.get<std::vector<double>>());
|
|
m.pae = std::move(rows);
|
|
} catch (const std::exception &) {
|
|
}
|
|
}
|
|
return m;
|
|
}
|
|
|
|
RugnuxModel ReadRugnuxModelFile(const std::string &path) {
|
|
const std::string pae_path = PaePathFor(path);
|
|
std::string pae_json;
|
|
if (!pae_path.empty() && std::ifstream(pae_path))
|
|
pae_json = ReadFile(pae_path);
|
|
return ParseRugnuxModel(ReadFile(path), path, pae_json);
|
|
}
|
|
|
|
std::string PaePathFor(const std::string &model_path) {
|
|
const size_t slash = model_path.find_last_of("/\\");
|
|
const std::string dir = slash == std::string::npos ? std::string() : model_path.substr(0, slash + 1);
|
|
const std::string name = model_path.substr(dir.size());
|
|
const size_t at = name.rfind("-model_v");
|
|
if (at == std::string::npos)
|
|
return {};
|
|
const size_t dot = name.find('.', at);
|
|
const std::string version = name.substr(at + 8, dot == std::string::npos ? std::string::npos : dot - at - 8);
|
|
return dir + name.substr(0, at) + "-predicted_aligned_error_v" + version + ".json";
|
|
}
|