diff --git a/docs/ACKNOWLEDGEMENT.md b/docs/ACKNOWLEDGEMENT.md index 8a2e79807..389b7bbd2 100644 --- a/docs/ACKNOWLEDGEMENT.md +++ b/docs/ACKNOWLEDGEMENT.md @@ -537,7 +537,8 @@ respect to every atom is taken by the FFT method - the derivatives of the target structure factors transformed into a map and summed against each atom's density derivatives: R. C. Agarwal, "A new least-squares refinement technique based on the fast Fourier transform algorithm" (1978), Acta Cryst. A34, 791-809 [doi:10.1107/S0567739478001618](https://doi.org/10.1107/S0567739478001618). -The positions and B factors are minimised by L-BFGS (Ceres): D. C. Liu and J. Nocedal, "On the limited +The positions and B factors are minimised by L-BFGS (rugnux's own, its memory kept from one macro-cycle to +the next, the B factors' steps scaled by their estimated curvature): D. C. Liu and J. Nocedal, "On the limited memory BFGS method for large scale optimization" (1989), Math. Program. 45, 503-528 [doi:10.1007/BF01589116](https://doi.org/10.1007/BF01589116). The geometry restraints are those of the CCP4 Monomer Library, a subset of which rugnux carries (see THIRD_PARTY_NOTICES.md): A. A. Vagin, diff --git a/image_analysis/structure_refinement/CMakeLists.txt b/image_analysis/structure_refinement/CMakeLists.txt index 44d65a19d..d9cf337f4 100644 --- a/image_analysis/structure_refinement/CMakeLists.txt +++ b/image_analysis/structure_refinement/CMakeLists.txt @@ -18,7 +18,7 @@ ADD_LIBRARY(JFJochStructureRefinement STATIC $<$:ModelScaleGPU.cu> ModelScaleGPU.h SigmaA.cpp SigmaA.h MonomerLibrary.cpp MonomerLibrary.h - RefinementMath.h + RefinementMath.h RefinementMinimiser.h RefinementRestraints.cpp RefinementRestraints.h RefinementXray.cpp RefinementXray.h Refinement.cpp Refinement.h diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index 3df00741f..c0e3e5770 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -18,6 +18,7 @@ #include #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, , 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 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 w(SHELLS, 0), s2_sum(SHELLS, 0); + std::vector 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(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::vector> f2; // per element and charge, per shell + std::vector shape(atoms.size(), 0.0); + for (size_t i = 0; i < atoms.size(); i++) { + const gemmi::Atom &a = *atoms[i]; + std::vector &f = f2[{static_cast(a.element.elem), a.charge}]; + if (f.empty()) { + const auto &coef = gemmi::IT92::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 ¶ms, const std::vector &shape) { + constexpr double EPS = 0.5; // A^2 + const std::vector b0 = B(); + std::mt19937 rng(7); + std::vector v(params.size()); + for (double &x : v) + x = rng() % 2 == 0 ? 1.0 : -1.0; + std::array>, 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(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> 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 ¶m_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(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 BDiagonal(const std::vector &xray, const std::vector &restraints, + const std::vector ¶m_atoms, double weight) { + std::vector h; + for (int i : param_atoms) + h.push_back(weight * xray[i] + restraints[i]); + std::vector 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 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 ¶m_atoms, double weight, - int restraint_set, int iterations) { + int restraint_set, int iterations, RefinementMinimiser &minimiser, std::vector 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 ¶m } else p.push_back(a.b_iso); } - ceres::GradientProblem problem(std::make_unique(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(p.size(), 1.0) : std::move(diagonal)); + minimiser.Minimise([&](const std::vector &x, double &f, std::vector &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 &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 &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; diff --git a/image_analysis/structure_refinement/RefinementMinimiser.h b/image_analysis/structure_refinement/RefinementMinimiser.h new file mode 100644 index 000000000..42bb3afe1 --- /dev/null +++ b/image_analysis/structure_refinement/RefinementMinimiser.h @@ -0,0 +1,153 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include +#include +#include + +// Limited-memory BFGS whose memory outlives one call, so that a refinement's macro-cycles - each a few dozen +// iterations on a target whose scale and likelihood parameters were refitted in between - go on from what the +// earlier ones learnt of the curvature instead of starting each from a steepest-descent step. The initial inverse +// Hessian is a diagonal the caller gives (SetDiagonal: an estimate of the Hessian's diagonal, so that a parameter +// with ten times the curvature of another takes a tenth of its step), scaled each iteration by the latest pair +// (Oren-Shanno in that metric). The line search backtracks from the full step to sufficient decrease. +// Following Liu & Nocedal (1989) Math. Program. 45, 503-528 +class RefinementMinimiser { +public: + // f(x, g): the target at x, and its gradient in g (sized as x); false where it cannot be evaluated there. + using Target = std::function &x, double &f, std::vector &g)>; + + explicit RefinementMinimiser(size_t memory = 20) : memory_(memory) {} + + // Forget the pairs (the target changed beyond what the memory still describes, e.g. another weight). + void Reset() { + s_.clear(); + y_.clear(); + } + // The estimated diagonal of the Hessian; every element > 0. + void SetDiagonal(std::vector h) { h_ = std::move(h); } + + struct Outcome { + int iterations = 0; + int evaluations = 0; + double f_start = 0, f_end = 0; + }; + + // At most `iterations` iterations from x (x is where it ended); stops early where an iteration lowers f by + // less than `f_tolerance` relative to |f|, or `stop` says so (checked before every iteration). `max_step`, in + // the parameters' units, is the most any parameter moves in the first trial of a line search. + Outcome Minimise(const Target &target, std::vector &x, int iterations, double f_tolerance, double max_step, + const std::function &stop = nullptr) { + const size_t n = x.size(); + if (h_.size() != n) + h_.assign(n, 1.0); + Outcome out; + double f = 0; + std::vector g(n), d(n), x_new(n), g_new(n); + if (!target(x, f, g)) + return out; + out.evaluations++; + out.f_start = out.f_end = f; + for (int it = 0; it < iterations; it++) { + if (stop && stop()) + break; + Direction(g, d); + double gd = Dot(g, d); + if (!(gd < 0)) { // not a descent direction: the memory no longer fits, start again from the diagonal + Reset(); + Direction(g, d); + gd = Dot(g, d); + if (!(gd < 0)) + break; + } + double d_max = 0; + for (double v : d) + d_max = std::max(d_max, std::fabs(v)); + double alpha = d_max > max_step ? max_step / d_max : 1.0; + double f_new = 0; + bool accepted = false; + for (int ls = 0; ls < 10; ls++) { + for (size_t k = 0; k < n; k++) + x_new[k] = x[k] + alpha * d[k]; + const bool ok = target(x_new, f_new, g_new); + out.evaluations++; + if (ok && f_new <= f + 1e-4 * alpha * gd) { + accepted = true; + break; + } + // the minimum of the quadratic through f, its slope and f_new, kept within [0.1, 0.5] of the step + const double q = ok ? -gd * alpha * alpha / (2.0 * (f_new - f - gd * alpha)) : 0.1 * alpha; + alpha = std::clamp(q, 0.1 * alpha, 0.5 * alpha); + } + if (!accepted) + break; + std::vector s(n), y(n); + for (size_t k = 0; k < n; k++) { + s[k] = x_new[k] - x[k]; + y[k] = g_new[k] - g[k]; + } + const double sy = Dot(s, y); + if (sy > 1e-10 * std::sqrt(Dot(s, s) * Dot(y, y))) { + s_.push_back(std::move(s)); + y_.push_back(std::move(y)); + if (s_.size() > memory_) { + s_.pop_front(); + y_.pop_front(); + } + } + const double decrease = f - f_new; + x.swap(x_new); + g.swap(g_new); + f = f_new; + out.iterations++; + out.f_end = f; + if (decrease <= f_tolerance * std::max(std::fabs(f), 1.0)) + break; + } + return out; + } + +private: + static double Dot(const std::vector &a, const std::vector &b) { + double s = 0; + for (size_t k = 0; k < a.size(); k++) + s += a[k] * b[k]; + return s; + } + + // d = -H g by the two-loop recursion, H0 = gamma diag(1/h). + void Direction(const std::vector &g, std::vector &d) const { + const size_t n = g.size(), m = s_.size(); + std::vector q = g, a(m), rho(m); + for (size_t i = m; i-- > 0;) { + rho[i] = 1.0 / Dot(s_[i], y_[i]); + a[i] = rho[i] * Dot(s_[i], q); + for (size_t k = 0; k < n; k++) + q[k] -= a[i] * y_[i][k]; + } + double gamma = 1.0; + if (m > 0) { + double yhy = 0; + for (size_t k = 0; k < n; k++) + yhy += y_[m - 1][k] * y_[m - 1][k] / h_[k]; + gamma = Dot(s_[m - 1], y_[m - 1]) / yhy; + } + for (size_t k = 0; k < n; k++) + q[k] *= gamma / h_[k]; + for (size_t i = 0; i < m; i++) { + const double b = rho[i] * Dot(y_[i], q); + for (size_t k = 0; k < n; k++) + q[k] += s_[i][k] * (a[i] - b); + } + for (size_t k = 0; k < n; k++) + d[k] = -q[k]; + } + + size_t memory_; + std::deque> s_, y_; + std::vector h_; +}; diff --git a/image_analysis/structure_refinement/RefinementRestraints.cpp b/image_analysis/structure_refinement/RefinementRestraints.cpp index ed5f22145..04e997305 100644 --- a/image_analysis/structure_refinement/RefinementRestraints.cpp +++ b/image_analysis/structure_refinement/RefinementRestraints.cpp @@ -324,6 +324,16 @@ double EvaluateRestraints(const RestraintSet &set, const std::vector BRestraintCurvature(const RestraintSet &set, const std::vector> &xyzb) { + std::vector h(xyzb.size(), 0.0); + for (const RestraintRecord &r : set.records) + if (r.kind == static_cast(RestraintKind::BSimilarity)) { + h[r.atom[0]] += 2.0 / (r.sigma * r.sigma); + h[r.atom[1]] += 2.0 / (r.sigma * r.sigma); + } + return h; +} + RefinementTopology BuildRefinementTopology(gemmi::Structure &st, bool hydrogens) { RefinementTopology out; gemmi::Model &model = st.models.at(0); diff --git a/image_analysis/structure_refinement/RefinementRestraints.h b/image_analysis/structure_refinement/RefinementRestraints.h index b4717c44b..5e4be92f9 100644 --- a/image_analysis/structure_refinement/RefinementRestraints.h +++ b/image_analysis/structure_refinement/RefinementRestraints.h @@ -46,6 +46,9 @@ double RestraintValue(const RestraintSet &set, const RestraintRecord &r, const s double EvaluateRestraints(const RestraintSet &set, const std::vector> &xyzb, std::vector> *grad); +// The diagonal of the Hessian of the set's B similarity restraints with respect to each atom's B, at `xyzb`. +std::vector BRestraintCurvature(const RestraintSet &set, const std::vector> &xyzb); + // What refinement knows of a model's chemistry. struct RefinementTopology { RestraintSet geometry; // bonds, angles, torsions, chiral volumes, planes diff --git a/tests/RefinementTest.cpp b/tests/RefinementTest.cpp index 9571e3f1b..7081ba73d 100644 --- a/tests/RefinementTest.cpp +++ b/tests/RefinementTest.cpp @@ -17,6 +17,7 @@ #include "../image_analysis/structure_refinement/MonomerLibrary.h" #include "../image_analysis/structure_refinement/Refinement.h" #include "../image_analysis/structure_refinement/RefinementMath.h" +#include "../image_analysis/structure_refinement/RefinementMinimiser.h" #include "../image_analysis/structure_refinement/RefinementRestraints.h" #include "../image_analysis/structure_refinement/RefinementXray.h" #ifdef JFJOCH_USE_CUDA @@ -555,6 +556,48 @@ TEST_CASE("RefinementXray_gradient_GPU", "[Refinement]") { } #endif +TEST_CASE("RefinementMinimiser_quadratic", "[Refinement]") { + // f = sum a_k (x_k - 1)^2 with curvatures spread over two decades + const size_t n = 50; + std::vector a(n); + for (size_t k = 0; k < n; k++) + a[k] = std::pow(100.0, static_cast(k) / (n - 1)); + const RefinementMinimiser::Target target = [&](const std::vector &x, double &f, std::vector &g) { + f = 0; + for (size_t k = 0; k < n; k++) { + f += a[k] * (x[k] - 1) * (x[k] - 1); + g[k] = 2 * a[k] * (x[k] - 1); + } + return true; + }; + auto distance = [&](const std::vector &x) { + double d = 0; + for (double v : x) + d = std::max(d, std::fabs(v - 1)); + return d; + }; + + // the exact diagonal: one step + RefinementMinimiser exact; + std::vector h(n); + for (size_t k = 0; k < n; k++) + h[k] = 2 * a[k]; + exact.SetDiagonal(h); + std::vector x(n, 0.0); + exact.Minimise(target, x, 1, 0.0, 10.0); + CHECK(distance(x) < 1e-9); + + // without it, the memory carried from one call to the next does what one long call does + RefinementMinimiser split, whole; + std::vector xs(n, 0.0), xw(n, 0.0); + for (int call = 0; call < 6; call++) + split.Minimise(target, xs, 10, 0.0, 10.0); + whole.Minimise(target, xw, 60, 0.0, 10.0); + CHECK(distance(xw) < 1e-3); + for (size_t k = 0; k < n; k++) + CHECK(xs[k] == Catch::Approx(xw[k]).margin(1e-12)); +} + TEST_CASE("Refinement_restrained_recovers_model", "[Refinement]") { const gemmi::UnitCell cell(36, 40, 44, 90, 90, 90); const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name("P 21 21 21");