leonarski_fandClaude Opus 5
13de0e7a87
The rigid-body target is an Fcalc + solvent mask + scale re-fit per evaluation, and its forward-difference Jacobian made six of them one after the other after the central one - 7 of the ~8 evaluations per LM iteration. The six shifted evaluations now run in parallel, column j on its own Evaluator over its own copy of the model (an evaluation moves every atom, so two cannot share one), set to the same zone and given the central evaluation's bulk-solvent pair - the pair the serial loop's shifted evaluations used, since the zone's solvent is fitted on the zone's first evaluation and every Evaluate starts with the central one. Each column's arithmetic is the serial one and the Jacobian is assembled in column order, so the refinement is bit-identical; the evaluation count is kept as the serial loop kept it (up to and including a failing column). The per-zone solvent fit (FitModelScale) gets the thread count as well. Inside a null replicate (a pool worker) the columns run inline, as before. To check: RIGID_BODY_* and MODEL_* keys and md5 of maps/.mtz/_model.cif identical with and without this commit on the audit set; rigid-body seconds on a large model. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013nW6FNRP1bBJJ8pfHiByAT
404 lines
20 KiB
C++
404 lines
20 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "RigidBodyRefine.h"
|
|
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <chrono>
|
|
#include <cmath>
|
|
#include <complex>
|
|
#include <vector>
|
|
|
|
#include <ceres/ceres.h>
|
|
#include <ceres/rotation.h>
|
|
|
|
#include "gemmi/dencalc.hpp" // DensityCalculator
|
|
#include "gemmi/it92.hpp" // IT92 x-ray form factors
|
|
#include "gemmi/scaling.hpp" // Scaling (bulk solvent + anisotropic B)
|
|
#include "gemmi/solmask.hpp" // SolventMasker
|
|
|
|
#include "ModelFFT.h" // MapToFPhi
|
|
#include "ModelScaling.h" // FitModelScale
|
|
#include "../common/JFJochMath.h" // PI
|
|
#include "../common/Logger.h"
|
|
#include "../common/ParallelFor.h"
|
|
|
|
namespace {
|
|
|
|
using Table = gemmi::IT92<float>;
|
|
|
|
// The ladder the placement is walked down. It starts coarse because the model arrives already placed
|
|
// but out by a cell's worth of non-isomorphism: at 6 A a few hundred reflections see the body as a
|
|
// blob and the target has one broad minimum, and each finer zone starts from the previous one's
|
|
// answer. It stops at 3.5 A, which is where rigid-body refinement is conventionally run (it is
|
|
// REFMAC's own default through dimple) - the movement being recovered is a few tenths of an
|
|
// angstrom, a tenth of that resolution, so it is well determined there, while a finer zone costs
|
|
// (1/d)^3 in grid points and reflections for a placement it cannot meaningfully sharpen.
|
|
constexpr double LADDER[] = {6.0, 4.5, 3.5};
|
|
|
|
// The step of the forward-difference Jacobian, as a fraction of the zone's resolution - so it is
|
|
// 0.06 A of atom displacement at 6 A and 0.035 A at 3.5 A. A step fixed in angstroms instead is far
|
|
// too small for the coarse zones, where a structure factor barely notices it and the derivative is
|
|
// swallowed by the jitter of the scale re-fit: measured, a fixed 0.02 A left the 6 A zone at 0.35
|
|
// degrees where this rule takes it to 2.79, which is most of the way to the answer.
|
|
constexpr double JACOBIAN_STEP_FRACTION = 0.01;
|
|
|
|
// Parameters are carried as six lengths in angstroms - the first three are the angle-axis rotation
|
|
// vector multiplied by the model's rms radius, so a unit of each of the six moves a typical atom by
|
|
// the same amount. That makes the Jacobian step isotropic in something physical, rather than mixing
|
|
// radians with angstroms.
|
|
struct Placement {
|
|
gemmi::Position centre; // the model centroid: rotating about it decorrelates R from t
|
|
double rms_radius = 1.0; // rms distance of the atoms from the centroid
|
|
|
|
void Apply(const double q[6], const std::vector<gemmi::Position> &base, gemmi::Model &model) const {
|
|
const double aa[3] = {q[0] / rms_radius, q[1] / rms_radius, q[2] / rms_radius};
|
|
size_t i = 0;
|
|
for (gemmi::Chain &ch : model.chains)
|
|
for (gemmi::Residue &r : ch.residues)
|
|
for (gemmi::Atom &a : r.atoms) {
|
|
const double p[3] = {base[i].x - centre.x, base[i].y - centre.y, base[i].z - centre.z};
|
|
double rp[3];
|
|
ceres::AngleAxisRotatePoint(aa, p, rp);
|
|
a.pos = gemmi::Position(rp[0] + centre.x + q[3],
|
|
rp[1] + centre.y + q[4],
|
|
rp[2] + centre.z + q[5]);
|
|
++i;
|
|
}
|
|
}
|
|
};
|
|
|
|
// One target evaluation: place the model, recompute Fcalc and the bulk-solvent mask to the zone's
|
|
// resolution, re-fit the scale, and hand back the amplitude residuals.
|
|
class Evaluator {
|
|
public:
|
|
Evaluator(gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
|
|
const std::vector<gemmi::Position> &base, const Placement &placement, size_t nthreads)
|
|
: model_(model), cell_(cell), sg_(sg), base_(base), placement_(placement), nthreads_(nthreads) {}
|
|
|
|
// The zone's observations, and the scale the residuals are expressed in.
|
|
void SetZone(const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs, double d_min) {
|
|
fobs_ = fobs;
|
|
d_min_ = d_min;
|
|
double sum = 0;
|
|
for (const auto &hv : fobs_.v)
|
|
sum += hv.value.value;
|
|
f_mean_ = fobs_.v.empty() ? 1.0 : sum / static_cast<double>(fobs_.v.size());
|
|
solvent_fitted_ = false;
|
|
}
|
|
|
|
// Hold the zone's bulk solvent at a pair another evaluator already fitted, as if this one had
|
|
// fitted it itself (the Jacobian's column evaluators take the central evaluation's).
|
|
void UseSolvent(double k, double b) {
|
|
k_sol = k;
|
|
b_sol = b;
|
|
solvent_fitted_ = true;
|
|
}
|
|
|
|
size_t NumObservations() const { return fobs_.v.size(); }
|
|
double JacobianStep() const { return JACOBIAN_STEP_FRACTION * d_min_; }
|
|
int evaluations = 0;
|
|
int unmatched = 0; // zone observations with no calculated amplitude to compare against
|
|
double k_sol = 0, b_sol = 0; // the bulk solvent the zone's target was evaluated with
|
|
|
|
bool Residuals(const double q[6], double *residuals) {
|
|
++evaluations;
|
|
placement_.Apply(q, base_, model_);
|
|
|
|
gemmi::DensityCalculator<Table, float> dc;
|
|
dc.d_min = d_min_;
|
|
dc.rate = 1.5;
|
|
dc.grid.unit_cell = cell_;
|
|
dc.grid.spacegroup = &sg_;
|
|
dc.set_refmac_compatible_blur(model_);
|
|
dc.put_model_density_on_grid(model_);
|
|
gemmi::AsuData<std::complex<float>> fcalc =
|
|
MapToFPhi(dc.grid).prepare_asu_data(dc.d_min, dc.blur, false, false, false);
|
|
|
|
gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
|
|
gemmi::Grid<float> mask_grid;
|
|
mask_grid.unit_cell = cell_;
|
|
mask_grid.spacegroup = &sg_;
|
|
mask_grid.set_size_from_spacing(dc.requested_grid_spacing(), gemmi::GridSizeRounding::Up);
|
|
masker.put_mask_on_grid(mask_grid, model_);
|
|
gemmi::AsuData<std::complex<float>> fmask =
|
|
MapToFPhi(mask_grid).prepare_asu_data(dc.d_min, 0);
|
|
if (fmask.size() != fcalc.size())
|
|
return false;
|
|
|
|
// Re-fitted at every evaluation: with the scale held at the starting placement's value the
|
|
// target would measure the scale as much as the placement, and the body would translate to
|
|
// repair a scale error instead of moving where the density is.
|
|
//
|
|
// The bulk solvent is not part of that scale. k_sol and b_sol describe the disordered solvent
|
|
// of the crystal rather than the fit of one placement, so they are fitted once per zone - by
|
|
// FitModelScale, inside the same physical box the reported fit is searched in - and then held
|
|
// while the overall scale and the anisotropic B follow the body. Leaving them free at every
|
|
// evaluation, which is what gemmi's unbounded Levenberg-Marquardt did here, puts a solvent
|
|
// term with no physical meaning inside the target that decides where the model goes: measured
|
|
// over a corpus of deposited models, 40% of the evaluations came out with b_sol outside
|
|
// 10-80 A^2, some of them negative, which is a solvent that GROWS with resolution.
|
|
gemmi::Scaling<float> scaling(cell_, &sg_);
|
|
scaling.use_solvent = true;
|
|
scaling.prepare_points(fcalc, fobs_, &fmask);
|
|
if (scaling.points.empty())
|
|
return false;
|
|
if (!solvent_fitted_) {
|
|
FitModelScale(scaling, {}, nthreads_);
|
|
k_sol = scaling.k_sol;
|
|
b_sol = scaling.b_sol;
|
|
solvent_fitted_ = true;
|
|
}
|
|
scaling.k_sol = k_sol;
|
|
scaling.b_sol = b_sol;
|
|
scaling.fix_k_sol = true;
|
|
scaling.fix_b_sol = true;
|
|
scaling.fit_isotropic_b_approximately();
|
|
scaling.fit_parameters();
|
|
scaling.scale_data(fcalc, &fmask);
|
|
|
|
// Both are sorted and in the same ASU, so one merge pass matches them. An observation with no
|
|
// calculated amplitude gets residual 0, which drops it from the target rather than scoring it
|
|
// as a perfect fit: its Jacobian row below comes out zero as well, and the set that matches is
|
|
// fixed by the cell, the group and the zone, so it does not move as the body does. It is
|
|
// counted and reported because a large count is a statement about the model rather than about
|
|
// this refinement - a group whose reflection conditions the data do not obey leaves half of
|
|
// them with nothing to compare against.
|
|
auto c = fcalc.v.begin();
|
|
unmatched = 0;
|
|
for (size_t i = 0; i < fobs_.v.size(); ++i) {
|
|
const gemmi::Miller &h = fobs_.v[i].hkl;
|
|
while (c != fcalc.v.end() && c->hkl < h)
|
|
++c;
|
|
const bool matched = c != fcalc.v.end() && c->hkl == h;
|
|
if (!matched)
|
|
++unmatched;
|
|
residuals[i] = matched ? (fobs_.v[i].value.value - std::abs(c->value)) / f_mean_ : 0.0;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
private:
|
|
gemmi::Model &model_;
|
|
const gemmi::UnitCell &cell_;
|
|
const gemmi::SpaceGroup &sg_;
|
|
const std::vector<gemmi::Position> &base_;
|
|
Placement placement_;
|
|
gemmi::AsuData<gemmi::ValueSigma<float>> fobs_;
|
|
double d_min_ = 0;
|
|
double f_mean_ = 1;
|
|
bool solvent_fitted_ = false;
|
|
size_t nthreads_ = 1;
|
|
};
|
|
|
|
// Ceres' own numeric differentiation steps by |x| * relative_step_size, which is zero at the start of
|
|
// every zone (the placement begins at no shift), so the Jacobian is supplied here instead, by
|
|
// forward differences at a step chosen in the parameters' units. Analytic dF/dp would need
|
|
// derivatives GEMMI's structure-factor path does not have, and at six parameters it is not worth it:
|
|
// a Jacobian costs seven evaluations, and the evaluations at 6-3.5 A are cheap.
|
|
//
|
|
// The six shifted evaluations are independent of each other, so they run in parallel, column j on
|
|
// `columns[j]`: an evaluator over its own copy of the model (an evaluation moves every atom to the
|
|
// placement it is asked about, so two cannot share one), set to the same zone and handed the solvent
|
|
// pair the central evaluation fitted - which is the pair the serial loop's shifted evaluations used,
|
|
// since the zone's solvent is fitted by the first evaluation of the zone and every Evaluate starts
|
|
// with the central one. Each column's arithmetic is the serial one, so the Jacobian is too.
|
|
class RigidBodyCost : public ceres::CostFunction {
|
|
public:
|
|
RigidBodyCost(Evaluator &ev, std::vector<Evaluator> &columns, size_t nthreads)
|
|
: ev_(ev), columns_(columns), nthreads_(nthreads) {
|
|
set_num_residuals(static_cast<int>(ev.NumObservations()));
|
|
mutable_parameter_block_sizes()->push_back(6);
|
|
}
|
|
|
|
bool Evaluate(double const *const *parameters, double *residuals, double **jacobians) const override {
|
|
const int n = num_residuals();
|
|
if (!ev_.Residuals(parameters[0], residuals))
|
|
return false;
|
|
if (jacobians != nullptr && jacobians[0] != nullptr) {
|
|
const double step = ev_.JacobianStep();
|
|
std::vector<std::vector<double>> shifted(6, std::vector<double>(n));
|
|
std::array<bool, 6> ok{};
|
|
ParallelFor(6, nthreads_, [&](int j) {
|
|
double q[6];
|
|
std::copy(parameters[0], parameters[0] + 6, q);
|
|
q[j] += step;
|
|
columns_[j].UseSolvent(ev_.k_sol, ev_.b_sol);
|
|
ok[j] = columns_[j].Residuals(q, shifted[j].data());
|
|
});
|
|
// Counted as the serial loop counted them: it stopped at the first column that failed.
|
|
for (int j = 0; j < 6; j++) {
|
|
++ev_.evaluations;
|
|
ev_.unmatched = columns_[j].unmatched;
|
|
if (!ok[j])
|
|
return false;
|
|
for (int i = 0; i < n; i++)
|
|
jacobians[0][i * 6 + j] = (shifted[j][i] - residuals[i]) / step;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
private:
|
|
Evaluator &ev_;
|
|
std::vector<Evaluator> &columns_;
|
|
size_t nthreads_;
|
|
};
|
|
|
|
// The directions in which this space group's origin is free. Translating the whole cell content along
|
|
// one of them multiplies every F by a phase and leaves every |F| EXACTLY unchanged, so the target
|
|
// cannot determine that component: all three directions in P1, the unique axis in a polar group. The
|
|
// R-free gate cannot stand in for this - it is a function of |F| too, so along such a direction it
|
|
// sees only grid noise and commits or not by coin flip, while the other five parameters carry the
|
|
// noise in with them. The free directions are the common fixed subspace of the group's rotation
|
|
// parts, and the projector onto it is simply their average.
|
|
gemmi::Mat33 GaugeProjector(const gemmi::SpaceGroup &sg, const gemmi::UnitCell &cell) {
|
|
const gemmi::GroupOps gops = sg.operations();
|
|
double m[3][3] = {};
|
|
const double n = static_cast<double>(gops.sym_ops.size()) * gemmi::Op::DEN;
|
|
for (const gemmi::Op &op : gops.sym_ops)
|
|
for (int i = 0; i < 3; i++)
|
|
for (int j = 0; j < 3; j++)
|
|
m[i][j] += static_cast<double>(op.rot[i][j]) / n;
|
|
const gemmi::Mat33 mean(m[0][0], m[0][1], m[0][2],
|
|
m[1][0], m[1][1], m[1][2],
|
|
m[2][0], m[2][1], m[2][2]);
|
|
// Fractional projector taken into orthogonal space, where the parameters live.
|
|
return cell.orth.mat.multiply(mean).multiply(cell.frac.mat);
|
|
}
|
|
|
|
} // namespace
|
|
|
|
std::vector<gemmi::Position> ModelPositions(const gemmi::Model &model) {
|
|
std::vector<gemmi::Position> pos;
|
|
for (const gemmi::Chain &ch : model.chains)
|
|
for (const gemmi::Residue &r : ch.residues)
|
|
for (const gemmi::Atom &a : r.atoms)
|
|
pos.push_back(a.pos);
|
|
return pos;
|
|
}
|
|
|
|
void SetModelPositions(gemmi::Model &model, const std::vector<gemmi::Position> &pos) {
|
|
size_t i = 0;
|
|
for (gemmi::Chain &ch : model.chains)
|
|
for (gemmi::Residue &r : ch.residues)
|
|
for (gemmi::Atom &a : r.atoms)
|
|
a.pos = pos[i++];
|
|
}
|
|
|
|
// One rigid body, not groups: a fragment-screening model arrives already solved and isomorphous, and
|
|
// the movement to recover is the crystal's, not the molecule's. Splitting it into domains or giving a
|
|
// bound ligand its own six parameters would refine against evidence this data does not separately
|
|
// carry, and the ligand is what the difference map is meant to show rather than model away.
|
|
//
|
|
// Some of the translation would be a gauge rather than a quantity - the origin is free in all three
|
|
// directions in P1 and along the unique axis in a polar group, and |F| does not change when the whole
|
|
// content moves along it - so that component is projected out after every zone. Neither of the two
|
|
// things that might look like they cover it actually does: the R-free gate is a function of |F| and
|
|
// therefore blind to exactly this, and the LM damping follows the gauge column of the Jacobian, which
|
|
// is not zero but noise divided by the difference step.
|
|
RigidBodyRefineResult RefineRigidBody(gemmi::Model &model,
|
|
const gemmi::UnitCell &cell,
|
|
const gemmi::SpaceGroup &sg,
|
|
const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs,
|
|
double d_min,
|
|
Logger &logger,
|
|
size_t nthreads) {
|
|
const auto t0 = std::chrono::steady_clock::now();
|
|
RigidBodyRefineResult result;
|
|
|
|
const std::vector<gemmi::Position> base = ModelPositions(model);
|
|
if (base.empty() || fobs.v.empty())
|
|
return result;
|
|
|
|
Placement placement;
|
|
for (const gemmi::Position &p : base)
|
|
placement.centre += p;
|
|
placement.centre *= 1.0 / static_cast<double>(base.size());
|
|
double r2 = 0;
|
|
for (const gemmi::Position &p : base)
|
|
r2 += placement.centre.dist_sq(p);
|
|
placement.rms_radius = std::sqrt(r2 / static_cast<double>(base.size()));
|
|
if (!(placement.rms_radius > 0))
|
|
return result;
|
|
|
|
std::vector<double> ladder;
|
|
for (double zone : LADDER)
|
|
if (zone >= d_min)
|
|
ladder.push_back(zone);
|
|
if (ladder.empty())
|
|
ladder.push_back(d_min);
|
|
|
|
const gemmi::Mat33 gauge = GaugeProjector(sg, cell);
|
|
Evaluator ev(model, cell, sg, base, placement, nthreads);
|
|
// The Jacobian's six column evaluators, each over its own copy of the model (see RigidBodyCost).
|
|
// The copies only ever hold probe placements; `model` is left at the answer below.
|
|
std::vector<gemmi::Model> column_models(6, model);
|
|
std::vector<Evaluator> columns;
|
|
columns.reserve(6);
|
|
for (int j = 0; j < 6; j++)
|
|
columns.emplace_back(column_models[j], cell, sg, base, placement, 1);
|
|
double q[6] = {0, 0, 0, 0, 0, 0};
|
|
bool any_zone_solved = false;
|
|
for (double zone : ladder) {
|
|
gemmi::AsuData<gemmi::ValueSigma<float>> zone_obs;
|
|
zone_obs.unit_cell_ = fobs.unit_cell_;
|
|
zone_obs.spacegroup_ = fobs.spacegroup_;
|
|
for (const auto &hv : fobs.v)
|
|
if (cell.calculate_d(hv.hkl) >= zone)
|
|
zone_obs.v.push_back(hv);
|
|
if (zone_obs.v.size() < 50)
|
|
continue;
|
|
result.zones.push_back(zone); // the ladder WALKED, which a thin zone drops out of
|
|
ev.SetZone(zone_obs, zone);
|
|
for (Evaluator &c : columns)
|
|
c.SetZone(zone_obs, zone);
|
|
|
|
ceres::Problem problem;
|
|
problem.AddResidualBlock(new RigidBodyCost(ev, columns, nthreads), nullptr, q);
|
|
ceres::Solver::Options options;
|
|
options.linear_solver_type = ceres::DENSE_QR;
|
|
options.max_num_iterations = 15;
|
|
options.function_tolerance = 1e-4;
|
|
options.parameter_tolerance = 1e-4;
|
|
options.logging_type = ceres::LoggingType::SILENT;
|
|
ceres::Solver::Summary summary;
|
|
const int evaluations_before = ev.evaluations;
|
|
const auto zone_t0 = std::chrono::steady_clock::now();
|
|
ceres::Solve(options, &problem, &summary);
|
|
any_zone_solved = any_zone_solved || summary.IsSolutionUsable();
|
|
const gemmi::Vec3 along = gauge.multiply(gemmi::Vec3(q[3], q[4], q[5]));
|
|
q[3] -= along.x; q[4] -= along.y; q[5] -= along.z;
|
|
logger.Debug("Rigid body zone {:.1f} A: {} reflections ({} without a calculated amplitude), "
|
|
"solvent k_sol {:.2f} b_sol {:.0f} A^2, {} iterations, {} evaluations, {:.2f} s, "
|
|
"rotation {:.3f} deg, translation {:.3f} A", zone, zone_obs.v.size(), ev.unmatched,
|
|
ev.k_sol, ev.b_sol,
|
|
summary.iterations.empty() ? 0 : summary.iterations.size() - 1,
|
|
ev.evaluations - evaluations_before,
|
|
std::chrono::duration<double>(std::chrono::steady_clock::now() - zone_t0).count(),
|
|
std::sqrt(q[0]*q[0] + q[1]*q[1] + q[2]*q[2]) / placement.rms_radius * 180.0 / PI,
|
|
std::sqrt(q[3]*q[3] + q[4]*q[4] + q[5]*q[5]));
|
|
}
|
|
|
|
placement.Apply(q, base, model); // Ceres left the model at a Jacobian probe; put it at the answer
|
|
result.evaluations = ev.evaluations;
|
|
result.converged = any_zone_solved;
|
|
result.k_sol = ev.k_sol;
|
|
result.b_sol = ev.b_sol;
|
|
const double aa = std::sqrt(q[0] * q[0] + q[1] * q[1] + q[2] * q[2]) / placement.rms_radius;
|
|
result.angle_deg = aa * 180.0 / PI;
|
|
result.shift_A = std::sqrt(q[3] * q[3] + q[4] * q[4] + q[5] * q[5]);
|
|
result.seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();
|
|
|
|
if (!result.zones.empty())
|
|
logger.Info("Model validation: rigid body over {} resolution zone(s) down to {:.1f} A, "
|
|
"{} evaluations in {:.2f} s: rotation {:.3f} deg, translation {:.3f} A",
|
|
result.zones.size(), result.zones.back(), result.evaluations, result.seconds,
|
|
result.angle_deg, result.shift_A);
|
|
else
|
|
logger.Info("Model validation: rigid body had no resolution zone with enough reflections to "
|
|
"run in; the model is left where it arrived");
|
|
return result;
|
|
}
|