leonarski_f
2c1df0ebfa
Deterministic, rounding changes. The per-bin ML fit took 4 rounds x (40 + 40) golden-section steps on the free set and a 50-step golden section for D on the working set - each step a Bessel I0/I1 per row, ~4.6M per Refit at 92k reflections, 27-37% of a trial's wall on the high-resolution sets. Each one-dimensional search is now a safeguarded Newton iteration on the analytic first and second derivatives (bracketed on the sign of the derivative, bisection where a step leaves the bracket or the curvature is not positive), within the same brackets; the free-set alternation keeps its 4 rounds from the moments, so beta is the one it was (a converged joint fit was tried: the 4 rounds are not converged in the low-resolution bins, and converging them moves 6iu9's FOM by -0.003, so it was not kept). D on the working set starts from the bin's D of the context's previous Refit (XrayRowModel::bin_d), restored with a trial's fit and cleared at a trial's first Refit, so a trial depends on nothing its slot ran before (side by side = one at a time, md5). 14 parity sets, score only: R-work/R-free identical to 4 digits, mean |FOM - REFMAC| 0.01186 -> 0.01186, D within 1e-5 and var within 1e-4 (relative) per reflection. refit.ml (TIMING sum, -N 8): 5reo 1.47 -> 0.11 s, 5f6m 3.88 -> 0.26, 9yl4 1.37 -> 0.11, 5f6m MD start 4.97 -> 0.46. Wall: 5reo 2.03 -> 1.75, 5f6m 3.45 -> 2.60, 9yl4 7.99 -> 7.36, 7d98 4.04 -> 3.78, 5f6m MD start 6.10 -> 4.82 s. 5f6m MD starts 01-04 median R-free 0.1693 -> 0.1692. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit a91511945e0ccd7145a93877c0c231e87f86fb44)
531 lines
25 KiB
C++
531 lines
25 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "RefinementXray.h"
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <map>
|
|
#include <mutex>
|
|
|
|
#include <fftw3.h>
|
|
|
|
#include "gemmi/dencalc.hpp"
|
|
#include "gemmi/it92.hpp"
|
|
#include "gemmi/math.hpp"
|
|
|
|
#include "../../common/Logger.h"
|
|
#include "../../common/ParallelFor.h"
|
|
#include "ModelFFT.h"
|
|
#include "ModelGrid.h"
|
|
#include "RefinementMath.h"
|
|
|
|
namespace {
|
|
|
|
using Table = gemmi::IT92<float>;
|
|
|
|
// The free reflections a likelihood bin is fitted on, at least.
|
|
constexpr size_t FREE_PER_BIN = 100;
|
|
constexpr size_t MAX_BINS = 30;
|
|
|
|
// The density calculator the refinement grids atoms with: DensityCalculator's settings for the grid of
|
|
// RigidBodyZoneGrid(), and the given blur.
|
|
gemmi::DensityCalculator<Table, float> Calculator(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
|
|
double d_min, double blur) {
|
|
gemmi::DensityCalculator<Table, float> dc;
|
|
dc.d_min = d_min;
|
|
dc.rate = 1.5; // RigidBodyRefine.cpp's GRID_RATE, which RigidBodyZoneGrid sizes the grid with
|
|
dc.grid.unit_cell = cell;
|
|
dc.grid.spacegroup = &sg;
|
|
dc.blur = blur;
|
|
return dc;
|
|
}
|
|
|
|
// The likelihood's variance of reflection i: epsilon beta, plus the measurement error of each of the structure
|
|
// factor's components that the observed amplitude measures - sigma^2 twice for an acentric reflection, once for a
|
|
// centric one (with_sigma = false leaves it out).
|
|
double Variance(const RefinementReflections &refl, size_t i, double beta, bool with_sigma) {
|
|
const double s2 = with_sigma ? double(refl.sigma[i]) * refl.sigma[i] : 0.0;
|
|
return refl.epsilon[i] * beta + (refl.centric[i] ? s2 : 2.0 * s2);
|
|
}
|
|
|
|
// The first and second derivatives of one row's -ln L (refine_math::AmplitudeNegLogLikelihood) in D and in the
|
|
// variance, at D and var.
|
|
struct RowDerivatives {
|
|
double d, dd, v, vv;
|
|
};
|
|
|
|
RowDerivatives LikelihoodDerivatives(double fo, double fc, double d, double var, bool centric) {
|
|
const double q = fo * fo + d * d * fc * fc;
|
|
if (!centric) {
|
|
// x = 2 D Fo Fc / var, m = I1(x) / I0(x) = d ln I0 / dx, m' = 1 - m / x - m^2
|
|
const double a = 2.0 * fo * fc / var, x = a * d;
|
|
double log_i0, m;
|
|
refine_math::BesselI0I1(x, log_i0, m);
|
|
const double m1 = x > 0 ? 1.0 - m / x - m * m : 0.5;
|
|
return {2.0 * d * fc * fc / var - a * m,
|
|
2.0 * fc * fc / var - a * a * m1,
|
|
1.0 / var - q / (var * var) + m * x / var,
|
|
(-1.0 + 2.0 * q / var - x * (m1 * x + 2.0 * m)) / (var * var)};
|
|
}
|
|
// x = D Fo Fc / var, t = tanh x = d ln cosh / dx
|
|
const double a = fo * fc / var, x = a * d;
|
|
const double t = std::tanh(x), t1 = 1.0 - t * t;
|
|
return {d * fc * fc / var - a * t,
|
|
fc * fc / var - a * a * t1,
|
|
0.5 / var - q / (2.0 * var * var) + t * x / var,
|
|
(-0.5 + q / var - x * (t1 * x + 2.0 * t)) / (var * var)};
|
|
}
|
|
|
|
// The minimum of a function over [lo, hi] from x, given its first and second derivative (derivatives(x, g, h)):
|
|
// Newton's method on the first derivative, kept inside the interval where that changes sign and bisecting it
|
|
// wherever a step would leave it or the curvature is not positive. A minimum on a bound is the bound.
|
|
template <typename F>
|
|
double NewtonMinimum(F derivatives, double x, double lo, double hi) {
|
|
x = std::clamp(x, lo, hi);
|
|
for (int iteration = 0; iteration < 200; iteration++) {
|
|
double g = 0, h = 0;
|
|
derivatives(x, g, h);
|
|
if (g < 0)
|
|
lo = x;
|
|
else
|
|
hi = x;
|
|
double next = h > 0 ? x - g / h : lo - 1.0;
|
|
if (!(next > lo && next < hi))
|
|
next = 0.5 * (lo + hi);
|
|
const bool converged = std::fabs(next - x) < 1e-10 || hi - lo < 1e-10;
|
|
x = next;
|
|
if (converged)
|
|
break;
|
|
}
|
|
return x;
|
|
}
|
|
|
|
class RefinementEngineCPU : public RefinementEngine {
|
|
public:
|
|
RefinementEngineCPU(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
|
|
const RefinementReflections &refl, size_t nthreads)
|
|
: cell_(cell), sg_(sg), d_min_(d_min), refl_(refl), nthreads_(nthreads),
|
|
grid_(RigidBodyZoneGrid(cell, sg, d_min)), composition_(grid_, d_min, refl.hkl),
|
|
orbit_leaders_(OrbitLeaders(grid_, nthreads)) {}
|
|
|
|
void SetBlur(double blur) override { blur_ = blur; }
|
|
void SetRowModel(const XrayRowModel &rows) override { rows_ = rows; }
|
|
bool OnGPU() const override { return false; }
|
|
|
|
void Fcalc(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
|
|
std::vector<std::complex<double>> fd;
|
|
Compose(model, fd);
|
|
f.assign(fd.begin(), fd.end());
|
|
}
|
|
|
|
void Fmask(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
|
|
gemmi::Grid<float> mask = grid_;
|
|
PutMaskOnGrid(mask, model, orbit_leaders_, nthreads_);
|
|
const gemmi::FPhiGrid<float> fm = MapToFPhi(mask);
|
|
f.clear();
|
|
for (const gemmi::Miller &h : refl_.hkl)
|
|
f.push_back(fm.get_value_by_hkl(h));
|
|
}
|
|
|
|
double Evaluate(const gemmi::Model &model, std::vector<std::array<double, 4>> *grad) override {
|
|
std::vector<std::complex<double>> f;
|
|
Compose(model, f);
|
|
const size_t n = f.size();
|
|
// The target and, per reflection, dT/dF of the atoms' structure factor (dT/dRe F + i dT/dIm F).
|
|
std::vector<std::complex<double>> dt_df(n, 0.0);
|
|
double sum = 0;
|
|
for (size_t m = 0; m < n; m++) {
|
|
if (refl_.free[m])
|
|
continue;
|
|
const double k = rows_.k_total[m];
|
|
const std::complex<double> fm = k * (f[m] + std::complex<double>(rows_.fmask[m]));
|
|
const double fc = std::abs(fm);
|
|
double dfc = 0;
|
|
sum += refine_math::AmplitudeNegLogLikelihood<double>(refl_.fo[m], fc, rows_.d[m], rows_.var[m],
|
|
refl_.centric[m], dfc);
|
|
if (fc > 0)
|
|
dt_df[m] = k * dfc * fm / fc;
|
|
}
|
|
if (grad == nullptr)
|
|
return sum;
|
|
|
|
// The FFT method's map, dT/drho(x) = (V/N) Re sum_m sum_ops conj(dT/dF_m) scale_m exp(2 pi i h.t)
|
|
// exp(2 pi i (hR).x), as the inverse real transform of a half-l coefficient grid: a term A at index k
|
|
// goes in as A/2 at k, or as conj(A)/2 at -k where k is in the half not stored; on the planes l = 0 and
|
|
// l = nw/2, which hold both, as both.
|
|
const int nu = grid_.nu, nv = grid_.nv, nw = grid_.nw, nwh = nw / 2 + 1;
|
|
const size_t ncomplex = static_cast<size_t>(nu) * nv * nwh;
|
|
std::complex<float> *in = reinterpret_cast<std::complex<float> *>(fftwf_alloc_complex(ncomplex));
|
|
float *out = fftwf_alloc_real(static_cast<size_t>(nu) * nv * nw);
|
|
std::fill(in, in + ncomplex, std::complex<float>(0, 0));
|
|
const std::vector<SymmetryComposition::Term> &terms = composition_.Terms();
|
|
const size_t ops = composition_.Ops();
|
|
auto add = [&](int ku, int kv, int kw, std::complex<double> value) {
|
|
in[gemmi::modulo(ku, nu) + static_cast<size_t>(nu) * (gemmi::modulo(kv, nv) + static_cast<size_t>(nv) * kw)] +=
|
|
std::complex<float>(value);
|
|
};
|
|
for (size_t m = 0; m < n; m++) {
|
|
if (dt_df[m] == 0.0)
|
|
continue;
|
|
const double scale = composition_.Centring() * std::exp(blur_ * 0.25 * composition_.InvD2()[m]);
|
|
const std::complex<double> c = std::conj(dt_df[m]) * scale;
|
|
for (size_t o = 0; o < ops; o++) {
|
|
const SymmetryComposition::Term &t = terms[m * ops + o];
|
|
const std::complex<double> a = 0.5 * c * t.phase;
|
|
const gemmi::Miller &k = t.k;
|
|
if (k[2] > 0 && 2 * k[2] != nw)
|
|
add(k[0], k[1], k[2], a);
|
|
else if (k[2] < 0 && -2 * k[2] != nw)
|
|
add(-k[0], -k[1], -k[2], std::conj(a));
|
|
else {
|
|
const int l = k[2] < 0 ? -k[2] : k[2];
|
|
add(k[0], k[1], l, a);
|
|
add(-k[0], -k[1], l, std::conj(a));
|
|
}
|
|
}
|
|
}
|
|
MapFromFPhiExecute(nu, nv, nw, in, out);
|
|
gemmi::Grid<float> map = grid_;
|
|
const float norm = static_cast<float>(cell_.volume / static_cast<double>(map.point_count()));
|
|
for (size_t i = 0; i < map.data.size(); i++)
|
|
map.data[i] = out[i] * norm;
|
|
fftwf_free(in);
|
|
fftwf_free(out);
|
|
AtomGradientsFromMap(map, model, blur_, nthreads_, *grad);
|
|
return sum;
|
|
}
|
|
|
|
private:
|
|
void Compose(const gemmi::Model &model, std::vector<std::complex<double>> &f) const {
|
|
gemmi::DensityCalculator<Table, float> dc = Calculator(cell_, sg_, d_min_, blur_);
|
|
PutModelDensityOnGrid(dc, model, {}, nthreads_); // one copy: the symmetry is composed below
|
|
composition_.Compose(MapToFPhi(dc.grid), blur_, f, nullptr, nthreads_);
|
|
}
|
|
|
|
gemmi::UnitCell cell_;
|
|
const gemmi::SpaceGroup &sg_;
|
|
double d_min_;
|
|
const RefinementReflections &refl_;
|
|
size_t nthreads_;
|
|
gemmi::Grid<float> grid_;
|
|
SymmetryComposition composition_;
|
|
std::vector<size_t> orbit_leaders_;
|
|
double blur_ = 0;
|
|
XrayRowModel rows_;
|
|
};
|
|
|
|
} // namespace
|
|
|
|
RefinementReflections PrepareRefinementReflections(const std::vector<MergedReflection> &merged,
|
|
const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
|
|
double d_min) {
|
|
const gemmi::GroupOps gops = sg.operations();
|
|
const gemmi::ReciprocalAsu asu(&sg);
|
|
// ASU, last one wins where a reflection is listed twice (both hands of an anomalous merge)
|
|
std::map<gemmi::Miller, const MergedReflection *> by_hkl;
|
|
for (const MergedReflection &r : merged) {
|
|
if (!std::isfinite(r.F) || r.F < 0)
|
|
continue;
|
|
gemmi::Miller h{{r.h, r.k, r.l}};
|
|
if (!asu.is_in(h))
|
|
h = asu.to_asu(h, gops).first;
|
|
by_hkl[h] = &r;
|
|
}
|
|
std::vector<gemmi::Miller> hkl;
|
|
for (const auto &[h, r] : by_hkl)
|
|
hkl.push_back(h);
|
|
const SymmetryComposition composition(RigidBodyZoneGrid(cell, sg, d_min), d_min, hkl);
|
|
RefinementReflections out;
|
|
for (size_t i = 0; i < hkl.size(); i++) {
|
|
if (composition.Row()[i] < 0)
|
|
continue;
|
|
const MergedReflection &r = *by_hkl[hkl[i]];
|
|
out.hkl.push_back(hkl[i]);
|
|
out.fo.push_back(r.F);
|
|
out.sigma.push_back(std::isfinite(r.sigmaF) ? r.sigmaF : 0.0f);
|
|
out.free.push_back(r.rfree_flag);
|
|
out.centric.push_back(gops.is_reflection_centric(hkl[i]));
|
|
out.epsilon.push_back(static_cast<float>(gops.epsilon_factor(hkl[i])));
|
|
out.inv_d2.push_back(cell.calculate_1_d2(hkl[i]));
|
|
}
|
|
return out;
|
|
}
|
|
|
|
BulkSolventMask::BulkSolventMask(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
|
|
size_t nthreads)
|
|
: grid_(MaskGrid(cell, sg, d_min)), nthreads_(nthreads),
|
|
spacing_(std::max(std::min(d_min / 3.0, MASK_SPACING_MAX), std::cbrt(cell.volume / MASK_POINTS_MAX))) {}
|
|
|
|
gemmi::Grid<float> BulkSolventMask::MaskGrid(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min) {
|
|
gemmi::Grid<float> grid;
|
|
grid.unit_cell = cell;
|
|
grid.spacegroup = &sg;
|
|
const double spacing = std::max(std::min(d_min / 3.0, MASK_SPACING_MAX), std::cbrt(cell.volume / MASK_POINTS_MAX));
|
|
// set_size_from_spacing()'s size, without allocating the values
|
|
const std::array<int, 3> n = gemmi::good_grid_size({{cell.a / spacing, cell.b / spacing, cell.c / spacing}},
|
|
gemmi::GridSizeRounding::Up, &sg);
|
|
grid.nu = n[0];
|
|
grid.nv = n[1];
|
|
grid.nw = n[2];
|
|
grid.calculate_spacing();
|
|
grid.axis_order = gemmi::AxisOrder::XYZ;
|
|
return grid;
|
|
}
|
|
|
|
void BulkSolventMask::Fmask(const gemmi::Model &model, const std::vector<gemmi::Miller> &hkl,
|
|
std::vector<std::complex<float>> &f) const {
|
|
if (orbit_leaders_.empty()) // at the first mask: a GPU engine's refinement never draws one here
|
|
orbit_leaders_ = OrbitLeaders(grid_, nthreads_);
|
|
gemmi::Grid<float> mask = grid_;
|
|
mask.data.resize(mask.point_count());
|
|
PutMaskOnGrid(mask, model, orbit_leaders_, nthreads_);
|
|
const gemmi::FPhiGrid<float> fm = MapToFPhi(mask);
|
|
f.clear();
|
|
for (const gemmi::Miller &h : hkl)
|
|
f.push_back(fm.get_value_by_hkl(h));
|
|
}
|
|
|
|
void EstimateLikelihoodParameters(const RefinementReflections &refl, const std::vector<std::complex<float>> &fmodel,
|
|
XrayRowModel &rows, size_t nthreads) {
|
|
const size_t n = refl.Size();
|
|
std::vector<double> fc(n);
|
|
for (size_t i = 0; i < n; i++)
|
|
fc[i] = std::abs(fmodel[i]);
|
|
std::vector<size_t> free, work;
|
|
for (size_t i = 0; i < n; i++)
|
|
if (fc[i] > 0)
|
|
(refl.free[i] ? free : work).push_back(i);
|
|
// Too few free reflections to bin: the working set, which says the model is better than it is, but
|
|
// still gives the likelihood a scale.
|
|
if (free.size() < FREE_PER_BIN)
|
|
free.insert(free.end(), work.begin(), work.end());
|
|
std::sort(free.begin(), free.end(), [&](size_t a, size_t b) { return refl.inv_d2[a] < refl.inv_d2[b]; });
|
|
const size_t bins = std::clamp<size_t>(free.size() / FREE_PER_BIN, 1, MAX_BINS);
|
|
// the working reflections of each bin: those from its first free reflection's 1/d^2 to the next bin's
|
|
std::vector<double> lower(bins);
|
|
for (size_t b = 0; b < bins; b++)
|
|
lower[b] = refl.inv_d2[free[b * free.size() / bins]];
|
|
std::vector<std::vector<size_t>> work_in_bin(bins);
|
|
for (size_t i : work) {
|
|
const size_t b = std::upper_bound(lower.begin(), lower.end(), refl.inv_d2[i]) - lower.begin();
|
|
work_in_bin[b > 0 ? b - 1 : 0].push_back(i);
|
|
}
|
|
|
|
std::vector<double> centre(bins), bin_d(bins), bin_beta(bins);
|
|
const bool warm = rows.bin_d.size() == bins;
|
|
// each bin is fitted on its own
|
|
ParallelFor(static_cast<int>(bins), nthreads, [&](int b) {
|
|
const std::vector<size_t> idx(free.begin() + b * free.size() / bins, free.begin() + (b + 1) * free.size() / bins);
|
|
double s2 = 0, sfo = 0, sfc2 = 0, sfofc = 0;
|
|
for (size_t i : idx) {
|
|
s2 += refl.inv_d2[i];
|
|
sfo += double(refl.fo[i]) * refl.fo[i] / refl.epsilon[i];
|
|
sfc2 += fc[i] * fc[i] / refl.epsilon[i];
|
|
sfofc += refl.fo[i] * fc[i] / refl.epsilon[i];
|
|
}
|
|
centre[b] = s2 / static_cast<double>(idx.size());
|
|
double d = sfc2 > 0 ? std::max(sfofc / sfc2, 0.05) : 1.0;
|
|
double beta = std::max((sfo - d * d * sfc2) / static_cast<double>(idx.size()), 1e-3 * sfo / idx.size());
|
|
// beta: the spread of the free amplitudes about D Fc, measurement error included, by maximum likelihood
|
|
// (alternating one-dimensional minimisations over ln D and ln beta)
|
|
for (int round = 0; round < 4; round++) {
|
|
const double ld = NewtonMinimum([&](double x, double &g, double &h) {
|
|
const double dx = std::exp(x);
|
|
for (size_t i : idx) {
|
|
const RowDerivatives r = LikelihoodDerivatives(refl.fo[i], fc[i], dx, refl.epsilon[i] * beta,
|
|
refl.centric[i]);
|
|
g += dx * r.d;
|
|
h += dx * dx * r.dd + dx * r.d;
|
|
}
|
|
}, std::log(d), std::log(d) - 2.0, std::log(std::min(d * 7.4, 3.0)));
|
|
d = std::exp(ld);
|
|
const double lb = NewtonMinimum([&](double x, double &g, double &h) {
|
|
for (size_t i : idx) {
|
|
const double var = refl.epsilon[i] * std::exp(x);
|
|
const RowDerivatives r = LikelihoodDerivatives(refl.fo[i], fc[i], d, var, refl.centric[i]);
|
|
g += var * r.v;
|
|
h += var * var * r.vv + var * r.v;
|
|
}
|
|
}, std::log(beta), std::log(beta) - 3.0, std::log(beta) + 3.0);
|
|
beta = std::exp(lb);
|
|
}
|
|
// D: by maximum likelihood on the working reflections, with the measurement error added to that beta, from
|
|
// where the last fit found it. Where the error is a large part of the spread, the variance counts it twice
|
|
// and D falls towards 0 - the high-resolution D and figure of merit REFMAC 5.8 reports.
|
|
// Following Murshudov et al. (2011) Acta Cryst. D67, 355-367 (D from the working set, the variance from the free)
|
|
if (!work_in_bin[b].empty()) {
|
|
std::vector<double> var;
|
|
for (size_t i : work_in_bin[b])
|
|
var.push_back(Variance(refl, i, beta, true));
|
|
d = NewtonMinimum([&](double x, double &g, double &h) {
|
|
for (size_t k = 0; k < var.size(); k++) {
|
|
const size_t i = work_in_bin[b][k];
|
|
const RowDerivatives r = LikelihoodDerivatives(refl.fo[i], fc[i], x, var[k], refl.centric[i]);
|
|
g += r.d;
|
|
h += r.dd;
|
|
}
|
|
}, warm ? rows.bin_d[b] : d, 0.0, 3.0);
|
|
}
|
|
bin_d[b] = d;
|
|
bin_beta[b] = beta;
|
|
});
|
|
rows.bin_d = bin_d;
|
|
rows.d.resize(n);
|
|
rows.var.resize(n);
|
|
for (size_t i = 0; i < n; i++) {
|
|
// linear in 1/d^2 between the bins' centres, flat beyond the outermost
|
|
const double s = refl.inv_d2[i];
|
|
double d = bin_d[0], beta = bin_beta[0];
|
|
if (s >= centre[bins - 1]) {
|
|
d = bin_d[bins - 1];
|
|
beta = bin_beta[bins - 1];
|
|
} else if (s > centre[0]) {
|
|
const size_t b = std::upper_bound(centre.begin(), centre.end(), s) - centre.begin();
|
|
const double t = (s - centre[b - 1]) / (centre[b] - centre[b - 1]);
|
|
d = bin_d[b - 1] + t * (bin_d[b] - bin_d[b - 1]);
|
|
beta = bin_beta[b - 1] + t * (bin_beta[b] - bin_beta[b - 1]);
|
|
}
|
|
rows.d[i] = static_cast<float>(d);
|
|
rows.var[i] = static_cast<float>(Variance(refl, i, beta, true));
|
|
}
|
|
}
|
|
|
|
RFactors ComputeRFactors(const RefinementReflections &refl, const std::vector<std::complex<float>> &fmodel) {
|
|
double nw = 0, dw = 0, nf = 0, df = 0;
|
|
RFactors r;
|
|
for (size_t i = 0; i < refl.Size(); i++) {
|
|
const double diff = std::fabs(refl.fo[i] - std::abs(fmodel[i]));
|
|
if (refl.free[i]) {
|
|
nf += diff;
|
|
df += refl.fo[i];
|
|
r.n_free++;
|
|
} else {
|
|
nw += diff;
|
|
dw += refl.fo[i];
|
|
r.n_work++;
|
|
}
|
|
}
|
|
r.r_work = dw > 0 ? nw / dw : 1;
|
|
r.r_free = df > 0 ? nf / df : 1;
|
|
return r;
|
|
}
|
|
|
|
std::vector<std::array<double, 4>> ModelXyzB(const gemmi::Model &model) {
|
|
std::vector<std::array<double, 4>> out;
|
|
for (const gemmi::Chain &ch : model.chains)
|
|
for (const gemmi::Residue &r : ch.residues)
|
|
for (const gemmi::Atom &a : r.atoms)
|
|
out.push_back({a.pos.x, a.pos.y, a.pos.z, a.b_iso});
|
|
return out;
|
|
}
|
|
|
|
namespace {
|
|
std::mutex stage_mutex;
|
|
std::map<std::string, std::pair<double, int>> stage_times;
|
|
} // namespace
|
|
|
|
RefinementStage::RefinementStage(const char *name)
|
|
: name_(name), start_(std::chrono::steady_clock::now()), nvtx_(NvtxRangeStart(name)) {}
|
|
|
|
RefinementStage::~RefinementStage() {
|
|
NvtxRangeEnd(nvtx_);
|
|
const double seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start_).count();
|
|
std::lock_guard lock(stage_mutex);
|
|
std::pair<double, int> &t = stage_times[name_];
|
|
t.first += seconds;
|
|
t.second++;
|
|
}
|
|
|
|
void RefinementStage::Log(Logger &logger) {
|
|
std::lock_guard lock(stage_mutex);
|
|
for (const auto &[name, t] : stage_times)
|
|
logger.Debug("TIMING refine.{} {:.3f} {}", name, t.first, t.second);
|
|
stage_times.clear();
|
|
}
|
|
|
|
double RefinementBlur(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min, double b_floor) {
|
|
// set_refmac_compatible_blur(): u_to_b / 1.1 * spacing^2 - b_min, not below 0, the spacing the one asked for
|
|
const gemmi::DensityCalculator<Table, float> dc = Calculator(cell, sg, d_min, 0);
|
|
const double spacing = dc.requested_grid_spacing();
|
|
return std::max(0.0, gemmi::u_to_b() / 1.1 * spacing * spacing - b_floor);
|
|
}
|
|
|
|
std::unique_ptr<RefinementEngine> MakeRefinementEngineCPU(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
|
|
double d_min, const RefinementReflections &refl,
|
|
size_t nthreads) {
|
|
return std::make_unique<RefinementEngineCPU>(cell, sg, d_min, refl, nthreads);
|
|
}
|
|
|
|
void AtomGradientsFromMap(const gemmi::Grid<float> &map, const gemmi::Model &model, double blur, size_t nthreads,
|
|
std::vector<std::array<double, 4>> &grad) {
|
|
const gemmi::DensityCalculator<Table, float> dc = Calculator(map.unit_cell, *map.spacegroup, 1.0, blur);
|
|
std::vector<const gemmi::Atom *> atoms;
|
|
for (const gemmi::Chain &ch : model.chains)
|
|
for (const gemmi::Residue &r : ch.residues)
|
|
for (const gemmi::Atom &a : r.atoms)
|
|
atoms.push_back(&a);
|
|
grad.assign(atoms.size(), {0, 0, 0, 0});
|
|
const int nu = map.nu, nv = map.nv, nw = map.nw;
|
|
// As ModelGrid.cpp walks a density box: the orthogonalization over the grid size in single precision.
|
|
const gemmi::UpperTriangularMat33 &om = map.orth_n;
|
|
const float m00 = static_cast<float>(om.a11), m01 = static_cast<float>(om.a12), m02 = static_cast<float>(om.a13),
|
|
m11 = static_cast<float>(om.a22), m12 = static_cast<float>(om.a23), m22 = static_cast<float>(om.a33);
|
|
constexpr double FOUR_PI2 = 4 * refine_math::PI * refine_math::PI;
|
|
ParallelChunks(static_cast<int>(atoms.size()), nthreads, [&](int lo, int hi) {
|
|
for (int i = lo; i < hi; i++) {
|
|
const gemmi::Atom &atom = *atoms[i];
|
|
const float b = static_cast<float>(atom.b_iso + blur);
|
|
const auto precal = Table::get(atom.element, atom.charge, atom.serial).precalculate_density_iso(b, 0.0f);
|
|
const double radius = dc.estimate_radius(precal, b);
|
|
const float max_r2 = static_cast<float>(radius * radius);
|
|
const gemmi::Fractional fpos = map.unit_cell.fractionalize(atom.pos);
|
|
const double n3[3] = {fpos.x * nu, fpos.y * nv, fpos.z * nw};
|
|
float nctr[3];
|
|
int c[3], d[3];
|
|
for (int k = 0; k < 3; k++) {
|
|
nctr[k] = static_cast<float>(n3[k]);
|
|
c[k] = gemmi::iround(n3[k]);
|
|
}
|
|
d[0] = std::min(static_cast<int>(std::ceil(radius / map.spacing[0])), nu - 1);
|
|
d[1] = std::min(static_cast<int>(std::ceil(radius / map.spacing[1])), nv - 1);
|
|
d[2] = std::min(static_cast<int>(std::ceil(radius / map.spacing[2])), nw - 1);
|
|
double gx = 0, gy = 0, gz = 0, gb = 0;
|
|
for (int w = c[2] - d[2]; w <= c[2] + d[2]; w++) {
|
|
const float fz = nctr[2] - static_cast<float>(w);
|
|
const float z = m22 * fz;
|
|
const int w_ = gemmi::modulo(w, nw);
|
|
for (int v = c[1] - d[1]; v <= c[1] + d[1]; v++) {
|
|
const float fy = nctr[1] - static_cast<float>(v);
|
|
const float y = m11 * fy + m12 * fz;
|
|
const float yz = y * y + z * z;
|
|
if (yz > max_r2)
|
|
continue;
|
|
const float x_yz = m01 * fy + m02 * fz;
|
|
const size_t row = static_cast<size_t>(nu) * (gemmi::modulo(v, nv) + static_cast<size_t>(nv) * w_);
|
|
for (int u = c[0] - d[0]; u <= c[0] + d[0]; u++) {
|
|
const float x = m00 * (nctr[0] - static_cast<float>(u)) + x_yz;
|
|
const float r2 = yz + x * x;
|
|
if (r2 > max_r2)
|
|
continue;
|
|
const double value = map.data[row + gemmi::modulo(u, nu)];
|
|
double dr2 = 0, db = 0;
|
|
for (size_t q = 0; q < std::size(precal.a); q++) {
|
|
const double e = precal.a[q] * gemmi::unsafe_expapprox(std::max(precal.b[q] * r2, -88.f));
|
|
dr2 += e * precal.b[q];
|
|
db += e * precal.b[q] * (1.5 + precal.b[q] * r2);
|
|
}
|
|
gx += value * dr2 * x;
|
|
gy += value * dr2 * y;
|
|
gz += value * dr2 * z;
|
|
gb += value * db;
|
|
}
|
|
}
|
|
}
|
|
// rho = occ sum_q a_q exp(b_q r^2): d rho / d(centre) = occ 2 (sum a_q b_q e_q) (centre - point), and
|
|
// with b_q = -4 pi^2 / (B_q + B + blur), d rho / dB = occ sum a_q e_q b_q (1.5 + b_q r^2) / (4 pi^2).
|
|
grad[i] = {2.0 * atom.occ * gx, 2.0 * atom.occ * gy, 2.0 * atom.occ * gz, atom.occ * gb / FOUR_PI2};
|
|
}
|
|
});
|
|
}
|