diff --git a/rugnux/RigidBodyGPU.cpp b/rugnux/RigidBodyGPU.cpp index 5206ea917..5bc07daf9 100644 --- a/rugnux/RigidBodyGPU.cpp +++ b/rugnux/RigidBodyGPU.cpp @@ -133,6 +133,7 @@ RigidBodyGPUZone MakeZone(const gemmi::Model &model, const gemmi::UnitCell &cell const gemmi::ValueSigma &vs = fobs->v[i].value; zone.obs_row.push_back(m); zone.obs_fobs.push_back(vs.value); + zone.obs_sigma.push_back(vs.sigma); sum += vs.value; // the points gemmi's prepare_points() takes if (m >= 0 && !std::isnan(vs.value) && !std::isnan(vs.sigma)) @@ -257,6 +258,7 @@ void RigidBodyTargetGPU::SetZone(const gemmi::AsuData> fobs_ = fobs; d_min_ = d_min; solvent_fitted_ = false; + host_scale_ = false; have_point_ = false; const gemmi::Grid grid = RigidBodyZoneGrid(cell_, sg_, d_min); orbit_leaders_ = OrbitLeaders(grid, nthreads_); @@ -272,6 +274,36 @@ void RigidBodyTargetGPU::SetZone(const gemmi::AsuData> } } +// The scale fitted on the host, exactly as RigidBodyTarget::Residuals() fits it, for a zone the device +// does not reproduce: one with five or fewer reflections for fit_isotropic_b_approximately(), where +// FitModelScale()'s grid fits chain from one to the next and every fit starts from the last. +void RigidBodyTargetGPU::HostScale() { + std::vector> fc, fm; + engine_.DownloadPoints(fc, fm); + gemmi::Scaling scaling(cell_, &sg_); + scaling.use_solvent = true; + for (size_t p = 0; p < zone_->point_obs.size(); p++) { + const auto &o = fobs_.v[zone_->point_obs[p]]; + const int m = zone_->obs_row[zone_->point_obs[p]]; + scaling.points.push_back({o.hkl, zone_->row_stol2[m], std::complex(fc[p][0], fc[p][1]), + std::complex(fm[p][0], fm[p][1]), o.value.value, o.value.sigma}); + } + if (!solvent_fitted_) { + FitModelScale(scaling, {}, nthreads_); + k_sol = scaling.k_sol; + b_sol = scaling.b_sol; + solvent_fitted_ = true; + } + scaling.k_sol = k_sol; + scaling.b_sol = b_sol; + scaling.fix_k_sol = true; + scaling.fix_b_sol = true; + scaling.fit_isotropic_b_approximately(); + scaling.fit_parameters(); + k_overall_ = scaling.k_overall; + b_star_ = scaling.b_star; +} + // RigidBodyTarget::Residuals() with the density, the composition, the transforms and the residuals on // the device. bool RigidBodyTargetGPU::Residuals(const double q[6], double *residuals) { @@ -289,42 +321,31 @@ bool RigidBodyTargetGPU::Residuals(const double q[6], double *residuals) { if (zone_->row_hkl.empty()) return false; - std::vector> fc, fm; - engine_.DownloadPoints(fc, fm); - gemmi::Scaling scaling(cell_, &sg_); - scaling.use_solvent = true; - for (size_t p = 0; p < zone_->point_obs.size(); p++) { - const auto &o = fobs_.v[zone_->point_obs[p]]; - const int m = zone_->obs_row[zone_->point_obs[p]]; - scaling.points.push_back({o.hkl, zone_->row_stol2[m], std::complex(fc[p][0], fc[p][1]), - std::complex(fm[p][0], fm[p][1]), o.value.value, o.value.sigma}); - } - if (scaling.points.empty()) + if (zone_->point_obs.empty()) return false; // The scale as RigidBodyTarget::Residuals() fits it: the solvent once per zone, then the // overall scale and the anisotropic B at every evaluation. - if (!solvent_fitted_) { - FitModelScale(scaling, {}, nthreads_); - k_sol = scaling.k_sol; - b_sol = scaling.b_sol; - solvent_fitted_ = true; + if (!solvent_fitted_ && !host_scale_) { + try { + engine_.FitSolvent(k_sol, b_sol); + solvent_fitted_ = true; + } catch (const std::logic_error &) { + host_scale_ = true; + } } - scaling.k_sol = k_sol; - scaling.b_sol = b_sol; - scaling.fix_k_sol = true; - scaling.fix_b_sol = true; - scaling.fit_isotropic_b_approximately(); - scaling.fit_parameters(); - - const gemmi::SMat33 &b = scaling.b_star; - const double b_star[6] = {b.u11, b.u22, b.u33, b.u12, b.u13, b.u23}; - engine_.Residuals(scaling.k_overall, b_star, k_sol, b_sol, residuals); + if (host_scale_) { + HostScale(); + } else { + double b[6]; + engine_.FitScale(k_sol, b_sol, k_overall_, b); + b_star_ = {b[0], b[1], b[2], b[3], b[4], b[5]}; + } + const double b_star[6] = {b_star_.u11, b_star_.u22, b_star_.u33, b_star_.u12, b_star_.u13, b_star_.u23}; + engine_.Residuals(k_overall_, b_star, k_sol, b_sol, residuals); unmatched = static_cast(std::count(zone_->obs_row.begin(), zone_->obs_row.end(), -1)); have_point_ = true; std::copy(q, q + 6, q_.begin()); - k_overall_ = scaling.k_overall; - b_star_ = scaling.b_star; return true; } catch (const JFJochException &e) { throw RigidBodyGPUFailure(e.what()); diff --git a/rugnux/RigidBodyGPU.cu b/rugnux/RigidBodyGPU.cu index c90c372b7..9866f4038 100644 --- a/rugnux/RigidBodyGPU.cu +++ b/rugnux/RigidBodyGPU.cu @@ -11,6 +11,7 @@ #include #include +#include "ModelScaleGPU.h" #include "../common/JFJochException.h" #include "../image_analysis/indexing/CUDAMemHelpers.h" @@ -502,6 +503,7 @@ struct RigidBodyGPUEngineImpl { CudaDevicePtr products, x; std::map, cufftHandle> plans; + std::unique_ptr scale; // The zone. int n_atoms = 0; @@ -556,6 +558,16 @@ struct RigidBodyGPUEngineImpl { point_fmask = CudaDevicePtr(no, sync); products = CudaDevicePtr(MAX_SCALE_PARAMS * (MAX_SCALE_PARAMS + 6), sync); x = CudaDevicePtr(MAX_SCALE_PARAMS * 6, sync); + scale = std::make_unique(stream, no); + } + + // The points' fcmol and fmask at this evaluation. + void PackPoints() { + if (n_points == 0) + return; + pack_points_kernel<<>>(point_obs, obs_row, static_cast(n_points), + f, fmask, point_fcmol, point_fmask); + cuda_err(cudaGetLastError()); } ~RigidBodyGPUEngineImpl() { @@ -706,6 +718,7 @@ size_t RigidBodyGPUEngine::DeviceBytes(const RigidBodyGPUCapacity &c) { bytes += nr * (3 * sizeof(int) + 2 * sizeof(double) + 7 * sizeof(double2) + sizeof(float2)); bytes += no * (2 * sizeof(int) + sizeof(float) + (1 + 6 + MAX_SCALE_PARAMS) * sizeof(double) + 2 * sizeof(float2)); bytes += CubBytes(na, c.pairs + 1); + bytes += ModelScaleGPU::DeviceBytes(no); return bytes; } @@ -795,6 +808,17 @@ void RigidBodyGPUEngine::SetZone(const RigidBodyGPUZone &zone) { up(e.obs_row, zone.obs_row.data(), zone.obs_row.size() * sizeof(int)); up(e.obs_fobs, zone.obs_fobs.data(), zone.obs_fobs.size() * sizeof(float)); up(e.point_obs, zone.point_obs.data(), zone.point_obs.size() * sizeof(int)); + std::vector> point_hkl; + std::vector point_stol2; + std::vector point_fobs, point_sigma; + for (int i : zone.point_obs) { + const int m = zone.obs_row[i]; + point_hkl.push_back(zone.row_hkl[m]); + point_stol2.push_back(zone.row_stol2[m]); + point_fobs.push_back(zone.obs_fobs[i]); + point_sigma.push_back(zone.obs_sigma[i]); + } + e.scale->SetPoints(point_hkl, point_stol2, point_fobs, point_sigma, zone.constraints, zone.frac); e.Plan(1); e.Plan(3); cuda_err(cudaStreamSynchronize(e.stream)); @@ -830,15 +854,30 @@ void RigidBodyGPUEngine::DownloadPoints(std::vector> &fcmol fmask.resize(e.n_points); if (e.n_points == 0) return; - pack_points_kernel<<>>(e.point_obs, e.obs_row, - static_cast(e.n_points), e.f, e.fmask, - e.point_fcmol, e.point_fmask); - cuda_err(cudaGetLastError()); + e.PackPoints(); cuda_err(cudaMemcpyAsync(fcmol.data(), e.point_fcmol, e.n_points * sizeof(float2), cudaMemcpyDeviceToHost, e.stream)); cuda_err(cudaMemcpyAsync(fmask.data(), e.point_fmask, e.n_points * sizeof(float2), cudaMemcpyDeviceToHost, e.stream)); cuda_err(cudaStreamSynchronize(e.stream)); } +void RigidBodyGPUEngine::FitSolvent(double &k_sol, double &b_sol) { + RigidBodyGPUEngineImpl &e = *impl_; + cuda_err(cudaSetDevice(e.device)); + e.PackPoints(); + const ModelSolventFit fit = e.scale->FitSolvent(e.point_fcmol, e.point_fmask); + k_sol = fit.k_sol; + b_sol = fit.b_sol; +} + +void RigidBodyGPUEngine::FitScale(double k_sol, double b_sol, double &k_overall, double b_star[6]) { + RigidBodyGPUEngineImpl &e = *impl_; + cuda_err(cudaSetDevice(e.device)); + e.PackPoints(); + const ModelScaleParams p = e.scale->Fit(e.point_fcmol, e.point_fmask, k_sol, b_sol); + k_overall = p.k_overall; + std::copy(p.b_star, p.b_star + 6, b_star); +} + void RigidBodyGPUEngine::Residuals(double k_overall, const double b_star[6], double k_sol, double b_sol, double *residuals) { RigidBodyGPUEngineImpl &e = *impl_; diff --git a/rugnux/RigidBodyGPU.h b/rugnux/RigidBodyGPU.h index a1de80d0d..9567796ab 100644 --- a/rugnux/RigidBodyGPU.h +++ b/rugnux/RigidBodyGPU.h @@ -75,6 +75,7 @@ public: private: // The placement at q as x -> R (x - centroid) + centroid + t. void Placement(const double q[6], double rotation[9], double translation[3]) const; + void HostScale(); RigidBodyGPUPool &pool_; RigidBodyGPUEngine &engine_; @@ -87,6 +88,7 @@ private: double d_min_ = 0; std::unique_ptr zone_; bool solvent_fitted_ = false; + bool host_scale_ = false; // the zone's scale is fitted on the host (HostScale) std::vector orbit_leaders_; // for the host mask bool have_point_ = false; diff --git a/rugnux/RigidBodyGPUEngine.h b/rugnux/RigidBodyGPUEngine.h index 35ca075b1..7119568d6 100644 --- a/rugnux/RigidBodyGPUEngine.h +++ b/rugnux/RigidBodyGPUEngine.h @@ -53,7 +53,7 @@ struct RigidBodyGPUZone { // The observations the residuals are over: each one's row (-1 without one) and amplitude. std::vector obs_row; - std::vector obs_fobs; + std::vector obs_fobs, obs_sigma; double f_mean = 1; // The scale's points: the observations gemmi's prepare_points() would take, in order. std::vector point_obs; @@ -101,7 +101,12 @@ public: // instead (the stub until ModelMaskGPU is in). void Fcalc(const double rotation[9], const double translation[3]); void Fmask(const float *host_mask); - // The scale's points: fcmol (Fcalc as complex) and fmask, downloaded (the stub). + // The scale at this evaluation's Fcalc and Fmask (ModelScaleGPU): FitModelScale()'s solvent grid, and + // the overall scale and anisotropic B at a fixed solvent. FitSolvent() throws std::logic_error where + // the grid's fits would chain (too few reflections for the isotropic start), which the host then does. + void FitSolvent(double &k_sol, double &b_sol); + void FitScale(double k_sol, double b_sol, double &k_overall, double b_star[6]); + // The scale's points: fcmol (Fcalc as complex) and fmask, downloaded for the host's fit. void DownloadPoints(std::vector> &fcmol, std::vector> &fmask); // The residuals at the scale given, NumObservations() of them, into `residuals` (host). void Residuals(double k_overall, const double b_star[6], double k_sol, double b_sol, double *residuals);