Jungfraujoch/image_analysis/structure_refinement/RigidBodyGPU.cpp

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "RigidBodyGPU.h"

#include <algorithm>
#include <cmath>

#include <Eigen/Dense>
#include <ceres/rotation.h>

#include "gemmi/dencalc.hpp"   // DensityCalculator
#include "gemmi/calculate.hpp"  // count_atom_sites
#include "gemmi/it92.hpp"      // IT92 x-ray form factors
#include "gemmi/scaling.hpp"   // Scaling

#include "ModelScaling.h"           // FitModelScale
#include "ModelStructureFactorsGPU.h"   // ModelDensityAtoms
#include "RigidBodyGPUEngine.h"
#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"
#include "../../common/Logger.h"

namespace {

using Table = gemmi::IT92<float>;

// What a zone needs of the model, the cell and the group: the grid, each atom's density and mask
// radius, the images and the scale's constraints. The same for every fit of a validation.
RigidBodyGPUZone ModelZone(const gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                           double d_min) {
    RigidBodyGPUZone zone;
    const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell, sg, d_min);
    zone.nu = grid.nu;
    zone.nv = grid.nv;
    zone.nw = grid.nw;
    for (int i = 0; i < 3; i++)
        for (int j = 0; j < 3; j++) {
            zone.orth[3 * i + j] = cell.orth.mat[i][j];
            zone.frac[3 * i + j] = cell.frac.mat[i][j];
        }
    zone.volume = cell.volume;

    gemmi::DensityCalculator<Table, float> dc;
    dc.d_min = d_min;
    dc.rate = 1.5;
    dc.grid.unit_cell = cell;
    dc.grid.spacegroup = &sg;
    dc.set_refmac_compatible_blur(model);
    zone.blur = dc.blur;
    ModelDensityAtoms(model, dc, zone.atoms, zone.mask_atom, zone.mask_radius);

    const gemmi::GroupOps gops = sg.operations();
    for (const gemmi::Op::Tran &cen : gops.cen_ops)
        for (const gemmi::Op &op : gops.sym_ops) {
            std::array<double, 12> image{};
            for (int i = 0; i < 3; i++) {
                for (int j = 0; j < 3; j++)
                    image[3 * i + j] = static_cast<double>(op.rot[i][j]) / gemmi::Op::DEN;
                image[9 + i] = static_cast<double>(op.tran[i] + cen[i]) / gemmi::Op::DEN;
            }
            zone.images.push_back(image);
        }
    for (const gemmi::Vec6 &c : gemmi::adp_symmetry_constraints(&sg))
        zone.constraints.push_back(c);
    return zone;
}

// The zone's observations and the composition of Fcalc at them.
void AddObservations(RigidBodyGPUZone &zone, const gemmi::UnitCell &cell, const SymmetryComposition &composition,
                     const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs) {
    zone.ops = composition.Ops();
    const size_t rows = composition.Hkl().size();
    zone.row_hkl.reserve(rows);
    zone.row_scale.reserve(rows);
    zone.row_stol2.reserve(rows);
    for (size_t m = 0; m < rows; m++) {
        const gemmi::Miller &h = composition.Hkl()[m];
        zone.row_hkl.push_back({h[0], h[1], h[2]});
        // prepare_asu_data()'s unblur, exp(B_blur |s|^2 / 4), as Compose() applies it
        zone.row_scale.push_back(composition.Centring() * std::exp(zone.blur * 0.25 * composition.InvD2()[m]));
        zone.row_stol2.push_back(cell.calculate_stol_sq(h));
    }
    zone.terms.reserve(composition.Terms().size());
    for (const SymmetryComposition::Term &t : composition.Terms()) {
        RigidBodyGPUTerm g{};
        // The half-u transform holds h >= 0, and F1(-k) = conj F1(k) for a real map.
        g.conj = t.k[0] < 0;
        const int ku = g.conj ? -t.k[0] : t.k[0], kv = g.conj ? -t.k[1] : t.k[1], kw = g.conj ? -t.k[2] : t.k[2];
        g.index = ku + (zone.nu / 2 + 1) * (gemmi::modulo(kv, zone.nv) + zone.nv * gemmi::modulo(kw, zone.nw));
        g.phase[0] = t.phase.real();
        g.phase[1] = t.phase.imag();
        g.s[0] = t.s.x;
        g.s[1] = t.s.y;
        g.s[2] = t.s.z;
        zone.terms.push_back(g);
    }
    double sum = 0;
    for (size_t i = 0; i < fobs.v.size(); i++) {
        const int m = composition.Row()[i];
        const gemmi::ValueSigma<float> &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)) {
            zone.point_obs.push_back(static_cast<int>(i));
            zone.point_hkl.push_back(zone.row_hkl[m]);
            zone.point_stol2.push_back(zone.row_stol2[m]);
            zone.point_fobs.push_back(vs.value);
            zone.point_sigma.push_back(vs.sigma);
        }
    }
    zone.f_mean = fobs.v.empty() ? 1.0 : sum / static_cast<double>(fobs.v.size());
}

bool SameObservations(const gemmi::AsuData<gemmi::ValueSigma<float>> &a,
                      const gemmi::AsuData<gemmi::ValueSigma<float>> &b) {
    if (a.v.size() != b.v.size())
        return false;
    for (size_t i = 0; i < a.v.size(); i++)
        if (a.v[i].hkl != b.v[i].hkl || !(a.v[i].value.value == b.v[i].value.value) ||
            !(a.v[i].value.sigma == b.v[i].value.sigma))
            return false;
    return true;
}

}  // namespace

std::unique_ptr<RigidBodyGPUPool> RigidBodyGPUPool::Plan(const gemmi::Model &model, const gemmi::UnitCell &cell,
                                                         const gemmi::SpaceGroup &sg, double d_min,
                                                         size_t max_observations, Logger &logger) {
    if (get_gpu_count() == 0)
        return nullptr;
    std::unique_ptr<RigidBodyGPUPool> pool(new RigidBodyGPUPool);
    pool->cell_ = cell;
    pool->d_min_ = d_min;
    pool->max_observations_ = max_observations;
    pool->cap_ = std::make_unique<RigidBodyGPUCapacity>();
    pool->cap_->observations = max_observations;
    pool->cap_->rows = max_observations;
    if (!pool->Cover(model, sg, logger))
        return nullptr;
    try {
        pool->device_ = RigidBodyGPUEngine::CurrentDevice();
    } catch (const JFJochException &e) {
        throw RigidBodyGPUFailure(e.what());
    }
    return pool;
}

bool RigidBodyGPUPool::Cover(const gemmi::Model &model, const gemmi::SpaceGroup &sg, Logger &logger) {
    try {
        // Sized for the largest zone of the ladder, which every fit walks a part of.
        RigidBodyGPUCapacity &cap = *cap_;
        for (double zone_d : RigidBodyLadder(d_min_)) {
            const RigidBodyGPUZone &zone = Zone(model, cell_, sg, zone_d);
            if (!RigidBodyGPUEngine::Supports(zone)) {
                logger.Info("Model validation: the rigid body runs on the CPU - the cell is too small for the "
                            "GPU's gridding at {:.1f} A", zone_d);
                return false;
            }
            cap.atoms = std::max(cap.atoms, zone.atoms.size());
            cap.grid_points = std::max(cap.grid_points, static_cast<size_t>(zone.nu) * zone.nv * zone.nw);
            cap.complex_points = std::max(cap.complex_points, static_cast<size_t>(zone.nu / 2 + 1) * zone.nv * zone.nw);
            cap.bricks = std::max(cap.bricks, RigidBodyGPUEngine::Bricks(zone.nu, zone.nv, zone.nw));
            cap.pairs = std::max(cap.pairs, RigidBodyGPUEngine::PairBound(zone));
            cap.fft_work_bytes = std::max(cap.fft_work_bytes, RigidBodyGPUEngine::FFTWorkBytes(zone.nu, zone.nv, zone.nw));
        }
        cap.terms = std::max(cap.terms, max_observations_ * sg.operations().sym_ops.size());
        return true;
    } catch (const JFJochException &e) {
        throw RigidBodyGPUFailure(e.what());
    }
}

size_t RigidBodyGPUPool::EngineBytes() const {
    return RigidBodyGPUEngine::DeviceBytes(*cap_);
}

void RigidBodyGPUPool::Reserve(const std::vector<size_t> &per_card) {
    // The engines give the same bits on any card of one model, so where they are changes only where they run.
    // A card that is merely busy is waited for by the engines' allocations, never left for the CPU.
    try {
        const int cards = get_gpu_count();
        const size_t rounds = per_card.empty() ? 0 : *std::max_element(per_card.begin(), per_card.end());
        for (size_t round = 0; round < rounds; round++)
            for (size_t k = 0; k < per_card.size(); k++)
                if (round < per_card[k])
                    engines_.push_back(std::make_unique<RigidBodyGPUEngine>(*cap_, (device_ + static_cast<int>(k)) % cards));
        set_gpu(device_);   // each engine was made on its own card
    } catch (const JFJochException &e) {
        throw RigidBodyGPUFailure(e.what());
    }
    for (auto &e : engines_)
        idle_.push_back(e.get());
}

RigidBodyGPUPool::~RigidBodyGPUPool() = default;

std::shared_ptr<const RigidBodyGPUZone> RigidBodyGPUPool::ObservedZone(
        const gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
        const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs, double d_min) {
    const RigidBodyGPUZone &model_zone = Zone(model, cell, sg, d_min);
    const size_t atoms = gemmi::count_atom_sites(model);
    std::lock_guard lock(observed_m_);
    for (const ObservedEntry &o : observed_)
        if (o.group == sg.xhm() && o.atoms == atoms && o.d_min == d_min && SameObservations(o.fobs, fobs))
            return o.zone;
    const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell, sg, d_min);
    std::vector<gemmi::Miller> hkl;
    hkl.reserve(fobs.v.size());
    for (const auto &hv : fobs.v)
        hkl.push_back(hv.hkl);
    const SymmetryComposition composition(grid, d_min, hkl);
    auto zone = std::make_shared<RigidBodyGPUZone>(model_zone);
    AddObservations(*zone, cell, composition, fobs);
    observed_.push_back({sg.xhm(), atoms, d_min, fobs, zone});
    return zone;
}

const RigidBodyGPUZone &RigidBodyGPUPool::Zone(const gemmi::Model &model, const gemmi::UnitCell &cell,
                                              const gemmi::SpaceGroup &sg, double d_min) {
    const std::tuple<std::string, size_t, double> key{sg.xhm(), gemmi::count_atom_sites(model), d_min};
    std::lock_guard lock(zones_m_);
    auto it = zones_.find(key);
    if (it == zones_.end())
        it = zones_.emplace(key, std::make_unique<RigidBodyGPUZone>(ModelZone(model, cell, sg, d_min))).first;
    return *it->second;
}

RigidBodyGPUEngine &RigidBodyGPUPool::Acquire() {
    // One on the calling thread's own card where one is idle - the threads that drive a validation are
    // pinned to the cards in turn - and any other otherwise.
    const int device = RigidBodyGPUEngine::CurrentDevice();
    std::unique_lock lock(m_);
    cv_.wait(lock, [this] { return !idle_.empty(); });
    size_t pick = idle_.size() - 1;
    for (size_t i = idle_.size(); i-- > 0;)
        if (idle_[i]->Device() == device) {
            pick = i;
            break;
        }
    RigidBodyGPUEngine *e = idle_[pick];
    idle_.erase(idle_.begin() + static_cast<std::ptrdiff_t>(pick));
    return *e;
}

void RigidBodyGPUPool::Release(RigidBodyGPUEngine &engine) {
    {
        std::lock_guard lock(m_);
        idle_.push_back(&engine);
    }
    cv_.notify_one();
}

RigidBodyTargetGPU::RigidBodyTargetGPU(RigidBodyGPUPool &pool, gemmi::Model &model, const gemmi::UnitCell &cell,
                                       const gemmi::SpaceGroup &sg, size_t nthreads)
    : RigidBodyTargetBase(model), pool_(pool), lease_(pool), engine_(lease_.Engine()), model_(model), cell_(cell),
      sg_(sg), nthreads_(nthreads) {
    std::vector<std::array<double, 3>> relative;
    for (const gemmi::Position &p : base_)
        relative.push_back({p.x - centre_.x, p.y - centre_.y, p.z - centre_.z});
    try {
        engine_.SetBody(relative);
    } catch (const JFJochException &e) {
        throw RigidBodyGPUFailure(e.what());
    }
}

RigidBodyTargetGPU::~RigidBodyTargetGPU() = default;

// Place()'s placement as a matrix: the columns are the rotated axes, rotated as Place() rotates a point.
void RigidBodyTargetGPU::Placement(const double q[6], double rotation[9], double translation[3]) const {
    const double aa[3] = {q[0] / rms_radius_, q[1] / rms_radius_, q[2] / rms_radius_};
    for (int j = 0; j < 3; j++) {
        const double e[3] = {j == 0 ? 1.0 : 0.0, j == 1 ? 1.0 : 0.0, j == 2 ? 1.0 : 0.0};
        double column[3];
        ceres::AngleAxisRotatePoint(aa, e, column);
        for (int i = 0; i < 3; i++)
            rotation[3 * i + j] = column[i];
    }
    translation[0] = centre_.x + q[3];
    translation[1] = centre_.y + q[4];
    translation[2] = centre_.z + q[5];
}

void RigidBodyTargetGPU::SetZone(const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs, double d_min) {
    // The body where this zone starts, for the mask it shares with fixed atoms: the last evaluation's.
    std::array<double, 6> q_start{};
    if (have_point_)
        q_start = q_;
    d_min_ = d_min;
    solvent_fitted_ = false;
    host_scale_ = false;
    have_point_ = false;
    zone_ = pool_.ObservedZone(model_, cell_, sg_, fobs, d_min);
    try {
        engine_.SetZone(*zone_);
        if (fixed != nullptr) {
            // The fixed atoms are copies of the body (RefineRigidBody): each one's placement of the body's
            // atoms, by the least-squares superposition of the two.
            // Following Kabsch (1976) Acta Cryst. A32, 922-923
            const std::vector<gemmi::Position> all = ModelPositions(*fixed);
            const size_t n = base_.size();
            std::vector<std::array<double, 12>> copies;
            for (size_t c = 0; c + n <= all.size(); c += n) {
                Eigen::Vector3d mean = Eigen::Vector3d::Zero();
                for (size_t i = 0; i < n; i++)
                    mean += Eigen::Vector3d(all[c + i].x, all[c + i].y, all[c + i].z);
                mean /= static_cast<double>(n);
                Eigen::Matrix3d h = Eigen::Matrix3d::Zero();
                for (size_t i = 0; i < n; i++) {
                    const Eigen::Vector3d p(base_[i].x - centre_.x, base_[i].y - centre_.y, base_[i].z - centre_.z);
                    const Eigen::Vector3d q(all[c + i].x - mean.x(), all[c + i].y - mean.y(), all[c + i].z - mean.z());
                    h += p * q.transpose();
                }
                const Eigen::JacobiSVD<Eigen::Matrix3d> svd(h, Eigen::ComputeFullU | Eigen::ComputeFullV);
                Eigen::Matrix3d d = Eigen::Matrix3d::Identity();
                d(2, 2) = (svd.matrixV() * svd.matrixU().transpose()).determinant() < 0 ? -1.0 : 1.0;
                const Eigen::Matrix3d r = svd.matrixV() * d * svd.matrixU().transpose();
                // The body's relative coordinates have their mean at zero, so the copy's centroid is t.
                std::array<double, 12> x{};
                for (int i = 0; i < 3; i++)
                    for (int j = 0; j < 3; j++)
                        x[3 * i + j] = r(i, j);
                x[9] = mean.x();
                x[10] = mean.y();
                x[11] = mean.z();
                copies.push_back(x);
            }
            std::array<double, 12> moving{};
            Placement(q_start.data(), moving.data(), moving.data() + 9);
            engine_.SetFixedCopies(copies, moving);
        }
    } catch (const JFJochException &e) {
        throw RigidBodyGPUFailure(e.what());
    }
}

// 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<std::array<float, 2>> fc, fm;
    engine_.DownloadPoints(fc, fm);
    gemmi::Scaling<float> scaling(cell_, &sg_);
    scaling.use_solvent = true;
    for (size_t p = 0; p < zone_->point_obs.size(); p++) {
        const std::array<int, 3> &h = zone_->point_hkl[p];
        scaling.points.push_back({{h[0], h[1], h[2]}, zone_->point_stol2[p], std::complex<float>(fc[p][0], fc[p][1]),
                                  std::complex<float>(fm[p][0], fm[p][1]), zone_->point_fobs[p], zone_->point_sigma[p]});
    }
    if (!solvent_fitted_) {
        FitModelScale(scaling, {.b_hi = solvent_b_max_}, 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) {
    try {
        ++evaluations;
        double rotation[9], translation[3];
        Placement(q, rotation, translation);
        if (zone_->row_hkl.empty())
            return false;
        engine_.Fcalc(rotation, translation);
        if (fixed == nullptr && !(hold_mask && have_point_))
            engine_.Fmask();

        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_ && !host_scale_) {
            try {
                engine_.FitSolvent(solvent_b_max_, k_sol, b_sol);
                solvent_fitted_ = true;
            } catch (const ModelScaleGPUTooFewReflections &) {
                host_scale_ = true;
            }
        }
        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<int>(std::count(zone_->obs_row.begin(), zone_->obs_row.end(), -1));

        have_point_ = true;
        std::copy(q, q + 6, q_.begin());
        return true;
    } catch (const JFJochException &e) {
        throw RigidBodyGPUFailure(e.what());
    }
}

// RigidBodyTarget::Jacobian(): the rotation columns by forward difference, three copies of the body
// gridded and transformed together, the translation columns exact, the scale re-fit projected out.
bool RigidBodyTargetGPU::Jacobian(const double q[6], double *jacobian) {
    if (!have_point_ || !std::equal(q, q + 6, q_.begin())) {
        std::vector<double> residuals(NumObservations());
        if (!Residuals(q, residuals.data()))
            return false;
    }
    try {
        ++jacobians;
        const double step = RigidBodyJacobianStep(d_min_);
        double rotation[3][9], translation[3][3];
        for (int j = 0; j < 3; j++) {
            double qj[6];
            std::copy(q_.begin(), q_.end(), qj);
            qj[j] += step;
            Placement(qj, rotation[j], translation[j]);
        }
        const double b_star[6] = {b_star_.u11, b_star_.u22, b_star_.u33, b_star_.u12, b_star_.u13, b_star_.u23};
        std::vector<double> jtj, jtq;
        engine_.Jacobian(rotation, translation, step, k_overall_, b_star, k_sol, b_sol, jtj, jtq);

        // Following Golub & Pereyra (1973) SIAM J. Numer. Anal. 10, 413-432, in the form of Kaufman (1975) BIT 15, 49-57
        const int p = static_cast<int>(zone_->constraints.size()) + 1;
        Eigen::MatrixXd a(p, p), c(p, 6);
        for (int i = 0; i < p; i++) {
            for (int j = 0; j < p; j++)
                a(i, j) = jtj[i * p + j];
            for (int j = 0; j < 6; j++)
                c(i, j) = jtq[i * 6 + j];
        }
        const Eigen::MatrixXd x = a.ldlt().solve(c);
        std::vector<double> xv(p * 6);
        for (int i = 0; i < p; i++)
            for (int j = 0; j < 6; j++)
                xv[i * 6 + j] = x(i, j);
        engine_.ProjectJacobian(xv, jacobian);
        return true;
    } catch (const JFJochException &e) {
        throw RigidBodyGPUFailure(e.what());
    }
}

size_t RigidBodyTargetGPU::NumObservations() const {
    return zone_ ? zone_->obs_row.size() : 0;
}