Jungfraujoch/image_analysis/structure_refinement/MolecularReplacementGPU.cu

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

#include "MolecularReplacementGPU.h"

#include <algorithm>
#include <cmath>
#include <string>

#include <cufft.h>

#include "../../common/JFJochException.h"
#include "../indexing/CUDAMemHelpers.h"

namespace {

void cuda_err(cudaError_t val) {
    if (val != cudaSuccess)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
}

void cufft_err(cufftResult val) {
    if (val != CUFFT_SUCCESS)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                              "cuFFT error " + std::to_string(static_cast<int>(val)));
}

constexpr int THREADS = 256;
constexpr int MAX_BINS = 16;   // SEARCH_BINS (MolecularReplacement.cpp) is 10
// The rotation function's per-reflection scratch is made for this many orientations at a time.
constexpr size_t ROTATION_BATCH = 2048;

struct Transform {
    const float2 *data;
    int nu, nv, nwh;
    float box, blur;
};

struct Rot {
    float m[9];   // R^T, row-major: the model turned by R has F_R(s) = F_model(R^T s)
};

__device__ __forceinline__ float2 transform_at(const Transform &g, int i, int j, int l) {
    bool conj = false;
    if (l < 0) {   // the half-l grid holds l >= 0; F(-h) = conj F(h)
        i = -i; j = -j; l = -l;
        conj = true;
    }
    i %= g.nu; if (i < 0) i += g.nu;
    j %= g.nv; if (j < 0) j += g.nv;
    float2 v = g.data[i + static_cast<size_t>(g.nu) * (j + static_cast<size_t>(g.nv) * l)];
    if (conj)
        v.y = -v.y;
    return v;
}

// MolecularTransform::At on the device.
__device__ float2 transform_interpolate(const Transform &g, const Rot &r, float3 s0) {
    const float s[3] = {r.m[0] * s0.x + r.m[1] * s0.y + r.m[2] * s0.z,
                        r.m[3] * s0.x + r.m[4] * s0.y + r.m[5] * s0.z,
                        r.m[6] * s0.x + r.m[7] * s0.y + r.m[8] * s0.z};
    float f[3];
    int i0[3];
    for (int a = 0; a < 3; a++) {
        const float u = s[a] * g.box;
        i0[a] = static_cast<int>(floorf(u));
        f[a] = u - i0[a];
    }
    float2 out = make_float2(0, 0);
    for (int dx = 0; dx < 2; dx++)
        for (int dy = 0; dy < 2; dy++)
            for (int dz = 0; dz < 2; dz++) {
                const float w = (dx ? f[0] : 1 - f[0]) * (dy ? f[1] : 1 - f[1]) * (dz ? f[2] : 1 - f[2]);
                const float2 v = transform_at(g, i0[0] + dx, i0[1] + dy, i0[2] + dz);
                out.x += w * v.x;
                out.y += w * v.y;
            }
    const float unblur = expf(g.blur * (s[0] * s[0] + s[1] * s[1] + s[2] * s[2]) / 4.0f);
    return make_float2(out.x * unblur, out.y * unblur);
}

// A block-wide sum in a fixed order: each thread's value into shared memory, then a tree.
__device__ double block_sum(double v, double *sh) {
    sh[threadIdx.x] = v;
    __syncthreads();
    for (int stride = THREADS / 2; stride > 0; stride /= 2) {
        if (threadIdx.x < stride)
            sh[threadIdx.x] += sh[threadIdx.x + stride];
        __syncthreads();
    }
    const double out = sh[0];
    __syncthreads();
    return out;
}

// One block per orientation: |F|^2 summed over the copies for every reflection, binned, normalised by
// its bin's mean, and correlated with E^2 - 1.
__global__ void rotation_function(Transform g, const Rot *rot, const float3 *s, const float *e2, const int *bin,
                                  const int *bin_n, int n, int nops, int bins, double w_mean, double w_ss,
                                  float *scratch, double *out) {
    __shared__ double sh[THREADS];
    __shared__ double bin_mean[MAX_BINS];
    const Rot r = rot[blockIdx.x];
    float *ic = scratch + static_cast<size_t>(blockIdx.x) * n;
    double part[MAX_BINS];
    for (int b = 0; b < bins; b++)
        part[b] = 0;
    for (int i = threadIdx.x; i < n; i += THREADS) {
        float sum = 0;
        for (int j = 0; j < nops; j++) {
            const float2 f = transform_interpolate(g, r, s[static_cast<size_t>(i) * nops + j]);
            sum += f.x * f.x + f.y * f.y;
        }
        ic[i] = sum;
        part[bin[i]] += sum;
    }
    for (int b = 0; b < bins; b++) {
        const double total = block_sum(part[b], sh);
        if (threadIdx.x == 0)
            bin_mean[b] = total / bin_n[b];
    }
    __syncthreads();
    double x_sum = 0, x_ss = 0, xy = 0;
    for (int i = threadIdx.x; i < n; i += THREADS) {
        const double x = ic[i] / bin_mean[bin[i]];
        x_sum += x;
        x_ss += x * x;
        xy += x * (e2[i] - w_mean);
    }
    x_sum = block_sum(x_sum, sh);
    x_ss = block_sum(x_ss, sh);
    xy = block_sum(xy, sh);
    if (threadIdx.x == 0) {
        const double sxx = x_ss - x_sum * x_sum / n;
        out[blockIdx.x] = sxx > 0 && w_ss > 0 ? xy / sqrt(sxx * w_ss) : 0.0;
    }
}

// g_j = G_j(h) exp(2 pi i h.t_j) for every term and operator, and each term's sum of |g_j|^2.
__global__ void translation_copies(Transform g, Rot r, const float3 *s, const float2 *phase, int terms, int nops,
                                   float2 *copies, float *term_sum) {
    const int t = blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= terms)
        return;
    float sum = 0;
    for (int j = 0; j < nops; j++) {
        const size_t e = static_cast<size_t>(t) * nops + j;
        const float2 f = transform_interpolate(g, r, s[e]);
        const float2 p = phase[e];
        const float2 v = make_float2(f.x * p.x - f.y * p.y, f.x * p.y + f.y * p.x);
        copies[e] = v;
        sum += v.x * v.x + v.y * v.y;
    }
    term_sum[t] = sum;
}

// Each term's weight over its bin's mean |g|^2 sum, the bins summed in a fixed order by one block.
__global__ void translation_weights(const float *term_sum, const float *weight, const int *bin, const int *bin_n,
                                    int terms, int bins, float *scaled) {
    __shared__ double sh[THREADS];
    __shared__ double bin_mean[MAX_BINS];
    double part[MAX_BINS];
    for (int b = 0; b < bins; b++)
        part[b] = 0;
    for (int t = threadIdx.x; t < terms; t += THREADS)
        part[bin[t]] += term_sum[t];
    for (int b = 0; b < bins; b++) {
        const double total = block_sum(part[b], sh);
        if (threadIdx.x == 0)
            bin_mean[b] = total / bin_n[b];
    }
    __syncthreads();
    for (int t = threadIdx.x; t < terms; t += THREADS)
        scaled[t] = static_cast<float>(weight[t] / bin_mean[bin[t]]);
}

// Every coefficient of the series, gathered: a grid point sums the (term, pair) entries that land on it,
// in the order they were listed.
__global__ void translation_gather(const int *offset, const int *entry, const float2 *copies, const float *scaled,
                                   const int *pair_j, const int *pair_k, int nops, int npairs, size_t cells,
                                   float2 *coef) {
    const size_t c = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
    if (c >= cells)
        return;
    float2 sum = make_float2(0, 0);
    for (int e = offset[c]; e < offset[c + 1]; e++) {
        const int t = entry[e] / npairs, p = entry[e] % npairs;
        const float2 a = copies[static_cast<size_t>(t) * nops + pair_j[p]];
        const float2 b = copies[static_cast<size_t>(t) * nops + pair_k[p]];
        const float w = scaled[t];
        sum.x += w * (a.x * b.x + a.y * b.y);   // a conj(b)
        sum.y += w * (a.y * b.x - a.x * b.y);
    }
    coef[c] = sum;
}

// The cross terms with a known partial structure, added to the coefficients: w conj(F_known) g_j, or its
// conjugate, for every entry landing on a grid point, in the order they were listed.
__global__ void translation_gather_cross(const int *offset, const int *entry, const float2 *copies,
                                         const float *scaled, const float2 *known, const int *cross_term,
                                         const int *cross_op, const char *cross_conj, int nops, size_t cells,
                                         float2 *coef) {
    const size_t c = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
    if (c >= cells)
        return;
    float2 sum = coef[c];
    for (int e = offset[c]; e < offset[c + 1]; e++) {
        const int x = entry[e], t = cross_term[x];
        const float2 k = known[t];
        const float2 g = copies[static_cast<size_t>(t) * nops + cross_op[x]];
        const float w = scaled[t];
        const float re = k.x * g.x + k.y * g.y, im = k.x * g.y - k.y * g.x;   // conj(k) g
        sum.x += w * re;
        sum.y += w * (cross_conj[x] ? -im : im);
    }
    coef[c] = sum;
}

// Per block: the sum and the sum of squares of the function, and its highest point (the lowest index
// on a tie).
__global__ void translation_stats(const float *map, size_t n, double *sum, double *sum2, float *best,
                                  long long *best_index) {
    __shared__ double sh[THREADS];
    __shared__ float sh_best[THREADS];
    __shared__ long long sh_index[THREADS];
    double s = 0, s2 = 0;
    float b = -INFINITY;
    long long bi = -1;
    for (size_t i = static_cast<size_t>(blockIdx.x) * THREADS + threadIdx.x; i < n;
         i += static_cast<size_t>(gridDim.x) * THREADS) {
        const float v = map[i];
        s += v;
        s2 += static_cast<double>(v) * v;
        if (v > b) {
            b = v;
            bi = static_cast<long long>(i);
        }
    }
    s = block_sum(s, sh);
    s2 = block_sum(s2, sh);
    sh_best[threadIdx.x] = b;
    sh_index[threadIdx.x] = bi;
    __syncthreads();
    for (int stride = THREADS / 2; stride > 0; stride /= 2) {
        if (threadIdx.x < stride) {
            const float o = sh_best[threadIdx.x + stride];
            const long long oi = sh_index[threadIdx.x + stride];
            if (o > sh_best[threadIdx.x] || (o == sh_best[threadIdx.x] && oi >= 0 && oi < sh_index[threadIdx.x])) {
                sh_best[threadIdx.x] = o;
                sh_index[threadIdx.x] = oi;
            }
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        sum[blockIdx.x] = s;
        sum2[blockIdx.x] = s2;
        best[blockIdx.x] = sh_best[0];
        best_index[blockIdx.x] = sh_index[0];
    }
}

Rot transposed(const gemmi::Mat33 &m) {
    Rot r;
    for (int i = 0; i < 3; i++)
        for (int j = 0; j < 3; j++)
            r.m[3 * i + j] = static_cast<float>(m.a[j][i]);
    return r;
}

template <typename T>
CudaDevicePtr<T> upload(const std::vector<T> &v, cudaStream_t stream) {
    CudaDevicePtr<T> d(std::max<size_t>(v.size(), 1));
    if (!v.empty())
        cuda_err(cudaMemcpyAsync(d.get(), v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice, stream));
    return d;
}

std::vector<float3> float3s(const std::vector<gemmi::Vec3> &v) {
    std::vector<float3> out(v.size());
    for (size_t i = 0; i < v.size(); i++)
        out[i] = make_float3(static_cast<float>(v[i].x), static_cast<float>(v[i].y), static_cast<float>(v[i].z));
    return out;
}

std::vector<int> bin_counts(const std::vector<int> &bin, int bins) {
    std::vector<int> n(bins, 0);
    for (int b : bin)
        n[b]++;
    return n;
}

}  // namespace

struct MolecularReplacementGPUState {
    CudaStream stream;
    CudaDevicePtr<float2> grid;
    Transform transform{};
};

MolecularReplacementGPU::MolecularReplacementGPU(const MolecularTransform &ft)
    : state_(std::make_unique<MolecularReplacementGPUState>()) {
    const gemmi::FPhiGrid<float> &g = ft.Grid();
    std::vector<float2> host(g.data.size());
    for (size_t i = 0; i < host.size(); i++)
        host[i] = make_float2(g.data[i].real(), g.data[i].imag());
    state_->grid = upload(host, state_->stream);
    state_->transform = {state_->grid.get(), g.nu, g.nv, g.nw, static_cast<float>(ft.Box()),
                         static_cast<float>(ft.Blur())};
    cuda_err(cudaStreamSynchronize(state_->stream));
}

MolecularReplacementGPU::~MolecularReplacementGPU() = default;

std::vector<double> MolecularReplacementGPU::RotationFunction(const SearchReflections &refl,
                                                              const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                                                              const std::vector<gemmi::Mat33> &rotations) {
    cudaStream_t stream = state_->stream;
    const int n = static_cast<int>(refl.hkl.size());
    const std::vector<gemmi::Vec3> s = CopyVectors(refl, cell, sg);
    const int nops = static_cast<int>(s.size() / std::max(n, 1));
    double w_mean = 0;
    for (double e2 : refl.e2)
        w_mean += e2;
    w_mean /= n;
    double w_ss = 0;
    for (double e2 : refl.e2)
        w_ss += (e2 - w_mean) * (e2 - w_mean);
    const std::vector<float> e2(refl.e2.begin(), refl.e2.end());
    auto d_s = upload(float3s(s), stream);
    auto d_e2 = upload(e2, stream);
    auto d_bin = upload(refl.bin, stream);
    auto d_bin_n = upload(bin_counts(refl.bin, refl.bins), stream);
    std::vector<Rot> rot(rotations.size());
    for (size_t i = 0; i < rotations.size(); i++)
        rot[i] = transposed(rotations[i]);
    auto d_rot = upload(rot, stream);
    const size_t batch = std::min(ROTATION_BATCH, std::max<size_t>(rotations.size(), 1));
    CudaDevicePtr<float> scratch(batch * std::max(n, 1));
    CudaDevicePtr<double> d_out(std::max<size_t>(rotations.size(), 1));
    for (size_t first = 0; first < rotations.size(); first += batch) {
        const int count = static_cast<int>(std::min(batch, rotations.size() - first));
        rotation_function<<<count, THREADS, 0, stream>>>(state_->transform, d_rot.get() + first, d_s, d_e2, d_bin,
                                                         d_bin_n, n, nops, refl.bins, w_mean, w_ss, scratch,
                                                         d_out.get() + first);
        cuda_err(cudaGetLastError());
    }
    std::vector<double> out(rotations.size());
    cuda_err(cudaMemcpyAsync(out.data(), d_out.get(), out.size() * sizeof(double), cudaMemcpyDeviceToHost, stream));
    cuda_err(cudaStreamSynchronize(stream));
    return out;
}

std::vector<TranslationPeak> MolecularReplacementGPU::TranslationFunction(
        const TranslationSetup &setup, const std::vector<gemmi::Mat33> &rotations,
        const std::vector<std::complex<double>> *known) {
    std::vector<TranslationPeak> out(rotations.size());
    const int terms = static_cast<int>(setup.Terms()), nops = setup.ops;
    const int npairs = static_cast<int>(setup.pair_j.size());
    if (terms == 0 || rotations.empty() || (nops < 2 && known == nullptr))
        return out;
    cudaStream_t stream = state_->stream;
    const int nu = setup.nu, nv = setup.nv, nw = setup.nw, nwh = nw / 2 + 1;
    const size_t cells = static_cast<size_t>(nu) * nv * nwh, points = static_cast<size_t>(nu) * nv * nw;

    // The gather's lists: for every point of the half grid in cuFFT's layout ((u nv + v) nwh + l), the
    // (term, pair) entries landing on it, in entry order.
    std::vector<int> offset(cells + 1, 0);
    for (size_t e = 0; e < setup.target_u.size(); e++)
        if (setup.target_u[e] >= 0)
            offset[(static_cast<size_t>(setup.target_u[e]) * nv + setup.target_v[e]) * nwh + setup.target_l[e] + 1]++;
    for (size_t c = 0; c < cells; c++)
        offset[c + 1] += offset[c];
    std::vector<int> entry(offset[cells]);
    {
        std::vector<int> fill(offset.begin(), offset.end() - 1);
        for (size_t e = 0; e < setup.target_u.size(); e++)
            if (setup.target_u[e] >= 0)
                entry[fill[(static_cast<size_t>(setup.target_u[e]) * nv + setup.target_v[e]) * nwh
                           + setup.target_l[e]]++] = static_cast<int>(e);
    }
    // The same for the cross terms.
    std::vector<int> cross_offset(cells + 1, 0), cross_entry;
    std::vector<float2> known_f;
    if (known != nullptr) {
        for (size_t e = 0; e < setup.cross_term.size(); e++)
            cross_offset[(static_cast<size_t>(setup.cross_u[e]) * nv + setup.cross_v[e]) * nwh + setup.cross_l[e] + 1]++;
        for (size_t c = 0; c < cells; c++)
            cross_offset[c + 1] += cross_offset[c];
        cross_entry.resize(cross_offset[cells]);
        std::vector<int> fill(cross_offset.begin(), cross_offset.end() - 1);
        for (size_t e = 0; e < setup.cross_term.size(); e++)
            cross_entry[fill[(static_cast<size_t>(setup.cross_u[e]) * nv + setup.cross_v[e]) * nwh
                             + setup.cross_l[e]]++] = static_cast<int>(e);
        for (const std::complex<double> &k : *known)
            known_f.push_back(make_float2(static_cast<float>(k.real()), static_cast<float>(k.imag())));
    }
    auto d_cross_offset = upload(cross_offset, stream);
    auto d_cross_entry = upload(cross_entry, stream);
    auto d_known = upload(known_f, stream);
    auto d_cross_term = upload(setup.cross_term, stream);
    auto d_cross_op = upload(setup.cross_op, stream);
    auto d_cross_conj = upload(setup.cross_conj, stream);
    std::vector<float2> phase(setup.phase.size());
    for (size_t i = 0; i < phase.size(); i++)
        phase[i] = make_float2(static_cast<float>(setup.phase[i].real()), static_cast<float>(setup.phase[i].imag()));
    const std::vector<float> weight(setup.weight.begin(), setup.weight.end());
    auto d_s = upload(float3s(setup.s), stream);
    auto d_phase = upload(phase, stream);
    auto d_weight = upload(weight, stream);
    auto d_bin = upload(setup.bin, stream);
    auto d_bin_n = upload(bin_counts(setup.bin, setup.bins), stream);
    auto d_pair_j = upload(setup.pair_j, stream);
    auto d_pair_k = upload(setup.pair_k, stream);
    auto d_offset = upload(offset, stream);
    auto d_entry = upload(entry, stream);
    CudaDevicePtr<float2> copies(static_cast<size_t>(terms) * nops);
    CudaDevicePtr<float> term_sum(terms), scaled(terms);
    CudaDevicePtr<float2> coef(cells);
    CudaDevicePtr<float> map(points);
    const int stat_blocks = static_cast<int>(std::min<size_t>(1024, (points + THREADS - 1) / THREADS));
    const size_t nstat = static_cast<size_t>(stat_blocks) * rotations.size();
    CudaDevicePtr<double> d_sum(nstat), d_sum2(nstat);
    CudaDevicePtr<float> d_best(nstat);
    CudaDevicePtr<long long> d_best_index(nstat);

    cufftHandle plan;
    cufft_err(cufftPlan3d(&plan, nu, nv, nw, CUFFT_C2R));
    try {
        cufft_err(cufftSetStream(plan, stream));
        for (size_t r = 0; r < rotations.size(); r++) {
            const Rot rot = transposed(rotations[r]);
            translation_copies<<<(terms + THREADS - 1) / THREADS, THREADS, 0, stream>>>(
                    state_->transform, rot, d_s, d_phase, terms, nops, copies, term_sum);
            translation_weights<<<1, THREADS, 0, stream>>>(term_sum, d_weight, d_bin, d_bin_n, terms, setup.bins,
                                                           scaled);
            translation_gather<<<static_cast<unsigned>((cells + THREADS - 1) / THREADS), THREADS, 0, stream>>>(
                    d_offset, d_entry, copies, scaled, d_pair_j, d_pair_k, nops, npairs, cells, coef);
            if (known != nullptr)
                translation_gather_cross<<<static_cast<unsigned>((cells + THREADS - 1) / THREADS), THREADS, 0, stream>>>(
                        d_cross_offset, d_cross_entry, copies, scaled, d_known, d_cross_term, d_cross_op, d_cross_conj,
                        nops, cells, coef);
            cufft_err(cufftExecC2R(plan, reinterpret_cast<cufftComplex *>(coef.get()), map.get()));
            translation_stats<<<stat_blocks, THREADS, 0, stream>>>(
                    map, points, d_sum.get() + r * stat_blocks, d_sum2.get() + r * stat_blocks,
                    d_best.get() + r * stat_blocks, d_best_index.get() + r * stat_blocks);
            cuda_err(cudaGetLastError());
        }
    } catch (...) {
        cufftDestroy(plan);
        throw;
    }
    std::vector<double> sum(nstat), sum2(nstat);
    std::vector<float> best(nstat);
    std::vector<long long> best_index(nstat);
    cuda_err(cudaMemcpyAsync(sum.data(), d_sum.get(), nstat * sizeof(double), cudaMemcpyDeviceToHost, stream));
    cuda_err(cudaMemcpyAsync(sum2.data(), d_sum2.get(), nstat * sizeof(double), cudaMemcpyDeviceToHost, stream));
    cuda_err(cudaMemcpyAsync(best.data(), d_best.get(), nstat * sizeof(float), cudaMemcpyDeviceToHost, stream));
    cuda_err(cudaMemcpyAsync(best_index.data(), d_best_index.get(), nstat * sizeof(long long),
                             cudaMemcpyDeviceToHost, stream));
    cuda_err(cudaStreamSynchronize(stream));
    cufftDestroy(plan);

    for (size_t r = 0; r < rotations.size(); r++) {
        double s = 0, s2 = 0;
        float b = -INFINITY;
        long long bi = 0;
        for (int k = 0; k < stat_blocks; k++) {
            const size_t i = r * stat_blocks + k;
            s += sum[i];
            s2 += sum2[i];
            if (best[i] > b) {
                b = best[i];
                bi = best_index[i];
            }
        }
        const double count = static_cast<double>(points);
        const double mean = s / count, sd = std::sqrt(std::max(0.0, s2 / count - mean * mean));
        out[r].z = sd > 0 ? (b - mean) / sd : 0.0;
        const long long w = bi % nw, v = (bi / nw) % nv, u = bi / (static_cast<long long>(nw) * nv);
        out[r].t = gemmi::Fractional(static_cast<double>(u) / nu, static_cast<double>(v) / nv,
                                     static_cast<double>(w) / nw);
    }
    return out;
}