mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 87 Wiki Activity

Refinement: an own L-BFGS whose memory outlives the macro-cycle, B steps scaled by their curvature

Browse Source 
Each macro-cycle of a weight trial used to start a fresh Ceres solve: the L-BFGS memory was dropped every
10 (positions) or 5 (B) iterations and every cycle began with a steepest-descent step. RefinementMinimiser
keeps one memory for the positions and one for the B factors across a trial's macro-cycles (the target's
scale and likelihood parameters are refitted in between, which the memory tolerates), with a backtracking
Armijo line search and Oren-Shanno scaling of a diagonal initial inverse Hessian.

The diagonal: for B, the X-ray curvature per atom - the likelihood's curvature in |F_model| per resolution
shell against the atom's (s^2/4)^2 f^2 exp(-B s^2/2), calibrated once per trial by a central-difference
probe along a random direction - plus the B restraints'. For the positions none: tried with the restraints'
Gauss-Newton diagonal (bonds and angles, ~1e4 per axis against the X-ray's ~1e2) it slowed the minimiser,
since that stiffness is in directions that move bonded atoms together.

Same budget (10 cycles x 10 + 5 iterations, same ladder), median R-free (REFMAC 0-cycle scorer) from the
refinement battery's MD starts, base -> this: 5f6m (8 starts) 0.250 -> 0.210, 8dqb (8) 0.387 -> 0.345,
7arr 0.383 -> 0.356, 11if 0.325 -> 0.295, 5reo 0.316 -> 0.326 (CA RMSD 0.48 -> 0.43), 7d98 0.416 -> 0.393;
REFMAC converged: 0.168 / 0.284 / 0.311 / 0.251 / 0.261 / 0.361. Deposited models (14 parity sets): R-work
lower, R-free +0.0035 on average (their free sets were fitted to; the minimiser now moves the bottom rung
further). Time unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
(cherry picked from commit dfba4c1bd3c9345f31b64fd1cadc805810960825)
This commit is contained in:
leonarski_f committed 2026-10-26 01:38:24 +01:00
1 parent 626afc49dd
commit 4424d8884d
7 files changed
+343 -18
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

image_analysis/structure_refinement/Refinement.cpp
+131 -16
View File
@@ -18,6 +18,7 @@
#include <ceres/rotation.h>
#include "gemmi/fourier.hpp" // get_size_for_hkl, get_f_phi_on_grid
#include "gemmi/it92.hpp"
#include "gemmi/modify.hpp" // remove_hydrogens
#include "gemmi/scaling.hpp"
#include "gemmi/stats.hpp" // calculate_data_statistics
@@ -28,6 +29,7 @@
#include "ModelFFT.h"
#include "ModelScaling.h"
#include "RefinementMath.h"
#include "RefinementMinimiser.h"
#include "RefinementXray.h"
#ifdef JFJOCH_USE_CUDA
#include "../../common/CUDAWrapper.h"
@@ -48,6 +50,10 @@ constexpr double REPULSION_MARGIN = 0.5; // A: pairs this much beyond their co
constexpr double MAX_RMSD_BONDS = 0.02, MAX_RMSD_ANGLES = 2.0;
// R-free differences smaller than this do not decide between two weights: the tighter geometry is taken.
constexpr double R_FREE_TIE = 0.001;
// The minimiser (RefinementMinimiser): a solve stops where an iteration lowers the target by less than this,
// relative; the most a position (A) or a B (A^2) moves in the first trial of a line search.
constexpr double MINIMISER_F_TOLERANCE = 1e-7;
constexpr double MAX_STEP_XYZ = 0.3, MAX_STEP_B = 10.0;
// The figure of merit of the likelihood, <cos(dphi)>, at X = 2 Fo D Fc / var: I1(X)/I0(X) for an acentric
// reflection, tanh(X / 2) for a centric one.
@@ -305,6 +311,77 @@ struct Context {
return sc;
}
// The X-ray target's curvature in each atom's B, up to one factor (CalibrateBCurvature): the likelihood's
// curvature in |F_model| of the cycle, 2 (k_total D)^2 / var, summed in resolution shells, against the atom's own
// part of it, (s^2 / 4)^2 f(s)^2 exp(-B s^2 / 2). A hydrogen's is its parent's.
std::vector<double> XrayBCurvatureShape() const {
constexpr int SHELLS = 32;
double s2_max = 0;
for (double s2 : refl.inv_d2)
s2_max = std::max(s2_max, s2);
std::vector<double> w(SHELLS, 0), s2_sum(SHELLS, 0);
std::vector<int> n(SHELLS, 0);
for (size_t m = 0; m < refl.Size(); m++) {
if (refl.free[m])
continue;
const int k = std::min(SHELLS - 1, static_cast<int>(refl.inv_d2[m] / s2_max * SHELLS));
const double kd = rows.k_total[m] * rows.d[m];
w[k] += 2.0 * kd * kd / rows.var[m];
s2_sum[k] += refl.inv_d2[m];
n[k]++;
}
std::map<std::pair<int, signed char>, std::vector<double>> f2; // per element and charge, per shell
std::vector<double> shape(atoms.size(), 0.0);
for (size_t i = 0; i < atoms.size(); i++) {
const gemmi::Atom &a = *atoms[i];
std::vector<double> &f = f2[{static_cast<int>(a.element.elem), a.charge}];
if (f.empty()) {
const auto &coef = gemmi::IT92<double>::get(a.element, a.charge, a.serial);
f.resize(SHELLS, 0);
for (int k = 0; k < SHELLS; k++)
if (n[k] > 0)
f[k] = std::pow(coef.calculate_sf(s2_sum[k] / n[k] / 4.0), 2);
}
for (int k = 0; k < SHELLS; k++)
if (n[k] > 0) {
const double s2 = s2_sum[k] / n[k];
shape[i] += w[k] * f[k] * a.occ * a.occ * std::exp(-a.b_iso * s2 / 2.0) * s2 * s2 / 16.0;
}
}
for (size_t i = 0; i < rider_parent.size(); i++)
if (rider_parent[i] >= 0)
shape[rider_parent[i]] += shape[i];
return shape;
}
// The factor that turns XrayBCurvatureShape() into the X-ray target's curvature: its second derivative along a
// random +-1 direction of the B factors of `params`, by central differences of the gradient, over what the shape
// predicts for it.
double CalibrateBCurvature(const std::vector<int> &params, const std::vector<double> &shape) {
constexpr double EPS = 0.5; // A^2
const std::vector<double> b0 = B();
std::mt19937 rng(7);
std::vector<double> v(params.size());
for (double &x : v)
x = rng() % 2 == 0 ? 1.0 : -1.0;
std::array<std::vector<std::array<double, 4>>, 2> g;
for (int side = 0; side < 2; side++) {
for (size_t k = 0; k < params.size(); k++)
atoms[params[k]]->b_iso = static_cast<float>(b0[params[k]] + (side == 0 ? EPS : -EPS) * v[k]);
MoveRiders();
engine->Evaluate(Model(), &g[side]);
FoldRiders(g[side]);
}
SetB(b0);
MoveRiders();
double num = 0, den = 0;
for (size_t k = 0; k < params.size(); k++) {
num += v[k] * (g[0][params[k]][3] - g[1][params[k]][3]) / (2 * EPS);
den += shape[params[k]];
}
return den > 0 ? std::max(num / den, 0.0) : 0.0;
}
// F_model with the cycle's scale and solvent, for the atoms as they are now.
std::vector<std::complex<float>> FModel() {
engine->Fcalc(Model(), fcalc);
@@ -343,19 +420,15 @@ private:
const Context &ctx_;
};
// The positions or the B factors of the refinable atoms, against w times the X-ray target plus a restraint set.
// Following Liu & Nocedal (1989) Math. Program. 45, 503-528 for the minimiser (Ceres' L-BFGS)
class AtomTarget : public ceres::FirstOrderFunction {
// The positions or the B factors of the refinable atoms, against w times the X-ray target plus a restraint set
// (minimised by RefinementMinimiser).
class AtomTarget {
public:
enum class Mode { Xyz, B };
AtomTarget(Context &ctx, Mode mode, const std::vector<int> &param_atoms, double weight, int restraint_set)
: ctx_(ctx), mode_(mode), param_atoms_(param_atoms), weight_(weight), restraint_set_(restraint_set) {}
int NumParameters() const override {
return static_cast<int>(param_atoms_.size()) * (mode_ == Mode::Xyz ? 3 : 1);
}
bool Evaluate(const double *p, double *cost, double *gradient) const override {
bool Evaluate(const double *p, double *cost, double *gradient) const {
const RefinementStage stage(mode_ == Mode::Xyz ? "eval.xyz" : "eval.b");
for (size_t k = 0; k < param_atoms_.size(); k++) {
gemmi::Atom &a = *ctx_.atoms[param_atoms_[k]];
@@ -402,8 +475,32 @@ private:
int restraint_set_;
};
// The minimiser's diagonal for the B factors of `param_atoms` (RefinementMinimiser::SetDiagonal): weight times
// the X-ray target's curvature in B plus the B restraints', each element at least a thousandth of the median. The
// positions take none: their restraints' curvature is that of stiff bonds and angles, which move atoms together,
// so as a diagonal it slows rather than scales the minimiser.
std::vector<double> BDiagonal(const std::vector<double> &xray, const std::vector<double> &restraints,
const std::vector<int> &param_atoms, double weight) {
std::vector<double> h;
for (int i : param_atoms)
h.push_back(weight * xray[i] + restraints[i]);
std::vector<double> sorted = h;
std::nth_element(sorted.begin(), sorted.begin() + sorted.size() / 2, sorted.end());
const double floor = std::max(1e-3 * sorted[sorted.size() / 2], 1e-12);
for (double &v : h)
v = std::max(v, floor);
return h;
}
// The curvature of the B factors the minimiser scales its steps by: the X-ray target's (XrayBCurvatureShape() times
// CalibrateBCurvature()) and the B restraints' (BRestraintCurvature()), per atom.
struct BCurvature {
std::vector<double> xray, restraints;
};
// `minimiser` carries its memory from the call before; `diagonal` (B only) the curvature it scales the steps by.
void Minimise(Context &ctx, AtomTarget::Mode mode, const std::vector<int> &param_atoms, double weight,
int restraint_set, int iterations) {
int restraint_set, int iterations, RefinementMinimiser &minimiser, std::vector<double> diagonal = {}) {
if (param_atoms.empty())
return;
const RefinementStage stage(mode == AtomTarget::Mode::Xyz ? "minimise.xyz" : "minimise.b");
@@ -417,10 +514,13 @@ void Minimise(Context &ctx, AtomTarget::Mode mode, const std::vector<int> &param
} else
p.push_back(a.b_iso);
}
ceres::GradientProblem problem(std::make_unique<AtomTarget>(ctx, mode, param_atoms, weight, restraint_set));
ceres::GradientProblemSolver::Summary summary;
const CancelCheck cancel_check(ctx);
ceres::Solve(cancel_check.Options(iterations), problem, p.data(), &summary);
const AtomTarget target(ctx, mode, param_atoms, weight, restraint_set);
minimiser.SetDiagonal(diagonal.empty() ? std::vector<double>(p.size(), 1.0) : std::move(diagonal));
minimiser.Minimise([&](const std::vector<double> &x, double &f, std::vector<double> &g) {
return target.Evaluate(x.data(), &f, g.data());
},
p, iterations, MINIMISER_F_TOLERANCE, mode == AtomTarget::Mode::Xyz ? MAX_STEP_XYZ : MAX_STEP_B,
[&]() { return ctx.Cancelled(); });
// the atoms where the solver ended, not where its last line search probed
for (size_t k = 0; k < param_atoms.size(); k++) {
gemmi::Atom &a = *ctx.atoms[param_atoms[k]];
@@ -867,6 +967,10 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const s
const RefinementStage stage(adp_w > 0 ? "rung.b" : "rung.xyz");
ctx.SetPositions(pos0);
ctx.SetB(b0);
// one memory for the positions and one for the B factors, kept from macro-cycle to macro-cycle
RefinementMinimiser xyz_minimiser, b_minimiser;
BCurvature curvature;
double xray_scale = 0;
for (int cycle = 0; cycle < settings.macro_cycles && !ctx.Cancelled(); cycle++) {
PlaceRiders(ctx, topo);
ctx.Refit();
@@ -875,9 +979,20 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const s
const RefinementStage s1("set_restraints");
ctx.engine->SetRestraints(0, set);
}
Minimise(ctx, AtomTarget::Mode::Xyz, refinable, w, 0, settings.xyz_iterations);
if (adp_w > 0)
Minimise(ctx, AtomTarget::Mode::B, refinable, adp_w, 1, settings.b_iterations);
Minimise(ctx, AtomTarget::Mode::Xyz, refinable, w, 0, settings.xyz_iterations, xyz_minimiser);
if (adp_w > 0) {
{
const RefinementStage s2("curvature");
curvature.xray = ctx.XrayBCurvatureShape();
if (cycle == 0)
xray_scale = ctx.CalibrateBCurvature(refinable, curvature.xray);
for (double &x : curvature.xray)
x *= xray_scale;
curvature.restraints = BRestraintCurvature(topo.adp, ctx.XyzB());
}
Minimise(ctx, AtomTarget::Mode::B, refinable, adp_w, 1, settings.b_iterations, b_minimiser,
BDiagonal(curvature.xray, curvature.restraints, refinable, adp_w));
}
}
PlaceRiders(ctx, topo);
TrialOutcome out;