leonarski_fandClaude Opus 5.5
9ad92b6bfe
A pure move. ModelValidation, RigidBodyRefine, RigidBodyGPU, ModelFFT, ModelGrid, ModelScaling, ModelMaskGPU, ModelScaleGPU and SigmaA - everything that works on an atomic model - become the JFJochStructureRefinement library, linked by JFJochImageAnalysis. WriteModel (the placed-model mmCIF/PDB writer) goes to writer/ as its own small JFJochModelWriter target, so JFJochWriter, which a writer-only build compiles, does not gain a gemmi dependency. Only include paths and CMake lists change. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
588 lines
29 KiB
C++
588 lines
29 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) : 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);
|
|
}
|
|
|
|
// 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 (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, {}, 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);
|
|
}));
|
|
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 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;
|
|
|
|
const std::vector<double> ladder = RigidBodyLadder(d_min);
|
|
const gemmi::Mat33 gauge = 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;
|
|
}
|