diff --git a/docs/ACKNOWLEDGEMENT.md b/docs/ACKNOWLEDGEMENT.md index 84a72f9b5..4f4dc67ff 100644 --- a/docs/ACKNOWLEDGEMENT.md +++ b/docs/ACKNOWLEDGEMENT.md @@ -341,6 +341,15 @@ random about its own centroid, and the rotations are drawn uniformly from SO(3) random unit quaternion. K. Shoemake, "Uniform Random Rotations", in *Graphics Gems III*, ed. D. Kirk, Academic Press (1992), 124-132 (no DOI). +**Variable projection** — the rigid-body placement of a supplied model re-fits the overall scale at +every step, and its Jacobian folds that re-fit in by projecting the scale's own derivatives out of the +placement's, in Kaufman's simplified form of Golub and Pereyra's derivative of the reduced problem. +G. H. Golub and V. Pereyra, "The Differentiation of Pseudo-Inverses and Nonlinear Least Squares +Problems Whose Variables Separate" (1973), SIAM J. Numer. Anal. 10, 413-432 +[doi:10.1137/0710036](https://doi.org/10.1137/0710036). L. Kaufman, "A variable projection method +for solving separable nonlinear least squares problems" (1975), BIT 15, 49-57 +[doi:10.1007/BF01932995](https://doi.org/10.1007/BF01932995). + ## Software and computing methods Decoding bitshuffle+LZ4 images on the GPU, rather than decompressing them on the host and uploading diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 87f46fe3e..b8624e08d 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -10,6 +10,7 @@ * `rugnux --model` draws its null's random placements away from every orientation equivalent to the model's under the space group or a twin law, so a random placement can no longer be refined onto the model's own solution and make the crystal's own model read as not fitting. * Rugnux adds beam-stop holder arms that let part of the beam through to the beam-stop mask, without changing the mask of a sweep that has none. * Rugnux's `--model` rigid-body refinement computes the model's density and solvent mask on all threads, with results unchanged. +* Rugnux's `--model` rigid-body refinement computes the model's structure factors from one copy of the model and derives the translation part of its Jacobian exactly, so it runs faster; placements can differ slightly. ### 1.0.0-rc.173 diff --git a/docs/CPU_DATA_ANALYSIS.md b/docs/CPU_DATA_ANALYSIS.md index e1ccfabe0..147c15f8d 100644 --- a/docs/CPU_DATA_ANALYSIS.md +++ b/docs/CPU_DATA_ANALYSIS.md @@ -61,6 +61,7 @@ The methods draw on, and in places reimplement, solutions from: - A. Fokine & A. Urzhumtsev, "Flat bulk-solvent model: obtaining optimal parameters", *Acta Cryst.* **D58** (2002), 1387-1392 (the flat bulk-solvent model, its optimal parameters and the range they are physically meaningful over, used when scaling a model to the data in §14). - P. V. Afonine, R. W. Grosse-Kunstleve & P. D. Adams, "A robust bulk-solvent correction and anisotropic scaling procedure", *Acta Cryst.* **D61** (2005), 850-855 (the grid search over that range that fits k_sol and b_sol, with the overall scale and anisotropic B refitted at each grid point). - K. Shoemake, "Uniform Random Rotations", in *Graphics Gems III*, ed. D. Kirk, Academic Press (1992), 124-132 (the uniform random rotations the model-fit null of §14.5 is built from). +- G. H. Golub & V. Pereyra, "The differentiation of pseudo-inverses and nonlinear least squares problems whose variables separate", *SIAM J. Numer. Anal.* **10** (1973), 413-432, and L. Kaufman, "A variable projection method for solving separable nonlinear least squares problems", *BIT* **15** (1975), 49-57 (the scale re-fit folded into the rigid-body Jacobian of §14.8). - Z. Otwinowski & W. Minor, "Processing of X-ray diffraction data collected in oscillation mode", *Methods Enzymol.* **276** (1997), 307-326 (reweighted, de-biased profile-fit variances). - G. Winter et al., "DIALS: implementation and evaluation of a new integration package", *Acta Cryst.* **D74** (2018), 85-97, and J. Beilsten-Edmands et al., *Acta Cryst.* **D76** (2020), 385-399 (CC1/2 resolution cutoff, merge outlier rejection, scaling error model). - R. H. Blessing, "An empirical correction for absorption anisotropy", *Acta Cryst.* **A51** (1995), 33-38 (absorption as spherical harmonics of the beam directions). diff --git a/docs/CPU_DATA_ANALYSIS_DECISIONS.md b/docs/CPU_DATA_ANALYSIS_DECISIONS.md index 6fe63891d..dbd77f5a1 100644 --- a/docs/CPU_DATA_ANALYSIS_DECISIONS.md +++ b/docs/CPU_DATA_ANALYSIS_DECISIONS.md @@ -244,7 +244,7 @@ Re-fractionalizing a model into the data cell puts it in the right box but not n It is **one** rigid body. A fragment-screening model arrives already solved and isomorphous, and what is being recovered is the crystal's movement, not the molecule's; splitting it into domains, or giving a bound ligand six parameters of its own, would refine against evidence these data do not separately carry — and the ligand is what the difference map is there to show, not to model away. -The refinement walks a coarse-to-fine ladder, 6 Å → 4.5 Å → 3.5 Å, each zone starting from the previous one's answer. It stops at 3.5 Å because that is where rigid-body refinement is conventionally run and because the movement being recovered is a few tenths of an ångström, a tenth of that resolution; a finer zone costs $(1/d)^3$ in grid points and reflections for a placement it cannot meaningfully sharpen. Each evaluation recomputes $F_\mathrm{calc}$ and the solvent mask for the moved model **and re-fits the scale of §14.2** — otherwise the target would measure the scale as much as the placement, and the body would translate to repair a scale error instead of moving to where the density is. The minimiser is Levenberg–Marquardt on the amplitude residuals with a forward-difference Jacobian (seven evaluations per Jacobian; GEMMI's structure-factor path carries no analytic $\partial F/\partial p$, and at six parameters it does not need one), the step being a fixed fraction of the zone's resolution. +The refinement walks a coarse-to-fine ladder, 6 Å → 4.5 Å → 3.5 Å, each zone starting from the previous one's answer. It stops at 3.5 Å because that is where rigid-body refinement is conventionally run and because the movement being recovered is a few tenths of an ångström, a tenth of that resolution; a finer zone costs $(1/d)^3$ in grid points and reflections for a placement it cannot meaningfully sharpen. Each evaluation recomputes $F_\mathrm{calc}$ and the solvent mask for the moved model **and re-fits the scale of §14.2** — otherwise the target would measure the scale as much as the placement, and the body would translate to repair a scale error instead of moving to where the density is. The minimiser is Levenberg–Marquardt on the amplitude residuals. $F_\mathrm{calc}$ is computed from **one copy** of the model, gridded without symmetrization and transformed once, and composed over the space group's operators, $F(\mathbf h) = n_\mathrm{cen}\sum_{(R,\mathbf t)} e^{2\pi i\,\mathbf h\cdot\mathbf t}\,F_1(\mathbf hR)$ — the transform of the symmetrized map, rearranged. That makes the derivative with respect to the translation exact and free ($F_1(\mathbf k)$ only picks up the phase $e^{2\pi i\,\mathbf s_\mathbf k\cdot\Delta}$). The rotation's three columns are forward differences, the step a fixed fraction of the zone's resolution, each needing only one copy gridded and transformed. The Jacobian holds the scale and the bulk-solvent mask at the evaluation's; the residuals keep both exact, and the scale re-fit is folded into the Jacobian by variable projection (Golub & Pereyra 1973, in Kaufman's 1975 form), so a step is judged on the target the evaluations actually compute. The refinement sees **only the working reflections**. The step is then committed only if it lowers **R-free**, computed at full resolution on the free set it never saw; otherwise the model is put back exactly where it was read and the maps are the ones it would have given. The whole addition costs about two seconds. Note that the origin is a gauge in some space groups — free in all three directions in $P1$, and along the unique axis in a polar group — so those components of the translation are undetermined; nothing is done about that beyond the Levenberg–Marquardt damping and the R-free gate, which between them make an undetermined direction harmless rather than unstable. diff --git a/rugnux/RigidBodyRefine.cpp b/rugnux/RigidBodyRefine.cpp index d0128686e..165c22f90 100644 --- a/rugnux/RigidBodyRefine.cpp +++ b/rugnux/RigidBodyRefine.cpp @@ -11,6 +11,7 @@ #include #include +#include #include #include @@ -22,6 +23,7 @@ #include "ModelGrid.h" // PutModelDensityOnGrid, PutMaskOnGrid #include "ModelScaling.h" // FitModelScale #include "../common/JFJochMath.h" // PI +#include "../common/ParallelFor.h" #include "../common/Logger.h" namespace { @@ -37,38 +39,13 @@ using Table = gemmi::IT92; // (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 +// 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; -// 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 &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; - } - } -}; - // 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; @@ -81,201 +58,34 @@ gemmi::Grid ZoneGrid(const gemmi::UnitCell &cell, const gemmi::SpaceGroup return grid; } -// 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 &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. `orbit_leaders` are those - // of the zone's grid, and must outlive the zone. - void SetZone(const gemmi::AsuData> &fobs, double d_min, - const std::vector &orbit_leaders) { - fobs_ = fobs; - d_min_ = d_min; - orbit_leaders_ = &orbit_leaders; - double sum = 0; - for (const auto &hv : fobs_.v) - sum += hv.value.value; - f_mean_ = fobs_.v.empty() ? 1.0 : sum / static_cast(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 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_, *orbit_leaders_, nthreads_); - gemmi::AsuData> fcalc = - MapToFPhi(dc.grid).prepare_asu_data(dc.d_min, dc.blur, false, false, false); - - gemmi::Grid mask_grid = ZoneGrid(cell_, sg_, d_min_); - PutMaskOnGrid(mask_grid, model_, *orbit_leaders_, nthreads_); - gemmi::AsuData> 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 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 &base_; - Placement placement_; - gemmi::AsuData> fobs_; - double d_min_ = 0; - const std::vector *orbit_leaders_ = nullptr; - 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. -// -// The columns run on threads of their own rather than on ParallelFor's pool, for the reason the null's -// replicates do (ModelValidation.cpp): an evaluation spreads its gridding over the pool, which it -// cannot do from a pool worker, where a parallel pass runs inline. On the pool, the six columns would -// each grid on one core and leave the rest of the machine idle for most of the fit. A run given one -// thread evaluates them one after the other, on its own. +// every zone (the placement begins at no shift), so the Jacobian is supplied by RigidBodyTarget. class RigidBodyCost : public ceres::CostFunction { public: - RigidBodyCost(Evaluator &ev, std::vector &columns, size_t nthreads) - : ev_(ev), columns_(columns), nthreads_(nthreads) { - set_num_residuals(static_cast(ev.NumObservations())); + explicit RigidBodyCost(RigidBodyTarget &target) : target_(target) { + set_num_residuals(static_cast(target.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(); // 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 (!ev_.Residuals(parameters[0], residuals)) + if (!target_.Residuals(parameters[0], residuals)) return false; std::copy(parameters[0], parameters[0] + 6, last_q_.begin()); - last_residuals_.assign(residuals, residuals + n); - } - if (jacobians != nullptr && jacobians[0] != nullptr) { - const double step = ev_.JacobianStep(); - std::vector> shifted(6, std::vector(n)); - std::array ok{}; - const std::launch policy = nthreads_ > 1 ? std::launch::async : std::launch::deferred; - std::vector> running; - for (int j = 0; j < 6; j++) - running.push_back(std::async(policy, [&, 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()); - })); - for (std::future &f : running) - f.get(); - // 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; - } + last_residuals_.assign(residuals, residuals + num_residuals()); } + if (jacobians != nullptr && jacobians[0] != nullptr) + return target_.Jacobian(parameters[0], jacobians[0]); return true; } private: - Evaluator &ev_; - std::vector &columns_; - size_t nthreads_; + RigidBodyTarget &target_; mutable std::array last_q_{}; mutable std::vector last_residuals_; // at last_q_; empty until an evaluation succeeds }; @@ -321,6 +131,315 @@ void SetModelPositions(gemmi::Model &model, const std::vector & a.pos = pos[i++]; } +SymmetryComposition::SymmetryComposition(const gemmi::Grid &grid, double d_min, + const std::vector &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(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(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(grid.nu) * (gemmi::modulo(kv, grid.nv) + static_cast(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])); + terms_.push_back(t); + } + } +} + +void SymmetryComposition::Compose(const gemmi::FPhiGrid &f1, double unblur, + std::vector> &f, + std::vector, 3>> *df_dt, + size_t nthreads) const { + const int n = static_cast(hkl_.size()); + f.assign(n, 0.0); + if (df_dt != nullptr) + df_dt->assign(n, {}); + const std::complex two_pi_i(0.0, 2 * PI); + ParallelChunks(n, nthreads, [&](int lo, int hi) { + for (int m = lo; m < hi; m++) { + std::complex sum = 0; + std::array, 3> d{}; + for (size_t o = 0; o < ops_; o++) { + const Term &t = terms_[m * ops_ + o]; + const std::complex value = f1.data[t.index]; + const std::complex 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]; + } + }); +} + +RigidBodyTarget::RigidBodyTarget(gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, + size_t nthreads) + : model_(model), cell_(cell), sg_(sg), nthreads_(nthreads), base_(ModelPositions(model)), + column_models_(3, model) { + if (base_.empty()) + return; + for (const gemmi::Position &p : base_) + centre_ += p; + centre_ *= 1.0 / static_cast(base_.size()); + double r2 = 0; + for (const gemmi::Position &p : base_) + r2 += centre_.dist_sq(p); + rms_radius_ = std::sqrt(r2 / static_cast(base_.size())); +} + +// 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 RigidBodyTarget::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> &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(fobs_.v.size()); + solvent_fitted_ = false; + have_point_ = false; + const gemmi::Grid grid = ZoneGrid(cell_, sg_, d_min); + orbit_leaders_ = OrbitLeaders(grid, nthreads_); + std::vector 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> &f, + std::vector, 3>> *df_dt) const { + gemmi::DensityCalculator 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> fc; + std::vector, 3>> dfc_dt; + CopyFcalc(model_, fc, &dfc_dt); + + const std::vector &hkl = composition_->Hkl(); + std::vector> fmask; + if (hold_mask && have_point_) { + fmask = fmask_; + } else { + // The mask is symmetric as gridded, so it is read at h directly. + gemmi::Grid mask_grid = ZoneGrid(cell_, sg_, d_min_); + PutMaskOnGrid(mask_grid, model_, orbit_leaders_, nthreads_); + const gemmi::FPhiGrid 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> fcalc, fmask_data; + for (size_t m = 0; m < hkl.size(); m++) { + fcalc.v.push_back({hkl[m], std::complex(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 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 &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 residuals(NumObservations()); + if (!Residuals(q, residuals.data())) + return false; + } + ++jacobians; + const double step = JACOBIAN_STEP_FRACTION * d_min_; + std::array>, 3> shifted; + const std::launch policy = nthreads_ > 1 ? std::launch::async : std::launch::deferred; + std::vector> 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 &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 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 jk(n * p, 0.0); + std::vector dy_da(p); + const std::vector &row = composition_->Row(); + const std::vector &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::Point point{hkl[m], cell_.calculate_stol_sq(hkl[m]), + std::complex(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 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 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 @@ -359,20 +478,10 @@ RigidBodyRefineResult RefineRigidBody(gemmi::Model &model, size_t nthreads) { const auto t0 = std::chrono::steady_clock::now(); RigidBodyRefineResult result; - - const std::vector base = ModelPositions(model); - if (base.empty() || fobs.v.empty()) + if (fobs.v.empty()) return result; - - Placement placement; - for (const gemmi::Position &p : base) - placement.centre += p; - placement.centre *= 1.0 / static_cast(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(base.size())); - if (!(placement.rms_radius > 0)) + RigidBodyTarget target(model, cell, sg, nthreads); + if (!target.Usable()) return result; std::vector ladder; @@ -383,14 +492,6 @@ RigidBodyRefineResult RefineRigidBody(gemmi::Model &model, 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 column_models(6, model); - std::vector columns; - columns.reserve(6); - for (int j = 0; j < 6; j++) - columns.emplace_back(column_models[j], cell, sg, base, placement, nthreads); double q[6] = {0, 0, 0, 0, 0, 0}; bool any_zone_solved = false; for (double zone : ladder) { @@ -403,13 +504,10 @@ RigidBodyRefineResult RefineRigidBody(gemmi::Model &model, if (zone_obs.v.size() < 50) continue; result.zones.push_back(zone); // the ladder WALKED, which a thin zone drops out of - const std::vector orbit_leaders = OrbitLeaders(ZoneGrid(cell, sg, zone), nthreads); - ev.SetZone(zone_obs, zone, orbit_leaders); - for (Evaluator &c : columns) - c.SetZone(zone_obs, zone, orbit_leaders); + target.SetZone(zone_obs, zone); ceres::Problem problem; - problem.AddResidualBlock(new RigidBodyCost(ev, columns, nthreads), nullptr, q); + problem.AddResidualBlock(new RigidBodyCost(target), nullptr, q); ceres::Solver::Options options; options.linear_solver_type = ceres::DENSE_QR; options.max_num_iterations = 15; @@ -417,38 +515,39 @@ RigidBodyRefineResult RefineRigidBody(gemmi::Model &model, options.parameter_tolerance = 1e-4; options.logging_type = ceres::LoggingType::SILENT; ceres::Solver::Summary summary; - const int evaluations_before = ev.evaluations; + 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, {:.2f} s, " - "rotation {:.3f} deg, translation {:.3f} A", zone, zone_obs.v.size(), ev.unmatched, - ev.k_sol, ev.b_sol, + "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, - ev.evaluations - evaluations_before, + target.evaluations - evaluations_before, target.jacobians - jacobians_before, std::chrono::duration(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[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])); } - placement.Apply(q, base, model); // Ceres left the model at a Jacobian probe; put it at the answer - result.evaluations = ev.evaluations; + 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 = 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.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; result.shift_A = std::sqrt(q[3] * q[3] + q[4] * q[4] + q[5] * q[5]); result.seconds = std::chrono::duration(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); + "{} evaluations and {} Jacobians in {:.2f} 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); else logger.Info("Model validation: rigid body had no resolution zone with enough reflections to " "run in; the model is left where it arrived"); diff --git a/rugnux/RigidBodyRefine.h b/rugnux/RigidBodyRefine.h index 6c15a3b3d..313a50cca 100644 --- a/rugnux/RigidBodyRefine.h +++ b/rugnux/RigidBodyRefine.h @@ -3,11 +3,16 @@ #pragma once +#include +#include #include +#include #include #include "gemmi/asudata.hpp" +#include "gemmi/grid.hpp" #include "gemmi/model.hpp" +#include "gemmi/recgrid.hpp" #include "gemmi/symmetry.hpp" #include "gemmi/unitcell.hpp" @@ -19,6 +24,7 @@ struct RigidBodyRefineResult { double shift_A = 0.0; // magnitude of the translation std::vector zones; // the resolution ladder actually walked, coarsest first int evaluations = 0; // structure-factor evaluations (each one re-fits the scale) + int jacobians = 0; // Jacobians, each three single-copy transforms (not counted above) double seconds = 0.0; double k_sol = 0.0, b_sol = 0.0; // the bulk solvent the last zone's target was scored through }; @@ -50,3 +56,100 @@ RigidBodyRefineResult RefineRigidBody(gemmi::Model &model, double d_min, Logger &logger, size_t nthreads = 1); + +// --- The pieces RefineRigidBody is built from, exposed for the tests --- + +// Fcalc of the whole crystal composed from the transform of ONE copy of its content, gridded on the +// zone's grid without symmetrization: F(h) = n_cen * sum over the group's operators x -> Rx + t of +// exp(+2 pi i h.t) F1(hR). gemmi's symmetrize() maps grid points onto grid points, so this is the +// transform of the symmetrized grid, rearranged - the same number up to rounding - without the +// symmetrization. Moving the copy by a Cartesian t multiplies F1(k) by exp(2 pi i s_k.t), so the +// derivative of F with respect to a translation of the body comes with it, exactly. +class SymmetryComposition { +public: + // `grid`: the zone's grid (size, cell and group). `hkl`: the zone's observed indices, sorted. Only + // those prepare_asu_data(d_min) would give a value on this grid are composed - in the group's + // reciprocal ASU, strictly inside d_min, inside the grid's Nyquist box, not systematically absent + // and not 000 - so an observation it does not match stays unmatched, as it was. + SymmetryComposition(const gemmi::Grid &grid, double d_min, const std::vector &hkl); + + // For each of `hkl`, its position among the composed indices, or -1. + const std::vector &Row() const { return row_; } + const std::vector &Hkl() const { return hkl_; } + + // F at the composed indices from `f1` = MapToFPhi() of the unsymmetrized grid, times the unblur + // factor prepare_asu_data() applies; with `df_dt`, also dF/dt for a Cartesian translation t. + void Compose(const gemmi::FPhiGrid &f1, double unblur, std::vector> &f, + std::vector, 3>> *df_dt, size_t nthreads) const; + +private: + struct Term { + size_t index; // of F1(hR) in the half-l grid, or of its Friedel mate + bool conj; // read as the Friedel mate: conj F1(-hR) + std::complex phase; // exp(+2 pi i h.t) + gemmi::Vec3 s; // hR as a Cartesian reciprocal vector + }; + std::vector row_; + std::vector hkl_; + std::vector inv_d2_; + size_t ops_ = 1; // terms per index + double centring_ = 1; // n_cen: every centring operator adds the same term + std::vector terms_; +}; + +// The target RefineRigidBody minimises over one resolution zone: the amplitude residuals of the model +// placed at q (the angle-axis rotation about the centroid times the rms radius, then the translation, +// all in angstroms), with the scale re-fitted at every evaluation. +class RigidBodyTarget { +public: + // q = 0 is the placement `model` has now. Every evaluation moves `model`. + RigidBodyTarget(gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, + size_t nthreads); + + bool Usable() const { return !base_.empty() && rms_radius_ > 0; } + double RmsRadius() const { return rms_radius_; } + void Place(const double q[6], gemmi::Model &model) const; + + void SetZone(const gemmi::AsuData> &fobs, double d_min); + size_t NumObservations() const { return fobs_.v.size(); } + + bool Residuals(const double q[6], double *residuals); + // d residuals / d q at q, NumObservations() x 6 row-major, with the scale re-fit folded in by + // projection and the bulk-solvent mask held. + bool Jacobian(const double q[6], double *jacobian); + + bool hold_mask = false; // for the tests: keep the last evaluation's bulk-solvent mask + int evaluations = 0; + int jacobians = 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 + +private: + void CopyFcalc(const gemmi::Model &model, std::vector> &f, + std::vector, 3>> *df_dt) const; + + gemmi::Model &model_; + const gemmi::UnitCell &cell_; + const gemmi::SpaceGroup &sg_; + size_t nthreads_; + std::vector base_; + gemmi::Position centre_; + double rms_radius_ = 0; + std::vector column_models_; // the rotation columns' own copies + + gemmi::AsuData> fobs_; + double d_min_ = 0; + double f_mean_ = 1; + std::vector orbit_leaders_; + std::optional composition_; // the zone's + bool solvent_fitted_ = false; + + // The last evaluation, which the Jacobian is taken at. + bool have_point_ = false; + std::array q_{}; + std::vector> fc_; + std::vector, 3>> dfc_dt_; + std::vector> fmask_; + double k_overall_ = 1; + gemmi::SMat33 b_star_{0, 0, 0, 0, 0, 0}; +}; diff --git a/tests/ModelValidationTest.cpp b/tests/ModelValidationTest.cpp index 60703b4b9..e34386a1d 100644 --- a/tests/ModelValidationTest.cpp +++ b/tests/ModelValidationTest.cpp @@ -880,3 +880,263 @@ TEST_CASE("ModelValidation_ParallelGriddingMatchesGemmi", "[ModelValidation]") { } } } + +namespace { + // The five groups of ModelValidation_ParallelGriddingMatchesGemmi: none, a centring, screws, a + // cubic body centring and a cubic face centring with a diamond glide. + const char *kRigidBodyCrysts[] = { + "CRYST1 40.000 50.000 60.000 90.00 90.00 90.00 P 1 1\n", + "CRYST1 40.000 50.000 60.000 90.00 100.00 90.00 C 1 2 1 4\n", + "CRYST1 40.000 50.000 60.000 90.00 90.00 90.00 P 21 21 21 4\n", + "CRYST1 60.000 60.000 60.000 90.00 90.00 90.00 I 2 3 24\n", + "CRYST1 80.000 80.000 80.000 90.00 90.00 90.00 F 41 3 2 96\n", + }; + + // ClusterPdb() in `cryst`, every third atom anisotropic. + gemmi::Structure AnisoCluster(const char *cryst) { + const auto path = WriteTemp("rigid_body_composition_test.pdb", ClusterPdb(cryst).c_str()); + gemmi::Structure st = gemmi::read_structure_gz(path, gemmi::CoorFormat::Detect); + std::filesystem::remove(path); + st.setup_cell_images(); + int i = 0; + for (gemmi::Chain &ch : st.models[0].chains) + for (gemmi::Residue &r : ch.residues) + for (gemmi::Atom &a : r.atoms) + if (i++ % 3 == 0) + a.aniso = {0.30f, 0.25f, 0.20f, 0.02f, -0.01f, 0.03f}; + return st; + } + + // The zone's density calculator, as the rigid body sets it up for `model`. + gemmi::DensityCalculator, float> ZoneDensity(const gemmi::Structure &st, + const gemmi::Model &model, double d_min) { + gemmi::DensityCalculator, float> dc; + dc.d_min = d_min; + dc.rate = 1.5; + dc.grid.unit_cell = st.cell; + dc.grid.spacegroup = st.find_spacegroup(); + dc.set_refmac_compatible_blur(model); + return dc; + } +} + +// The rigid body's Fcalc is composed from the transform of one copy of the model instead of being +// taken from the symmetrized grid. The two are the same sum rearranged, so they must agree to +// rounding, and must give a value to exactly the reflections prepare_asu_data() does - the +// systematic absences of the group and a reflection sitting exactly on d_min (the (10 0 0) of the +// 60 A cubic cell at 6 A) are left out by both. +TEST_CASE("ModelValidation_RigidBodyP1FcalcMatchesSymmetrized", "[ModelValidation]") { + for (const char *cryst : kRigidBodyCrysts) { + const gemmi::Structure st = AnisoCluster(cryst); + const gemmi::Model &model = st.models[0]; + for (double d_min : {6.0, 3.5}) { + auto sym = ZoneDensity(st, model, d_min); + sym.put_model_density_on_grid(model); + gemmi::FPhiGrid sym_f = MapToFPhi(sym.grid); + const auto ref = sym_f.prepare_asu_data(d_min, sym.blur, false, false, false); + // Every index of the ASU in the grid's box, whatever its resolution, absences and 000 too. + std::vector candidates; + for (const auto &hv : sym_f.prepare_asu_data(0, 0, true, true, false).v) + candidates.push_back(hv.hkl); + + auto copy = ZoneDensity(st, model, d_min); + PutModelDensityOnGrid(copy, model, {}, 4); + const SymmetryComposition composition(copy.grid, d_min, candidates); + std::vector> f; + composition.Compose(MapToFPhi(copy.grid), copy.blur, f, nullptr, 4); + + REQUIRE(composition.Hkl().size() == ref.v.size()); + double mean = 0, worst = 0; + for (size_t i = 0; i < ref.v.size(); i++) { + REQUIRE(composition.Hkl()[i] == ref.v[i].hkl); + mean += std::abs(ref.v[i].value) / static_cast(ref.v.size()); + worst = std::max(worst, std::abs(f[i] - std::complex(ref.v[i].value))); + } + INFO(cryst << " at " << d_min << " A: worst " << worst / mean << " of the mean |F|"); + CHECK(worst <= 1e-4 * mean); + } + } +} + +// The derivative of the composed Fcalc with respect to a translation of the body is a phase factor +// per term, exact - against a central difference of the regridded, recomposed Fcalc. Along a +// direction the origin is free in (all three in P1, b in C2) the amplitude does not change at all. +TEST_CASE("ModelValidation_RigidBodyTranslationDerivativeIsExact", "[ModelValidation]") { + for (const char *cryst : kRigidBodyCrysts) { + const gemmi::Structure st = AnisoCluster(cryst); + const double d_min = 3.5; + const gemmi::DensityCalculator, float> zone = ZoneDensity(st, st.models[0], d_min); + gemmi::Grid grid; + grid.unit_cell = zone.grid.unit_cell; + grid.spacegroup = zone.grid.spacegroup; + grid.set_size_from_spacing(zone.requested_grid_spacing(), gemmi::GridSizeRounding::Up); + std::vector candidates; + const gemmi::ReciprocalAsu asu(grid.spacegroup); + for (int h = -25; h <= 25; h++) + for (int k = -25; k <= 25; k++) + for (int l = -25; l <= 25; l++) + if (asu.is_in({{h, k, l}})) + candidates.push_back({{h, k, l}}); + const SymmetryComposition composition(grid, d_min, candidates); + REQUIRE(composition.Hkl().size() > 200); + + auto fcalc = [&](const gemmi::Vec3 &shift, std::vector, 3>> *df_dt) { + gemmi::Model model = st.models[0]; + for (gemmi::Chain &ch : model.chains) + for (gemmi::Residue &r : ch.residues) + for (gemmi::Atom &a : r.atoms) + a.pos += gemmi::Position(shift); + auto dc = ZoneDensity(st, model, d_min); + PutModelDensityOnGrid(dc, model, {}, 4); + std::vector> f; + composition.Compose(MapToFPhi(dc.grid), dc.blur, f, df_dt, 4); + return f; + }; + std::vector, 3>> df_dt; + const std::vector> f = fcalc({}, &df_dt); + // A five-point difference, whose truncation error goes as eps^4: the grid is float, and at the + // 1e-3 A of a plain central difference its rounding alone reads as 3e-3 of the derivative. What + // is left, about 5e-4, is the sampling - an atom gridded after a move is the moved atom only to + // the accuracy of the grid - which the exact derivative does not have and the difference does. + const double eps = 0.05; + const bool p1 = st.find_spacegroup()->number == 1, c2 = st.find_spacegroup()->number == 5; + for (int k = 0; k < 3; k++) { + gemmi::Vec3 e; + e.at(k) = eps; + const auto p_1 = fcalc(e, nullptr), m_1 = fcalc(-e, nullptr); + const auto p_2 = fcalc(2 * e, nullptr), m_2 = fcalc(-2 * e, nullptr); + double diff = 0, norm = 0, amplitude = 0; + for (size_t m = 0; m < f.size(); m++) { + const std::complex numeric = (8.0 * (p_1[m] - m_1[m]) - (p_2[m] - m_2[m])) / (12 * eps); + diff += std::norm(numeric - df_dt[m][k]); + norm += std::norm(df_dt[m][k]); + amplitude += gemmi::sq(std::real(std::conj(f[m]) * df_dt[m][k]) / std::abs(f[m])); + } + INFO(cryst << " axis " << k << ": relative error " << std::sqrt(diff / norm) + << ", amplitude part " << std::sqrt(amplitude / norm)); + CHECK(std::sqrt(diff / norm) <= 1e-3); + if (p1 || (c2 && k == 1)) + CHECK(std::sqrt(amplitude / norm) <= 1e-5); + } + } +} + +namespace { + // The rigid body's Jacobian at q against a central difference of its own residuals - which + // re-fit the scale at every evaluation - with the bulk-solvent mask held, as the Jacobian holds + // it. Per column (those in `columns`): the cosine between the two and the ratio of their norms. + // Also logged, not checked: how far the difference moves once the mask is let move with the body, + // which is what holding it costs. + void CheckRigidBodyJacobian(const char *cryst, const double q0[6], const std::vector &columns, + double min_cosine, double max_norm_error) { + Logger logger("CheckRigidBodyJacobian"); + const gemmi::Structure st = AnisoCluster(cryst); + const gemmi::SpaceGroup *sg = st.find_spacegroup(); + // "Observed" amplitudes: the model's own, as placed at q = 0. + const auto path = WriteTemp("rigid_body_jacobian_test.pdb", ClusterPdb(cryst).c_str()); + const auto ref = ModelReferenceIntensities(path, {}, {}, 3.0, logger); + std::filesystem::remove(path); + REQUIRE_FALSE(ref.empty()); + + for (double zone : {6.0, 3.5}) { + gemmi::AsuData> fobs; + fobs.unit_cell_ = st.cell; + fobs.spacegroup_ = sg; + for (const auto &r : ref) + if (st.cell.calculate_d({{r.h, r.k, r.l}}) >= zone) + fobs.v.push_back({{{r.h, r.k, r.l}}, {std::sqrt(r.I), 1.0f}}); + fobs.ensure_sorted(); + + gemmi::Model model = st.models[0]; + RigidBodyTarget target(model, st.cell, *sg, 4); + target.SetZone(fobs, zone); + const size_t n = target.NumObservations(); + std::vector r(n), jacobian(n * 6); + REQUIRE(target.Residuals(q0, r.data())); + REQUIRE(target.Jacobian(q0, jacobian.data())); + + auto difference = [&](int j, double h) { + double qp[6], qm[6]; + std::copy(q0, q0 + 6, qp); + std::copy(q0, q0 + 6, qm); + qp[j] += h; + qm[j] -= h; + std::vector rp(n), rm(n), d(n); + REQUIRE(target.Residuals(qp, rp.data())); + REQUIRE(target.Residuals(qm, rm.data())); + for (size_t i = 0; i < n; i++) + d[i] = (rp[i] - rm[i]) / (2 * h); + return d; + }; + std::vector> held(6), moving(6); + target.hold_mask = true; + REQUIRE(target.Residuals(q0, r.data())); // the mask the difference holds is q0's + for (int j : columns) + held[j] = difference(j, 0.005); + // The mask is binary on the grid, so with it moving the difference is taken at the step the + // Jacobian had when every column moved it (0.01 of the zone's resolution): at 0.005 A it + // mostly measures which grid points a few atoms happened to cross. + target.hold_mask = false; + for (int j : columns) + moving[j] = difference(j, 0.01 * zone); + + double largest = 0; + for (int j : columns) { + double s = 0; + for (size_t i = 0; i < n; i++) + s += gemmi::sq(held[j][i]); + largest = std::max(largest, std::sqrt(s)); + } + for (int j : columns) { + double dot = 0, ours = 0, theirs = 0, dot_moving = 0, norm_moving = 0; + for (size_t i = 0; i < n; i++) { + dot += jacobian[i * 6 + j] * held[j][i]; + ours += gemmi::sq(jacobian[i * 6 + j]); + theirs += gemmi::sq(held[j][i]); + dot_moving += moving[j][i] * held[j][i]; + norm_moving += gemmi::sq(moving[j][i]); + } + ours = std::sqrt(ours); + theirs = std::sqrt(theirs); + norm_moving = std::sqrt(norm_moving); + const double cosine = dot / (ours * theirs); + logger.Info("{} at {:.1f} A, column {}: cosine {:.5f}, norm {:.4f} of the difference's; " + "with the mask moving, cosine {:.4f} and norm {:.4f}", + std::string(cryst).substr(55, 11), zone, j, cosine, ours / theirs, + dot_moving / (norm_moving * theirs), norm_moving / theirs); + INFO(cryst << " at " << zone << " A, column " << j << ": cosine " << cosine + << ", norm ratio " << ours / theirs); + if (theirs < 1e-3 * largest) { + // A direction the origin is free in: |F| does not change, so neither may the column. + CHECK(ours < 1e-2 * largest); + continue; + } + CHECK(cosine >= min_cosine); + CHECK(std::fabs(ours / theirs - 1) <= max_norm_error); + } + } + } +} + +// The rotation columns at a placement already rotated (about 1.4 deg, so the step is taken from a +// rotated body and not from the model as read) with anisotropic atoms, whose U the placement does not +// turn with the body - the columns must differentiate exactly the function the residuals evaluate. +TEST_CASE("ModelValidation_RigidBodyRotationDerivativeMatchesDifferences", "[ModelValidation]") { + const double q0[6] = {0.12, -0.09, 0.07, 0.15, -0.10, 0.05}; + for (const char *cryst : kRigidBodyCrysts) + CheckRigidBodyJacobian(cryst, q0, {0, 1, 2}, 0.99, 0.03); +} + +// The whole Jacobian - analytic translation, forward-difference rotation, the scale re-fit folded in by +// projection - against the central difference of the residuals it describes. The projection is +// Kaufman's, which leaves out a term that grows with the residuals (the scale's derivatives depend on +// the placement too, weighted by how badly the model fits): with the body 0.07 A off it agrees with the +// difference to about 2e-3 in cosine, with it 0.35 A off to about 2e-2, and the bounds follow. +TEST_CASE("ModelValidation_RigidBodyJacobianMatchesNumericScaleRefit", "[ModelValidation]") { + const double close[6] = {0.0, 0.0, 0.0, 0.05, -0.04, 0.03}; + const double far[6] = {0.0, 0.0, 0.0, 0.25, -0.20, 0.15}; + for (const char *cryst : kRigidBodyCrysts) { + CheckRigidBodyJacobian(cryst, close, {0, 1, 2, 3, 4, 5}, 0.99, 0.02); + CheckRigidBodyJacobian(cryst, far, {0, 1, 2, 3, 4, 5}, 0.975, 0.03); + } +}