bioem.cpp 32.2 KB
Newer Older
Pilar Cossio's avatar
License  
Pilar Cossio committed
1 2 3 4
/* ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
        < BioEM software for Bayesian inference of Electron Microscopy images>
            Copyright (C) 2014 Pilar Cossio, David Rohr and Gerhard Hummer.
            Max Planck Institute of Biophysics, Frankfurt, Germany.
5

Pilar Cossio's avatar
License  
Pilar Cossio committed
6 7 8 9
                See license statement for terms of distribution.

   ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++*/

10 11 12 13 14 15 16 17
#include <mpi.h>

#define MPI_CHK(expr) \
	if (expr != MPI_SUCCESS) \
	{ \
		fprintf(stderr, "Error in MPI function %s: %d\n", __FILE__, __LINE__); \
	}

18 19 20 21 22 23 24 25 26
#include <fstream>
#include <boost/program_options.hpp>
#include <iostream>
#include <algorithm>
#include <iterator>
#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <cmath>
27

28
#ifdef WITH_OPENMP
29
#include <omp.h>
30
#endif
31 32 33 34 35 36 37 38 39 40

#include <fftw3.h>
#include <math.h>
#include "cmodules/timer.h"

#include "param.h"
#include "bioem.h"
#include "model.h"
#include "map.h"

41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63
#ifdef BIOEM_USE_NVTX
#include "nvToolsExt.h"

const uint32_t colors[] = { 0x0000ff00, 0x000000ff, 0x00ffff00, 0x00ff00ff, 0x0000ffff, 0x00ff0000, 0x00ffffff };
const int num_colors = sizeof(colors)/sizeof(colors[0]);

#define cuda_custom_timeslot(name,cid) { \
	int color_id = cid; \
	color_id = color_id%num_colors;\
	nvtxEventAttributes_t eventAttrib = {0}; \
	eventAttrib.version = NVTX_VERSION; \
	eventAttrib.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE; \
	eventAttrib.colorType = NVTX_COLOR_ARGB; \
	eventAttrib.color = colors[color_id]; \
	eventAttrib.messageType = NVTX_MESSAGE_TYPE_ASCII; \
	eventAttrib.message.ascii = name; \
	nvtxRangePushEx(&eventAttrib); \
}
#define cuda_custom_timeslot_end nvtxRangePop();
#else
#define cuda_custom_timeslot(name,cid)
#define cuda_custom_timeslot_end
#endif
64

65 66 67 68 69 70 71 72 73 74 75
#include "bioem_algorithm.h"

using namespace boost;
namespace po = boost::program_options;

using namespace std;

// A helper function of Boost
template<class T>
ostream& operator<<(ostream& os, const vector<T>& v)
{
76 77
	copy(v.begin(), v.end(), ostream_iterator<T>(os, " "));
	return os;
78 79 80 81
}

bioem::bioem()
{
82
	FFTAlgo = getenv("FFTALGO") == NULL ? 1 : atoi(getenv("FFTALGO"));
83
	DebugOutput = getenv("BIOEM_DEBUG_OUTPUT") == NULL ? 2 : atoi(getenv("BIOEM_DEBUG_OUTPUT"));
84 85 86 87 88 89 90 91
}

bioem::~bioem()
{
}

int bioem::configure(int ac, char* av[])
{
David Rohr's avatar
David Rohr committed
92 93 94 95 96
	// **************************************************************************************
	// **** Configuration Routine using boost for extracting parameters, models and maps ****
	// **************************************************************************************
	// ****** And Precalculating necessary grids, map crosscorrelations and kernels  ********
	// *************************************************************************************
97

David Rohr's avatar
David Rohr committed
98 99
	HighResTimer timer;

David Rohr's avatar
David Rohr committed
100 101 102 103 104
	if (mpi_rank == 0)
	{
		// *** Inizialzing default variables ***
		std::string infile, modelfile, mapfile;
		Model.readPDB = false;
105
		param.param_device.writeAngles = false;
David Rohr's avatar
David Rohr committed
106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205
		param.dumpMap = false;
		param.loadMap = false;
		RefMap.readMRC = false;
		RefMap.readMultMRC = false;

		// *************************************************************************************
		cout << " ++++++++++++ FROM COMMAND LINE +++++++++++\n\n";
		// *************************************************************************************

		// ********************* Command line reading input with BOOST ************************

		try {
			po::options_description desc("Command line inputs");
			desc.add_options()
			("Inputfile", po::value<std::string>(), "(Mandatory) Name of input parameter file")
			("Modelfile", po::value< std::string>() , "(Mandatory) Name of model file")
			("Particlesfile", po::value< std::string>(), "(Mandatory) Name of paricles file")
			("ReadPDB", "(Optional) If reading model file in PDB format")
			("ReadMRC", "(Optional) If reading particle file in MRC format")
			("ReadMultipleMRC", "(Optional) If reading Multiple MRCs")
			("DumpMaps", "(Optional) Dump maps after they were red from maps file")
			("LoadMapDump", "(Optional) Read Maps from dump instead of maps file")
			("help", "(Optional) Produce help message")
			;


			po::positional_options_description p;
			p.add("Inputfile", -1);
			p.add("Modelfile", -1);
			p.add("Particlesfile", -1);
			p.add("ReadPDB", -1);
			p.add("ReadMRC", -1);
			p.add("ReadMultipleMRC", -1);
			p.add("DumpMaps", -1);
			p.add("LoadMapDump", -1);

			po::variables_map vm;
			po::store(po::command_line_parser(ac, av).
					  options(desc).positional(p).run(), vm);
			po::notify(vm);

			if((ac < 6)) {
				std::cout << desc << std::endl;
				return 1;
			}
			if (vm.count("help")) {
				cout << "Usage: options_description [options]\n";
				cout << desc;
				return 1;
			}

			if (vm.count("Inputfile"))
			{
				cout << "Input file is: ";
				cout << vm["Inputfile"].as< std::string >() << "\n";
				infile = vm["Inputfile"].as< std::string >();
			}
			if (vm.count("Modelfile"))
			{
				cout << "Model file is: "
					 << vm["Modelfile"].as<  std::string  >() << "\n";
				modelfile = vm["Modelfile"].as<  std::string  >();
			}

			if (vm.count("ReadPDB"))
			{
				cout << "Reading model file in PDB format.\n";
				Model.readPDB = true;
			}

			if (vm.count("ReadMRC"))
			{
				cout << "Reading particle file in MRC format.\n";
				RefMap.readMRC=true;
			}

			if (vm.count("ReadMultipleMRC"))
			{
				cout << "Reading Multiple MRCs.\n";
				RefMap.readMultMRC=true;
			}

			if (vm.count("DumpMaps"))
			{
				cout << "Dumping Maps after reading from file.\n";
				param.dumpMap = true;
			}

			if (vm.count("LoadMapDump"))
			{
				cout << "Loading Map dump.\n";
				param.loadMap = true;
			}

			if (vm.count("Particlesfile"))
			{
				cout << "Paricle file is: "
					 << vm["Particlesfile"].as< std::string >() << "\n";
				mapfile = vm["Particlesfile"].as< std::string >();
			}
David Rohr's avatar
David Rohr committed
206
		}
David Rohr's avatar
David Rohr committed
207
		catch(std::exception& e)
208
		{
David Rohr's avatar
David Rohr committed
209 210
			cout << e.what() << "\n";
			return 1;
211
		}
David Rohr's avatar
David Rohr committed
212 213 214 215 216 217 218 219 220

		//check for consitency in multiple MRCs
		if(RefMap.readMultMRC && not(RefMap.readMRC))
		{
			cout << "For Multiple MRCs command --ReadMRC is necesary too";
			exit(1);
		}

		if (DebugOutput >= 2 && mpi_rank == 0) timer.ResetStart();
David Rohr's avatar
David Rohr committed
221
		// ********************* Reading Parameter Input ***************************
222
		param.readParameters(infile.c_str());
David Rohr's avatar
David Rohr committed
223 224

		// ********************* Reading Model Input ******************************
225
		Model.readModel(modelfile.c_str());
David Rohr's avatar
David Rohr committed
226 227

		// ********************* Reading Particle Maps Input **********************
228
		RefMap.readRefMaps(param, mapfile.c_str());
David Rohr's avatar
David Rohr committed
229
		if (DebugOutput >= 2 && mpi_rank == 0) printf("Reading Input Data %f\n", timer.GetCurrentElapsedTime());
David Rohr's avatar
David Rohr committed
230
	}
231

David Rohr's avatar
David Rohr committed
232
#ifdef WITH_MPI
233 234 235 236 237 238 239 240 241 242 243 244 245
	if (mpi_size > 1)
	{
		if (DebugOutput >= 2 && mpi_rank == 0) timer.ResetStart();
		MPI_Bcast(&param, sizeof(param), MPI_BYTE, 0, MPI_COMM_WORLD);
		//refCtf, CtfParam, angles automatically filled by precalculare function below

		MPI_Bcast(&Model, sizeof(Model), MPI_BYTE, 0, MPI_COMM_WORLD);
		if (mpi_rank != 0) Model.points = (bioem_model::bioem_model_point*) mallocchk(sizeof(bioem_model::bioem_model_point) * Model.nPointsModel);
		MPI_Bcast(Model.points, sizeof(bioem_model::bioem_model_point) * Model.nPointsModel, MPI_BYTE, 0, MPI_COMM_WORLD);

		MPI_Bcast(&RefMap, sizeof(RefMap), MPI_BYTE, 0, MPI_COMM_WORLD);
		if (mpi_rank != 0) RefMap.maps = (myfloat_t*) mallocchk(RefMap.refMapSize * sizeof(myfloat_t) * RefMap.ntotRefMap);
		MPI_Bcast(RefMap.maps, RefMap.refMapSize * sizeof(myfloat_t) * RefMap.ntotRefMap, MPI_BYTE, 0, MPI_COMM_WORLD);
David Rohr's avatar
David Rohr committed
246
		if (DebugOutput >= 2 && mpi_rank == 0) printf("MPI Broadcast of Input Data %f\n", timer.GetCurrentElapsedTime());
247 248
	}
	#endif
249

David Rohr's avatar
David Rohr committed
250
	// ****************** Precalculating Necessary Stuff *********************
251 252

	if (DebugOutput >= 2 && mpi_rank == 0) timer.ResetStart();
253
	param.PrepareFFTs();
254 255 256 257 258
	if (DebugOutput >= 2 && mpi_rank == 0)
	{
		printf("Time Prepare FFTs %f\n", timer.GetCurrentElapsedTime());
		timer.ResetStart();
	}
259
	precalculate();
David Rohr's avatar
David Rohr committed
260

261 262
	if (getenv("BIOEM_DEBUG_BREAK"))
	{
263 264 265
		const int cut = atoi(getenv("BIOEM_DEBUG_BREAK"));
		if (param.nTotGridAngles > cut) param.nTotGridAngles = cut;
		if (param.nTotCTFs > cut) param.nTotCTFs = cut;
266
	}
David Rohr's avatar
David Rohr committed
267

268 269 270 271 272
	if (DebugOutput >= 2 && mpi_rank == 0)
	{
		printf("Time Precalculate %f\n", timer.GetCurrentElapsedTime());
		timer.ResetStart();
	}
273
	pProb.init(RefMap.ntotRefMap, param.nTotGridAngles, param.nTotCC, *this);
274

275 276 277 278 279
	if (DebugOutput >= 2 && mpi_rank == 0)
	{
		printf("Time Init Probabilities %f\n", timer.GetCurrentElapsedTime());
		timer.ResetStart();
	}
280
	deviceInit();
281
	if (DebugOutput >= 2 && mpi_rank == 0) printf("Time Device Init %f\n", timer.GetCurrentElapsedTime());
282

283
	return(0);
284 285
}

286 287 288
void bioem::cleanup()
{
	//Deleting allocated pointers
289
	free_device_host(pProb.ptr);
290 291 292
	RefMap.freePointers();
}

293 294
int bioem::precalculate()
{
David Rohr's avatar
David Rohr committed
295
	// **************************************************************************************
David Rohr's avatar
David Rohr committed
296
	// **Precalculating Routine of Orientation grids, Map crosscorrelations and CTF Kernels**
David Rohr's avatar
David Rohr committed
297
	// **************************************************************************************
298
	HighResTimer timer;
299

300
	// Generating Grids of orientations
301
	if (DebugOutput >= 3) timer.ResetStart();
302
	param.CalculateGridsParam();
303

304 305 306 307 308
	if (DebugOutput >= 3)
	{
		printf("\tTime Precalculate Grids Param: %f\n", timer.GetCurrentElapsedTime());
		timer.ResetStart();
	}
309 310
	// Precalculating CTF Kernels stored in class Param
	param.CalculateRefCTF();
311

312 313 314 315 316
	if (DebugOutput >= 3)
	{
		printf("\tTime Precalculate CTFs: %f\n", timer.GetCurrentElapsedTime());
		timer.ResetStart();
	}
317 318
	//Precalculate Maps
	RefMap.precalculate(param, *this);
319
	if (DebugOutput >= 3) printf("\tTime Precalculate Maps: %f\n", timer.GetCurrentElapsedTime());
320

321
	return(0);
322 323 324 325
}

int bioem::run()
{
David Rohr's avatar
David Rohr committed
326 327 328
	// **************************************************************************************
	// **** Main BioEM routine, projects, convolutes and compares with Map using OpenMP ****
	// **************************************************************************************
David Rohr's avatar
David Rohr committed
329

330 331
	if (param.printModel)
	{
David Rohr's avatar
David Rohr committed
332

333 334 335
		//....
		return(0);
	}
336

David Rohr's avatar
David Rohr committed
337 338
	// **** If we want to control the number of threads -> omp_set_num_threads(XX); ******
	// ****************** Declarying class of Probability Pointer  *************************
339

David Rohr's avatar
David Rohr committed
340
	if (mpi_rank == 0) printf("\tInitializing Probabilities\n");
341 342 343
	// Inizialzing Probabilites to zero and constant to -Infinity
	for (int iRefMap = 0; iRefMap < RefMap.ntotRefMap; iRefMap ++)
	{
344 345 346 347
		bioem_Probability_map& pProbMap = pProb.getProbMap(iRefMap);

		pProbMap.Total = 0.0;
		pProbMap.Constoadd = -9999999;
348
		if (param.param_device.writeAngles)
349
		{
350 351 352
			for (int iOrient = 0; iOrient < param.nTotGridAngles; iOrient ++)
			{
				bioem_Probability_angle& pProbAngle = pProb.getProbAngle(iRefMap, iOrient);
353

354 355 356
				pProbAngle.forAngles = 0.0;
				pProbAngle.ConstAngle = -99999999;
			}
357 358
		}
	}
David Rohr's avatar
David Rohr committed
359
	// **************************************************************************************
360 361
	deviceStartRun();

David Rohr's avatar
David Rohr committed
362
	// ******************************** MAIN CYCLE ******************************************
David Rohr's avatar
David Rohr committed
363

364
	mycomplex_t* proj_mapFFT;
365
	myfloat_t* conv_map = NULL;
366
	mycomplex_t* conv_mapFFT;
367
	myfloat_t sumCONV, sumsquareCONV;
368 369

	//allocating fftw_complex vector
370 371
	proj_mapFFT = (mycomplex_t *) myfftw_malloc(sizeof(mycomplex_t) * param.param_device.NumberPixels * param.param_device.NumberFFTPixels1D);
	conv_mapFFT = (mycomplex_t *) myfftw_malloc(sizeof(mycomplex_t) * param.param_device.NumberPixels * param.param_device.NumberFFTPixels1D);
372
	if (!FFTAlgo) conv_map = (myfloat_t*) myfftw_malloc(sizeof(myfloat_t) * param.param_device.NumberPixels * param.param_device.NumberPixels);
373

David Rohr's avatar
David Rohr committed
374
	HighResTimer timer, timer2;
375

376
	if (DebugOutput >= 1 && mpi_rank == 0) printf("\tMain Loop GridAngles %d, CTFs %d, RefMaps %d, Shifts (%d/%d)², Pixels %d², OMP Threads %d, MPI Ranks %d\n", param.nTotGridAngles, param.nTotCTFs, RefMap.ntotRefMap, 2 * param.param_device.maxDisplaceCenter + param.param_device.GridSpaceCenter, param.param_device.GridSpaceCenter, param.param_device.NumberPixels, omp_get_max_threads(), mpi_size);
David Rohr's avatar
David Rohr committed
377 378 379 380

	const int iOrientStart = (int) ((long long int) mpi_rank * param.nTotGridAngles / mpi_size);
	int iOrientEnd = (int) ((long long int) (mpi_rank + 1) * param.nTotGridAngles / mpi_size);
	if (iOrientEnd > param.nTotGridAngles) iOrientEnd = param.nTotGridAngles;
David Rohr's avatar
David Rohr committed
381

David Rohr's avatar
David Rohr committed
382
	for (int iOrient = iOrientStart; iOrient < iOrientEnd; iOrient++)
383
	{
David Rohr's avatar
David Rohr committed
384 385
		// ***************************************************************************************
		// ***** Creating Projection for given orientation and transforming to Fourier space *****
David Rohr's avatar
David Rohr committed
386 387
		if (DebugOutput >= 1) timer2.ResetStart();
		if (DebugOutput >= 2) timer.ResetStart();
388
		createProjection(iOrient, proj_mapFFT);
David Rohr's avatar
David Rohr committed
389
		if (DebugOutput >= 2) printf("Time Projection %d: %f (rank %d)\n", iOrient, timer.GetCurrentElapsedTime(), mpi_rank);
390

David Rohr's avatar
David Rohr committed
391 392
		// ***************************************************************************************
		// ***** **** Internal Loop over convolutions **** *****
393 394
		for (int iConv = 0; iConv < param.nTotCTFs; iConv++)
		{
David Rohr's avatar
David Rohr committed
395
			// *** Calculating convolutions of projection map and crosscorrelations ***
396

397
			if (DebugOutput >= 2) timer.ResetStart();
398
			createConvolutedProjectionMap(iOrient, iConv, proj_mapFFT, conv_map, conv_mapFFT, sumCONV, sumsquareCONV);
David Rohr's avatar
David Rohr committed
399
			if (DebugOutput >= 2) printf("Time Convolution %d %d: %f (rank %d)\n", iOrient, iConv, timer.GetCurrentElapsedTime(), mpi_rank);
400

David Rohr's avatar
David Rohr committed
401 402
			// ***************************************************************************************
			// *** Comparing each calculated convoluted map with all experimental maps ***
403
			if (DebugOutput >= 2) timer.ResetStart();
404
			compareRefMaps(iOrient, iConv, conv_map, conv_mapFFT, sumCONV, sumsquareCONV);
405

406 407 408 409 410 411 412 413 414 415
			if (DebugOutput >= 2)
			{
				const double compTime = timer.GetCurrentElapsedTime();
				const int nShifts = 2 * param.param_device.maxDisplaceCenter / param.param_device.GridSpaceCenter + 1;
				const double nFlops = (double) RefMap.ntotRefMap * (double) nShifts * (double) nShifts *
									  (((double) param.param_device.NumberPixels - (double) param.param_device.maxDisplaceCenter / 2.) * ((double) param.param_device.NumberPixels - (double) param.param_device.maxDisplaceCenter / 2.) * 5. + 25.) / compTime;
				const double nGBs = (double) RefMap.ntotRefMap * (double) nShifts * (double) nShifts *
									(((double) param.param_device.NumberPixels - (double) param.param_device.maxDisplaceCenter / 2.) * ((double) param.param_device.NumberPixels - (double) param.param_device.maxDisplaceCenter / 2.) * 2. + 8.) * (double) sizeof(myfloat_t) / compTime;
				const double nGBs2 = (double) RefMap.ntotRefMap * ((double) param.param_device.NumberPixels * (double) param.param_device.NumberPixels + 8.) * (double) sizeof(myfloat_t) / compTime;

David Rohr's avatar
David Rohr committed
416
				printf("Time Comparison %d %d: %f sec (%f GFlops, %f GB/s (cached), %f GB/s) (rank %d)\n", iOrient, iConv, compTime, nFlops / 1000000000., nGBs / 1000000000., nGBs2 / 1000000000., mpi_rank);
417
			}
418
		}
David Rohr's avatar
David Rohr committed
419
		if (DebugOutput >= 1) printf("Total time for orientation %d: %f (rank %d)\n", iOrient, timer2.GetCurrentElapsedTime(), mpi_rank);
420 421 422 423
	}
	//deallocating fftw_complex vector
	myfftw_free(proj_mapFFT);
	myfftw_free(conv_mapFFT);
424
	if (!FFTAlgo) myfftw_free(conv_map);
David Rohr's avatar
David Rohr committed
425

426 427
	deviceFinishRun();

David Rohr's avatar
David Rohr committed
428
	// ************* Writing Out Probabilities ***************
429

David Rohr's avatar
David Rohr committed
430
	// *** Angular Probability ***
David Rohr's avatar
David Rohr committed
431

David Rohr's avatar
David Rohr committed
432
#ifdef WITH_MPI
433
	if (mpi_size > 1)
David Rohr's avatar
David Rohr committed
434
	{
435 436
		if (DebugOutput >= 1 && mpi_rank == 0) timer.ResetStart();
		//Reduce Constant and summarize probabilities
David Rohr's avatar
David Rohr committed
437
		{
438 439 440
			myfloat_t* tmp1 = new myfloat_t[RefMap.ntotRefMap];
			myfloat_t* tmp2 = new myfloat_t[RefMap.ntotRefMap];
			myfloat_t* tmp3 = new myfloat_t[RefMap.ntotRefMap];
David Rohr's avatar
David Rohr committed
441 442
			for (int i = 0;i < RefMap.ntotRefMap;i++)
			{
443
					tmp1[i] = pProb.getProbMap(i).Constoadd;
David Rohr's avatar
David Rohr committed
444
			}
445
			MPI_Allreduce(tmp1, tmp2, RefMap.ntotRefMap, MY_MPI_FLOAT, MPI_MAX, MPI_COMM_WORLD);
David Rohr's avatar
David Rohr committed
446 447
			for (int i = 0;i < RefMap.ntotRefMap;i++)
			{
448 449 450 451 452 453 454 455 456 457 458
				bioem_Probability_map& pProbMap = pProb.getProbMap(i);
				tmp1[i] = pProbMap.Total * exp(pProbMap.Constoadd - tmp2[i]);
			}
			MPI_Reduce(tmp1, tmp3, RefMap.ntotRefMap, MY_MPI_FLOAT, MPI_SUM, 0, MPI_COMM_WORLD);

			//Find MaxProb
			MPI_Status mpistatus;
			{
				int* tmpi1 = new int[RefMap.ntotRefMap];
				int* tmpi2 = new int[RefMap.ntotRefMap];
				for (int i = 0;i < RefMap.ntotRefMap;i++)
David Rohr's avatar
David Rohr committed
459
				{
460 461
					bioem_Probability_map& pProbMap = pProb.getProbMap(i);
					tmpi1[i] = tmp2[i] <= pProbMap.Constoadd ? mpi_rank : -1;
David Rohr's avatar
David Rohr committed
462
				}
463 464
				MPI_Allreduce(tmpi1, tmpi2, RefMap.ntotRefMap, MPI_INT, MPI_MAX, MPI_COMM_WORLD);
				for (int i = 0;i < RefMap.ntotRefMap;i++)
David Rohr's avatar
David Rohr committed
465
				{
466
					if (tmpi2[i] == -1)
David Rohr's avatar
David Rohr committed
467
					{
468
						if (mpi_rank == 0) printf("Error: Could not find highest probability\n");
David Rohr's avatar
David Rohr committed
469
					}
470
					else if (tmpi2[i] != 0) //Skip if rank 0 already has highest probability
David Rohr's avatar
David Rohr committed
471
					{
472 473 474 475 476 477 478 479
						if (mpi_rank == 0)
						{
							MPI_Recv(&pProb.getProbMap(i).max, sizeof(pProb.getProbMap(i).max), MPI_BYTE, tmpi2[i], i, MPI_COMM_WORLD, &mpistatus);
						}
						else if (mpi_rank == tmpi2[i])
						{
							MPI_Send(&pProb.getProbMap(i).max, sizeof(pProb.getProbMap(i).max), MPI_BYTE, 0, i, MPI_COMM_WORLD);
						}
David Rohr's avatar
David Rohr committed
480 481
					}
				}
482 483
				delete[] tmpi1;
				delete[] tmpi2;
David Rohr's avatar
David Rohr committed
484
			}
David Rohr's avatar
David Rohr committed
485

486
			if (mpi_rank == 0)
David Rohr's avatar
David Rohr committed
487
			{
488 489 490 491 492 493
				for (int i = 0;i < RefMap.ntotRefMap;i++)
				{
						bioem_Probability_map& pProbMap = pProb.getProbMap(i);
						pProbMap.Total = tmp3[i];
						pProbMap.Constoadd = tmp2[i];
				}
David Rohr's avatar
David Rohr committed
494
			}
David Rohr's avatar
David Rohr committed
495

496 497 498 499
			delete[] tmp1;
			delete[] tmp2;
			delete[] tmp3;
			if (DebugOutput >= 1 && mpi_rank == 0 && mpi_size > 1) printf("Time MPI Reduction: %f\n", timer.GetCurrentElapsedTime());
David Rohr's avatar
David Rohr committed
500
		}
501 502 503

		//Angle Reduction and Probability summation for individual angles
		if (param.param_device.writeAngles)
David Rohr's avatar
David Rohr committed
504
		{
505 506 507 508 509
			const int count = RefMap.ntotRefMap * param.nTotGridAngles;
			myfloat_t* tmp1 = new myfloat_t[count];
			myfloat_t* tmp2 = new myfloat_t[count];
			myfloat_t* tmp3 = new myfloat_t[count];
			for (int i = 0;i < RefMap.ntotRefMap;i++)
David Rohr's avatar
David Rohr committed
510
			{
511
					tmp1[i] = pProb.getProbMap(i).Constoadd;
David Rohr's avatar
David Rohr committed
512
			}
513
			MPI_Allreduce(tmp1, tmp2, count, MY_MPI_FLOAT, MPI_MAX, MPI_COMM_WORLD);
David Rohr's avatar
David Rohr committed
514 515 516 517 518
			for (int i = 0;i < RefMap.ntotRefMap;i++)
			{
				for (int j = 0;j < param.nTotGridAngles;j++)
				{
					bioem_Probability_angle& pProbAngle = pProb.getProbAngle(i, j);
519 520 521 522 523 524 525 526 527 528 529 530 531 532
					tmp1[i * param.nTotGridAngles + j] = pProbAngle.forAngles * exp(pProbAngle.ConstAngle - tmp2[i * param.nTotGridAngles + j]);
				}
			}
			MPI_Reduce(tmp1, tmp3, count, MY_MPI_FLOAT, MPI_SUM, 0, MPI_COMM_WORLD);
			if (mpi_rank == 0)
			{
				for (int i = 0;i < RefMap.ntotRefMap;i++)
				{
					for (int j = 0;j < param.nTotGridAngles;j++)
					{
						bioem_Probability_angle& pProbAngle = pProb.getProbAngle(i, j);
						pProbAngle.forAngles = tmp3[i * param.nTotGridAngles + j];
						pProbAngle.ConstAngle = tmp2[i * param.nTotGridAngles + j];
					}
David Rohr's avatar
David Rohr committed
533 534
				}
			}
535 536 537
			delete[] tmp1;
			delete[] tmp2;
			delete[] tmp3;
David Rohr's avatar
David Rohr committed
538
		}
539
	}
David Rohr's avatar
David Rohr committed
540
#endif
541

David Rohr's avatar
David Rohr committed
542
	if (mpi_rank == 0)
543
	{
David Rohr's avatar
David Rohr committed
544 545 546 547 548
		ofstream angProbfile;
		if(param.param_device.writeAngles)
		{
			angProbfile.open ("ANG_PROB");
		}
549 550 551 552 553
		ofstream ccProbfile;
                if(param.param_device.writeCC)
                {
                        ccProbfile.open ("CROSS_CORRELATION");
                }
David Rohr's avatar
David Rohr committed
554 555 556 557 558 559 560

		ofstream outputProbFile;
		outputProbFile.open ("Output_Probabilities");
		for (int iRefMap = 0; iRefMap < RefMap.ntotRefMap; iRefMap ++)
		{
			// **** Total Probability ***
			bioem_Probability_map& pProbMap = pProb.getProbMap(iRefMap);
561

David Rohr's avatar
David Rohr committed
562
			outputProbFile << "RefMap " << iRefMap << " Probability  "  << log(pProbMap.Total) + pProbMap.Constoadd + 0.5 * log(M_PI) + (1 - param.param_device.Ntotpi * 0.5)*(log(2 * M_PI) + 1) + log(param.param_device.volu) << " Constant " << pProbMap.Constoadd  << "\n";
563

David Rohr's avatar
David Rohr committed
564
			outputProbFile << "RefMap " << iRefMap << " Maximizing Param: ";
565

David Rohr's avatar
David Rohr committed
566 567 568 569 570 571 572 573 574
			// *** Param that maximize probability****
			outputProbFile << (pProbMap.Constoadd + 0.5 * log(M_PI) + (1 - param.param_device.Ntotpi * 0.5) * (log(2 * M_PI) + 1) + log(param.param_device.volu)) << " ";
			outputProbFile << param.angles[pProbMap.max.max_prob_orient].pos[0] << " ";
			outputProbFile << param.angles[pProbMap.max.max_prob_orient].pos[1] << " ";
			outputProbFile << param.angles[pProbMap.max.max_prob_orient].pos[2] << " ";
			outputProbFile << param.CtfParam[pProbMap.max.max_prob_conv].pos[0] << " ";
			outputProbFile << param.CtfParam[pProbMap.max.max_prob_conv].pos[1] << " ";
			outputProbFile << param.CtfParam[pProbMap.max.max_prob_conv].pos[2] << " ";
			outputProbFile << pProbMap.max.max_prob_cent_x << " ";
Pilar Cossio's avatar
Pilar Cossio committed
575 576 577 578
			outputProbFile << pProbMap.max.max_prob_cent_y << " " ;
	                outputProbFile << pProbMap.max.max_prob_norm  << " " ;
        	        outputProbFile << pProbMap.max.max_prob_mu ;
                	outputProbFile << "\n";
579

David Rohr's avatar
David Rohr committed
580
			// *** For individual files*** //angProbfile.open ("ANG_PROB_"iRefMap);
581

David Rohr's avatar
David Rohr committed
582
			if(param.param_device.writeAngles)
583
			{
David Rohr's avatar
David Rohr committed
584 585 586
				for (int iOrient = 0; iOrient < param.nTotGridAngles; iOrient++)
				{
					bioem_Probability_angle& pProbAngle = pProb.getProbAngle(iRefMap, iOrient);
587

David Rohr's avatar
David Rohr committed
588 589
					angProbfile << " " << iRefMap << " " << param.angles[iOrient].pos[0] << " " << param.angles[iOrient].pos[1] << " " << param.angles[iOrient].pos[2] << " " << log(pProbAngle.forAngles) + pProbAngle.ConstAngle + 0.5 * log(M_PI) + (1 - param.param_device.Ntotpi * 0.5)*(log(2 * M_PI) + 1) + log(param.param_device.volu) << " " << log(param.param_device.volu) << "\n";
				}
590
			}
591 592 593 594 595 596 597 598 599 600 601 602 603
                        if(param.param_device.writeCC)
                        {
				 int  cc=0;
		           for (int cent_x = 0; cent_x < param.param_device.NumberPixels; cent_x = cent_x + param.param_device.CCdisplace)
                	   {
	                        for (int cent_y = 0; cent_y < param.param_device.NumberPixels; cent_y = cent_y + param.param_device.CCdisplace)
        	                {
	      	                bioem_Probability_cc& pProbCC = pProb.getProbCC(iRefMap, cc);
                                ccProbfile << " " << iRefMap << " " << cent_x << " " << cent_y << " " << log(pProbCC.forCC) <<"\n";
                        	}
                	   }
                             
			}
604
		}
605

David Rohr's avatar
David Rohr committed
606 607 608 609
		if(param.param_device.writeAngles)
		{
			angProbfile.close();
		}
610 611 612 613 614 615
                if(param.param_device.writeCC)
                {
                        ccProbfile.close();
                }


David Rohr's avatar
David Rohr committed
616
		outputProbFile.close();
617
	}
618

619
	return(0);
620 621
}

622
int bioem::compareRefMaps(int iOrient, int iConv, const myfloat_t* conv_map, mycomplex_t* localmultFFT, myfloat_t sumC, myfloat_t sumsquareC, const int startMap)
623
{
David Rohr's avatar
David Rohr committed
624 625
	//***************************************************************************************
	//***** BioEM routine for comparing reference maps to convoluted maps *****
626
	if (FFTAlgo)
627
	{
David Rohr's avatar
David Rohr committed
628
		//With FFT Algorithm
629 630
		#pragma omp parallel for
		for (int iRefMap = startMap; iRefMap < RefMap.ntotRefMap; iRefMap ++)
631
		{
632
			const int num = omp_get_thread_num();
633
			calculateCCFFT(iRefMap, iOrient, iConv, sumC, sumsquareC, localmultFFT, param.fft_scratch_complex[num], param.fft_scratch_real[num]);
634 635 636
		}
	}
	else
637
	{
David Rohr's avatar
David Rohr committed
638
		//Without FFT Algorithm
639
		#pragma omp parallel for
640
		for (int iRefMap = startMap; iRefMap < RefMap.ntotRefMap; iRefMap ++)
641
		{
642
			compareRefMapShifted < -1 > (iRefMap, iOrient, iConv, conv_map, pProb, param.param_device, RefMap);
643 644 645 646 647
		}
	}
	return(0);
}

648
inline void bioem::calculateCCFFT(int iRefMap, int iOrient, int iConv, myfloat_t sumC, myfloat_t sumsquareC, mycomplex_t* localConvFFT, mycomplex_t* localCCT, myfloat_t* lCC)
649
{
David Rohr's avatar
David Rohr committed
650 651
	//***************************************************************************************
	//***** Calculating cross correlation in FFTALGOrithm *****
Pilar Cossio's avatar
Pilar Cossio committed
652

653
	const mycomplex_t* RefMapFFT = &RefMap.RefMapsFFT[iRefMap * param.FFTMapSize];
654
	for(int i = 0; i < param.param_device.NumberPixels * param.param_device.NumberFFTPixels1D; i++)
655
	{
656 657
		localCCT[i][0] = localConvFFT[i][0] * RefMapFFT[i][0] + localConvFFT[i][1] * RefMapFFT[i][1];
		localCCT[i][1] = localConvFFT[i][1] * RefMapFFT[i][0] - localConvFFT[i][0] * RefMapFFT[i][1];
658 659
	}

660
	myfftw_execute_dft_c2r(param.fft_plan_c2r_backward, localCCT, lCC);
661

662
	doRefMapFFT(iRefMap, iOrient, iConv, lCC, sumC, sumsquareC, pProb, param.param_device, RefMap);
663
}
664

665
int bioem::createProjection(int iMap, mycomplex_t* mapFFT)
666
{
David Rohr's avatar
David Rohr committed
667
	// **************************************************************************************
David Rohr's avatar
David Rohr committed
668 669
	// ****  BioEM Create Projection routine in Euler angle predefined grid******************
	// ********************* and turns projection into Fourier space ************************
David Rohr's avatar
David Rohr committed
670
	// **************************************************************************************
671

672 673
	cuda_custom_timeslot("Projection", 0);

674 675
	myfloat3_t RotatedPointsModel[Model.nPointsModel];
	myfloat_t rotmat[3][3];
676
	myfloat_t alpha, gam, beta;
677
	myfloat_t* localproj;
678

679
	localproj = param.fft_scratch_real[omp_get_thread_num()];
680
	memset(localproj, 0, param.param_device.NumberPixels * param.param_device.NumberPixels * sizeof(*localproj));
681

682 683 684
	alpha = param.angles[iMap].pos[0];
	beta = param.angles[iMap].pos[1];
	gam = param.angles[iMap].pos[2];
685

David Rohr's avatar
David Rohr committed
686
	// **** To see how things are going: cout << "Id " << omp_get_thread_num() <<  " Angs: " << alpha << " " << beta << " " << gam << "\n"; ***
687

David Rohr's avatar
David Rohr committed
688
	// ********** Creat Rotation with pre-defiend grid of orientations**********
689 690 691 692 693 694 695 696 697 698 699
	rotmat[0][0] = cos(gam) * cos(alpha) - cos(beta) * sin(alpha) * sin(gam);
	rotmat[0][1] = cos(gam) * sin(alpha) + cos(beta) * cos(alpha) * sin(gam);
	rotmat[0][2] = sin(gam) * sin(beta);
	rotmat[1][0] = -sin(gam) * cos(alpha) - cos(beta) * sin(alpha) * cos(gam);
	rotmat[1][1] = -sin(gam) * sin(alpha) + cos(beta) * cos(alpha) * cos(gam);
	rotmat[1][2] = cos(gam) * sin(beta);
	rotmat[2][0] = sin(beta) * sin(alpha);
	rotmat[2][1] = -sin(beta) * cos(alpha);
	rotmat[2][2] = cos(beta);

	for(int n = 0; n < Model.nPointsModel; n++)
700
	{
701 702 703
		RotatedPointsModel[n].pos[0] = 0.0;
		RotatedPointsModel[n].pos[1] = 0.0;
		RotatedPointsModel[n].pos[2] = 0.0;
704
	}
705
	for(int n = 0; n < Model.nPointsModel; n++)
706
	{
707
		for(int k = 0; k < 3; k++)
708
		{
709
			for(int j = 0; j < 3; j++)
710
			{
711
				RotatedPointsModel[n].pos[k] += rotmat[k][j] * Model.points[n].point.pos[j];
712 713 714 715 716 717
			}
		}
	}

	int i, j;

Pilar Cossio's avatar
Pilar Cossio committed
718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748
	//********** Projection with radius ***************
	if(param.doaaradius){

	        int irad;
        	myfloat_t dist, rad2;

          	for(int n = 0; n < Model.nPointsModel; n++)
                {
                	//Getting Centers of Sphere
               		i = floor(RotatedPointsModel[n].pos[0] / param.pixelSize + (myfloat_t) param.param_device.NumberPixels / 2.0f + 0.5f);
		        j = floor(RotatedPointsModel[n].pos[1] / param.pixelSize + (myfloat_t) param.param_device.NumberPixels / 2.0f + 0.5f);
			//Getting the radius
			irad=int( Model.points[n].radius / param.pixelSize );
                	rad2= Model.points[n].radius * Model.points[n].radius;

               		//Projecting over the radius
                	for(int ii= i - irad; ii < i + irad; ii++)
			{	
				for(int jj = j - irad; jj < j + irad; jj++)
				{
                       			 dist= ( (myfloat_t) (ii-i)*(ii-i)+(jj-j)*(jj-j) ) *  param.pixelSize ;
	                       		 if( dist < rad2 )
					{
					 localproj[ii * param.param_device.NumberPixels + jj] += param.pixelSize * 2 * sqrt( rad2 - dist ) * Model.points[n].density
					 / Model.NormDen * 3 / (4 * M_PI * rad2 *  Model.points[n].radius); 	
					}
                 		}
                	}
		}
	}
	else 
749
	{
Pilar Cossio's avatar
Pilar Cossio committed
750 751 752 753
	// ************ Simple Projection over the Z axis********************
		for(int n = 0; n < Model.nPointsModel; n++)
		{
 		//Getting pixel that represents coordinates & shifting the start at to Numpix/2,Numpix/2 )
754 755
		i = floor(RotatedPointsModel[n].pos[0] / param.pixelSize + (myfloat_t) param.param_device.NumberPixels / 2.0f + 0.5f);
		j = floor(RotatedPointsModel[n].pos[1] / param.pixelSize + (myfloat_t) param.param_device.NumberPixels / 2.0f + 0.5f);
756

757 758 759 760 761 762
		if (i < 0 || j < 0 || i >= param.param_device.NumberPixels || j >= param.param_device.NumberPixels)
		{
			if (DebugOutput >= 3) cout << "Model Point out of map: " << i << ", " << j << "\n";
			continue;
		}

763
		localproj[i * param.param_device.NumberPixels + j] += Model.points[n].density / Model.NormDen;
Pilar Cossio's avatar
Pilar Cossio committed
764
         	}
765 766
	}

David Rohr's avatar
David Rohr committed
767
	// **** Output Just to check****
768
#ifdef PILAR_DEBUG
769
	if(iMap == 10)
770 771 772 773 774 775
	{
		ofstream myexamplemap;
		ofstream myexampleRot;
		myexamplemap.open ("MAP_i10");
		myexampleRot.open ("Rot_i10");
		myexamplemap << "ANGLES " << alpha << " " << beta << " " << gam << "\n";
776
		for(int k = 0; k < param.param_device.NumberPixels; k++)
777
		{
778
			for(int j = 0; j < param.param_device.NumberPixels; j++) myexamplemap << "\nMAP " << k << " " << j << " " << localproj[k * param.param_device.NumberPixels + j];
779 780
		}
		myexamplemap << " \n";
781
		for(int n = 0; n < Model.nPointsModel; n++)myexampleRot << "\nCOOR " << RotatedPointsModel[n].pos[0] << " " << RotatedPointsModel[n].pos[1] << " " << RotatedPointsModel[n].pos[2];
782 783 784
		myexamplemap.close();
		myexampleRot.close();
	}
785
#endif
786

David Rohr's avatar
David Rohr committed
787 788
	// ***** Converting projection to Fourier Space for Convolution later with kernel****
	// ********** Omp Critical is necessary with FFTW*******
789
	myfftw_execute_dft_r2c(param.fft_plan_r2c_forward, localproj, mapFFT);
790

791 792
	cuda_custom_timeslot_end;

793 794 795
	return(0);
}

796
int bioem::createConvolutedProjectionMap(int iMap, int iConv, mycomplex_t* lproj, myfloat_t* Mapconv, mycomplex_t* localmultFFT, myfloat_t& sumC, myfloat_t& sumsquareC)
797
{
David Rohr's avatar
David Rohr committed
798 799
	// **************************************************************************************
	// ****  BioEM Create Convoluted Projection Map routine, multiplies in Fourier **********
David Rohr's avatar
David Rohr committed
800 801
	// **************** calculated Projection with convoluted precalculated Kernel***********
	// *************** and Backtransforming it to real Space ********************************
David Rohr's avatar
David Rohr committed
802
	// **************************************************************************************
803

804 805
	cuda_custom_timeslot("Convolution", 1);

806
	mycomplex_t* tmp = param.fft_scratch_complex[omp_get_thread_num()];
807

David Rohr's avatar
David Rohr committed
808
	// **** Multiplying FFTmap with corresponding kernel ****
809
	const mycomplex_t* refCTF = &param.refCTF[iConv * param.FFTMapSize];
810
	for(int i = 0; i < param.param_device.NumberPixels * param.param_device.NumberFFTPixels1D; i++)
811
	{
812 813 814
		localmultFFT[i][0] = lproj[i][0] * refCTF[i][0] + lproj[i][1] * refCTF[i][1];
		localmultFFT[i][1] = lproj[i][1] * refCTF[i][0] - lproj[i][0] * refCTF[i][1];
		// cout << "GG " << i << " " << j << " " << refCTF[i][0] << " " << refCTF[i][1] <<" " <<lproj[i][0] <<" " <<lproj[i][1] << "\n";
815 816
	}

David Rohr's avatar
David Rohr committed
817
	// *** Calculating Cross-correlations of cal-convoluted map with its self *****
818
	sumC = localmultFFT[0][0];
David Rohr's avatar
David Rohr committed
819

820
	sumsquareC = 0;
821
	if (FFTAlgo)
822
	{
823 824
		int jloopend = param.param_device.NumberFFTPixels1D;
		if ((param.param_device.NumberPixels & 1) == 0) jloopend--;
825 826
		for(int i = 0; i < param.param_device.NumberPixels; i++)
		{
827
			for (int j = 1;j < jloopend;j++)
828 829 830 831 832 833
			{
				int k = i * param.param_device.NumberFFTPixels1D + j;
				sumsquareC += (localmultFFT[k][0] * localmultFFT[k][0] + localmultFFT[k][1] * localmultFFT[k][1]) * 2;
			}
			int k = i * param.param_device.NumberFFTPixels1D;
			sumsquareC += localmultFFT[k][0] * localmultFFT[k][0] + localmultFFT[k][1] * localmultFFT[k][1];
834 835
			if ((param.param_device.NumberPixels & 1) == 0)
			{
836
				k += param.param_device.NumberFFTPixels1D - 1;
837 838
				sumsquareC += localmultFFT[k][0] * localmultFFT[k][0] + localmultFFT[k][1] * localmultFFT[k][1];
			}
839
		}
David Rohr's avatar
David Rohr committed
840

841 842
		myfloat_t norm2 = (myfloat_t) (param.param_device.NumberPixels * param.param_device.NumberPixels);
		sumsquareC = sumsquareC / norm2;
843
	}
844 845 846 847 848 849 850
	else
	{
		//FFTW_C2R will destroy the input array, so we have to work on a copy here
		memcpy(tmp, localmultFFT, sizeof(mycomplex_t) * param.param_device.NumberPixels * param.param_device.NumberFFTPixels1D);

		// **** Bringing convoluted Map to real Space ****
		myfftw_execute_dft_c2r(param.fft_plan_c2r_backward, tmp, Mapconv);
851

852 853 854 855 856 857 858 859 860
		for(int i = 0; i < param.param_device.NumberPixels * param.param_device.NumberPixels; i++)
		{
			sumsquareC += Mapconv[i] * Mapconv[i];
		}

		myfloat_t norm2 = (myfloat_t) (param.param_device.NumberPixels * param.param_device.NumberPixels);
		myfloat_t norm4 = norm2 * norm2;
		sumsquareC = sumsquareC / norm4;