leonarski_f
46ad2c10ed
The score's robust fine pass (pseudo-Huber IRLS of ModelScaleBox::robust) on the device, so that refinement on a GPU no longer fits its scale on the host every cycle - with the trials side by side that was a robust grid search on one thread per slot (13-133 s per refinement on the seven timing sets, against 1.7-8 s before the score commit). - ModelScaleGPU: the weights multiply the WSSR and J^T J / J^T r in scale_sums (a weight row per fit, reserved at the first robust fit); robust_residuals / robust_weights kernels; the median of each fit's first residuals taken on the host as FitModelScale takes it (nth_element at n/2, on the fits' threads); LevMarBatch split out of FitBatch so that the rounds restart from where a fit stands. FitSolvent(..., robust) refits the coarse winner in the fine pass, as FitModelScale's robust fit does; the plain fit is unchanged. - ModelScaleGPUPool::Fit(..., robust); its streams at the highest priority: a fit is a chain of small launches the host waits on, and on a card shared with the trials' long kernels each wait queued behind them (refit.scale 58-87 ms -> 14-45 ms per cycle, the result the same to the bit). - Refinement: the main context and every trial slot fit on their pool (robust); the CPU engine keeps FitModelScale's robust fit. The card's fit is FitModelScale's to rounding (float residuals and weights, the steps' anisotropic factor in float), deterministic run to run and whichever engine a fit takes: R-work/R-free of the 14 parity sets on the GPU equal the CPU engine's to 4 digits (largest difference 1.4e-6 on the 3.6 A set). Test: ModelScaleGPUPool_RobustMatchesFitModelScale (one reflection in fifty a 4x outlier, five cells). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit d15f048cbd7ed2d6b430f9f9cb5018ca284da8f0)
618 lines
27 KiB
Plaintext
618 lines
27 KiB
Plaintext
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "ModelScaleGPU.h"
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <utility>
|
|
|
|
#include "../../common/JFJochException.h"
|
|
#include "../../common/ParallelFor.h"
|
|
#include "RigidBodyGPUEngine.h" // ModelScaleGPUTooFewReflections
|
|
|
|
namespace {
|
|
|
|
void cuda_err(cudaError_t val) {
|
|
if (val != cudaSuccess)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
|
|
}
|
|
|
|
// A fixed launch shape, so every point is summed by the same thread in the same order whatever the
|
|
// card or the stream: the result is the same bit for bit run to run.
|
|
constexpr int THREADS = 128;
|
|
constexpr int BLOCKS = 48;
|
|
constexpr int SLOTS = 36; // 1 + 28 + 7: the WSSR, the lower half of J^T J and J^T r for 7 parameters
|
|
constexpr int MAX_FITS = 88; // FitModelScale's coarse pass, the larger of its two
|
|
constexpr int MAX_ROBUST_FITS = 32; // FitModelScale's fine pass, at most 5 x 5 points
|
|
constexpr int ROBUST_ROUNDS = 4; // as ModelScaling.cpp
|
|
|
|
// What a launch sums over the points for one fit.
|
|
constexpr int MODE_ISOTROPIC = 0; // fit_isotropic_b_approximately(): sx, sy, sxx, sxy, n
|
|
constexpr int MODE_WSSR = 1; // compute_wssr()
|
|
constexpr int MODE_MATRICES = 2; // compute_lm_matrices(): WSSR, J^T J, J^T r
|
|
constexpr int MODE_R = 3; // FitModelScale's R: sum |fobs - value|, sum fobs
|
|
|
|
struct Constraints {
|
|
float row[6][6];
|
|
};
|
|
|
|
// |Fcalc + k_sol exp(-b_sol s^2) Fmask| of every point for each fit, in gemmi's arithmetic: the solvent
|
|
// scale cast to float and the sum and its modulus in float, as a complex<float> is. The solvent pair is
|
|
// fixed for the whole of a fit, so this is done once per fit rather than at every step.
|
|
__global__ void solvent_amplitudes(const double *__restrict__ stol2, const float2 *__restrict__ fcmol,
|
|
const float2 *__restrict__ fmask, int n,
|
|
const ModelScaleFitState *__restrict__ fits, float *__restrict__ f_abs) {
|
|
const ModelScaleFitState f = fits[blockIdx.y];
|
|
for (int i = blockIdx.x * THREADS + threadIdx.x; i < n; i += BLOCKS * THREADS) {
|
|
const float solvent = static_cast<float>(f.k_sol * exp(-f.b_sol * stol2[i]));
|
|
const float2 fc = fcmol[i], fm = fmask[i];
|
|
f_abs[static_cast<size_t>(f.column) * n + i] = hypotf(fc.x + solvent * fm.x, fc.y + solvent * fm.y);
|
|
}
|
|
}
|
|
|
|
// The sums of gemmi Scaling<float> for one fit (blockIdx.y) over the points.
|
|
template <int NA>
|
|
__global__ void scale_sums(const int *__restrict__ hkl, const double *__restrict__ stol2,
|
|
const float *__restrict__ fobs, const unsigned char *__restrict__ strong,
|
|
const float *__restrict__ f_abs_all, const float *__restrict__ weight_all, int n,
|
|
const ModelScaleFitState *__restrict__ fits, Constraints c, double *__restrict__ partial) {
|
|
constexpr int NS = 1 + NA * (NA + 1) / 2 + NA; // WSSR, J^T J lower half, J^T r
|
|
const ModelScaleFitState f = fits[blockIdx.y];
|
|
const float *f_abs_fit = f_abs_all + static_cast<size_t>(f.column) * n;
|
|
float uf[6];
|
|
for (int k = 0; k < 6; k++)
|
|
uf[k] = static_cast<float>(f.b_star[k]);
|
|
double acc[NS];
|
|
for (int s = 0; s < NS; s++)
|
|
acc[s] = 0;
|
|
|
|
for (int i = blockIdx.x * THREADS + threadIdx.x; i < n; i += BLOCKS * THREADS) {
|
|
const float f_abs = f_abs_fit[i];
|
|
const double fo = fobs[i];
|
|
|
|
if (f.mode == MODE_ISOTROPIC) {
|
|
// fit_isotropic_b_approximately() (local change in gemmi_gph): no log ratio at |F| = 0
|
|
if (!strong[i] || f_abs == 0)
|
|
continue;
|
|
const double x = stol2[i];
|
|
const double y = logf(static_cast<float>(fo / f_abs));
|
|
acc[0] += x;
|
|
acc[1] += y;
|
|
acc[2] += x * x;
|
|
acc[3] += x * y;
|
|
acc[4] += 1;
|
|
continue;
|
|
}
|
|
|
|
if (f.mode == MODE_R) {
|
|
const double hx = hkl[3 * i], hy = hkl[3 * i + 1], hz = hkl[3 * i + 2];
|
|
const double *u = f.b_star;
|
|
const double r_u_r = hx * hx * u[0] + hy * hy * u[1] + hz * hz * u[2]
|
|
+ 2 * (hx * hy * u[3] + hx * hz * u[4] + hy * hz * u[5]);
|
|
// compute_value(): |F| times the scale cast to float, in float
|
|
const float value = f_abs * static_cast<float>(f.k_overall * exp(-0.25 * r_u_r));
|
|
acc[0] += fabs(fo - value);
|
|
acc[1] += fo;
|
|
continue;
|
|
}
|
|
|
|
// Below, the anisotropic factor is taken in float rather than in double as gemmi takes it: double
|
|
// arithmetic is most of the cost of a step on a card with little double throughput. In
|
|
// compute_value() gemmi casts the scale to float anyway, so the WSSR moves by the last bit of some
|
|
// points; the derivatives only set the direction of the next step, which the WSSR then accepts or
|
|
// not.
|
|
const float fx = hkl[3 * i], fy = hkl[3 * i + 1], fz = hkl[3 * i + 2];
|
|
const float r_u_r_f = fx * fx * uf[0] + fy * fy * uf[1] + fz * fz * uf[2]
|
|
+ 2 * (fx * fy * uf[3] + fx * fz * uf[4] + fy * fz * uf[5]);
|
|
const float k_aniso = expf(-0.25f * r_u_r_f);
|
|
// the robust fit's weight (ModelScaling.cpp, ComputeLMMatrices and ComputeWSSR multiply by it)
|
|
const double w = f.weighted ? weight_all[static_cast<size_t>(f.column) * n + i] : 1.0;
|
|
if (f.mode == MODE_WSSR) {
|
|
const float value = f_abs * static_cast<float>(f.k_overall * k_aniso);
|
|
const double dy = fo - value;
|
|
acc[0] += dy * dy * w;
|
|
continue;
|
|
}
|
|
|
|
// compute_value_and_derivatives() with k_sol and b_sol fixed, per point in float, summed in double
|
|
const float fe = f_abs * k_aniso;
|
|
const float y = static_cast<float>(f.k_overall) * fe;
|
|
const float du[6] = {-0.25f * y * (fx * fx), -0.25f * y * (fy * fy), -0.25f * y * (fz * fz),
|
|
-0.5f * y * (fx * fy), -0.5f * y * (fx * fz), -0.5f * y * (fy * fz)};
|
|
double dy_da[NA];
|
|
dy_da[0] = fe;
|
|
for (int j = 1; j < NA; j++) {
|
|
const float *r = c.row[j - 1];
|
|
dy_da[j] = r[0] * du[0] + r[1] * du[1] + r[2] * du[2] + r[3] * du[3] + r[4] * du[4] + r[5] * du[5];
|
|
}
|
|
const double dy = fo - y;
|
|
acc[0] += dy * dy * w;
|
|
int s = 1;
|
|
for (int j = 0; j < NA; j++)
|
|
for (int k = 0; k <= j; k++)
|
|
acc[s++] += dy_da[j] * dy_da[k] * w;
|
|
for (int j = 0; j < NA; j++)
|
|
acc[s++] += dy * dy_da[j] * w;
|
|
}
|
|
|
|
// The block's sum of each slot the mode fills: a fixed shuffle tree within each warp, then the warps
|
|
// in order.
|
|
const int used = f.mode == MODE_MATRICES ? NS : f.mode == MODE_ISOTROPIC ? 5 : f.mode == MODE_R ? 2 : 1;
|
|
__shared__ double warp_sum[THREADS / 32][NS];
|
|
const int lane = threadIdx.x % 32, warp = threadIdx.x / 32;
|
|
for (int s = 0; s < used; s++) {
|
|
double v = acc[s];
|
|
for (int offset = 16; offset > 0; offset /= 2)
|
|
v += __shfl_down_sync(0xffffffffu, v, offset);
|
|
if (lane == 0)
|
|
warp_sum[warp][s] = v;
|
|
}
|
|
__syncthreads();
|
|
if (threadIdx.x < SLOTS) {
|
|
double t = 0;
|
|
if (threadIdx.x < used)
|
|
for (int w = 0; w < THREADS / 32; w++)
|
|
t += warp_sum[w][threadIdx.x];
|
|
partial[(static_cast<size_t>(blockIdx.y) * BLOCKS + blockIdx.x) * SLOTS + threadIdx.x] = t;
|
|
}
|
|
}
|
|
|
|
// The robust fit's residuals |Fo - value| of each fit (blockIdx.y) at its parameters, compute_value() as in MODE_R.
|
|
__global__ void robust_residuals(const int *__restrict__ hkl, const float *__restrict__ fobs,
|
|
const float *__restrict__ f_abs_all, int n,
|
|
const ModelScaleFitState *__restrict__ fits, float *__restrict__ residual) {
|
|
const ModelScaleFitState f = fits[blockIdx.y];
|
|
const float *f_abs_fit = f_abs_all + static_cast<size_t>(f.column) * n;
|
|
for (int i = blockIdx.x * THREADS + threadIdx.x; i < n; i += BLOCKS * THREADS) {
|
|
const double hx = hkl[3 * i], hy = hkl[3 * i + 1], hz = hkl[3 * i + 2];
|
|
const double *u = f.b_star;
|
|
const double r_u_r = hx * hx * u[0] + hy * hy * u[1] + hz * hz * u[2]
|
|
+ 2 * (hx * hy * u[3] + hx * hz * u[4] + hy * hz * u[5]);
|
|
const float value = f_abs_fit[i] * static_cast<float>(f.k_overall * exp(-0.25 * r_u_r));
|
|
residual[static_cast<size_t>(f.column) * n + i] = fabsf(fobs[i] - value);
|
|
}
|
|
}
|
|
|
|
// The robust fit's weights 1 / sqrt(1 + (r / c)^2) from the residuals (ModelScaling.cpp, RobustReweight); a fit
|
|
// whose c is not positive keeps its weights.
|
|
__global__ void robust_weights(const float *__restrict__ residual, int n, const ModelScaleFitState *__restrict__ fits,
|
|
float *__restrict__ weight) {
|
|
const ModelScaleFitState f = fits[blockIdx.y];
|
|
if (!(f.c > 0))
|
|
return;
|
|
const size_t row = static_cast<size_t>(f.column) * n;
|
|
for (int i = blockIdx.x * THREADS + threadIdx.x; i < n; i += BLOCKS * THREADS) {
|
|
const double q = residual[row + i] / f.c;
|
|
weight[row + i] = static_cast<float>(1.0 / sqrt(1.0 + q * q));
|
|
}
|
|
}
|
|
|
|
// Each fit's block partials summed in block order.
|
|
__global__ void sum_blocks(const double *__restrict__ partial, double *__restrict__ sums) {
|
|
const int s = threadIdx.x;
|
|
double t = 0;
|
|
for (int b = 0; b < BLOCKS; b++)
|
|
t += partial[(static_cast<size_t>(blockIdx.x) * BLOCKS + b) * SLOTS + s];
|
|
sums[blockIdx.x * SLOTS + s] = t;
|
|
}
|
|
|
|
// gemmi's jordan_solve (levmar.hpp): Gauss-Jordan elimination with partial pivoting, x returned in b.
|
|
void JordanSolve(double *a, double *b, int n) {
|
|
for (int i = 0; i < n; i++) {
|
|
int maxnr = -1;
|
|
double amax = 0;
|
|
for (int j = i; j < n; j++) {
|
|
const double aji = std::fabs(a[n * j + i]);
|
|
if (aji > amax) {
|
|
maxnr = j;
|
|
amax = aji;
|
|
}
|
|
}
|
|
if (maxnr == -1) {
|
|
for (int j = i; j < n; j++)
|
|
if (a[n * i + j] != 0. || b[i] != 0.)
|
|
throw std::runtime_error("Trying to reverse singular matrix. Column " + std::to_string(i) +
|
|
" is zeroed.");
|
|
continue;
|
|
}
|
|
if (maxnr != i) {
|
|
for (int j = i; j < n; j++)
|
|
std::swap(a[n * maxnr + j], a[n * i + j]);
|
|
std::swap(b[i], b[maxnr]);
|
|
}
|
|
const double c = 1.0 / a[i * n + i];
|
|
for (int j = i; j < n; j++)
|
|
a[i * n + j] *= c;
|
|
b[i] *= c;
|
|
for (int k = 0; k < n; k++)
|
|
if (k != i) {
|
|
const double d = a[k * n + i];
|
|
for (int j = i; j < n; j++)
|
|
a[k * n + j] -= a[i * n + j] * d;
|
|
b[k] -= b[i] * d;
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
// gemmi's LevMar::fit (levmar.hpp), unrolled into the requests it makes of the target, so that many
|
|
// fits can share one launch per step. Each call of compute_lm_matrices() or compute_wssr() there is one
|
|
// request here; everything between two of them is the same code in the same order.
|
|
struct ModelScaleGPU::LevMarRun {
|
|
// gemmi's defaults
|
|
const int eval_limit = 100;
|
|
const double lambda_limit = 1e+15;
|
|
const double stop_rel_change = 1e-5;
|
|
const double lambda_up_factor = 10;
|
|
const double lambda_down_factor = 0.1;
|
|
|
|
double lambda = 0.001;
|
|
std::vector<double> initial_a, best_a, trial;
|
|
std::vector<double> alpha, beta;
|
|
double initial_wssr = NAN, wssr = NAN;
|
|
int eval_count = 0;
|
|
int small_change_counter = 0;
|
|
bool done = false;
|
|
|
|
int mode = MODE_MATRICES; // the next request, at `at`
|
|
std::vector<double> at;
|
|
|
|
explicit LevMarRun(const std::vector<double> &a) : initial_a(a), best_a(a), at(a) {
|
|
alpha.resize(a.size() * a.size());
|
|
beta.resize(a.size());
|
|
}
|
|
|
|
void Take(const double *sums) {
|
|
const size_t na = beta.size();
|
|
if (mode == MODE_MATRICES) {
|
|
int s = 1;
|
|
for (size_t j = 0; j < na; j++)
|
|
for (size_t k = 0; k <= j; k++)
|
|
alpha[na * j + k] = alpha[na * k + j] = sums[s++];
|
|
for (size_t j = 0; j < na; j++)
|
|
beta[j] = sums[s++];
|
|
if (eval_count == 0) {
|
|
initial_wssr = wssr = sums[0];
|
|
eval_count = 1;
|
|
} else {
|
|
++eval_count;
|
|
lambda *= lambda_down_factor;
|
|
}
|
|
Propose();
|
|
return;
|
|
}
|
|
const double new_wssr = sums[0];
|
|
++eval_count;
|
|
if (new_wssr < wssr) {
|
|
const double rel_change = (wssr - new_wssr) / wssr;
|
|
wssr = new_wssr;
|
|
best_a = trial;
|
|
if (wssr == 0) {
|
|
done = true;
|
|
return;
|
|
}
|
|
if (rel_change < stop_rel_change) {
|
|
if (++small_change_counter >= 2) {
|
|
done = true;
|
|
return;
|
|
}
|
|
} else {
|
|
small_change_counter = 0;
|
|
}
|
|
mode = MODE_MATRICES;
|
|
at = best_a;
|
|
} else {
|
|
if (lambda > lambda_limit) {
|
|
done = true;
|
|
return;
|
|
}
|
|
lambda *= lambda_up_factor;
|
|
Propose();
|
|
}
|
|
}
|
|
|
|
// The damped normal equations J^T J + lambda diag(J^T J), solved for the next trial.
|
|
void Propose() {
|
|
if (eval_limit > 0 && eval_count >= eval_limit) {
|
|
done = true;
|
|
return;
|
|
}
|
|
const size_t na = beta.size();
|
|
std::vector<double> temp_alpha = alpha;
|
|
for (size_t j = 0; j < na; j++)
|
|
temp_alpha[na * j + j] *= (1.0 + lambda);
|
|
trial = beta;
|
|
JordanSolve(temp_alpha.data(), trial.data(), static_cast<int>(na));
|
|
for (size_t i = 0; i < na; i++)
|
|
trial[i] += best_a[i];
|
|
mode = MODE_WSSR;
|
|
at = trial;
|
|
}
|
|
|
|
const std::vector<double> &Result() const {
|
|
return wssr < initial_wssr ? best_a : initial_a;
|
|
}
|
|
};
|
|
|
|
size_t ModelScaleGPU::RobustBytes(size_t max_points) {
|
|
return 2 * MAX_ROBUST_FITS * max_points * sizeof(float);
|
|
}
|
|
|
|
size_t ModelScaleGPU::DeviceBytes(size_t max_points) {
|
|
return max_points * (3 * sizeof(int) + sizeof(double) + sizeof(float) + sizeof(unsigned char))
|
|
+ MAX_FITS * max_points * sizeof(float)
|
|
+ MAX_FITS * (sizeof(ModelScaleFitState) + (BLOCKS + 1) * SLOTS * sizeof(double));
|
|
}
|
|
|
|
ModelScaleGPU::ModelScaleGPU(cudaStream_t stream, size_t max_points)
|
|
: stream_(stream), max_points_(max_points),
|
|
hkl_(3 * max_points), stol2_(max_points), fobs_(max_points), strong_(max_points),
|
|
f_abs_(MAX_FITS * max_points), fits_(MAX_FITS), partial_(static_cast<size_t>(MAX_FITS) * BLOCKS * SLOTS),
|
|
sums_(static_cast<size_t>(MAX_FITS) * SLOTS),
|
|
host_fits_(MAX_FITS), host_sums_(static_cast<size_t>(MAX_FITS) * SLOTS) {}
|
|
|
|
void ModelScaleGPU::SetPoints(const std::vector<std::array<int, 3>> &hkl,
|
|
const std::vector<double> &stol2,
|
|
const std::vector<float> &fobs,
|
|
const std::vector<float> &sigma,
|
|
const std::vector<std::array<double, 6>> &constraints,
|
|
const double frac[9]) {
|
|
if (hkl.size() > max_points_ || stol2.size() != hkl.size() || fobs.size() != hkl.size() ||
|
|
sigma.size() != hkl.size() || constraints.empty() || constraints.size() > 6)
|
|
throw std::invalid_argument("ModelScaleGPU::SetPoints: inconsistent points");
|
|
n_ = static_cast<int>(hkl.size());
|
|
n_params_ = 1 + static_cast<int>(constraints.size());
|
|
for (size_t j = 0; j < constraints.size(); j++)
|
|
for (int k = 0; k < 6; k++)
|
|
constraints_[j][k] = constraints[j][k];
|
|
std::copy(frac, frac + 9, frac_);
|
|
|
|
// fit_isotropic_b_approximately()'s filter: it skips weak reflections
|
|
std::vector<unsigned char> strong(n_);
|
|
n_strong_ = 0;
|
|
for (int i = 0; i < n_; i++) {
|
|
strong[i] = !(fobs[i] < 1 || fobs[i] < sigma[i]);
|
|
n_strong_ += strong[i];
|
|
}
|
|
if (n_ == 0)
|
|
return;
|
|
cuda_err(cudaMemcpyAsync(hkl_, hkl.data(), n_ * 3 * sizeof(int), cudaMemcpyHostToDevice, stream_));
|
|
cuda_err(cudaMemcpyAsync(stol2_, stol2.data(), n_ * sizeof(double), cudaMemcpyHostToDevice, stream_));
|
|
cuda_err(cudaMemcpyAsync(fobs_, fobs.data(), n_ * sizeof(float), cudaMemcpyHostToDevice, stream_));
|
|
cuda_err(cudaMemcpyAsync(strong_, strong.data(), n_, cudaMemcpyHostToDevice, stream_));
|
|
cuda_err(cudaStreamSynchronize(stream_)); // the host vectors go out of scope
|
|
}
|
|
|
|
void ModelScaleGPU::UploadFits(const std::vector<ModelScaleFitState> &fits) {
|
|
if (fits.size() > static_cast<size_t>(MAX_FITS))
|
|
throw std::logic_error("ModelScaleGPU: too many fits in one launch");
|
|
std::copy(fits.begin(), fits.end(), host_fits_.get());
|
|
cuda_err(cudaMemcpyAsync(fits_, host_fits_, fits.size() * sizeof(ModelScaleFitState), cudaMemcpyHostToDevice,
|
|
stream_));
|
|
}
|
|
|
|
const double *ModelScaleGPU::Sums(const std::vector<ModelScaleFitState> &fits) {
|
|
UploadFits(fits);
|
|
return Reduce(static_cast<int>(fits.size()));
|
|
}
|
|
|
|
const double *ModelScaleGPU::Reduce(int nfits) {
|
|
Constraints c{};
|
|
for (int j = 0; j < 6; j++)
|
|
for (int k = 0; k < 6; k++)
|
|
c.row[j][k] = static_cast<float>(constraints_[j][k]);
|
|
const dim3 grid(BLOCKS, nfits);
|
|
switch (n_params_) {
|
|
case 2: scale_sums<2><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, weight_, n_, fits_, c, partial_); break;
|
|
case 3: scale_sums<3><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, weight_, n_, fits_, c, partial_); break;
|
|
case 4: scale_sums<4><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, weight_, n_, fits_, c, partial_); break;
|
|
case 5: scale_sums<5><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, weight_, n_, fits_, c, partial_); break;
|
|
case 6: scale_sums<6><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, weight_, n_, fits_, c, partial_); break;
|
|
case 7: scale_sums<7><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, weight_, n_, fits_, c, partial_); break;
|
|
default: throw std::logic_error("ModelScaleGPU: unexpected number of parameters");
|
|
}
|
|
cuda_err(cudaGetLastError());
|
|
sum_blocks<<<nfits, SLOTS, 0, stream_>>>(partial_, sums_);
|
|
cuda_err(cudaGetLastError());
|
|
cuda_err(cudaMemcpyAsync(host_sums_, sums_, nfits * SLOTS * sizeof(double), cudaMemcpyDeviceToHost, stream_));
|
|
cuda_err(cudaStreamSynchronize(stream_));
|
|
return host_sums_;
|
|
}
|
|
|
|
// fit_isotropic_b_approximately() then fit_parameters() for every fit in `fits`, which start from the
|
|
// k_overall and b_star they hold, and are left holding the answer.
|
|
void ModelScaleGPU::FitBatch(const float2 *d_fcmol, const float2 *d_fmask, std::vector<ModelScaleFitState> &fits) {
|
|
const size_t nfits = fits.size();
|
|
for (size_t i = 0; i < nfits; i++) {
|
|
fits[i].column = static_cast<int>(i);
|
|
fits[i].mode = MODE_ISOTROPIC;
|
|
}
|
|
UploadFits(fits);
|
|
solvent_amplitudes<<<dim3(BLOCKS, static_cast<unsigned>(nfits)), THREADS, 0, stream_>>>(stol2_, d_fcmol, d_fmask,
|
|
n_, fits_, f_abs_);
|
|
cuda_err(cudaGetLastError());
|
|
const double *iso = Reduce(static_cast<int>(nfits));
|
|
for (size_t i = 0; i < nfits; i++) {
|
|
const double *s = iso + i * SLOTS;
|
|
const double sx = s[0], sy = s[1], sxx = s[2], sxy = s[3], n = s[4];
|
|
if (n <= 5)
|
|
continue;
|
|
const double slope = (n * sxy - sx * sy) / (n * sxx - sx * sx);
|
|
const double intercept = (sy - slope * sx) / n;
|
|
const double b_iso = -slope;
|
|
fits[i].k_overall = std::exp(intercept);
|
|
// set_b_overall({b_iso, b_iso, b_iso, 0, 0, 0}): b_star = F B F^T (SMat33::transformed_by)
|
|
const double *m = frac_;
|
|
auto elem = [&](int r, int q) {
|
|
return m[3 * r] * (m[3 * q] * b_iso) + m[3 * r + 1] * (m[3 * q + 1] * b_iso) + m[3 * r + 2] * (m[3 * q + 2] * b_iso);
|
|
};
|
|
const double b[6] = {elem(0, 0), elem(1, 1), elem(2, 2), elem(0, 1), elem(0, 2), elem(1, 2)};
|
|
std::copy(b, b + 6, fits[i].b_star);
|
|
}
|
|
LevMarBatch(fits);
|
|
}
|
|
|
|
void ModelScaleGPU::LevMarBatch(std::vector<ModelScaleFitState> &fits) {
|
|
const size_t nfits = fits.size();
|
|
// get_parameters() and set_parameters() with k_sol and b_sol fixed
|
|
const int na = n_params_;
|
|
auto get_parameters = [&](const ModelScaleFitState &f) {
|
|
std::vector<double> a(na);
|
|
a[0] = f.k_overall;
|
|
for (int j = 1; j < na; j++) {
|
|
const double *r = constraints_[j - 1];
|
|
const double *u = f.b_star;
|
|
a[j] = r[0] * u[0] + r[1] * u[1] + r[2] * u[2] + r[3] * u[3] + r[4] * u[4] + r[5] * u[5];
|
|
}
|
|
return a;
|
|
};
|
|
auto set_parameters = [&](ModelScaleFitState &f, const std::vector<double> &a) {
|
|
f.k_overall = a[0];
|
|
std::fill(f.b_star, f.b_star + 6, 0.0);
|
|
for (int j = 1; j < na; j++)
|
|
for (int k = 0; k < 6; k++)
|
|
f.b_star[k] += constraints_[j - 1][k] * a[j];
|
|
};
|
|
|
|
std::vector<LevMarRun> runs;
|
|
for (const auto &f : fits)
|
|
runs.emplace_back(get_parameters(f));
|
|
while (true) {
|
|
std::vector<size_t> live;
|
|
std::vector<ModelScaleFitState> request;
|
|
for (size_t i = 0; i < nfits; i++)
|
|
if (!runs[i].done) {
|
|
ModelScaleFitState f = fits[i];
|
|
set_parameters(f, runs[i].at);
|
|
f.mode = runs[i].mode;
|
|
live.push_back(i);
|
|
request.push_back(f);
|
|
}
|
|
if (live.empty())
|
|
break;
|
|
const double *sums = Sums(request);
|
|
for (size_t j = 0; j < live.size(); j++)
|
|
runs[live[j]].Take(sums + j * SLOTS);
|
|
}
|
|
for (size_t i = 0; i < nfits; i++)
|
|
set_parameters(fits[i], runs[i].Result());
|
|
}
|
|
|
|
// ModelScaling.cpp's robust rounds of each fit: the residuals at its parameters, their median taken once as c,
|
|
// the weights, and gemmi's Levenberg-Marquardt again from where the fit stands. The median is the element
|
|
// std::nth_element puts at n / 2, as there.
|
|
void ModelScaleGPU::RobustBatch(std::vector<ModelScaleFitState> &fits, size_t nthreads) {
|
|
const size_t nfits = fits.size();
|
|
if (nfits > static_cast<size_t>(MAX_ROBUST_FITS))
|
|
throw std::logic_error("ModelScaleGPU: too many robust fits in one launch");
|
|
if (!weight_) {
|
|
weight_ = CudaDevicePtr<float>(MAX_ROBUST_FITS * max_points_);
|
|
residual_ = CudaDevicePtr<float>(MAX_ROBUST_FITS * max_points_);
|
|
host_residual_ = CudaHostPtr<float>(MAX_ROBUST_FITS * max_points_);
|
|
}
|
|
for (size_t i = 0; i < nfits; i++) {
|
|
fits[i].column = static_cast<int>(i);
|
|
fits[i].weighted = 0;
|
|
fits[i].c = 0;
|
|
}
|
|
const dim3 grid(BLOCKS, static_cast<unsigned>(nfits));
|
|
for (int round = 0; round < ROBUST_ROUNDS; round++) {
|
|
UploadFits(fits);
|
|
robust_residuals<<<grid, THREADS, 0, stream_>>>(hkl_, fobs_, f_abs_, n_, fits_, residual_);
|
|
cuda_err(cudaGetLastError());
|
|
if (round == 0) {
|
|
cuda_err(cudaMemcpyAsync(host_residual_, residual_, nfits * n_ * sizeof(float), cudaMemcpyDeviceToHost,
|
|
stream_));
|
|
cuda_err(cudaStreamSynchronize(stream_));
|
|
ParallelFor(static_cast<int>(nfits), nthreads, [&](int i) {
|
|
std::vector<float> r(host_residual_.get() + static_cast<size_t>(i) * n_,
|
|
host_residual_.get() + static_cast<size_t>(i + 1) * n_);
|
|
std::nth_element(r.begin(), r.begin() + r.size() / 2, r.end());
|
|
fits[i].c = r[r.size() / 2];
|
|
});
|
|
UploadFits(fits);
|
|
}
|
|
robust_weights<<<grid, THREADS, 0, stream_>>>(residual_, n_, fits_, weight_);
|
|
cuda_err(cudaGetLastError());
|
|
cuda_err(cudaStreamSynchronize(stream_)); // host_fits_ is written again below
|
|
for (ModelScaleFitState &f : fits)
|
|
f.weighted = f.c > 0;
|
|
LevMarBatch(fits);
|
|
}
|
|
for (ModelScaleFitState &f : fits)
|
|
f.weighted = 0;
|
|
}
|
|
|
|
ModelScaleParams ModelScaleGPU::Fit(const float2 *d_fcmol, const float2 *d_fmask, double k_sol, double b_sol) {
|
|
ModelScaleParams out;
|
|
if (n_ == 0)
|
|
return out;
|
|
std::vector<ModelScaleFitState> fits(1, ModelScaleFitState{1.0, {0, 0, 0, 0, 0, 0}, k_sol, b_sol, 0, 0});
|
|
FitBatch(d_fcmol, d_fmask, fits);
|
|
out.k_overall = fits[0].k_overall;
|
|
std::copy(fits[0].b_star, fits[0].b_star + 6, out.b_star);
|
|
return out;
|
|
}
|
|
|
|
// FitModelScale's grid (ModelScaling.cpp, default ModelScaleBox with b_sol up to b_hi): a coarse pass
|
|
// over the box and a finer one around its winner, every grid point its own fit, the winner the lowest
|
|
// finite R in grid order. The fits of a pass run together, one launch per Levenberg-Marquardt step.
|
|
ModelSolventFit ModelScaleGPU::FitSolvent(const float2 *d_fcmol, const float2 *d_fmask, double b_hi, bool robust,
|
|
size_t nthreads) {
|
|
ModelSolventFit out;
|
|
if (n_ < 20)
|
|
return out;
|
|
if (n_strong_ <= 5)
|
|
throw ModelScaleGPUTooFewReflections();
|
|
|
|
const double k_lo = 0.10, k_hi = 0.60, b_lo = 10.0; // ModelScaleBox{}
|
|
const double b_step = std::max(10.0, (b_hi - b_lo) / 7.0);
|
|
double best_r = -1;
|
|
auto try_points = [&](std::vector<ModelScaleFitState> &fits, bool robust_pass) {
|
|
FitBatch(d_fcmol, d_fmask, fits);
|
|
if (robust_pass)
|
|
RobustBatch(fits, nthreads);
|
|
for (auto &f : fits)
|
|
f.mode = MODE_R;
|
|
const double *r = Sums(fits);
|
|
for (size_t i = 0; i < fits.size(); i++) {
|
|
++out.n_grid;
|
|
const double den = r[i * SLOTS + 1];
|
|
const double ri = den > 0 ? r[i * SLOTS] / den : 1.0;
|
|
if (std::isfinite(ri) && (best_r < 0 || ri < best_r)) {
|
|
best_r = ri;
|
|
out.k_sol = fits[i].k_sol;
|
|
out.b_sol = fits[i].b_sol;
|
|
out.scale.k_overall = fits[i].k_overall;
|
|
std::copy(fits[i].b_star, fits[i].b_star + 6, out.scale.b_star);
|
|
}
|
|
}
|
|
};
|
|
auto point = [](double ks, double bs) {
|
|
return ModelScaleFitState{1.0, {0, 0, 0, 0, 0, 0}, ks, bs, 0, 0};
|
|
};
|
|
|
|
std::vector<ModelScaleFitState> coarse;
|
|
for (double ks = k_lo; ks <= k_hi + 1e-9; ks += 0.05)
|
|
for (double bs = b_lo; bs <= b_hi + 1e-9; bs += b_step)
|
|
coarse.push_back(point(ks, bs));
|
|
try_points(coarse, false);
|
|
const double k0 = out.k_sol, b0 = out.b_sol;
|
|
const double k_hi2 = std::min(k_hi, k0 + 0.05);
|
|
const double b_hi2 = std::min(b_hi, b0 + b_step);
|
|
std::vector<ModelScaleFitState> fine;
|
|
for (double ks = std::max(k_lo, k0 - 0.05); ks <= k_hi2 + 1e-9; ks += 0.025)
|
|
for (double bs = std::max(b_lo, b0 - b_step); bs <= b_hi2 + 1e-9; bs += b_step / 2.0)
|
|
fine.push_back(point(ks, bs));
|
|
// a robust fit takes the fine pass robustly, the coarse winner included (as FitModelScale)
|
|
if (robust)
|
|
best_r = -1;
|
|
try_points(fine, robust);
|
|
out.r = best_r;
|
|
return out;
|
|
}
|