A pure move. ModelValidation, RigidBodyRefine, RigidBodyGPU, ModelFFT, ModelGrid, ModelScaling, ModelMaskGPU, ModelScaleGPU and SigmaA - everything that works on an atomic model - become the JFJochStructureRefinement library, linked by JFJochImageAnalysis. WriteModel (the placed-model mmCIF/PDB writer) goes to writer/ as its own small JFJochModelWriter target, so JFJochWriter, which a writer-only build compiles, does not gain a gemmi dependency. Only include paths and CMake lists change. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
522 lines
22 KiB
Plaintext
522 lines
22 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
|
|
|
|
// 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, 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) {
|
|
if (!strong[i])
|
|
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);
|
|
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;
|
|
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;
|
|
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];
|
|
for (int j = 0; j < NA; j++)
|
|
acc[s++] += dy * dy_da[j];
|
|
}
|
|
|
|
// 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;
|
|
}
|
|
}
|
|
|
|
// 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::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_, n_, fits_, c, partial_); break;
|
|
case 3: scale_sums<3><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break;
|
|
case 4: scale_sums<4><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break;
|
|
case 5: scale_sums<5><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break;
|
|
case 6: scale_sums<6><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break;
|
|
case 7: scale_sums<7><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, 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);
|
|
}
|
|
|
|
// 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());
|
|
}
|
|
|
|
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): 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) {
|
|
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, b_hi = 80.0; // ModelScaleBox{}
|
|
double best_r = -1;
|
|
auto try_points = [&](std::vector<ModelScaleFitState> &fits) {
|
|
FitBatch(d_fcmol, d_fmask, 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};
|
|
};
|
|
|
|
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 += 10.0)
|
|
coarse.push_back(point(ks, bs));
|
|
try_points(coarse);
|
|
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 + 10.0);
|
|
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 - 10.0); bs <= b_hi2 + 1e-9; bs += 5.0)
|
|
fine.push_back(point(ks, bs));
|
|
try_points(fine);
|
|
out.r = best_r;
|
|
return out;
|
|
}
|