leonarski_fandClaude Fable 5.1
f724bd1117
A Levenberg-Marquardt step is a few short launches the host then reads the sums of, a hundred or so steps per fit, and the process has every synchronisation sleep until the device raises the wake-up (set_gpu_blocking_sync): for these launches the wake-up costs more than the work. The fit's waits poll the stream instead. Outputs are byte-identical; the refinement's scale refit counter drops 20-27% and the molecular replacement's scale fits gain a little on the demo-sized sets. Summing the block partials on the host instead of in a kernel was tried alongside and dropped: the coarse solvent grid of 88 fits made it a megabyte of read-back per step, slower than the kernel. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
664 lines
29 KiB
Plaintext
664 lines
29 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 "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));
|
|
}
|
|
}
|
|
|
|
// The robust fit's c of each fit (blockIdx.x): the median of its residuals, the element std::nth_element puts at
|
|
// n / 2, by radix selection - the residuals are not negative, so their order is that of their bits as unsigned
|
|
// integers, and four passes find those bits a byte at a time. The counts are of integers: the same answer
|
|
// whatever order the threads take.
|
|
__global__ void residual_median(const float *__restrict__ residual, int n, const ModelScaleFitState *__restrict__ fits,
|
|
double *__restrict__ c) {
|
|
__shared__ unsigned int count[256];
|
|
__shared__ unsigned int prefix, rank;
|
|
const float *r = residual + static_cast<size_t>(fits[blockIdx.x].column) * n;
|
|
if (threadIdx.x == 0) {
|
|
prefix = 0;
|
|
rank = n / 2;
|
|
}
|
|
for (int shift = 24; shift >= 0; shift -= 8) {
|
|
for (int b = threadIdx.x; b < 256; b += blockDim.x)
|
|
count[b] = 0;
|
|
__syncthreads();
|
|
// the bits above this byte, which the residuals counted must share with the median
|
|
const unsigned int high = shift == 24 ? 0 : 0xffffffffu << (shift + 8);
|
|
for (int i = threadIdx.x; i < n; i += blockDim.x) {
|
|
const unsigned int u = __float_as_uint(r[i]);
|
|
if ((u & high) == prefix)
|
|
atomicAdd(&count[(u >> shift) & 255], 1u);
|
|
}
|
|
__syncthreads();
|
|
if (threadIdx.x == 0) {
|
|
unsigned int b = 0;
|
|
while (rank >= count[b]) {
|
|
rank -= count[b];
|
|
b++;
|
|
}
|
|
prefix |= b << shift;
|
|
}
|
|
__syncthreads();
|
|
}
|
|
if (threadIdx.x == 0)
|
|
c[blockIdx.x] = __uint_as_float(prefix);
|
|
}
|
|
|
|
// 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) + MAX_ROBUST_FITS * sizeof(double);
|
|
}
|
|
|
|
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_));
|
|
Wait();
|
|
return host_sums_;
|
|
}
|
|
|
|
void ModelScaleGPU::Wait() {
|
|
cudaError_t err;
|
|
while ((err = cudaStreamQuery(stream_)) == cudaErrorNotReady) {}
|
|
cuda_err(err);
|
|
}
|
|
|
|
// 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) {
|
|
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_);
|
|
median_ = CudaDevicePtr<double>(MAX_ROBUST_FITS);
|
|
host_median_ = CudaHostPtr<double>(MAX_ROBUST_FITS);
|
|
}
|
|
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) {
|
|
residual_median<<<static_cast<unsigned>(nfits), 1024, 0, stream_>>>(residual_, n_, fits_, median_);
|
|
cuda_err(cudaGetLastError());
|
|
cuda_err(cudaMemcpyAsync(host_median_, median_, nfits * sizeof(double), cudaMemcpyDeviceToHost, stream_));
|
|
Wait();
|
|
for (size_t i = 0; i < nfits; i++)
|
|
fits[i].c = host_median_[i];
|
|
UploadFits(fits);
|
|
}
|
|
robust_weights<<<grid, THREADS, 0, stream_>>>(residual_, n_, fits_, weight_);
|
|
cuda_err(cudaGetLastError());
|
|
Wait(); // 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,
|
|
double k_start, double b_start) {
|
|
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);
|
|
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};
|
|
};
|
|
|
|
double k0 = k_start, b0 = b_start;
|
|
if (k0 < 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);
|
|
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;
|
|
}
|