leonarski_f
37c3661b33
An integrated rotation frame is fitted twice, each fit independent of the other frames. The first fit is free and sets the frame's residual scale. The second starts again from the sweep's geometry with a Gaussian prior on how far the orientation and the beam centre move. The prior's two widths are measured once per run, on the geometry pre-pass, as the rms of the seven-frame running median of the frames' free shifts. This is a pass-level input, like the pre-pass mosaicity. A beam centre that does not wander in the pre-pass is therefore held near the sweep, and the data decide this per run without a switch. The orientation prior acts on the rotation away from the sweep orientation and not on the angle-axis parameters. Applied to the parameters it was up to 2.4x too tight, depending on how the crystal sits in the laboratory. Both fits use XtalOptimizer's back-rotating form with the spindle given, so the beam-centre component along the spindle, which on a single frame nearly trades against the orientation, gets the same restraint it gets over a sweep. The frames the first pass scores lattices on, the pre-pass and stills keep the free fit. The residual scale counts 3N components and has no parameter correction (measured as immaterial). Squashes the sparse-eb prototype (2a54995fd, 8fb05679a, 51ae67467) onto twin-subgroup-check, simplified as described above. On sparse small molecules SHELXL R1 improved by a geometric mean of 3.943 -> 3.768 over 11 well-measured sets. Protein REFMAC R-free and anomalous peaks stayed within noise. The rotation-scale walk is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HVeDQsU5a8AJUdJPgWevAK (cherry picked from commit 045ae107d1c6b6ab6dc40773924426134dee69bf) (cherry picked from commit 1f3c11f61c693e6eb5debbbb881245c69efe7472)
353 lines
15 KiB
C++
353 lines
15 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "XtalRefine.h"
|
|
|
|
#include <array>
|
|
|
|
#include "Dual.h"
|
|
#include "../../common/ParallelFor.h"
|
|
|
|
namespace {
|
|
// Ambient layout of the parameter vector and the order of the blocks in it.
|
|
constexpr int OFF_BEAM = 0, OFF_ROT = 2, OFF_AXIS = 4, OFF_P0 = 7, OFF_LEN = 10, OFF_ANG = 13, N_AMBIENT = 16;
|
|
|
|
// Derivative lanes. The observed half of a residual depends on beam, detector angles and spindle
|
|
// (OBS lanes), the predicted half on orientation and cell (PRED lanes); each half is carried on a
|
|
// dual number of its own width and the two are put side by side only in the sums.
|
|
constexpr int OBS = 6, PRED = 9, LANES = OBS + PRED;
|
|
using DO = Dual<OBS>;
|
|
using DP = Dual<PRED>;
|
|
|
|
// The reduced problem: beam (2) and orientation (3) only.
|
|
constexpr int R_OBS = 2, R_PRED = 3, R_LANES = R_OBS + R_PRED;
|
|
|
|
// Sums over one block of residuals, in lane coordinates.
|
|
template<int L>
|
|
struct Sums {
|
|
double cost = 0.0;
|
|
double g[L] = {};
|
|
double H[L][L] = {}; // upper triangle
|
|
|
|
void Add(double r, const double *J, double w2) {
|
|
cost += 0.5 * w2 * r * r;
|
|
for (int i = 0; i < L; i++) {
|
|
const double wj = w2 * J[i];
|
|
g[i] += wj * r;
|
|
for (int j = i; j < L; j++)
|
|
H[i][j] += wj * J[j];
|
|
}
|
|
}
|
|
|
|
void Add(const Sums &o) {
|
|
cost += o.cost;
|
|
for (int i = 0; i < L; i++) {
|
|
g[i] += o.g[i];
|
|
for (int j = i; j < L; j++)
|
|
H[i][j] += o.H[i][j];
|
|
}
|
|
}
|
|
};
|
|
|
|
std::vector<LMBlock> MakeBlocks(const XtalRefineProblem &p) {
|
|
std::vector<LMBlock> blocks(6);
|
|
blocks[0].offset = OFF_BEAM;
|
|
blocks[0].size = 2;
|
|
blocks[0].constant = p.beam_constant;
|
|
|
|
blocks[1].offset = OFF_ROT;
|
|
blocks[1].size = 2;
|
|
blocks[1].constant = p.beam_and_orientation_only || p.detector_rot_constant;
|
|
for (int j = 0; j < 2; j++) {
|
|
blocks[1].lower[j] = p.detector_rot_lower[j];
|
|
blocks[1].upper[j] = p.detector_rot_upper[j];
|
|
}
|
|
|
|
blocks[2].offset = OFF_AXIS;
|
|
blocks[2].size = 3;
|
|
blocks[2].constant = p.beam_and_orientation_only || p.rot_vec_constant;
|
|
blocks[2].sphere = true;
|
|
|
|
blocks[3].offset = OFF_P0;
|
|
blocks[3].size = 3;
|
|
|
|
blocks[4].offset = OFF_LEN;
|
|
blocks[4].size = 3;
|
|
blocks[4].constant = p.beam_and_orientation_only || p.latt_vec1_constant;
|
|
|
|
blocks[5].offset = OFF_ANG;
|
|
blocks[5].size = 3;
|
|
blocks[5].constant = p.beam_and_orientation_only || p.latt_vec2_constant;
|
|
|
|
for (int j = 0; j < 3; j++) {
|
|
blocks[4].lower[j] = p.latt_vec1_lower[j];
|
|
blocks[4].upper[j] = p.latt_vec1_upper[j];
|
|
blocks[5].lower[j] = p.latt_vec2_lower[j];
|
|
blocks[5].upper[j] = p.latt_vec2_upper[j];
|
|
}
|
|
return blocks;
|
|
}
|
|
|
|
// Where each lane lands among the tangent coordinates of the free blocks; -1 for a held one.
|
|
template<size_t L>
|
|
std::array<int, L> LaneToTangent(const std::vector<LMBlock> &blocks, const std::array<int, L> &lane_block,
|
|
const std::array<int, L> &lane_index) {
|
|
std::array<int, L> map{};
|
|
std::array<int, 6> tangent_offset{};
|
|
int t = 0;
|
|
for (size_t b = 0; b < blocks.size(); b++) {
|
|
tangent_offset[b] = t;
|
|
t += blocks[b].TangentSize();
|
|
}
|
|
for (size_t l = 0; l < L; l++)
|
|
map[l] = blocks[lane_block[l]].constant ? -1 : tangent_offset[lane_block[l]] + lane_index[l];
|
|
return map;
|
|
}
|
|
|
|
template<int L>
|
|
void ToTangent(const Sums<L> &s, const std::array<int, L> &map, Eigen::VectorXd &g, Eigen::MatrixXd &H) {
|
|
for (int i = 0; i < L; i++) {
|
|
if (map[i] < 0)
|
|
continue;
|
|
g[map[i]] += s.g[i];
|
|
for (int j = i; j < L; j++) {
|
|
if (map[j] < 0)
|
|
continue;
|
|
H(map[i], map[j]) += s.H[i][j];
|
|
if (map[i] != map[j])
|
|
H(map[j], map[i]) += s.H[i][j];
|
|
}
|
|
}
|
|
}
|
|
|
|
void AddPriors(const XtalRefineProblem &p, const double *x, double &cost, Eigen::VectorXd *g,
|
|
Eigen::MatrixXd *H, int beam_tangent, int rot_tangent, int orientation_tangent) {
|
|
if (p.orientation_prior_weight > 0.0) {
|
|
const double w = p.orientation_prior_weight;
|
|
const double *J = p.orientation_prior_jacobian;
|
|
double d[3], r[3];
|
|
for (int k = 0; k < 3; k++)
|
|
d[k] = x[OFF_P0 + k] - p.orientation_prior[k];
|
|
for (int i = 0; i < 3; i++) {
|
|
r[i] = w * (J[3 * i] * d[0] + J[3 * i + 1] * d[1] + J[3 * i + 2] * d[2]);
|
|
cost += 0.5 * r[i] * r[i];
|
|
}
|
|
if (g && orientation_tangent >= 0)
|
|
for (int k = 0; k < 3; k++)
|
|
for (int i = 0; i < 3; i++) {
|
|
(*g)[orientation_tangent + k] += w * J[3 * i + k] * r[i];
|
|
for (int l = 0; l < 3; l++)
|
|
(*H)(orientation_tangent + k, orientation_tangent + l) += w * w * J[3 * i + k] * J[3 * i + l];
|
|
}
|
|
}
|
|
for (const auto &prior: p.priors) {
|
|
const bool on_beam = prior.block == XtalRefinePrior::Block::Beam;
|
|
const double *v = x + (on_beam ? OFF_BEAM : OFF_ROT);
|
|
const double r = prior.weight * (prior.gx * v[0] + prior.gy * v[1] - prior.p0);
|
|
cost += 0.5 * r * r;
|
|
const int t = on_beam ? beam_tangent : rot_tangent;
|
|
if (!g || t < 0)
|
|
continue;
|
|
const double J[2] = {prior.weight * prior.gx, prior.weight * prior.gy};
|
|
for (int i = 0; i < 2; i++) {
|
|
(*g)[t + i] += J[i] * r;
|
|
for (int j = 0; j < 2; j++)
|
|
(*H)(t + i, t + j) += J[i] * J[j];
|
|
}
|
|
}
|
|
}
|
|
|
|
template<class D>
|
|
D Seed(double value, int lane, bool free) {
|
|
return free ? D::Variable(value, lane) : D(value);
|
|
}
|
|
|
|
// Residual blocks of at least this many residuals; the cut depends on the count alone.
|
|
constexpr int MIN_RESIDUALS_PER_BLOCK = 256;
|
|
|
|
template<int L, class Fn>
|
|
Sums<L> SumResiduals(const XtalRefineProblem &p, int num_threads, Fn &&residual) {
|
|
const int n = static_cast<int>(p.residuals.size());
|
|
std::vector<Sums<L>> partial(ReductionBlocks(n, MIN_RESIDUALS_PER_BLOCK));
|
|
ParallelBlocks(n, std::max(1, num_threads), [&](int b, int lo, int hi) {
|
|
for (int i = lo; i < hi; i++)
|
|
residual(i, partial[b]);
|
|
}, MIN_RESIDUALS_PER_BLOCK);
|
|
Sums<L> total;
|
|
for (const auto &s: partial)
|
|
total.Add(s);
|
|
return total;
|
|
}
|
|
|
|
double WeightSq(const XtalRefineProblem &p, int i) {
|
|
return p.weight_sq.empty() ? 1.0 : p.weight_sq[i];
|
|
}
|
|
|
|
bool AllFinite(double cost, const Eigen::VectorXd *g, const Eigen::MatrixXd *H) {
|
|
return std::isfinite(cost) && (!g || (g->allFinite() && H->allFinite()));
|
|
}
|
|
|
|
// Seven-block residual (XtalResidualFixedDistance), with every block that depends on parameters
|
|
// alone - detector-angle sines and cosines, the spindle back-rotation of each frame, the reciprocal
|
|
// basis of the cell, the orientation's rotation - worked out once per evaluation.
|
|
bool EvaluateGeneral(const XtalRefineProblem &p, const std::array<int, LANES> &map, int num_threads,
|
|
const double *x, double &cost, Eigen::VectorXd *g, Eigen::MatrixXd *H) {
|
|
const bool beam_free = map[0] >= 0, rot_free = map[2] >= 0, axis_free = map[4] >= 0;
|
|
const bool p0_free = map[OBS] >= 0, len_free = map[OBS + 3] >= 0, ang_free = map[OBS + 6] >= 0;
|
|
|
|
const DO beam[2] = {Seed<DO>(x[OFF_BEAM], 0, beam_free),
|
|
Seed<DO>(x[OFF_BEAM + 1], 1, beam_free)};
|
|
const DO rot1 = Seed<DO>(x[OFF_ROT], 2, rot_free);
|
|
const DO rot2 = Seed<DO>(x[OFF_ROT + 1], 3, rot_free);
|
|
const DO c1 = cos(rot1), s1 = sin(rot1), c2 = cos(rot2), s2 = sin(rot2);
|
|
|
|
DO axis[3] = {x[OFF_AXIS], x[OFF_AXIS + 1], x[OFF_AXIS + 2]};
|
|
if (axis_free) {
|
|
double J[3][2];
|
|
SpherePlusJacobian3(x + OFF_AXIS, J);
|
|
for (int k = 0; k < 3; k++) {
|
|
axis[k].v[4] = J[k][0];
|
|
axis[k].v[5] = J[k][1];
|
|
}
|
|
}
|
|
std::vector<AngleAxisRotator<DO>> rot_back;
|
|
rot_back.reserve(p.frame_angle_rad.size());
|
|
for (const double angle: p.frame_angle_rad) {
|
|
const DO aa_back[3] = {angle * axis[0], angle * axis[1], angle * axis[2]};
|
|
rot_back.emplace_back(aa_back);
|
|
}
|
|
|
|
DP p0[3], len[3], ang[3];
|
|
for (int k = 0; k < 3; k++) {
|
|
p0[k] = Seed<DP>(x[OFF_P0 + k], k, p0_free);
|
|
len[k] = Seed<DP>(x[OFF_LEN + k], 3 + k, len_free);
|
|
ang[k] = Seed<DP>(x[OFF_ANG + k], 6 + k, ang_free);
|
|
}
|
|
Eigen::Matrix<DP, 3, 1> bxc, cxa, axb;
|
|
DP invV;
|
|
XtalResidual::ReciprocalBasis(len, ang, p.crystal_system, bxc, cxa, axb, invV);
|
|
const AngleAxisRotator<DP> rot_p0(p0);
|
|
|
|
const Sums<LANES> s = SumResiduals<LANES>(p, num_threads, [&](int i, Sums<LANES> &acc) {
|
|
const XtalResidual &res = p.residuals[i];
|
|
DO obs[3];
|
|
res.ObservedRecipCore(beam, p.distance_mm, c1, s1, c2, s2, rot_back[p.frame[i]], obs);
|
|
DP unrot[3], pred[3];
|
|
res.CombineRecipUnrot(bxc, cxa, axb, invV, unrot);
|
|
rot_p0.Rotate(unrot, pred);
|
|
const double w2 = WeightSq(p, i);
|
|
for (int k = 0; k < 3; k++) {
|
|
double J[LANES];
|
|
for (int l = 0; l < OBS; l++)
|
|
J[l] = obs[k].v[l];
|
|
for (int l = 0; l < PRED; l++)
|
|
J[OBS + l] = -pred[k].v[l];
|
|
acc.Add(obs[k].a - pred[k].a, J, w2);
|
|
}
|
|
});
|
|
|
|
cost = s.cost;
|
|
if (g)
|
|
ToTangent<LANES>(s, map, *g, *H);
|
|
AddPriors(p, x, cost, g, H, map[0], map[2], map[OBS]);
|
|
return AllFinite(cost, g, H);
|
|
}
|
|
|
|
// The reduced residual (XtalResidualBeamOrientation): detector, spindle and cell held, so the
|
|
// back-rotation of each frame and the unrotated prediction of each residual are constants.
|
|
bool EvaluateBeamOrientation(const XtalRefineProblem &p, const std::array<int, R_LANES> &map,
|
|
const std::vector<AngleAxisRotator<double>> &rot_back,
|
|
const std::vector<std::array<double, 3>> &unrot, double c1, double s1,
|
|
double c2, double s2, int num_threads,
|
|
const double *x, double &cost, Eigen::VectorXd *g, Eigen::MatrixXd *H) {
|
|
using DB = Dual<R_OBS>;
|
|
using DR = Dual<R_PRED>;
|
|
const bool beam_free = map[0] >= 0;
|
|
const DB beam[2] = {beam_free ? DB::Variable(x[OFF_BEAM], 0) : DB(x[OFF_BEAM]),
|
|
beam_free ? DB::Variable(x[OFF_BEAM + 1], 1) : DB(x[OFF_BEAM + 1])};
|
|
const DR p0[3] = {DR::Variable(x[OFF_P0], 0), DR::Variable(x[OFF_P0 + 1], 1),
|
|
DR::Variable(x[OFF_P0 + 2], 2)};
|
|
const AngleAxisRotator<DR> rot_p0(p0);
|
|
|
|
const Sums<R_LANES> s = SumResiduals<R_LANES>(p, num_threads, [&](int i, Sums<R_LANES> &acc) {
|
|
const XtalResidual &res = p.residuals[i];
|
|
DB obs[3];
|
|
res.ObservedRecipCore(beam, p.distance_mm, c1, s1, c2, s2, rot_back[p.frame[i]], obs);
|
|
DR pred[3];
|
|
rot_p0.Rotate(unrot[i].data(), pred);
|
|
const double w2 = WeightSq(p, i);
|
|
for (int k = 0; k < 3; k++) {
|
|
double J[R_LANES];
|
|
for (int l = 0; l < R_OBS; l++)
|
|
J[l] = obs[k].v[l];
|
|
for (int l = 0; l < R_PRED; l++)
|
|
J[R_OBS + l] = -pred[k].v[l];
|
|
acc.Add(obs[k].a - pred[k].a, J, w2);
|
|
}
|
|
});
|
|
|
|
cost = s.cost;
|
|
if (g)
|
|
ToTangent<R_LANES>(s, map, *g, *H);
|
|
AddPriors(p, x, cost, g, H, map[0], -1, map[R_OBS]);
|
|
return AllFinite(cost, g, H);
|
|
}
|
|
}
|
|
|
|
LMSummary SolveXtalRefine(XtalRefineProblem &p, int num_threads) {
|
|
const std::vector<LMBlock> blocks = MakeBlocks(p);
|
|
|
|
std::vector<double> x(N_AMBIENT);
|
|
const auto put = [&](int off, const double *v, int n) { for (int i = 0; i < n; i++) x[off + i] = v[i]; };
|
|
put(OFF_BEAM, p.beam, 2);
|
|
put(OFF_ROT, p.detector_rot, 2);
|
|
put(OFF_AXIS, p.rot_vec, 3);
|
|
put(OFF_P0, p.latt_vec0, 3);
|
|
put(OFF_LEN, p.latt_vec1, 3);
|
|
put(OFF_ANG, p.latt_vec2, 3);
|
|
|
|
LMSummary summary;
|
|
if (p.beam_and_orientation_only) {
|
|
const std::array<int, R_LANES> lane_block = {0, 0, 3, 3, 3};
|
|
const std::array<int, R_LANES> lane_index = {0, 1, 0, 1, 2};
|
|
const auto map = LaneToTangent<R_LANES>(blocks, lane_block, lane_index);
|
|
|
|
std::vector<AngleAxisRotator<double>> rot_back;
|
|
rot_back.reserve(p.frame_angle_rad.size());
|
|
for (const double angle: p.frame_angle_rad)
|
|
rot_back.push_back(XtalFrameConstants::BackRotator(angle, p.rot_vec));
|
|
Eigen::Matrix<double, 3, 1> bxc, cxa, axb;
|
|
double invV;
|
|
XtalResidual::ReciprocalBasis(p.latt_vec1, p.latt_vec2, p.crystal_system, bxc, cxa, axb, invV);
|
|
std::vector<std::array<double, 3>> unrot(p.residuals.size());
|
|
for (size_t i = 0; i < p.residuals.size(); i++)
|
|
p.residuals[i].CombineRecipUnrot(bxc, cxa, axb, invV, unrot[i].data());
|
|
const double c1 = std::cos(p.detector_rot[0]), s1 = std::sin(p.detector_rot[0]);
|
|
const double c2 = std::cos(p.detector_rot[1]), s2 = std::sin(p.detector_rot[1]);
|
|
|
|
summary = SolveLM(x, blocks, p.options, [&](const double *at, double &cost, Eigen::VectorXd *g,
|
|
Eigen::MatrixXd *H) {
|
|
return EvaluateBeamOrientation(p, map, rot_back, unrot, c1, s1, c2, s2, num_threads, at, cost, g, H);
|
|
});
|
|
} else {
|
|
const std::array<int, LANES> lane_block = {0, 0, 1, 1, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5};
|
|
const std::array<int, LANES> lane_index = {0, 1, 0, 1, 0, 1, 0, 1, 2, 0, 1, 2, 0, 1, 2};
|
|
const auto map = LaneToTangent<LANES>(blocks, lane_block, lane_index);
|
|
summary = SolveLM(x, blocks, p.options, [&](const double *at, double &cost, Eigen::VectorXd *g,
|
|
Eigen::MatrixXd *H) {
|
|
return EvaluateGeneral(p, map, num_threads, at, cost, g, H);
|
|
});
|
|
}
|
|
|
|
if (summary.IsSolutionUsable()) {
|
|
const auto get = [&](int off, double *v, int n) { for (int i = 0; i < n; i++) v[i] = x[off + i]; };
|
|
get(OFF_BEAM, p.beam, 2);
|
|
get(OFF_ROT, p.detector_rot, 2);
|
|
get(OFF_AXIS, p.rot_vec, 3);
|
|
get(OFF_P0, p.latt_vec0, 3);
|
|
get(OFF_LEN, p.latt_vec1, 3);
|
|
get(OFF_ANG, p.latt_vec2, 3);
|
|
}
|
|
return summary;
|
|
}
|