leonarski_f
a8d3a8dbc2
BulkSolventMask::MaskGrid, documented as "sizes and spacing set, values not allocated", allocated and zeroed the whole fine mask grid (up to 2^26 floats) - in BulkSolventMask's constructor and in every GPU engine's, i.e. once more for each weight-trial slot, only for the grid's size. It now takes good_grid_size() as set_size_from_spacing() does, and the CPU engine's Fmask() allocates its copy. rugnux_refine_test, RTX 5080, -N 8: 9yl4 8.6 -> 7.8 s, 7d98 4.85 -> 4.0 s (engine 0.66 -> 0.16 s, slots 0.60 -> 0.23 s on 9yl4); the high-resolution sets, whose mask grid is small, are unchanged. RefinementMinimiser keeps s.y of every pair it stores and Direction() takes rho and gamma from it instead of two more dot products per pair per call (the same Dot() of the same vectors): Direction's CPU time -22% on 9yl4; no wall change with the trials side by side, where other slots fill the card meanwhile. Bitwise: refined CIF identical on 5reo, 5f6m, 9yl4, 7d98, a 5f6m MD start (GPU) and 6o2h (CPU engine). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 6f5caf5e1b1183a752acc5cca78bfb7ef002b12a)
528 lines
25 KiB
C++
528 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);
|
|
}
|
|
|
|
// -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));
|
|
// 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);
|
|
// 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};
|
|
}
|
|
});
|
|
}
|