leonarski_f
6075239be0
The likelihood's variance per bin (beta) is fitted by maximum likelihood on the free reflections as the whole spread of Fo about D Fc, and D is then fitted on the working reflections with the variance eps beta + 2 sigma^2 (acentric; + sigma^2 centric), which is also the target's variance. Murshudov et al. (2011) Acta D67 355: D from the working set, the variance from the free set. The measurement error is thereby counted twice, so D and the figure of merit fall where sigma is a large part of the spread - at high resolution. Measured on REFMAC 5.8.0431 0-cycle output (black box) of the 14 parity sets: its FOM inverted per reflection gives var = eps a + 2 sigma^2 (acentric, fitted coefficient 2.00) and eps a + sigma^2 (centric); its D (FC_ALL/FC_ALL_LS = the log's SigmaA_Fc1) and a are reproduced by this estimate (on REFMAC's own Fc, one set to rms 0.009 in D over all 20 bins). The previous per-bin free-set ML put D at 0.39-0.48 where REFMAC has 0.06-0.13 in the outermost bins. rugnux score vs REFMAC 0-cycle, 14 sets: mean |FOM difference| 0.056 -> 0.012 (max 0.112 -> 0.042), per-bin D rms 0.107 -> 0.068; R-work/R-free unchanged (D does not enter them). The score's FOM is now the mean over the working reflections, as REFMAC reports it, and the refined model's maps are weighted by the refinement's own m and D rather than a separate sigma_A estimate. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit d00d635e2874bff6cd216b7f24367362d0cc2e56) (cherry picked from commit b93e0ea2fc1ffcc6bb40dae9b77816a321f4d6ec)
481 lines
22 KiB
C++
481 lines
22 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;
|
|
}
|
|
|
|
// 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 double ld = GoldenMinimum([&](double x) { return BinNegLogLikelihood(refl, fc, idx, std::exp(x), beta, false); },
|
|
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())
|
|
d = GoldenMinimum([&](double x) { return BinNegLogLikelihood(refl, fc, work_in_bin[b], x, beta, true); },
|
|
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};
|
|
}
|
|
});
|
|
}
|