mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases 88 Wiki Activity
Files
rc181
Jungfraujoch/image_analysis/structure_refinement/RefinementADP.cpp
T
leonarski_fandClaude Opus 5.5 6472ce3ff1 RESULT-CHANGING: Refinement: anisotropic atoms and TLS groups, the displacement model decided on the free set 
S1 - the model's displacements kept: an atom's ANISOU is no longer made isotropic; where the model has TLS
groups whose atoms bring no U, their U is made from the groups and the B column read as the total or as the
residual to them, whichever (or the isotropic B) scores the lowest R-free as given. Refined models are written
with ANISOU and TLS groups (mmCIF _pdbx_refine_tls*, and PDB REMARK 3, which gemmi's writer leaves out).

S2/S3 - both engines grid an atom with its full U (U = dev + B/(8 pi^2) I, so the B machinery works on it with
its shape held) and give the target's derivatives in U by the FFT method, CPU and GPU (one warp per atom, no
atomics). After the weights are chosen, the refinement tries, as trials side by side: the B factors refined on;
TLS groups with a residual B (the model's own, one per chain, and chains cut into 2/4/8/16 TLSMD segments by
dynamic programming over a quadratic B(r) fit); and a U per atom with rigid-bond (RIGU form), similarity and
isotropy restraints at three weights. A trial displaces another only within the geometry bounds and with both
R-free (beyond the tie) and the free -ln L lower. The water-build rounds hold the chosen U shapes.

An anisotropic atom's density is now cut at the radius of its largest eigenvalue rather than its largest
diagonal element (also in the validation's density and the GPU gather), which left skewed atoms truncated.

Refine arm + open-arm subset (11 sets, vs 51cbe0ef7), refined R-free: 1exr 0.2064 -> 0.1377 (SHELXL U kept),
193l 0.1689 -> 0.1322 (aniso), 2r8s 0.2262 -> 0.2054 (TLS 47), 1hr2 0.2463 -> 0.2316 (TLS 30), 1k4c 0.2110 ->
0.1881 (TLS 42), 36gk 0.2445 -> 0.2221, 7l6j 0.1766 -> 0.1553, 7os3 0.2040 -> 0.2014, 8agq 0.1948 -> 0.1595,
8sa8 0.1830 -> 0.1548, 6f3p 0.1821 -> 0.1453. Every output U positive definite. The first refinement takes
about 2-3x as long (the hypothesis trials; measured on a shared GPU).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HVeDQqN9qRD57vnShzTC37
2026-10-27 17:34:27 +01:00

320 lines
14 KiB
C++
Raw Permalink Blame History

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "RefinementADP.h"
#include <algorithm>
#include <cmath>
#include <limits>
#include <Eigen/Dense>
namespace {
constexpr double DEG = 3.14159265358979323846 / 180.0;
} // namespace
gemmi::SMat33<double> TlsBody::U(const gemmi::Position &pos) const {
const double x = pos.x - origin.x, y = pos.y - origin.y, z = pos.z - origin.z;
const gemmi::Mat33 a(0, z, -y, -z, 0, x, y, -x, 0);
const gemmi::Mat33 ala = a.multiply(L.as_mat33()).multiply(a.transpose());
const gemmi::Mat33 as = a.multiply(S);
gemmi::SMat33<double> u = T;
u.u11 += ala[0][0] + 2 * as[0][0];
u.u22 += ala[1][1] + 2 * as[1][1];
u.u33 += ala[2][2] + 2 * as[2][2];
u.u12 += ala[0][1] + as[0][1] + as[1][0];
u.u13 += ala[0][2] + as[0][2] + as[2][0];
u.u23 += ala[1][2] + as[1][2] + as[2][1];
return u;
}
std::array<double, TlsBody::PARAMETERS> TlsBody::Get() const {
return {T.u11, T.u22, T.u33, T.u12, T.u13, T.u23, L.u11, L.u22, L.u33, L.u12, L.u13, L.u23,
S[0][0], S[1][1], S[0][1], S[0][2], S[1][0], S[1][2], S[2][0], S[2][1]};
}
void TlsBody::Set(const std::array<double, PARAMETERS> &p) {
T = {p[0], p[1], p[2], p[3], p[4], p[5]};
L = {p[6], p[7], p[8], p[9], p[10], p[11]};
S = gemmi::Mat33(p[12], p[14], p[15], p[16], p[13], p[17], p[18], p[19], -(p[12] + p[13]));
}
std::array<std::array<double, 6>, TlsBody::PARAMETERS> TlsBody::Derivatives(const gemmi::Position &pos) const {
std::array<std::array<double, 6>, PARAMETERS> d{};
TlsBody unit;
unit.origin = origin;
for (int k = 0; k < PARAMETERS; k++) {
std::array<double, PARAMETERS> p{};
p[k] = 1;
unit.Set(p);
const gemmi::SMat33<double> u = unit.U(pos);
d[k] = {u.u11, u.u22, u.u33, u.u12, u.u13, u.u23};
}
return d;
}
std::array<double, TlsBody::PARAMETERS> TlsBody::Gradient(const std::vector<gemmi::Atom *> &atom_list,
const std::vector<std::array<double, 6>> &grad_u) const {
gemmi::Mat33 mt(0), ml(0), ms(0);
for (int i : atoms) {
const std::array<double, 6> &g = grad_u[i];
const gemmi::Mat33 gm(g[0], g[3] / 2, g[4] / 2, g[3] / 2, g[1], g[5] / 2, g[4] / 2, g[5] / 2, g[2]);
const gemmi::Position &pos = atom_list[i]->pos;
const double x = pos.x - origin.x, y = pos.y - origin.y, z = pos.z - origin.z;
const gemmi::Mat33 a(0, z, -y, -z, 0, x, y, -x, 0);
const gemmi::Mat33 atg = a.transpose().multiply(gm);
const gemmi::Mat33 atga = atg.multiply(a);
for (int u = 0; u < 3; u++)
for (int v = 0; v < 3; v++) {
mt[u][v] += gm[u][v];
ml[u][v] += atga[u][v];
ms[u][v] += 2 * atg[u][v];
}
}
// symmetric T and L: an off-diagonal parameter takes both of its elements; S33 = -(S11 + S22)
return {mt[0][0], mt[1][1], mt[2][2], 2 * mt[0][1], 2 * mt[0][2], 2 * mt[1][2],
ml[0][0], ml[1][1], ml[2][2], 2 * ml[0][1], 2 * ml[0][2], 2 * ml[1][2],
ms[0][0] - ms[2][2], ms[1][1] - ms[2][2], ms[0][1], ms[0][2], ms[1][0], ms[1][2], ms[2][0], ms[2][1]};
}
std::vector<TlsBody> TlsBodiesFromModel(const gemmi::Model &model, const std::vector<gemmi::TlsGroup> &groups) {
std::vector<TlsBody> bodies;
for (const gemmi::TlsGroup &g : groups) {
TlsBody b;
b.id = g.id;
b.origin = g.origin;
b.T = g.T;
b.L = g.L.scaled(DEG * DEG);
for (int i = 0; i < 3; i++)
for (int j = 0; j < 3; j++)
b.S[i][j] = g.S[i][j] * DEG;
// the trace of S does not enter U; S33 is taken as -(S11 + S22)
const double t = (b.S[0][0] + b.S[1][1] + b.S[2][2]) / 3.0;
for (int i = 0; i < 3; i++)
b.S[i][i] -= t;
bool finite = std::isfinite(b.origin.x) && std::isfinite(b.origin.y) && std::isfinite(b.origin.z);
for (double v : b.Get())
finite = finite && std::isfinite(v);
if (finite)
bodies.push_back(b);
}
// the atoms of each group's residues, an atom to the first group that names it
int i = 0;
for (const gemmi::Chain &ch : model.chains)
for (const gemmi::Residue &r : ch.residues) {
int body = -1;
if (!r.is_water())
for (size_t k = 0; k < groups.size() && body < 0; k++) {
for (const gemmi::TlsGroup::Selection &sel : groups[k].selections)
if (sel.chain == ch.name && !(r.seqid < sel.res_begin) && !(sel.res_end < r.seqid))
body = static_cast<int>(k);
}
// the group index among those kept
int kept = -1;
if (body >= 0)
for (size_t k = 0; k < bodies.size(); k++)
if (bodies[k].id == groups[body].id)
kept = static_cast<int>(k);
for (size_t a = 0; a < r.atoms.size(); a++, i++)
if (kept >= 0)
bodies[kept].atoms.push_back(i);
}
std::erase_if(bodies, [](const TlsBody &b) { return b.atoms.empty(); });
return bodies;
}
std::vector<gemmi::TlsGroup> TlsBodiesToGroups(const gemmi::Model &model, const std::vector<TlsBody> &bodies) {
// per atom its body, per residue the body of its first atom in one
std::vector<int> atom_body;
for (const gemmi::Chain &ch : model.chains)
for (const gemmi::Residue &r : ch.residues)
atom_body.insert(atom_body.end(), r.atoms.size(), -1);
for (size_t b = 0; b < bodies.size(); b++)
for (int a : bodies[b].atoms)
atom_body[a] = static_cast<int>(b);
std::vector<gemmi::TlsGroup> groups(bodies.size());
for (size_t b = 0; b < bodies.size(); b++) {
gemmi::TlsGroup &g = groups[b];
g.id = std::to_string(b + 1);
g.num_id = static_cast<short>(b + 1);
g.origin = bodies[b].origin;
g.T = bodies[b].T;
g.L = bodies[b].L.scaled(1.0 / (DEG * DEG));
for (int i = 0; i < 3; i++)
for (int j = 0; j < 3; j++)
g.S[i][j] = bodies[b].S[i][j] / DEG;
}
size_t i = 0;
for (const gemmi::Chain &ch : model.chains) {
int open = -1; // the body whose run of residues is open in this chain
for (const gemmi::Residue &r : ch.residues) {
int body = -1;
for (size_t a = 0; a < r.atoms.size(); a++)
if (atom_body[i + a] >= 0 && body < 0)
body = atom_body[i + a];
i += r.atoms.size();
if (body >= 0 && body == open)
groups[body].selections.back().res_end = r.seqid;
else if (body >= 0) {
groups[body].selections.push_back({ch.name, r.seqid, r.seqid, ""});
}
open = body;
}
}
return groups;
}
URestraints MakeURestraints(const std::vector<std::pair<int, int>> &b_similarity_pairs, const std::vector<char> &refines_u) {
URestraints r;
std::vector<char> paired(refines_u.size(), 0);
for (const auto &[i, j] : b_similarity_pairs)
if (refines_u[i] && refines_u[j]) {
r.pairs.emplace_back(i, j);
paired[i] = paired[j] = 1;
}
for (size_t i = 0; i < refines_u.size(); i++)
if (refines_u[i] && !paired[i])
r.isotropic.push_back(static_cast<int>(i));
return r;
}
double EvaluateURestraints(const URestraints &r, const std::vector<gemmi::Position> &pos,
const std::vector<std::array<double, 6>> &u, std::vector<std::array<double, 6>> *grad) {
// a symmetric matrix as its six elements: the off-diagonal ones count twice in a Frobenius norm, and a derivative
// with respect to U_ab (a != b) is the sum of the matrix elements ab and ba
double e = 0;
const double wr = 1.0 / (URestraints::SIGMA_RIGID * URestraints::SIGMA_RIGID);
const double ws = 1.0 / (URestraints::SIGMA_SIMILAR * URestraints::SIGMA_SIMILAR);
const double wi = 1.0 / (URestraints::SIGMA_ISOTROPIC * URestraints::SIGMA_ISOTROPIC);
for (const auto &[i, j] : r.pairs) {
std::array<double, 6> d;
for (int c = 0; c < 6; c++)
d[c] = u[i][c] - u[j][c];
gemmi::Vec3 n = pos[j] - pos[i];
const double len = n.length();
std::array<double, 6> g{};
if (len > 0) {
n = n / len;
// v = (U_i - U_j) n
const double v[3] = {d[0] * n.x + d[3] * n.y + d[4] * n.z, d[3] * n.x + d[1] * n.y + d[5] * n.z,
d[4] * n.x + d[5] * n.y + d[2] * n.z};
const double nn[3] = {n.x, n.y, n.z};
e += wr * (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]);
g[0] += 2 * wr * v[0] * nn[0];
g[1] += 2 * wr * v[1] * nn[1];
g[2] += 2 * wr * v[2] * nn[2];
g[3] += 2 * wr * (v[0] * nn[1] + v[1] * nn[0]);
g[4] += 2 * wr * (v[0] * nn[2] + v[2] * nn[0]);
g[5] += 2 * wr * (v[1] * nn[2] + v[2] * nn[1]);
}
e += ws * (d[0] * d[0] + d[1] * d[1] + d[2] * d[2] + 2 * (d[3] * d[3] + d[4] * d[4] + d[5] * d[5]));
for (int c = 0; c < 3; c++)
g[c] += 2 * ws * d[c];
for (int c = 3; c < 6; c++)
g[c] += 4 * ws * d[c];
if (grad != nullptr)
for (int c = 0; c < 6; c++) {
(*grad)[i][c] += g[c];
(*grad)[j][c] -= g[c];
}
}
for (int i : r.isotropic) {
const double t = (u[i][0] + u[i][1] + u[i][2]) / 3.0;
const double d[6] = {u[i][0] - t, u[i][1] - t, u[i][2] - t, u[i][3], u[i][4], u[i][5]};
e += wi * (d[0] * d[0] + d[1] * d[1] + d[2] * d[2] + 2 * (d[3] * d[3] + d[4] * d[4] + d[5] * d[5]));
if (grad != nullptr) {
// d/dU_aa of the traceless part's norm: 2 d_aa (the trace's own derivative sums to zero against d)
for (int c = 0; c < 3; c++)
(*grad)[i][c] += 2 * wi * d[c];
for (int c = 3; c < 6; c++)
(*grad)[i][c] += 4 * wi * d[c];
}
}
return e;
}
std::vector<int> TlsSegments(const std::vector<std::vector<int>> &residue_atoms, const std::vector<gemmi::Position> &pos,
const std::vector<double> &b, int k, int min_residues) {
using Mat = Eigen::Matrix<double, 10, 10>;
using Vec = Eigen::Matrix<double, 10, 1>;
const int n = static_cast<int>(residue_atoms.size());
k = std::min(k, n / std::max(min_residues, 1));
if (k <= 1)
return {0};
// The chain in blocks of consecutive residues (cuts between blocks only), so that a long chain stays a few hundred
// blocks: the dynamic programming is quadratic in them.
const int per_block = std::max(1, (n + 299) / 300);
const int min_blocks = std::max(1, (min_residues + per_block - 1) / per_block);
std::vector<int> block_first;
for (int r = 0; r < n; r += per_block)
block_first.push_back(r);
const int nb = static_cast<int>(block_first.size());
k = std::min(k, nb / min_blocks);
if (k <= 1)
return {0};
// the normal equations of each block, about the chain's centroid in units of 10 A (for their conditioning)
gemmi::Position c;
int count = 0;
for (const std::vector<int> &ra : residue_atoms)
for (int a : ra) {
c += pos[a];
count++;
}
if (count == 0)
return {0};
c /= count;
std::vector<Mat> ata(nb + 1, Mat::Zero());
std::vector<Vec> atb(nb + 1, Vec::Zero());
std::vector<double> bb(nb + 1, 0.0);
for (int q = 0; q < nb; q++) {
ata[q + 1] = ata[q];
atb[q + 1] = atb[q];
bb[q + 1] = bb[q];
const int last = q + 1 < nb ? block_first[q + 1] : n;
for (int r = block_first[q]; r < last; r++)
for (int a : residue_atoms[r]) {
const double x = (pos[a].x - c.x) / 10, y = (pos[a].y - c.y) / 10, z = (pos[a].z - c.z) / 10;
Vec f;
f << 1, x, y, z, x * x, y * y, z * z, x * y, x * z, y * z;
ata[q + 1] += f * f.transpose();
atb[q + 1] += b[a] * f;
bb[q + 1] += b[a] * b[a];
}
}
// the residual of the fit of blocks [i, j)
auto cost = [&](int i, int j) {
const Mat m = ata[j] - ata[i] + 1e-6 * Mat::Identity();
const Vec v = atb[j] - atb[i];
const Vec x = m.ldlt().solve(v);
return std::max(0.0, bb[j] - bb[i] - v.dot(x));
};
const double INF = std::numeric_limits<double>::infinity();
// best[s][j]: blocks [0, j) in s + 1 segments; from[s][j]: where the last of them starts
std::vector<std::vector<double>> best(k, std::vector<double>(nb + 1, INF));
std::vector<std::vector<int>> from(k, std::vector<int>(nb + 1, 0));
std::vector<std::vector<double>> seg(nb + 1, std::vector<double>(nb + 1, INF));
for (int i = 0; i < nb; i++)
for (int j = i + min_blocks; j <= nb; j++)
seg[i][j] = cost(i, j);
for (int j = min_blocks; j <= nb; j++)
best[0][j] = seg[0][j];
for (int s = 1; s < k; s++)
for (int j = (s + 1) * min_blocks; j <= nb; j++)
for (int i = s * min_blocks; i <= j - min_blocks; i++)
if (best[s - 1][i] + seg[i][j] < best[s][j]) {
best[s][j] = best[s - 1][i] + seg[i][j];
from[s][j] = i;
}
std::vector<int> starts;
int j = nb;
for (int s = k - 1; s > 0; s--) {
j = from[s][j];
starts.push_back(block_first[j]);
}
starts.push_back(0);
std::reverse(starts.begin(), starts.end());
return starts;
}
Reference in new issue View Git Blame Copy Permalink