leonarski_f
626afc49dd
EstimateLikelihoodParameters' golden-section searches over D at a fixed beta - four per bin on the free reflections, and the one on the working reflections the D/FOM change added, 50 evaluations over 95 % of the data - evaluated -ln L reading every row through the index and taking log(var) at every evaluation. The rows are now gathered per search with their variance and its logarithm; the sum is the same expressions in the same order, so D, beta and everything after are the same to the bit (refined model and report identical on the seven timing sets). refit.ml 200 -> 140 ms per cycle on the 1.1 A set, one thread per trial slot. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 23ad3bfc254d10af310f17dbb8f85616f666a783)
520 lines
24 KiB
C++
520 lines
24 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);
|
|
}
|
|
|
|
// -ln L summed over rows `idx` at D and beta.
|
|
double BinNegLogLikelihood(const RefinementReflections &refl, const std::vector<double> &fc,
|
|
const std::vector<size_t> &idx, double d, double beta, bool with_sigma) {
|
|
double sum = 0, dfc = 0;
|
|
for (size_t i : idx)
|
|
sum += refine_math::AmplitudeNegLogLikelihood<double>(refl.fo[i], fc[i], d, Variance(refl, i, beta, with_sigma),
|
|
refl.centric[i], dfc);
|
|
return sum;
|
|
}
|
|
|
|
// Rows `idx` at a fixed beta, gathered for searches over D alone: the variance and its logarithm taken once rather
|
|
// than at every evaluation. FixedVarianceNegLogLikelihood() at D is BinNegLogLikelihood() at D and that beta, to
|
|
// the bit: the same expressions in the same order, only log(var) hoisted.
|
|
struct FixedVarianceRows {
|
|
std::vector<double> fo, fc, var, log_var;
|
|
std::vector<char> centric;
|
|
FixedVarianceRows(const RefinementReflections &refl, const std::vector<double> &fc_all,
|
|
const std::vector<size_t> &idx, double beta, bool with_sigma) {
|
|
for (size_t i : idx) {
|
|
fo.push_back(refl.fo[i]);
|
|
fc.push_back(fc_all[i]);
|
|
var.push_back(Variance(refl, i, beta, with_sigma));
|
|
log_var.push_back(std::log(var.back()));
|
|
centric.push_back(refl.centric[i]);
|
|
}
|
|
}
|
|
};
|
|
|
|
// refine_math::AmplitudeNegLogLikelihood() summed over the rows at D, with log(var) given.
|
|
double FixedVarianceNegLogLikelihood(const FixedVarianceRows &rows, double d) {
|
|
double sum = 0;
|
|
for (size_t k = 0; k < rows.fo.size(); k++) {
|
|
const double fo = rows.fo[k], fc = rows.fc[k], var = rows.var[k];
|
|
if (!rows.centric[k]) {
|
|
const double x = 2.0 * d * fo * fc / var;
|
|
double log_i0, ratio;
|
|
refine_math::BesselI0I1(x, log_i0, ratio);
|
|
sum += rows.log_var[k] + (fo * fo + d * d * fc * fc) / var - log_i0;
|
|
} else {
|
|
const double x = d * fo * fc / var;
|
|
const double log_cosh = x + std::log(1.0 + std::exp(-2.0 * x)) - 0.6931471805599453;
|
|
sum += 0.5 * rows.log_var[k] + (fo * fo + d * d * fc * fc) / (2.0 * var) - log_cosh;
|
|
}
|
|
}
|
|
return sum;
|
|
}
|
|
|
|
// The minimum of f over [lo, hi] by golden-section search.
|
|
template <typename F>
|
|
double GoldenMinimum(F f, double lo, double hi, int iterations) {
|
|
const double g = 0.5 * (std::sqrt(5.0) - 1.0);
|
|
double a = lo, b = hi;
|
|
double c = b - g * (b - a), d = a + g * (b - a);
|
|
double fc = f(c), fd = f(d);
|
|
for (int i = 0; i < iterations; i++) {
|
|
if (fc < fd) {
|
|
b = d;
|
|
d = c;
|
|
fd = fc;
|
|
c = b - g * (b - a);
|
|
fc = f(c);
|
|
} else {
|
|
a = c;
|
|
c = d;
|
|
fc = fd;
|
|
d = a + g * (b - a);
|
|
fd = f(d);
|
|
}
|
|
}
|
|
return 0.5 * (a + b);
|
|
}
|
|
|
|
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));
|
|
grid.set_size_from_spacing(spacing, gemmi::GridSizeRounding::Up);
|
|
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_;
|
|
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);
|
|
// 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 FixedVarianceRows at_beta(refl, fc, idx, beta, false);
|
|
const double ld = GoldenMinimum([&](double x) { return FixedVarianceNegLogLikelihood(at_beta, std::exp(x)); },
|
|
std::log(d) - 2.0, std::log(std::min(d * 7.4, 3.0)), 40);
|
|
d = std::exp(ld);
|
|
const double lb = GoldenMinimum([&](double x) { return BinNegLogLikelihood(refl, fc, idx, d, std::exp(x), false); },
|
|
std::log(beta) - 3.0, std::log(beta) + 3.0, 40);
|
|
beta = std::exp(lb);
|
|
}
|
|
// D: by maximum likelihood on the working reflections, with the measurement error added to that beta.
|
|
// 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()) {
|
|
const FixedVarianceRows working(refl, fc, work_in_bin[b], beta, true);
|
|
d = GoldenMinimum([&](double x) { return FixedVarianceNegLogLikelihood(working, x); }, 0.0, 3.0, 50);
|
|
}
|
|
bin_d[b] = d;
|
|
bin_beta[b] = beta;
|
|
});
|
|
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};
|
|
}
|
|
});
|
|
}
|