Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ZART: Symmetry and random sorting #962

Open
wants to merge 4 commits into
base: devel
Choose a base branch
from
Open

ZART: Symmetry and random sorting #962

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
48be7d0
Symmetry and random sorting for ZART
DavidHerreros Dec 12, 2024
3d49a93
Merge branch 'devel' into zart_cuda
DavidHerreros Dec 20, 2024
01f9d8f
Merge remote-tracking branch 'origin/devel' into zart_cuda
DavidHerreros Jan 25, 2025
0153fba
Merge remote-tracking branch 'origin/zart_cuda' into zart_cuda
DavidHerreros Jan 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 60 additions & 19 deletions src/xmipp/libraries/reconstruction_adapt_cuda11/forward_art_zernike3d_gpu.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
#include <fstream>
#include <iterator>
#include <numeric>
#include <random>
#include <stdexcept>
#include "data/cpu.h"

Expand Down Expand Up @@ -76,6 +77,8 @@ void ProgForwardArtZernike3DGPU::readParams()
debug_iter = checkParam("--debug_iter");
save_iter = getIntParam("--save_iter");
sort_last_N = getIntParam("--sort_last");
sort_random = checkParam("--sort_random");
fnSym = getParam("--sym");
FileName outPath = getParam("-o");
outPath = outPath.afterLastOf("/");
fnVolO = fnOutDir + "/" + outPath;
Expand Down Expand Up @@ -107,6 +110,7 @@ void ProgForwardArtZernike3DGPU::show() const
<< "Zernike Degree: " << L1 << std::endl
<< "SH Degree: " << L2 << std::endl
<< "Step: " << loop_step << std::endl
<< "Symmetry group: " << fnSym << std::endl
<< "Correct CTF: " << useCTF << std::endl
<< "Correct heretogeneity: " << useZernike << std::endl
<< "Remove negative values: " << removeNegValues << std::endl
Expand Down Expand Up @@ -142,6 +146,7 @@ void ProgForwardArtZernike3DGPU::defineParams()
addParamsLine(" [--ltk <ltv=1e-4>] : Tikhonov regualrization");
addParamsLine(" [--ll1 <ll1=1e-4>] : L1 regualrization");
addParamsLine(" [--lst <ll1=1e-4>] : Soft threshold regualrization");
addParamsLine(" [--sym <sym=c1>] : Symmetry to be considered during the reconstruction");
addParamsLine(" [--useZernike] : Correct heterogeneity with Zernike3D coefficients");
addParamsLine(" [--useCTF] : Correct CTF during ART reconstruction");
addParamsLine(" [--phaseFlipped] : Input images have been phase flipped");
Expand All @@ -152,6 +157,7 @@ void ProgForwardArtZernike3DGPU::defineParams()
addParamsLine(" [--save_iter <s=0>] : Save intermidiate volume after #save_iter iterations");
addParamsLine(
" [--sort_last <N=2>] : The algorithm sorts projections in the most orthogonally possible way. ");
addParamsLine(" [--sort_random] : Random sort of projections");
addParamsLine(
" : The most orthogonal way is defined as choosing the projection which "
"maximizes the ");
Expand Down Expand Up @@ -277,8 +283,8 @@ void ProgForwardArtZernike3DGPU::preProcess()
mask2D.setXmippOrigin();

vecSize = 0;
numCoefficients(L1, L2, vecSize);
fillVectorTerms(L1, L2, vL1, vN, vL2, vM);
numCoefficients(L1, L2);
fillVectorTerms(L1, L2);

initX = STARTINGX(Vrefined());
endX = FINISHINGX(Vrefined());
Expand Down Expand Up @@ -336,6 +342,20 @@ void ProgForwardArtZernike3DGPU::preProcess()
}


// Symmetry list
Matrix2D<double> L(3, 3), R(3, 3);
L.initIdentity(3);
R.initIdentity(3);
SymList SL_aux;
SL_aux.readSymmetryFile(fnSym);
LV.push_back(L);
RV.push_back(R);
for (int isym = 0; isym < SL_aux.symsNo(); isym++) {
SL_aux.getMatrices(isym, L, R, false);
LV.push_back(L);
RV.push_back(R);
}

// Create GPU interface
const cuda_forward_art_zernike3D::Program<PrecisionType>::ConstantParameters parameters = {
.VRecMaskF = VRecMaskF,
Expand Down Expand Up @@ -415,14 +435,29 @@ void ProgForwardArtZernike3DGPU::processImage(const FileName &fnImg,
I.read(fnImg);
I().setXmippOrigin();

// Forward Model
artModel<Direction::Forward>();
auto rot_dbl = static_cast<double>(rot);
auto tilt_dbl = static_cast<double>(tilt);
auto psi_dbl = static_cast<double>(psi);
double new_rot, new_tilt, new_psi;

for (int isym = 0; isym < RV.size(); isym++)
{

Euler_apply_transf(LV[isym], RV[isym], rot_dbl, tilt_dbl, psi_dbl, new_rot, new_tilt, new_psi);

rot = static_cast<PrecisionType>(new_rot);
tilt = static_cast<PrecisionType>(new_tilt);
psi = static_cast<PrecisionType>(new_psi);

// ART update
artModel<Direction::Backward>();
// Forward Model
artModel<Direction::Forward>();

// ART update
artModel<Direction::Backward>();
}
}

void ProgForwardArtZernike3DGPU::numCoefficients(int l1, int l2, int &vecSize)
void ProgForwardArtZernike3DGPU::numCoefficients(int l1, int l2)
{
for (int h = 0; h <= l2; h++) {
int numSPH = 2 * h + 1;
Expand All @@ -436,12 +471,7 @@ void ProgForwardArtZernike3DGPU::numCoefficients(int l1, int l2, int &vecSize)
}
}

void ProgForwardArtZernike3DGPU::fillVectorTerms(int l1,
int l2,
Matrix1D<int> &vL1,
Matrix1D<int> &vN,
Matrix1D<int> &vL2,
Matrix1D<int> &vM)
void ProgForwardArtZernike3DGPU::fillVectorTerms(int l1, int l2)
{
int idx = 0;
vL1.initZeros(vecSize);
Expand Down Expand Up @@ -497,7 +527,18 @@ void ProgForwardArtZernike3DGPU::run()

startProcessing();

sortOrthogonal();
if (sort_random)
{
std::vector<size_t> imgs_id(getInputMd()->size());
std::iota(imgs_id.begin(), imgs_id.end(), 0);
auto rng = std::default_random_engine {};
std::shuffle(std::begin(imgs_id), std::end(imgs_id), rng);
ordered_list.initZeros(getInputMd()->size());
for (size_t i=0; i < XSIZE(ordered_list); i++)
A1D_ELEM(ordered_list, i) = imgs_id[i];
}
else
sortOrthogonal();

if (!oroot.empty()) {
if (oext.empty())
Expand Down Expand Up @@ -593,12 +634,12 @@ void ProgForwardArtZernike3DGPU::sortOrthogonal()
double min_prod = MAXFLOAT;
;
int min_prod_proj = 0;
std::vector<double> rot;
std::vector<double> tilt;
std::vector<double> rotV;
std::vector<double> tiltV;
Matrix2D<double> v(numIMG, 3);
Matrix2D<double> euler;
getInputMd()->getColumnValues(MDL_ANGLE_ROT, rot);
getInputMd()->getColumnValues(MDL_ANGLE_TILT, tilt);
getInputMd()->getColumnValues(MDL_ANGLE_ROT, rotV);
getInputMd()->getColumnValues(MDL_ANGLE_TILT, tiltV);

// Initialization
ordered_list.resize(numIMG);
Expand All @@ -609,7 +650,7 @@ void ProgForwardArtZernike3DGPU::sortOrthogonal()

// Compute the Euler matrix for each image and keep only
// the third row of each one
Euler_angles2matrix(rot[i], tilt[i], 0., euler);
Euler_angles2matrix(rotV[i], tiltV[i], 0., euler);
euler.getRow(2, z);
v.setRow(i, z);
}
Expand Down
25 changes: 14 additions & 11 deletions src/xmipp/libraries/reconstruction_adapt_cuda11/forward_art_zernike3d_gpu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,13 +28,17 @@
#define _PROG_FORWARD_ART_ZERNIKE3D_GPU

#include <ctpl_stl.h>
#include <exception>
#include <fstream>
#include <numeric>
#include <core/matrix1d.h>
#include <core/xmipp_error.h>
#include <core/xmipp_image.h>
#include <core/xmipp_metadata_program.h>
#include <data/blobs.h>
#include <data/fourier_filter.h>
#include <data/fourier_projection.h>
#include <core/symmetries.h>
#include <reconstruction_cuda11/cuda_forward_art_zernike3d.h>

#include <memory>
Expand Down Expand Up @@ -85,10 +89,14 @@ class ProgForwardArtZernike3DGPU : public XmippMetadataProgram {
double ltv, ltk, ll1, lst, lmr;
// Remove negative values
bool removeNegValues;
// Symmetry
FileName fnSym;
// Symmetry list
std::vector<Matrix2D<double>> LV, RV;

public:
/** Resume computations */
bool resume;
bool resume = false;
// Number of ART iterations
int niter;
// Sort last N projections
Expand Down Expand Up @@ -121,6 +129,8 @@ class ProgForwardArtZernike3DGPU : public XmippMetadataProgram {
bool flip;
// CTF Check
bool hasCTF;
// Random image sorting
bool sort_random;
// Original defocus
double defocusU, defocusV, defocusAngle;
// CTF
Expand All @@ -134,7 +144,7 @@ class ProgForwardArtZernike3DGPU : public XmippMetadataProgram {
// Vector containing the degree of the spherical harmonics
std::vector<PrecisionType> clnm;
// Show optimization
bool showOptimization;
bool showOptimization = false;
// Row ids ordered in a orthogonal fashion
MultidimArray<size_t> ordered_list;
// Save iter counter
Expand Down Expand Up @@ -182,23 +192,16 @@ class ProgForwardArtZernike3DGPU : public XmippMetadataProgram {
An exception is thrown if any of the files is not found*/
void preProcess();

/** Create the processing working files.
* The working files are:
* nmaTodo.xmd for images to process (nmaTodo = mdIn - nmaDone)
* nmaDone.xmd image already processed (could exists from a previous run)
*/
// virtual void createWorkFiles();

/** Predict angles and shift.
At the input the pose parameters must have an initial guess of the
parameters. At the output they have the estimated pose.*/
void processImage(const FileName &fnImg, const FileName &fnImgOut, const MDRow &rowIn, MDRow &rowOut);

/// Length of coefficients vector
void numCoefficients(int l1, int l2, int &vecSize);
void numCoefficients(int l1, int l2);

/// Zernike and SPH coefficients allocation
void fillVectorTerms(int l1, int l2, Matrix1D<int> &vL1, Matrix1D<int> &vN, Matrix1D<int> &vL2, Matrix1D<int> &vM);
void fillVectorTerms(int l1, int l2);

///Deform a volumen using Zernike-Spherical harmonic basis
void deformVol(MultidimArray<double> &mP,
Expand Down