Jungfraujoch/image_analysis/structure_refinement/ModelStructureFactorsGPUEngine.h

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

#pragma once

// The device half of ModelStructureFactorsGPU (ModelStructureFactorsGPU.h). The host half works out with
// gemmi - which stays out of nvcc - everything that depends on the cell, the group and the resolution, and
// per evaluation the atoms' densities and positions; this half grids, transforms and composes. CUDA builds
// only.

#include <array>
#include <cstddef>
#include <memory>
#include <vector>

#include "ModelDensityGPU.h"
#include "ModelMaskGPU.h"

// Everything an engine needs that does not depend on the model's coordinates.
struct ModelStructureFactorsGPUSetup {
    ModelDensityGrid grid;                       // the structure-factor grid, u fastest
    ModelDensityGrid mask_grid;                  // the bulk-solvent mask's (SolventMaskGrid, ModelGrid.h): `grid`, or finer
    double volume = 0;
    // The group's operators as gemmi holds them: rot (row-major) and tran, in units of 1 / den.
    std::vector<std::array<int, 12>> sym_ops;
    int den = 24;
    std::vector<ModelMaskOp> mask_ops;           // every operator with every centring vector
    std::vector<std::array<int, 3>> rows;        // the reflections computed, in the order they come back
    std::vector<double> row_scale;               // n_cen * prepare_asu_data()'s unblur, per row
    size_t max_atoms = 0, max_pairs = 0;
    // The largest map MapFromFPhi() will be given: points of the real grid, and of its half-l transform.
    size_t map_points = 0, map_complex_points = 0;
};

struct ModelStructureFactorsGPUEngineImpl;

class ModelStructureFactorsGPUEngine {
public:
    // Every call names its device and leaves the calling thread's current device as it found it.

    // The device's total memory.
    static size_t TotalMemory(int device);
    // The cuFFT work area the structure factors' and the mask's r2c, and a map of (nu, nv, nw)'s c2r, take on `device`.
    static size_t FFTWorkBytes(int device, const ModelStructureFactorsGPUSetup &setup, int map_nu, int map_nv,
                               int map_nw);
    // The bytes an engine for `setup` reserves on the device, given its cuFFT work area.
    static size_t DeviceBytes(const ModelStructureFactorsGPUSetup &setup, size_t fft_work_bytes);

    ModelStructureFactorsGPUEngine(int device, const ModelStructureFactorsGPUSetup &setup, size_t fft_work_bytes);
    ~ModelStructureFactorsGPUEngine();

    // F_calc and F_mask at the rows for one set of atoms: each atom's density, its grid position
    // (fractional times the grid size, wrapped into [0, n); the 4th number is not read) and the atoms of
    // the bulk-solvent mask. Both come back as (re, im) per row.
    void Compute(const std::vector<ModelDensityAtom> &atoms, const std::vector<std::array<float, 4>> &pos,
                 const std::vector<ModelMaskAtom> &mask_atoms, std::vector<std::array<float, 2>> &fcalc,
                 std::vector<std::array<float, 2>> &fmask);

    // The real-space map rho(x) = 1/V sum F exp(-2 pi i h.x) of a half-l coefficient grid in gemmi's ZYX
    // order - h slowest, l fastest and halved: nu * nv * (nw / 2 + 1) complex numbers, (re, im)
    // interleaved, NaN read as zero - into `map`, nu * nv * nw floats with x fastest.
    void Map(int nu, int nv, int nw, double volume, const float *coefficients, float *map);

private:
    std::unique_ptr<ModelStructureFactorsGPUEngineImpl> impl_;
};