ModelScaleGPU: FitSolvent 3x faster; a test on a model's own points
- |Fcalc + solvent| is taken once per fit instead of at every step: the solvent pair is fixed for the whole of a fit. - The per-step anisotropic factor and the derivatives are computed per point in float, and summed in double as before. Double arithmetic was most of the cost on a card with little double throughput. The final R of the grid stays in gemmi's double arithmetic. - Each mode reduces only the slots it fills; 48 blocks per fit instead of 128. - The first upload of a batch no longer adds a stream synchronisation. FitSolvent per zone at 3.5 A (RTX 5080, real zone hkl sets, synthetic amplitudes): 3.2 ms at 3.3k points, 9.2 ms at 72k and 21.6 ms at 213k, against 9.7 / 29.6 / 62.9 ms before. Fit() is 0.25-0.66 ms. Against gemmi, k_overall and b* agree to 1e-9 - 1.4e-5 relative, and the grid winner is the same. New test ModelScaleGPU_MatchesGemmiOnAModelsPoints: the rigid body's own points (the ClusterPdb fixture with anisotropic atoms, its own amplitudes, a displaced placement, the rigid body's Fcalc and mask), five groups at 6 and 3.5 A. The R tolerance of the synthetic cases is now 1e-5. That is the resolution of gemmi's own fit: its Levenberg-Marquardt stops at a WSSR change of 1e-5, and perturbing gemmi's own Fcalc by 1e-5 moves its answer by up to 1.2e-4 of <Fobs>. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C
This commit is contained in:
1 parent
e39489dfe0
commit
274d42a2e7
3 files changed
+237
-59
No files matched your search
+95
-53
@@ -21,9 +21,9 @@ void cuda_err(cudaError_t 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 = 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 = 128; // FitModelScale's coarse pass is 88 fits
|
||||
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
|
||||
@@ -32,26 +32,41 @@ 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 {
|
||||
double row[6][6];
|
||||
float row[6][6];
|
||||
};
|
||||
|
||||
// gemmi Scaling<float>'s arithmetic for one point: the solvent term and |F| in float, as a
|
||||
// complex<float> is, the anisotropic scale in double.
|
||||
// |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 float2 *__restrict__ fcmol, const float2 *__restrict__ fmask, int n,
|
||||
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 solvent = static_cast<float>(f.k_sol * exp(-f.b_sol * stol2[i]));
|
||||
const float2 fc = fcmol[i], fm = fmask[i];
|
||||
const float f_abs = hypotf(fc.x + solvent * fm.x, fc.y + solvent * fm.y);
|
||||
const float f_abs = f_abs_fit[i];
|
||||
const double fo = fobs[i];
|
||||
|
||||
if (f.mode == MODE_ISOTROPIC) {
|
||||
@@ -67,34 +82,43 @@ __global__ void scale_sums(const int *__restrict__ hkl, const double *__restrict
|
||||
continue;
|
||||
}
|
||||
|
||||
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 double k_aniso = exp(-0.25 * r_u_r);
|
||||
|
||||
if (f.mode == MODE_WSSR || f.mode == MODE_R) {
|
||||
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 * k_aniso);
|
||||
const double dy = fo - value;
|
||||
if (f.mode == MODE_WSSR) {
|
||||
acc[0] += dy * dy;
|
||||
} else {
|
||||
acc[0] += fabs(dy);
|
||||
acc[1] += fo;
|
||||
}
|
||||
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;
|
||||
}
|
||||
|
||||
// compute_value_and_derivatives() with k_sol and b_sol fixed
|
||||
const double fe = f_abs * k_aniso;
|
||||
const double y = f.k_overall * fe;
|
||||
// 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;
|
||||
const double du[6] = {-0.25 * y * (hx * hx), -0.25 * y * (hy * hy), -0.25 * y * (hz * hz),
|
||||
-0.5 * y * (hx * hy), -0.5 * y * (hx * hz), -0.5 * y * (hy * hz)};
|
||||
for (int j = 1; j < NA; j++) {
|
||||
const double *r = c.row[j - 1];
|
||||
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;
|
||||
@@ -107,10 +131,12 @@ __global__ void scale_sums(const int *__restrict__ hkl, const double *__restrict
|
||||
acc[s++] += dy * dy_da[j];
|
||||
}
|
||||
|
||||
// The block's sum of each slot: a fixed shuffle tree within each warp, then the warps in order.
|
||||
// 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 < NS; s++) {
|
||||
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);
|
||||
@@ -120,7 +146,7 @@ __global__ void scale_sums(const int *__restrict__ hkl, const double *__restrict
|
||||
__syncthreads();
|
||||
if (threadIdx.x < SLOTS) {
|
||||
double t = 0;
|
||||
if (threadIdx.x < NS)
|
||||
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;
|
||||
@@ -277,13 +303,14 @@ struct ModelScaleGPU::LevMarRun {
|
||||
|
||||
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),
|
||||
fits_(MAX_FITS), partial_(static_cast<size_t>(MAX_FITS) * BLOCKS * SLOTS),
|
||||
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) {}
|
||||
|
||||
@@ -319,23 +346,32 @@ void ModelScaleGPU::SetPoints(const std::vector<std::array<int, 3>> &hkl,
|
||||
cuda_err(cudaStreamSynchronize(stream_)); // the host vectors go out of scope
|
||||
}
|
||||
|
||||
const double *ModelScaleGPU::Sums(const float2 *d_fcmol, const float2 *d_fmask,
|
||||
const std::vector<ModelScaleFitState> &fits) {
|
||||
const int nfits = static_cast<int>(fits.size());
|
||||
if (nfits > MAX_FITS)
|
||||
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_, nfits * sizeof(ModelScaleFitState), cudaMemcpyHostToDevice, stream_));
|
||||
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{};
|
||||
std::copy(&constraints_[0][0], &constraints_[0][0] + 36, &c.row[0][0]);
|
||||
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_, d_fcmol, d_fmask, n_, fits_, c, partial_); break;
|
||||
case 3: scale_sums<3><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break;
|
||||
case 4: scale_sums<4><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break;
|
||||
case 5: scale_sums<5><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break;
|
||||
case 6: scale_sums<6><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break;
|
||||
case 7: scale_sums<7><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break;
|
||||
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());
|
||||
@@ -350,9 +386,15 @@ const double *ModelScaleGPU::Sums(const float2 *d_fcmol, const float2 *d_fmask,
|
||||
// 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 (auto &f : fits)
|
||||
f.mode = MODE_ISOTROPIC;
|
||||
const double *iso = Sums(d_fcmol, d_fmask, fits);
|
||||
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];
|
||||
@@ -407,7 +449,7 @@ void ModelScaleGPU::FitBatch(const float2 *d_fcmol, const float2 *d_fmask, std::
|
||||
}
|
||||
if (live.empty())
|
||||
break;
|
||||
const double *sums = Sums(d_fcmol, d_fmask, request);
|
||||
const double *sums = Sums(request);
|
||||
for (size_t j = 0; j < live.size(); j++)
|
||||
runs[live[j]].Take(sums + j * SLOTS);
|
||||
}
|
||||
@@ -419,7 +461,7 @@ ModelScaleParams ModelScaleGPU::Fit(const float2 *d_fcmol, const float2 *d_fmask
|
||||
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});
|
||||
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);
|
||||
@@ -442,7 +484,7 @@ ModelSolventFit ModelScaleGPU::FitSolvent(const float2 *d_fcmol, const float2 *d
|
||||
FitBatch(d_fcmol, d_fmask, fits);
|
||||
for (auto &f : fits)
|
||||
f.mode = MODE_R;
|
||||
const double *r = Sums(d_fcmol, d_fmask, fits);
|
||||
const double *r = Sums(fits);
|
||||
for (size_t i = 0; i < fits.size(); i++) {
|
||||
++out.n_grid;
|
||||
const double den = r[i * SLOTS + 1];
|
||||
@@ -457,7 +499,7 @@ ModelSolventFit ModelScaleGPU::FitSolvent(const float2 *d_fcmol, const float2 *d
|
||||
}
|
||||
};
|
||||
auto point = [](double ks, double bs) {
|
||||
return ModelScaleFitState{1.0, {0, 0, 0, 0, 0, 0}, ks, bs, 0};
|
||||
return ModelScaleFitState{1.0, {0, 0, 0, 0, 0, 0}, ks, bs, 0, 0};
|
||||
};
|
||||
|
||||
std::vector<ModelScaleFitState> coarse;
|
||||
|
||||
+11
-4
@@ -14,8 +14,8 @@
|
||||
// The rigid body's scale fit on the GPU: what RigidBodyTarget::Residuals does with a gemmi
|
||||
// Scaling<float> (use_solvent, k_sol and b_sol fixed) - fit_isotropic_b_approximately() followed by
|
||||
// gemmi's Levenberg-Marquardt, and FitModelScale's k_sol/b_sol grid. The sums over the reflections run
|
||||
// on the device, in double, in a fixed order; the Levenberg-Marquardt control (the damping, the 7x7
|
||||
// solve, the stop rules) is gemmi's own, ported line for line and run on the host.
|
||||
// on the device, accumulated in double in a fixed order; the Levenberg-Marquardt control (the damping,
|
||||
// the 7x7 solve, the stop rules) is gemmi's own, ported line for line and run on the host.
|
||||
|
||||
struct ModelScaleParams {
|
||||
double k_overall = 1.0;
|
||||
@@ -35,6 +35,7 @@ struct ModelScaleFitState {
|
||||
double b_star[6];
|
||||
double k_sol, b_sol;
|
||||
int mode;
|
||||
int column; // the fit's row of |Fcalc + solvent| on the device
|
||||
};
|
||||
|
||||
class ModelScaleGPU {
|
||||
@@ -63,8 +64,13 @@ public:
|
||||
private:
|
||||
struct LevMarRun;
|
||||
|
||||
// The sums of every fit in `fits` over all points, one row of SLOTS doubles per fit, on the host.
|
||||
const double *Sums(const float2 *d_fcmol, const float2 *d_fmask, const std::vector<ModelScaleFitState> &fits);
|
||||
// Queues the copy of `fits` to the device. host_fits_ is written only after the stream has been
|
||||
// synchronized, which every Reduce() ends with.
|
||||
void UploadFits(const std::vector<ModelScaleFitState> &fits);
|
||||
// The sums of the first `nfits` uploaded fits over all points, one row of SLOTS doubles per fit, on
|
||||
// the host.
|
||||
const double *Reduce(int nfits);
|
||||
const double *Sums(const std::vector<ModelScaleFitState> &fits);
|
||||
void FitBatch(const float2 *d_fcmol, const float2 *d_fmask, std::vector<ModelScaleFitState> &fits);
|
||||
|
||||
cudaStream_t stream_;
|
||||
@@ -79,6 +85,7 @@ private:
|
||||
CudaDevicePtr<double> stol2_;
|
||||
CudaDevicePtr<float> fobs_;
|
||||
CudaDevicePtr<unsigned char> strong_;
|
||||
CudaDevicePtr<float> f_abs_; // per fit of a batch, per point
|
||||
CudaDevicePtr<ModelScaleFitState> fits_;
|
||||
CudaDevicePtr<double> partial_; // per fit, per block, per slot
|
||||
CudaDevicePtr<double> sums_; // per fit, per slot
|
||||
|
||||
+131
-2
@@ -13,13 +13,22 @@
|
||||
#include <random>
|
||||
#include <vector>
|
||||
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
|
||||
#include "gemmi/asumask.hpp"
|
||||
#include "gemmi/mmread_gz.hpp"
|
||||
#include "gemmi/scaling.hpp"
|
||||
#include "gemmi/symmetry.hpp"
|
||||
#include "gemmi/unitcell.hpp"
|
||||
|
||||
#include "../common/Logger.h"
|
||||
#include "../rugnux/ModelFFT.h"
|
||||
#include "../rugnux/ModelGrid.h"
|
||||
#include "../rugnux/ModelScaling.h"
|
||||
#include "../rugnux/ModelScaleGPU.h"
|
||||
#include "../rugnux/ModelValidation.h"
|
||||
#include "../rugnux/RigidBodyRefine.h"
|
||||
|
||||
namespace {
|
||||
|
||||
@@ -76,6 +85,82 @@ gemmi::Scaling<float> MakeScaling(const gemmi::UnitCell &cell, const gemmi::Spac
|
||||
return scaling;
|
||||
}
|
||||
|
||||
// A synthetic "protein" of 150 carbons, every third anisotropic, in `cryst` - ModelValidationTest.cpp's
|
||||
// rigid-body fixture - and the Scaling points RigidBodyTarget::Residuals fits at a displaced
|
||||
// placement q0: the model's own amplitudes as Fobs, and the Fcalc and bulk-solvent Fmask of the
|
||||
// displaced model, gridded and composed as the rigid body does.
|
||||
gemmi::Scaling<float> ModelPoints(const char *cryst, double zone) {
|
||||
std::string pdb = cryst;
|
||||
std::mt19937 rng(20260902);
|
||||
std::uniform_real_distribution<double> x(2, 14), y(2, 16), z(2, 18);
|
||||
char line[96];
|
||||
for (int i = 1; i <= 150; i++) {
|
||||
std::snprintf(line, sizeof line, "ATOM %5d C UNK A 1 %8.3f%8.3f%8.3f 1.00 20.00 C\n",
|
||||
i, x(rng), y(rng), z(rng));
|
||||
pdb += line;
|
||||
}
|
||||
pdb += "END\n";
|
||||
const std::string path = "model_scale_gpu_test.pdb";
|
||||
std::ofstream(path) << pdb;
|
||||
Logger logger("ModelScaleGPUTest");
|
||||
const auto reference = ModelReferenceIntensities(path, {}, {}, 3.0, logger);
|
||||
gemmi::Structure st = gemmi::read_structure_gz(path, gemmi::CoorFormat::Detect);
|
||||
std::filesystem::remove(path);
|
||||
st.setup_cell_images();
|
||||
int i = 0;
|
||||
for (gemmi::Chain &ch : st.models[0].chains)
|
||||
for (gemmi::Residue &r : ch.residues)
|
||||
for (gemmi::Atom &a : r.atoms)
|
||||
if (i++ % 3 == 0)
|
||||
a.aniso = {0.30f, 0.25f, 0.20f, 0.02f, -0.01f, 0.03f};
|
||||
const gemmi::SpaceGroup &sg = *st.find_spacegroup();
|
||||
|
||||
gemmi::AsuData<gemmi::ValueSigma<float>> fobs;
|
||||
fobs.unit_cell_ = st.cell;
|
||||
fobs.spacegroup_ = &sg;
|
||||
for (const auto &r : reference)
|
||||
if (st.cell.calculate_d({{r.h, r.k, r.l}}) >= zone)
|
||||
fobs.v.push_back({{{r.h, r.k, r.l}}, {std::sqrt(r.I), 1.0f}});
|
||||
fobs.ensure_sorted();
|
||||
|
||||
gemmi::Model model = st.models[0];
|
||||
RigidBodyTarget target(model, st.cell, sg, 4);
|
||||
const double q0[6] = {0.12, -0.09, 0.07, 0.15, -0.10, 0.05};
|
||||
target.Place(q0, model);
|
||||
|
||||
gemmi::Grid<float> grid;
|
||||
grid.unit_cell = st.cell;
|
||||
grid.spacegroup = &sg;
|
||||
grid.set_size_from_spacing(zone / 3.0, gemmi::GridSizeRounding::Up);
|
||||
std::vector<gemmi::Miller> hkl;
|
||||
for (const auto &hv : fobs.v)
|
||||
hkl.push_back(hv.hkl);
|
||||
const SymmetryComposition composition(grid, zone, hkl);
|
||||
|
||||
gemmi::DensityCalculator<gemmi::IT92<float>, float> dc;
|
||||
dc.d_min = zone;
|
||||
dc.rate = 1.5;
|
||||
dc.grid.unit_cell = st.cell;
|
||||
dc.grid.spacegroup = &sg;
|
||||
dc.set_refmac_compatible_blur(model);
|
||||
PutModelDensityOnGrid(dc, model, {}, 4);
|
||||
std::vector<std::complex<double>> fc;
|
||||
composition.Compose(MapToFPhi(dc.grid), dc.blur, fc, nullptr, 4);
|
||||
|
||||
gemmi::Grid<float> mask = grid;
|
||||
PutMaskOnGrid(mask, model, OrbitLeaders(mask, 4), 4);
|
||||
const gemmi::FPhiGrid<float> fm = MapToFPhi(mask);
|
||||
gemmi::AsuData<std::complex<float>> fcalc, fmask;
|
||||
for (size_t m = 0; m < composition.Hkl().size(); m++) {
|
||||
fcalc.v.push_back({composition.Hkl()[m], std::complex<float>(fc[m])});
|
||||
fmask.v.push_back({composition.Hkl()[m], fm.get_value_by_hkl(composition.Hkl()[m])});
|
||||
}
|
||||
gemmi::Scaling<float> scaling(st.cell, &sg);
|
||||
scaling.use_solvent = true;
|
||||
scaling.prepare_points(fcalc, fobs, &fmask);
|
||||
return scaling;
|
||||
}
|
||||
|
||||
struct GpuPoints {
|
||||
CudaStream stream;
|
||||
ModelScaleGPU scale;
|
||||
@@ -113,6 +198,12 @@ double MaxAbs(const double b[6]) {
|
||||
return m;
|
||||
}
|
||||
|
||||
// The tolerances are those of gemmi's fit rather than of the arithmetic. Its Levenberg-Marquardt stops
|
||||
// when the WSSR has changed by less than 1e-5 twice, which leaves the scale short of the minimum by
|
||||
// about sqrt(1e-5) of the residual, so where a step is accepted or the fit stops on a near-tie the last
|
||||
// bit decides it: measured, gemmi moves by up to 1e-4 of <Fobs> when its own Fcalc changes by 1e-5.
|
||||
constexpr double R_TOLERANCE = 1e-5;
|
||||
|
||||
void CheckSameScale(const ModelScaleParams &gpu, const gemmi::Scaling<float> &cpu) {
|
||||
const double b_cpu[6] = {cpu.b_star.u11, cpu.b_star.u22, cpu.b_star.u33,
|
||||
cpu.b_star.u12, cpu.b_star.u13, cpu.b_star.u23};
|
||||
@@ -163,7 +254,7 @@ TEST_CASE("ModelScaleGPU_FitMatchesGemmi", "[ModelValidation][gpu]") {
|
||||
gemmi::Scaling<float> with_gpu = ref;
|
||||
with_gpu.k_overall = p.k_overall;
|
||||
with_gpu.b_star = {p.b_star[0], p.b_star[1], p.b_star[2], p.b_star[3], p.b_star[4], p.b_star[5]};
|
||||
CHECK(std::fabs(with_gpu.calculate_r_factor() - ref.calculate_r_factor()) < 1e-6);
|
||||
CHECK(std::fabs(with_gpu.calculate_r_factor() - ref.calculate_r_factor()) < R_TOLERANCE);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -183,11 +274,49 @@ TEST_CASE("ModelScaleGPU_SolventGridMatchesFitModelScale", "[ModelValidation][gp
|
||||
CHECK(fit.n_grid == report.n_grid);
|
||||
CHECK(fit.k_sol == cpu.k_sol); // the same grid point, so the same double
|
||||
CHECK(fit.b_sol == cpu.b_sol);
|
||||
CHECK(std::fabs(fit.r - report.r_work_fit) < 1e-6);
|
||||
CHECK(std::fabs(fit.r - report.r_work_fit) < R_TOLERANCE);
|
||||
CheckSameScale(fit.scale, cpu);
|
||||
}
|
||||
}
|
||||
|
||||
// The rigid body's own sequence on a model's points: FitModelScale once for the zone, then the
|
||||
// per-evaluation fit at that solvent pair.
|
||||
TEST_CASE("ModelScaleGPU_MatchesGemmiOnAModelsPoints", "[ModelValidation][gpu]") {
|
||||
if (get_gpu_count() == 0)
|
||||
return;
|
||||
const char *crysts[] = {
|
||||
"CRYST1 40.000 50.000 60.000 90.00 90.00 90.00 P 1 1\n",
|
||||
"CRYST1 40.000 50.000 60.000 90.00 100.00 90.00 C 1 2 1 4\n",
|
||||
"CRYST1 40.000 50.000 60.000 90.00 90.00 90.00 P 21 21 21 4\n",
|
||||
"CRYST1 60.000 60.000 60.000 90.00 90.00 90.00 I 2 3 24\n",
|
||||
"CRYST1 80.000 80.000 80.000 90.00 90.00 90.00 F 41 3 2 96\n",
|
||||
};
|
||||
for (const char *cryst : crysts)
|
||||
for (const double zone : {6.0, 3.5}) {
|
||||
INFO(cryst << " at " << zone << " A");
|
||||
gemmi::Scaling<float> cpu = ModelPoints(cryst, zone);
|
||||
REQUIRE(cpu.points.size() > 50);
|
||||
GpuPoints gpu(cpu);
|
||||
|
||||
const ModelScaleReport report = FitModelScale(cpu, {}, 4);
|
||||
const ModelSolventFit solvent = gpu.scale.FitSolvent(gpu.fcmol, gpu.fmask);
|
||||
CHECK(solvent.k_sol == cpu.k_sol);
|
||||
CHECK(solvent.b_sol == cpu.b_sol);
|
||||
CHECK(std::fabs(solvent.r - report.r_work_fit) < R_TOLERANCE);
|
||||
CheckSameScale(solvent.scale, cpu);
|
||||
|
||||
gemmi::Scaling<float> ref = cpu;
|
||||
ref.fix_k_sol = true;
|
||||
ref.fix_b_sol = true;
|
||||
ref.k_overall = 1;
|
||||
ref.b_star = {0, 0, 0, 0, 0, 0};
|
||||
ref.fit_isotropic_b_approximately();
|
||||
ref.fit_parameters();
|
||||
const ModelScaleParams p = gpu.scale.Fit(gpu.fcmol, gpu.fmask, cpu.k_sol, cpu.b_sol);
|
||||
CheckSameScale(p, ref);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_CASE("ModelScaleGPU_Deterministic", "[ModelValidation][gpu]") {
|
||||
if (get_gpu_count() == 0)
|
||||
return;
|
||||
|
||||
Reference in new issue
Block a user