leonarski_fandClaude Opus 5.5
48d895d715
From the twin-multi prototype (af742b889, d197507c7), rebased onto twin-gate rule B and cleaned up. The twin arm takes a list of laws: RefinementReflections::mates per law and mates_from (the converse, which the GPU gathers the gradient by, as a three-fold is not an involution), one fraction per law, the twinned amplitude k sqrt(a0 |U(h)|^2 + sum_d a_d |U(T_d h)|^2) on the CPU and the GPU, and the fraction fit as the same intercept regression with one column per law (a law whose coefficient comes out negative taken out). RefinementTwinDomains: every coset of find_twin_laws, of any order. DecideTwin keeps the single-law rung as it was (null at 3 SE, or 2 with the L-test). Where that law is taken and the lattice has more domains, all of them are refined at once and taken where the added domains beat their own null (their mates unrelated, the adopted law's kept) by 3 SE, summed per twin group, each at a fraction of 0.05 or more (domains under it dropped and the rest refined again). TwinNull generalised from pairs to twin groups: kept adjacent, partners (Laws+1)(k+1) apart; for one law it is the null it was. CloseFreeSet, in RefineModel before both arms: every orbit under the domain operators in one set. An unsplit orbit stays; a split one goes free with the probability of its own free share (a salted hash of the orbit), so the free fraction is kept - the prototype's majority took a C-centred pseudo-hexagonal merge's free set from 4020 to 891, an unsalted hash inflated free sets by 60%. Open arm, against twin-grid on 23 twin/pseudo-symmetric sets (runs twinfix-grid2 / twinfix-multi / twinfix-multi2): the second rung ran on 9 sets and was taken on one, 9hnc (P2 on a tetragonal metric, deposited four domains at 0.06 each): domains 0.058/0.058/0.055, z 13.1, R-free 0.2174 -> 0.2063. On the other 8 the added domains refine to 0.000-0.016 and the single law stands. The closure moves reflections on 5ebi and the two C222 pseudo-hexagonal sets 6fvz/6fwc (their law, FREE_SET_SPLIT before, is now tested and refused); 5ebi's R-free 0.347 -> 0.377 on the closed free set. Single-law sets move by up to 0.003 (the regression and GPU gather round differently), all others <= 0.0002. Wall time +10-50% on the sets that reach the second rung. 5ebi deposited data (--mode refine): the closure moves 1 reflection and its strongest law is tested: R-free 0.390 -> 0.364. Report: REFINE_TWIN_MULTI*, REFINE_TWIN_FREE_CLOSED. Tests: RefinementTwin_domains_gradient, RefinementTwin_free_set_closed. Docs: §15.12, RUGNUX_REPORT.md, ACKNOWLEDGEMENT.md. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UqXYdL4qBGYRiSXSGXfCLu
914 lines
42 KiB
C++
914 lines
42 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 <numeric>
|
|
#include <set>
|
|
|
|
#include <fftw3.h>
|
|
|
|
#include "gemmi/dencalc.hpp"
|
|
#include "gemmi/it92.hpp"
|
|
#include "gemmi/math.hpp"
|
|
#include "gemmi/twin.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;
|
|
dc.cutoff = REFINEMENT_DENSITY_CUTOFF;
|
|
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.RowHkl()),
|
|
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(WithAniso(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_.RowHkl())
|
|
f.push_back(fm.get_value_by_hkl(h));
|
|
}
|
|
|
|
double Evaluate(const gemmi::Model &model_in, std::vector<std::array<double, 4>> *grad) override {
|
|
const gemmi::Model &model = WithAniso(model_in);
|
|
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;
|
|
const size_t laws = refl_.Laws();
|
|
double own_fraction = 1.0;
|
|
for (float a : rows_.twin_fractions)
|
|
own_fraction -= a;
|
|
for (size_t m = 0; m < refl_.Size(); m++) {
|
|
if (refl_.free[m])
|
|
continue;
|
|
const double k = rows_.k_total[m];
|
|
if (!refl_.Twinned()) {
|
|
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;
|
|
continue;
|
|
}
|
|
// twinned: |F_model| = k sqrt(a_0 |U(h)|^2 + sum_d a_d |U(T_d h)|^2), U = F + fmask, so that
|
|
// dT/dU(T_d h) = dT/d|F_model| k^2 a_d U(T_d h) / |F_model|
|
|
const std::complex<double> u1 = f[m] + std::complex<double>(rows_.fmask[m]);
|
|
double i_model = own_fraction * std::norm(u1);
|
|
for (size_t d = 0; d < laws; d++) {
|
|
const int j = refl_.mates[d][m];
|
|
i_model += rows_.twin_fractions[d] * std::norm(f[j] + std::complex<double>(rows_.fmask[j]));
|
|
}
|
|
const double fc = k * std::sqrt(std::max(i_model, 0.0));
|
|
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) {
|
|
const double g = dfc * k * k / fc;
|
|
dt_df[m] += g * own_fraction * u1;
|
|
for (size_t d = 0; d < laws; d++) {
|
|
const int j = refl_.mates[d][m];
|
|
dt_df[j] += g * double(rows_.twin_fractions[d]) * (f[j] + std::complex<double>(rows_.fmask[j]));
|
|
}
|
|
}
|
|
}
|
|
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);
|
|
const bool aniso = AnyAniso();
|
|
AtomGradientsFromMap(map, model, blur_, nthreads_, *grad, aniso ? &grad_u_ : nullptr);
|
|
if (!aniso)
|
|
grad_u_.clear();
|
|
return sum;
|
|
}
|
|
|
|
private:
|
|
// The model as the engine grids it: where some atom is anisotropic, a copy with every atom's U set (ApplyAniso).
|
|
const gemmi::Model &WithAniso(const gemmi::Model &model) {
|
|
if (!AnyAniso())
|
|
return model;
|
|
aniso_model_ = model;
|
|
ApplyAniso(aniso_model_, *aniso_);
|
|
return aniso_model_;
|
|
}
|
|
|
|
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_, true, true); // 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_;
|
|
gemmi::Model aniso_model_{1};
|
|
};
|
|
|
|
} // 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;
|
|
}
|
|
|
|
std::vector<gemmi::Miller> RefinementReflections::RowHkl() const {
|
|
std::vector<gemmi::Miller> out = hkl;
|
|
out.insert(out.end(), mate_hkl.begin(), mate_hkl.end());
|
|
return out;
|
|
}
|
|
|
|
std::vector<gemmi::Op> RefinementTwinLaws(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg) {
|
|
const gemmi::GroupOps gops = sg.operations();
|
|
std::vector<gemmi::Op> laws;
|
|
std::vector<gemmi::Op> found;
|
|
try {
|
|
found = gemmi::find_twin_laws(cell, &sg, 3.0, false);
|
|
} catch (const std::exception &) {
|
|
return laws; // a setting whose centring the lattice search does not reproduce: no law is looked for
|
|
}
|
|
for (const gemmi::Op &op : found) {
|
|
const gemmi::Op law{op.rot, {0, 0, 0}, 'x'};
|
|
const gemmi::Op square = law * law;
|
|
if (gops.find_by_rotation(square.rot) != nullptr || gops.find_by_rotation(square.negated_rot()) != nullptr)
|
|
laws.push_back(law);
|
|
}
|
|
return laws;
|
|
}
|
|
|
|
std::vector<gemmi::Op> RefinementTwinDomains(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg) {
|
|
std::vector<gemmi::Op> laws;
|
|
try {
|
|
for (const gemmi::Op &op : gemmi::find_twin_laws(cell, &sg, 3.0, false))
|
|
laws.push_back(gemmi::Op{op.rot, {0, 0, 0}, 'x'});
|
|
} catch (const std::exception &) {
|
|
}
|
|
return laws;
|
|
}
|
|
|
|
int SetTwinLaws(RefinementReflections &refl, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
|
|
const std::vector<gemmi::Op> &laws) {
|
|
const gemmi::GroupOps gops = sg.operations();
|
|
const gemmi::ReciprocalAsu asu(&sg);
|
|
const size_t n = refl.Size();
|
|
std::map<gemmi::Miller, int> row;
|
|
for (size_t m = 0; m < n; m++)
|
|
row[refl.hkl[m]] = static_cast<int>(m);
|
|
// each observed reflection's mates in the ASU, and those not observed - where the grid composes them - as rows
|
|
// of their own
|
|
std::vector<std::vector<gemmi::Miller>> mate_of(laws.size(), std::vector<gemmi::Miller>(n));
|
|
std::set<gemmi::Miller> unobserved;
|
|
for (size_t d = 0; d < laws.size(); d++)
|
|
for (size_t m = 0; m < n; m++) {
|
|
mate_of[d][m] = asu.to_asu(laws[d].apply_to_hkl(refl.hkl[m]), gops).first;
|
|
if (row.find(mate_of[d][m]) == row.end())
|
|
unobserved.insert(mate_of[d][m]);
|
|
}
|
|
const std::vector<gemmi::Miller> candidates(unobserved.begin(), unobserved.end());
|
|
const SymmetryComposition composition(RigidBodyZoneGrid(cell, sg, d_min), d_min, candidates);
|
|
refl.twin_laws = laws;
|
|
refl.mate_hkl.clear();
|
|
for (size_t i = 0; i < candidates.size(); i++)
|
|
if (composition.Row()[i] >= 0) {
|
|
row[candidates[i]] = static_cast<int>(n + refl.mate_hkl.size());
|
|
refl.mate_hkl.push_back(candidates[i]);
|
|
}
|
|
const size_t rows = refl.Rows();
|
|
refl.mates.assign(laws.size(), std::vector<int>(rows));
|
|
refl.mates_from.assign(laws.size(), std::vector<int>(rows, -1));
|
|
int split = 0;
|
|
for (size_t d = 0; d < laws.size(); d++) {
|
|
std::vector<int> &mate = refl.mates[d];
|
|
std::iota(mate.begin(), mate.end(), 0);
|
|
for (size_t m = 0; m < n; m++) {
|
|
const auto it = row.find(mate_of[d][m]);
|
|
if (it != row.end() && it->second != static_cast<int>(m)) {
|
|
mate[m] = it->second;
|
|
refl.mates_from[d][it->second] = static_cast<int>(m);
|
|
}
|
|
}
|
|
for (size_t m = 0; m < n; m++) {
|
|
const int j = mate[m];
|
|
if (j != static_cast<int>(m) && j < static_cast<int>(n) && refl.free[m] != refl.free[j] &&
|
|
(refl.mates[d][j] != static_cast<int>(m) || j > static_cast<int>(m)))
|
|
split++;
|
|
}
|
|
}
|
|
return split;
|
|
}
|
|
|
|
std::vector<std::complex<float>> TwinnedModel(const RefinementReflections &refl, const XrayRowModel &rows,
|
|
const std::vector<std::complex<float>> &f) {
|
|
const size_t n = refl.Size();
|
|
std::vector<std::complex<float>> out(n);
|
|
if (!refl.Twinned()) {
|
|
for (size_t m = 0; m < n; m++)
|
|
out[m] = rows.k_total[m] * (f[m] + rows.fmask[m]);
|
|
return out;
|
|
}
|
|
double own_fraction = 1.0;
|
|
for (float a : rows.twin_fractions)
|
|
own_fraction -= a;
|
|
for (size_t m = 0; m < n; m++) {
|
|
const std::complex<double> u1 = std::complex<double>(f[m]) + std::complex<double>(rows.fmask[m]);
|
|
const double i1 = std::norm(u1);
|
|
double i_model = own_fraction * i1;
|
|
for (size_t d = 0; d < refl.Laws(); d++) {
|
|
const int j = refl.mates[d][m];
|
|
i_model += rows.twin_fractions[d] *
|
|
std::norm(std::complex<double>(f[j]) + std::complex<double>(rows.fmask[j]));
|
|
}
|
|
const double amplitude = rows.k_total[m] * std::sqrt(std::max(i_model, 0.0));
|
|
out[m] = i1 > 0 ? std::complex<float>(u1 * (amplitude / std::sqrt(i1))) : std::complex<float>(amplitude, 0);
|
|
}
|
|
return out;
|
|
}
|
|
|
|
std::vector<float> FitTwinFractions(const RefinementReflections &refl, const XrayRowModel &rows,
|
|
const std::vector<std::complex<float>> &f) {
|
|
const size_t laws = refl.Laws();
|
|
std::vector<float> out(laws, 0.0f);
|
|
if (laws == 0)
|
|
return out;
|
|
// the working reflections with a mate that is another reflection: Fo^2 and the model's intensity at h and at
|
|
// each T_d h
|
|
const size_t nx = laws + 1;
|
|
struct Term {
|
|
double inv_d2, io;
|
|
std::vector<double> x; // the model's intensity at h, then at T_d h
|
|
};
|
|
std::vector<Term> terms;
|
|
for (size_t m = 0; m < refl.Size(); m++) {
|
|
if (refl.free[m])
|
|
continue;
|
|
bool related = false;
|
|
for (size_t d = 0; d < laws; d++)
|
|
related |= refl.mates[d][m] != static_cast<int>(m);
|
|
if (!related)
|
|
continue;
|
|
const double k2 = double(rows.k_total[m]) * rows.k_total[m];
|
|
Term t{refl.inv_d2[m], double(refl.fo[m]) * refl.fo[m], std::vector<double>(nx)};
|
|
t.x[0] = k2 * std::norm(std::complex<double>(f[m]) + std::complex<double>(rows.fmask[m]));
|
|
for (size_t d = 0; d < laws; d++) {
|
|
const int j = refl.mates[d][m];
|
|
t.x[d + 1] = k2 * std::norm(std::complex<double>(f[j]) + std::complex<double>(rows.fmask[j]));
|
|
}
|
|
terms.push_back(std::move(t));
|
|
}
|
|
if (terms.size() < nx + 2)
|
|
return out;
|
|
// normalised per resolution shell of TWIN_FRACTION_SHELL reflections: Fo^2 by its mean, the model's by the mean
|
|
// of its own intensity
|
|
std::stable_sort(terms.begin(), terms.end(), [](const Term &a, const Term &b) { return a.inv_d2 < b.inv_d2; });
|
|
const size_t shells = std::max<size_t>(1, terms.size() / TWIN_FRACTION_SHELL);
|
|
// the normal equations of y = b0 + sum_p b_p x_p, unknowns 0 (the intercept) to nx
|
|
const size_t nu = nx + 1;
|
|
std::vector<double> nm(nu * nu, 0.0), r(nu, 0.0);
|
|
std::vector<double> x(nu);
|
|
for (size_t k = 0; k < shells; k++) {
|
|
const size_t lo = k * terms.size() / shells, hi = (k + 1) * terms.size() / shells;
|
|
double so = 0, sc = 0;
|
|
for (size_t i = lo; i < hi; i++) {
|
|
so += terms[i].io;
|
|
sc += terms[i].x[0];
|
|
}
|
|
if (so <= 0 || sc <= 0)
|
|
continue;
|
|
for (size_t i = lo; i < hi; i++) {
|
|
x[0] = 1.0;
|
|
for (size_t p = 0; p < nx; p++)
|
|
x[p + 1] = terms[i].x[p] / sc;
|
|
const double y = terms[i].io / so;
|
|
for (size_t p = 0; p < nu; p++) {
|
|
r[p] += x[p] * y;
|
|
for (size_t q = 0; q < nu; q++)
|
|
nm[p * nu + q] += x[p] * x[q];
|
|
}
|
|
}
|
|
}
|
|
// solved over the unknowns still in, by Gaussian elimination; a law whose coefficient is negative is taken out
|
|
// (the most negative first) and the rest solved again
|
|
std::vector<char> in(nu, 1);
|
|
std::vector<double> b(nu, 0.0);
|
|
for (size_t round = 0; round < nu; round++) {
|
|
std::vector<size_t> idx;
|
|
for (size_t p = 0; p < nu; p++)
|
|
if (in[p])
|
|
idx.push_back(p);
|
|
const size_t k = idx.size();
|
|
std::vector<double> a(k * (k + 1));
|
|
for (size_t p = 0; p < k; p++) {
|
|
for (size_t q = 0; q < k; q++)
|
|
a[p * (k + 1) + q] = nm[idx[p] * nu + idx[q]];
|
|
a[p * (k + 1) + k] = r[idx[p]];
|
|
}
|
|
bool singular = false;
|
|
for (size_t c = 0; c < k && !singular; c++) {
|
|
size_t piv = c;
|
|
for (size_t p = c + 1; p < k; p++)
|
|
if (std::fabs(a[p * (k + 1) + c]) > std::fabs(a[piv * (k + 1) + c]))
|
|
piv = p;
|
|
if (!(std::fabs(a[piv * (k + 1) + c]) > 0)) {
|
|
singular = true;
|
|
break;
|
|
}
|
|
for (size_t q = 0; q <= k; q++)
|
|
std::swap(a[c * (k + 1) + q], a[piv * (k + 1) + q]);
|
|
for (size_t p = 0; p < k; p++)
|
|
if (p != c) {
|
|
const double factor = a[p * (k + 1) + c] / a[c * (k + 1) + c];
|
|
for (size_t q = c; q <= k; q++)
|
|
a[p * (k + 1) + q] -= factor * a[c * (k + 1) + q];
|
|
}
|
|
}
|
|
if (singular)
|
|
return out;
|
|
std::fill(b.begin(), b.end(), 0.0);
|
|
for (size_t p = 0; p < k; p++)
|
|
b[idx[p]] = a[p * (k + 1) + k] / a[p * (k + 1) + p];
|
|
size_t worst = 0;
|
|
for (size_t p = 2; p < nu; p++)
|
|
if (in[p] && b[p] < 0 && (worst == 0 || b[p] < b[worst]))
|
|
worst = p;
|
|
if (worst == 0)
|
|
break;
|
|
in[worst] = 0;
|
|
}
|
|
double total = 0;
|
|
for (size_t p = 1; p < nu; p++)
|
|
total += b[p];
|
|
if (!(total > 0))
|
|
return out;
|
|
for (size_t d = 0; d < laws; d++)
|
|
out[d] = static_cast<float>(std::max(0.0, b[d + 2] / total));
|
|
if (laws == 1)
|
|
out[0] = std::min(out[0], 0.5f);
|
|
return out;
|
|
}
|
|
|
|
double FitTwinFraction(const RefinementReflections &refl, const XrayRowModel &rows,
|
|
const std::vector<std::complex<float>> &f) {
|
|
double sum = 0;
|
|
for (float a : FitTwinFractions(refl, rows, f))
|
|
sum += a;
|
|
return sum;
|
|
}
|
|
|
|
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_(MaskSpacing(cell, d_min)) {}
|
|
|
|
double BulkSolventMask::MaskSpacing(const gemmi::UnitCell &cell, double d_min) {
|
|
const double spacing = std::max(std::min(d_min / 3.0, MASK_SPACING_MAX), std::cbrt(cell.volume / MASK_POINTS_MAX));
|
|
return std::min({spacing, MASK_SPACING_SHRINK, d_min / 3.0});
|
|
}
|
|
|
|
gemmi::Grid<float> BulkSolventMask::MaskGrid(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min) {
|
|
return SolventMaskGrid(cell, sg, d_min, MaskSpacing(cell, d_min));
|
|
}
|
|
|
|
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;
|
|
}
|
|
|
|
gemmi::SMat33<float> AnisoU(const AtomAniso &a, double b) {
|
|
const double u = b / gemmi::u_to_b();
|
|
gemmi::SMat33<double> m{a.dev[0] + u, a.dev[1] + u, a.dev[2] + u, a.dev[3], a.dev[4], a.dev[5]};
|
|
const std::array<double, 3> ev = m.calculate_eigenvalues();
|
|
const double lowest = std::min({ev[0], ev[1], ev[2]});
|
|
if (!(lowest >= ANISO_U_MIN))
|
|
m = m.added_kI(ANISO_U_MIN - (std::isfinite(lowest) ? lowest : 0.0));
|
|
return {static_cast<float>(m.u11), static_cast<float>(m.u22), static_cast<float>(m.u33),
|
|
static_cast<float>(m.u12), static_cast<float>(m.u13), static_cast<float>(m.u23)};
|
|
}
|
|
|
|
void ApplyAniso(gemmi::Model &model, const std::vector<AtomAniso> &aniso) {
|
|
size_t i = 0;
|
|
for (gemmi::Chain &ch : model.chains)
|
|
for (gemmi::Residue &r : ch.residues)
|
|
for (gemmi::Atom &a : r.atoms) {
|
|
if (i < aniso.size() && aniso[i].on)
|
|
a.aniso = AnisoU(aniso[i], a.b_iso);
|
|
else
|
|
a.aniso = {0, 0, 0, 0, 0, 0};
|
|
i++;
|
|
}
|
|
}
|
|
|
|
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);
|
|
}
|
|
|
|
namespace {
|
|
|
|
// AtomGradientsFromMap() for an anisotropic atom, gridded as PutModelDensityOnGrid() grids it: rho = occ sum_q a_q
|
|
// exp(r^T M_q r), M_q = -4 pi^2 (B_q I + 8 pi^2 U + blur I)^-1 and r = centre - point. Then d rho / d(centre) =
|
|
// occ sum_q 2 e_q M_q r, and, as d a_q / dU = -4 pi^2 a_q (B..)^-1 = a_q M_q and d(r^T M_q r) / dU = 2 (M_q r)(M_q r)^T,
|
|
// d rho / dU = occ sum_q e_q (M_q + 2 (M_q r)(M_q r)^T) - an off-diagonal U_ij twice its matrix element.
|
|
void AnisoGradient(const gemmi::Grid<float> &map, const gemmi::DensityCalculator<Table, float> &dc,
|
|
const gemmi::Atom &atom, double blur, std::array<double, 4> &grad, std::array<double, 6> *grad_u) {
|
|
using CReal = Table::Coef::coef_type;
|
|
const auto &coef = IonFormFactor<float>(atom.element.elem, atom.charge);
|
|
const auto aniso_b = atom.aniso.scaled(CReal(gemmi::u_to_b())).added_kI(CReal(blur));
|
|
const CReal b_max = AnisoRadiusB(aniso_b);
|
|
const double radius = dc.estimate_radius(coef.precalculate_density_iso(b_max, 0.0f), b_max);
|
|
const auto precal = coef.precalculate_density_aniso_b(aniso_b, 0.0f);
|
|
const float max_r2 = static_cast<float>(radius * radius);
|
|
const int nu = map.nu, nv = map.nv, nw = map.nw;
|
|
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);
|
|
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 g[3] = {0, 0, 0}, gu[6] = {0, 0, 0, 0, 0, 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;
|
|
if (yz + x * x > max_r2)
|
|
continue;
|
|
const double value = map.data[row + gemmi::modulo(u, nu)];
|
|
for (size_t q = 0; q < std::size(precal.a); q++) {
|
|
const gemmi::SMat33<float> &m = precal.b[q];
|
|
const gemmi::Vec3f r(x, y, z);
|
|
const double e = value * precal.a[q] * gemmi::unsafe_expapprox(std::max(m.r_u_r(r), -88.f));
|
|
const double mr[3] = {m.u11 * x + m.u12 * y + m.u13 * z, m.u12 * x + m.u22 * y + m.u23 * z,
|
|
m.u13 * x + m.u23 * y + m.u33 * z};
|
|
for (int k = 0; k < 3; k++)
|
|
g[k] += e * mr[k];
|
|
gu[0] += e * (m.u11 + 2 * mr[0] * mr[0]);
|
|
gu[1] += e * (m.u22 + 2 * mr[1] * mr[1]);
|
|
gu[2] += e * (m.u33 + 2 * mr[2] * mr[2]);
|
|
gu[3] += e * (m.u12 + 2 * mr[0] * mr[1]);
|
|
gu[4] += e * (m.u13 + 2 * mr[0] * mr[2]);
|
|
gu[5] += e * (m.u23 + 2 * mr[1] * mr[2]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
const double occ = atom.occ;
|
|
grad = {2 * occ * g[0], 2 * occ * g[1], 2 * occ * g[2], occ * (gu[0] + gu[1] + gu[2]) / gemmi::u_to_b()};
|
|
if (grad_u != nullptr)
|
|
*grad_u = {occ * gu[0], occ * gu[1], occ * gu[2], 2 * occ * gu[3], 2 * occ * gu[4], 2 * occ * gu[5]};
|
|
}
|
|
|
|
} // namespace
|
|
|
|
void AtomGradientsFromMap(const gemmi::Grid<float> &map, const gemmi::Model &model, double blur, size_t nthreads,
|
|
std::vector<std::array<double, 4>> &grad, std::vector<std::array<double, 6>> *grad_u) {
|
|
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});
|
|
if (grad_u != nullptr)
|
|
grad_u->assign(atoms.size(), {0, 0, 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];
|
|
if (atom.aniso.nonzero()) {
|
|
AnisoGradient(map, dc, atom, blur, grad[i], grad_u != nullptr ? &(*grad_u)[i] : nullptr);
|
|
continue;
|
|
}
|
|
const float b = static_cast<float>(atom.b_iso + blur);
|
|
const auto precal = IonFormFactor<float>(atom.element.elem, atom.charge).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};
|
|
}
|
|
});
|
|
}
|