leonarski_fandClaude Opus 5.5
dcb01100e7
An optional `fixed` model: other atoms of the crystal held where they are - the other copies of a molecular-replacement search - whose Fcalc is added to the body's at every evaluation and whose atoms the bulk-solvent mask covers together with the body's (computed once per zone, with the body where the zone starts, and held). With them the origin is no longer free, so the gauge projection is skipped. On the CPU the fixed part is gridded once per zone; on the GPU the fixed atoms must be copies of the body, and each one's placement is read off by Kabsch superposition and composed on the device with the engine's own density, transform and mask (SetFixedCopies), so a fit beside three copies of a 90 kDa model costs what a fit alone does. Without `fixed` nothing changes (the composition kernel's extra term is null). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3
625 lines
30 KiB
C++
625 lines
30 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 <future>
|
|
#include <memory>
|
|
#include <vector>
|
|
|
|
#include <Eigen/Dense>
|
|
#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 "ModelFFT.h" // MapToFPhi
|
|
#include "ModelGrid.h" // PutModelDensityOnGrid, PutMaskOnGrid
|
|
#include "ModelScaling.h" // FitModelScale
|
|
#ifdef JFJOCH_USE_CUDA
|
|
#include "RigidBodyGPU.h" // RigidBodyTargetGPU
|
|
#endif
|
|
#include "../../common/JFJochMath.h" // PI
|
|
#include "../../common/ParallelFor.h"
|
|
#include "../../common/Logger.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 rotation columns of the 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 rounding: 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;
|
|
|
|
// How finely a zone's maps are sampled: DensityCalculator's rate, for a spacing of d_min / (2 * rate).
|
|
constexpr double GRID_RATE = 1.5;
|
|
|
|
// 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 by RigidBodyTarget.
|
|
class RigidBodyCost : public ceres::CostFunction {
|
|
public:
|
|
explicit RigidBodyCost(RigidBodyTargetBase &target) : target_(target) {
|
|
set_num_residuals(static_cast<int>(target.NumObservations()));
|
|
mutable_parameter_block_sizes()->push_back(6);
|
|
}
|
|
|
|
bool Evaluate(double const *const *parameters, double *residuals, double **jacobians) const override {
|
|
// Ceres evaluates every step it accepts twice: the residuals alone to decide whether to take
|
|
// it, then residuals and Jacobian together at the same point. The residuals are a function of
|
|
// the placement alone, so the second time they are the first time's, copied.
|
|
if (!last_residuals_.empty() && std::equal(parameters[0], parameters[0] + 6, last_q_.begin())) {
|
|
std::copy(last_residuals_.begin(), last_residuals_.end(), residuals);
|
|
} else {
|
|
if (!target_.Residuals(parameters[0], residuals))
|
|
return false;
|
|
std::copy(parameters[0], parameters[0] + 6, last_q_.begin());
|
|
last_residuals_.assign(residuals, residuals + num_residuals());
|
|
}
|
|
if (jacobians != nullptr && jacobians[0] != nullptr)
|
|
return target_.Jacobian(parameters[0], jacobians[0]);
|
|
return true;
|
|
}
|
|
|
|
private:
|
|
RigidBodyTargetBase &target_;
|
|
mutable std::array<double, 6> last_q_{};
|
|
mutable std::vector<double> last_residuals_; // at last_q_; empty until an evaluation succeeds
|
|
};
|
|
|
|
// 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.
|
|
} // namespace
|
|
|
|
gemmi::Mat33 RigidBodyGaugeProjector(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);
|
|
}
|
|
|
|
// An empty grid of the zone's size, the size DensityCalculator gives its own at GRID_RATE.
|
|
gemmi::Grid<float> RigidBodyZoneGrid(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min) {
|
|
gemmi::Grid<float> grid;
|
|
grid.unit_cell = cell;
|
|
grid.spacegroup = &sg;
|
|
grid.set_size_from_spacing(d_min / (2 * GRID_RATE), gemmi::GridSizeRounding::Up);
|
|
return grid;
|
|
}
|
|
|
|
double RigidBodyJacobianStep(double d_min) {
|
|
return JACOBIAN_STEP_FRACTION * d_min;
|
|
}
|
|
|
|
std::vector<double> RigidBodyLadder(double d_min) {
|
|
std::vector<double> ladder;
|
|
for (double zone : LADDER)
|
|
if (zone >= d_min)
|
|
ladder.push_back(zone);
|
|
if (ladder.empty())
|
|
ladder.push_back(d_min);
|
|
return ladder;
|
|
}
|
|
|
|
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++];
|
|
}
|
|
|
|
SymmetryComposition::SymmetryComposition(const gemmi::Grid<float> &grid, double d_min,
|
|
const std::vector<gemmi::Miller> &hkl) {
|
|
const gemmi::UnitCell &cell = grid.unit_cell;
|
|
const gemmi::GroupOps gops = grid.spacegroup->operations();
|
|
const gemmi::ReciprocalAsu asu(grid.spacegroup);
|
|
// prepare_asu_data()'s box on the half-l transform of this grid, (nu, nv, nw/2 + 1): up to the
|
|
// Nyquist frequency but not on it along u and v, and to nw/2 along w.
|
|
const gemmi::Miller lim = cell.get_hkl_limits(d_min);
|
|
const int max_h = std::min((grid.nu - 1) / 2, lim[0]);
|
|
const int max_k = std::min((grid.nv - 1) / 2, lim[1]);
|
|
const int max_l = std::min(grid.nw / 2, lim[2]);
|
|
const double max_1_d2 = 1.0 / (d_min * d_min);
|
|
ops_ = gops.sym_ops.size();
|
|
centring_ = static_cast<double>(gops.cen_ops.size());
|
|
row_.assign(hkl.size(), -1);
|
|
for (size_t i = 0; i < hkl.size(); i++) {
|
|
const gemmi::Miller &h = hkl[i];
|
|
const double inv_d2 = cell.calculate_1_d2(h);
|
|
if (std::abs(h[0]) > max_h || std::abs(h[1]) > max_k || std::abs(h[2]) > max_l ||
|
|
!(inv_d2 < max_1_d2) || !asu.is_in(h) || gops.is_systematically_absent(h) ||
|
|
(h[0] == 0 && h[1] == 0 && h[2] == 0))
|
|
continue;
|
|
row_[i] = static_cast<int>(hkl_.size());
|
|
hkl_.push_back(h);
|
|
inv_d2_.push_back(inv_d2);
|
|
for (const gemmi::Op &op : gops.sym_ops) {
|
|
const gemmi::Miller k = op.apply_to_hkl(h);
|
|
Term t;
|
|
t.conj = k[2] < 0; // the half-l grid holds l >= 0, and F1(-k) = conj F1(k) for a real map
|
|
const int ku = t.conj ? -k[0] : k[0], kv = t.conj ? -k[1] : k[1], kw = t.conj ? -k[2] : k[2];
|
|
t.index = gemmi::modulo(ku, grid.nu) +
|
|
static_cast<size_t>(grid.nu) * (gemmi::modulo(kv, grid.nv) + static_cast<size_t>(grid.nv) * kw);
|
|
t.phase = std::polar(1.0, -op.phase_shift(h)); // gemmi's phase_shift is -2 pi h.t
|
|
t.s = cell.frac.mat.left_multiply(gemmi::Vec3(k[0], k[1], k[2]));
|
|
t.k = k;
|
|
terms_.push_back(t);
|
|
}
|
|
}
|
|
}
|
|
|
|
void SymmetryComposition::Compose(const gemmi::FPhiGrid<float> &f1, double unblur,
|
|
std::vector<std::complex<double>> &f,
|
|
std::vector<std::array<std::complex<double>, 3>> *df_dt,
|
|
size_t nthreads) const {
|
|
const int n = static_cast<int>(hkl_.size());
|
|
f.assign(n, 0.0);
|
|
if (df_dt != nullptr)
|
|
df_dt->assign(n, {});
|
|
const std::complex<double> two_pi_i(0.0, 2 * PI);
|
|
ParallelChunks(n, nthreads, [&](int lo, int hi) {
|
|
for (int m = lo; m < hi; m++) {
|
|
std::complex<double> sum = 0;
|
|
std::array<std::complex<double>, 3> d{};
|
|
for (size_t o = 0; o < ops_; o++) {
|
|
const Term &t = terms_[m * ops_ + o];
|
|
const std::complex<double> value = f1.data[t.index];
|
|
const std::complex<double> term = t.phase * (t.conj ? std::conj(value) : value);
|
|
sum += term;
|
|
d[0] += term * t.s.x;
|
|
d[1] += term * t.s.y;
|
|
d[2] += term * t.s.z;
|
|
}
|
|
// prepare_asu_data()'s unblur: exp(B_blur |s|^2 / 4)
|
|
const double scale = centring_ * std::exp(unblur * 0.25 * inv_d2_[m]);
|
|
f[m] = scale * sum;
|
|
if (df_dt != nullptr)
|
|
for (int k = 0; k < 3; k++)
|
|
(*df_dt)[m][k] = scale * two_pi_i * d[k];
|
|
}
|
|
});
|
|
}
|
|
|
|
RigidBodyTargetBase::RigidBodyTargetBase(const gemmi::Model &model)
|
|
: solvent_b_max_(SolventBMax(model)), base_(ModelPositions(model)) {
|
|
if (base_.empty())
|
|
return;
|
|
for (const gemmi::Position &p : base_)
|
|
centre_ += p;
|
|
centre_ *= 1.0 / static_cast<double>(base_.size());
|
|
double r2 = 0;
|
|
for (const gemmi::Position &p : base_)
|
|
r2 += centre_.dist_sq(p);
|
|
rms_radius_ = std::sqrt(r2 / static_cast<double>(base_.size()));
|
|
}
|
|
|
|
RigidBodyTarget::RigidBodyTarget(gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
|
|
size_t nthreads)
|
|
: RigidBodyTargetBase(model), model_(model), cell_(cell), sg_(sg), nthreads_(nthreads),
|
|
column_models_(3, model) {}
|
|
|
|
// 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. The rotation is about the model centroid, which decorrelates it from the
|
|
// translation, and is applied to the positions only: an anisotropic U is not turned with the body.
|
|
void RigidBodyTargetBase::Place(const double q[6], 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;
|
|
}
|
|
}
|
|
|
|
void RigidBodyTarget::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;
|
|
have_point_ = false;
|
|
const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell_, sg_, d_min);
|
|
orbit_leaders_ = OrbitLeaders(grid, nthreads_);
|
|
std::vector<gemmi::Miller> hkl;
|
|
for (const auto &hv : fobs_.v)
|
|
hkl.push_back(hv.hkl);
|
|
composition_.emplace(grid, d_min, hkl);
|
|
if (fixed != nullptr)
|
|
fixed_part_ = RigidBodyFixed(model_, *fixed, cell_, sg_, d_min_, *composition_, nthreads_);
|
|
}
|
|
|
|
RigidBodyFixedPart RigidBodyFixed(const gemmi::Model &moving, const gemmi::Model &fixed, const gemmi::UnitCell &cell,
|
|
const gemmi::SpaceGroup &sg, double d_min, const SymmetryComposition &composition,
|
|
size_t nthreads) {
|
|
RigidBodyFixedPart out;
|
|
gemmi::DensityCalculator<Table, float> dc;
|
|
dc.d_min = d_min;
|
|
dc.rate = GRID_RATE;
|
|
dc.grid.unit_cell = cell;
|
|
dc.grid.spacegroup = &sg;
|
|
dc.set_refmac_compatible_blur(fixed);
|
|
PutModelDensityOnGrid(dc, fixed, {}, nthreads);
|
|
composition.Compose(MapToFPhi(dc.grid), dc.blur, out.fcalc, nullptr, nthreads);
|
|
|
|
gemmi::Model all = moving;
|
|
all.chains.insert(all.chains.end(), fixed.chains.begin(), fixed.chains.end());
|
|
gemmi::Grid<float> mask_grid = RigidBodyZoneGrid(cell, sg, d_min);
|
|
PutMaskOnGrid(mask_grid, all, OrbitLeaders(mask_grid, nthreads), nthreads);
|
|
const gemmi::FPhiGrid<float> fm = MapToFPhi(mask_grid);
|
|
for (const gemmi::Miller &h : composition.Hkl())
|
|
out.fmask.push_back(fm.get_value_by_hkl(h));
|
|
return out;
|
|
}
|
|
|
|
// The model's Fcalc at the zone's composed indices, from one copy of it on the zone's grid.
|
|
void RigidBodyTarget::CopyFcalc(const gemmi::Model &model, std::vector<std::complex<double>> &f,
|
|
std::vector<std::array<std::complex<double>, 3>> *df_dt) const {
|
|
gemmi::DensityCalculator<Table, float> dc;
|
|
dc.d_min = d_min_;
|
|
dc.rate = GRID_RATE;
|
|
dc.grid.unit_cell = cell_;
|
|
dc.grid.spacegroup = &sg_;
|
|
dc.set_refmac_compatible_blur(model);
|
|
PutModelDensityOnGrid(dc, model, {}, nthreads_); // no orbits: the copy is not symmetrized
|
|
composition_->Compose(MapToFPhi(dc.grid), dc.blur, f, df_dt, nthreads_);
|
|
}
|
|
|
|
// 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.
|
|
bool RigidBodyTarget::Residuals(const double q[6], double *residuals) {
|
|
++evaluations;
|
|
Place(q, model_);
|
|
std::vector<std::complex<double>> fc;
|
|
std::vector<std::array<std::complex<double>, 3>> dfc_dt;
|
|
CopyFcalc(model_, fc, &dfc_dt);
|
|
|
|
const std::vector<gemmi::Miller> &hkl = composition_->Hkl();
|
|
std::vector<std::complex<float>> fmask;
|
|
if (fixed != nullptr) {
|
|
for (size_t m = 0; m < fc.size(); m++)
|
|
fc[m] += fixed_part_.fcalc[m];
|
|
fmask = fixed_part_.fmask;
|
|
} else if (hold_mask && have_point_) {
|
|
fmask = fmask_;
|
|
} else {
|
|
// The mask is symmetric as gridded, so it is read at h directly.
|
|
gemmi::Grid<float> mask_grid = RigidBodyZoneGrid(cell_, sg_, d_min_);
|
|
PutMaskOnGrid(mask_grid, model_, orbit_leaders_, nthreads_);
|
|
const gemmi::FPhiGrid<float> fm = MapToFPhi(mask_grid);
|
|
for (const gemmi::Miller &h : hkl)
|
|
fmask.push_back(fm.get_value_by_hkl(h));
|
|
}
|
|
if (hkl.empty())
|
|
return false;
|
|
gemmi::AsuData<std::complex<float>> fcalc, fmask_data;
|
|
for (size_t m = 0; m < hkl.size(); m++) {
|
|
fcalc.v.push_back({hkl[m], std::complex<float>(fc[m])});
|
|
fmask_data.v.push_back({hkl[m], fmask[m]});
|
|
}
|
|
|
|
// 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_data);
|
|
if (scaling.points.empty())
|
|
return false;
|
|
if (!solvent_fitted_) {
|
|
FitModelScale(scaling, {.b_hi = solvent_b_max_}, 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_data);
|
|
|
|
// 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 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.
|
|
const std::vector<int> &row = composition_->Row();
|
|
unmatched = 0;
|
|
for (size_t i = 0; i < fobs_.v.size(); ++i) {
|
|
if (row[i] < 0)
|
|
++unmatched;
|
|
residuals[i] = row[i] < 0 ? 0.0 : (fobs_.v[i].value.value - std::abs(fcalc.v[row[i]].value)) / f_mean_;
|
|
}
|
|
|
|
have_point_ = true;
|
|
std::copy(q, q + 6, q_.begin());
|
|
fc_ = std::move(fc);
|
|
dfc_dt_ = std::move(dfc_dt);
|
|
fmask_ = std::move(fmask);
|
|
k_overall_ = scaling.k_overall;
|
|
b_star_ = scaling.b_star;
|
|
return true;
|
|
}
|
|
|
|
// The Jacobian at the last evaluation's placement, without repeating it. The derivative of Fcalc
|
|
// with respect to the translation came with that evaluation, exact. The rotation's is a forward
|
|
// difference, one step per axis, each needing only one copy of the model gridded and transformed:
|
|
// the three run in parallel, column j on its own copy of the model, on threads of their own rather
|
|
// than on ParallelFor's pool, for the reason the null's replicates do (ModelValidation.cpp) - a
|
|
// gridding spreads over the pool, which it cannot do from a pool worker, where a parallel pass runs
|
|
// inline. The bulk-solvent mask is held at the evaluation's: the residuals see it move with the
|
|
// body, the Jacobian does not, which makes the steps Levenberg-Marquardt proposes slightly different
|
|
// but not the cost it accepts them on.
|
|
bool RigidBodyTarget::Jacobian(const double q[6], double *jacobian) {
|
|
if (!have_point_ || !std::equal(q, q + 6, q_.begin())) {
|
|
std::vector<double> residuals(NumObservations());
|
|
if (!Residuals(q, residuals.data()))
|
|
return false;
|
|
}
|
|
++jacobians;
|
|
const double step = RigidBodyJacobianStep(d_min_);
|
|
std::array<std::vector<std::complex<double>>, 3> shifted;
|
|
const std::launch policy = nthreads_ > 1 ? std::launch::async : std::launch::deferred;
|
|
std::vector<std::future<void>> running;
|
|
for (int j = 0; j < 3; j++)
|
|
running.push_back(std::async(policy, [&, j] {
|
|
double qj[6];
|
|
std::copy(q_.begin(), q_.end(), qj);
|
|
qj[j] += step;
|
|
Place(qj, column_models_[j]);
|
|
CopyFcalc(column_models_[j], shifted[j], nullptr);
|
|
if (fixed != nullptr)
|
|
for (size_t m = 0; m < shifted[j].size(); m++)
|
|
shifted[j][m] += fixed_part_.fcalc[m];
|
|
}));
|
|
for (std::future<void> &f : running)
|
|
f.get();
|
|
|
|
// d|F_scaled|/dq at the evaluation's scale: F_scaled = K(h) (Fcalc + k_sol exp(-b_sol s^2) Fmask),
|
|
// so the derivative of its amplitude is K Re(conj(F_total) dFcalc/dq) / |F_total|. The scale's
|
|
// own derivatives are gemmi's, for the parameters it re-fits: k_overall and the constrained B*.
|
|
gemmi::Scaling<float> scaling(cell_, &sg_);
|
|
scaling.use_solvent = true;
|
|
scaling.fix_k_sol = true;
|
|
scaling.fix_b_sol = true;
|
|
scaling.k_sol = k_sol;
|
|
scaling.b_sol = b_sol;
|
|
scaling.k_overall = k_overall_;
|
|
scaling.b_star = b_star_;
|
|
const size_t n = NumObservations();
|
|
const size_t p = 1 + scaling.constraint_matrix.size();
|
|
std::vector<double> jk(n * p, 0.0);
|
|
std::vector<double> dy_da(p);
|
|
const std::vector<int> &row = composition_->Row();
|
|
const std::vector<gemmi::Miller> &hkl = composition_->Hkl();
|
|
std::fill(jacobian, jacobian + n * 6, 0.0);
|
|
for (size_t i = 0; i < n; i++) {
|
|
const int m = row[i];
|
|
if (m < 0)
|
|
continue;
|
|
const gemmi::Scaling<float>::Point point{hkl[m], cell_.calculate_stol_sq(hkl[m]),
|
|
std::complex<float>(fc_[m]), fmask_[m],
|
|
fobs_.v[i].value.value, fobs_.v[i].value.sigma};
|
|
scaling.compute_value_and_derivatives(point, dy_da);
|
|
for (size_t a = 0; a < p; a++)
|
|
jk[i * p + a] = -dy_da[a] / f_mean_;
|
|
const std::complex<double> f_total(scaling.get_fcalc(point));
|
|
const double k = scaling.get_overall_scale_factor(hkl[m]) / (std::abs(f_total) * f_mean_);
|
|
for (int j = 0; j < 6; j++) {
|
|
const std::complex<double> dfc = j < 3 ? (shifted[j][m] - fc_[m]) / step : dfc_dt_[m][j - 3];
|
|
jacobian[i * 6 + j] = -k * std::real(std::conj(f_total) * dfc);
|
|
}
|
|
}
|
|
|
|
// The residuals are taken at the scale re-fitted to each placement, so their Jacobian is the
|
|
// fixed-scale one with the part a scale change absorbs projected out: J - J_k (J_k^T J_k)^-1 J_k^T J,
|
|
// J_k being the residuals' derivatives with respect to the scale parameters. Without it the steps
|
|
// would be those of a target whose scale stays put, overstating the curvature along any direction
|
|
// a scale change could follow.
|
|
// Following Golub & Pereyra (1973) SIAM J. Numer. Anal. 10, 413-432, in the form of Kaufman (1975) BIT 15, 49-57
|
|
Eigen::MatrixXd jtj = Eigen::MatrixXd::Zero(p, p), jtq = Eigen::MatrixXd::Zero(p, 6);
|
|
for (size_t i = 0; i < n; i++)
|
|
for (size_t a = 0; a < p; a++) {
|
|
for (size_t b = 0; b < p; b++)
|
|
jtj(a, b) += jk[i * p + a] * jk[i * p + b];
|
|
for (int j = 0; j < 6; j++)
|
|
jtq(a, j) += jk[i * p + a] * jacobian[i * 6 + j];
|
|
}
|
|
const Eigen::MatrixXd x = jtj.ldlt().solve(jtq);
|
|
for (size_t i = 0; i < n; i++)
|
|
for (int j = 0; j < 6; j++)
|
|
for (size_t a = 0; a < p; a++)
|
|
jacobian[i * 6 + j] -= jk[i * p + a] * x(a, j);
|
|
return true;
|
|
}
|
|
|
|
// 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.
|
|
double RigidBodyReachDeg(const gemmi::Model &model, double d_min) {
|
|
const std::vector<gemmi::Position> pos = ModelPositions(model);
|
|
if (pos.empty())
|
|
return 0.0;
|
|
gemmi::Position centre;
|
|
for (const gemmi::Position &p : pos)
|
|
centre += p;
|
|
centre *= 1.0 / static_cast<double>(pos.size());
|
|
double r2 = 0;
|
|
for (const gemmi::Position &p : pos)
|
|
r2 += centre.dist_sq(p);
|
|
const double rms_radius = std::sqrt(r2 / static_cast<double>(pos.size()));
|
|
// The first zone RefineRigidBody walks, which is d_min itself where d_min is coarser than all
|
|
// of the ladder.
|
|
const double first_zone = std::max(LADDER[0], d_min);
|
|
return first_zone / rms_radius * 180.0 / PI;
|
|
}
|
|
|
|
//
|
|
// 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,
|
|
RigidBodyGPUPool *gpu,
|
|
const gemmi::Model *fixed) {
|
|
const auto t0 = std::chrono::steady_clock::now();
|
|
RigidBodyRefineResult result;
|
|
if (fobs.v.empty())
|
|
return result;
|
|
std::unique_ptr<RigidBodyTargetBase> target_backend;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (gpu != nullptr)
|
|
target_backend = std::make_unique<RigidBodyTargetGPU>(*gpu, model, cell, sg, nthreads);
|
|
#endif
|
|
if (!target_backend)
|
|
target_backend = std::make_unique<RigidBodyTarget>(model, cell, sg, nthreads);
|
|
RigidBodyTargetBase &target = *target_backend;
|
|
if (!target.Usable())
|
|
return result;
|
|
target.fixed = fixed;
|
|
|
|
const std::vector<double> ladder = RigidBodyLadder(d_min);
|
|
// With fixed atoms beside the body the origin is theirs, not free.
|
|
const gemmi::Mat33 gauge = fixed != nullptr ? gemmi::Mat33(0, 0, 0, 0, 0, 0, 0, 0, 0)
|
|
: RigidBodyGaugeProjector(sg, cell);
|
|
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
|
|
target.SetZone(zone_obs, zone);
|
|
|
|
ceres::Problem problem;
|
|
problem.AddResidualBlock(new RigidBodyCost(target), 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 = target.evaluations, jacobians_before = target.jacobians;
|
|
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, {} Jacobians, "
|
|
"{:.2f} s, rotation {:.3f} deg, translation {:.3f} A", zone, zone_obs.v.size(),
|
|
target.unmatched, target.k_sol, target.b_sol,
|
|
summary.iterations.empty() ? 0 : summary.iterations.size() - 1,
|
|
target.evaluations - evaluations_before, target.jacobians - jacobians_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]) / target.RmsRadius() * 180.0 / PI,
|
|
std::sqrt(q[3]*q[3] + q[4]*q[4] + q[5]*q[5]));
|
|
}
|
|
|
|
target.Place(q, model); // Ceres left the model at its last trial; put it at the answer
|
|
result.evaluations = target.evaluations;
|
|
result.jacobians = target.jacobians;
|
|
result.converged = any_zone_solved;
|
|
result.k_sol = target.k_sol;
|
|
result.b_sol = target.b_sol;
|
|
const double aa = std::sqrt(q[0] * q[0] + q[1] * q[1] + q[2] * q[2]) / target.RmsRadius();
|
|
result.angle_deg = aa * 180.0 / PI;
|
|
const double axis_angle[3] = {q[0] / target.RmsRadius(), q[1] / target.RmsRadius(), q[2] / target.RmsRadius()};
|
|
double column[3][3];
|
|
for (int j = 0; j < 3; j++) {
|
|
const double e[3] = {j == 0 ? 1.0 : 0.0, j == 1 ? 1.0 : 0.0, j == 2 ? 1.0 : 0.0};
|
|
ceres::AngleAxisRotatePoint(axis_angle, e, column[j]);
|
|
}
|
|
result.rotation = gemmi::Mat33(column[0][0], column[1][0], column[2][0],
|
|
column[0][1], column[1][1], column[2][1],
|
|
column[0][2], column[1][2], column[2][2]);
|
|
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 and {} Jacobians in {:.3f} s: rotation {:.3f} deg, translation {:.3f} A{}",
|
|
result.zones.size(), result.zones.back(), result.evaluations, result.jacobians,
|
|
result.seconds, result.angle_deg, result.shift_A, gpu != nullptr ? " (GPU)" : "");
|
|
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;
|
|
}
|