bioem.cpp 33.5 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
	nProjectionsAtOnce = getenv("BIOEM_PROJECTIONS_AT_ONCE") == NULL ? 1 : atoi(getenv("BIOEM_PROJECTIONS_AT_ONCE"));
85
86
87
88
89
90
91
92
}

bioem::~bioem()
{
}

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

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

David Rohr's avatar
David Rohr committed
101
102
103
104
105
	if (mpi_rank == 0)
	{
		// *** Inizialzing default variables ***
		std::string infile, modelfile, mapfile;
		Model.readPDB = false;
106
		param.param_device.writeAngles = false;
David Rohr's avatar
David Rohr committed
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
206
		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
207
		}
David Rohr's avatar
David Rohr committed
208
		catch(std::exception& e)
209
		{
David Rohr's avatar
David Rohr committed
210
211
			cout << e.what() << "\n";
			return 1;
212
		}
David Rohr's avatar
David Rohr committed
213
214
215
216
217
218
219
220
221

		//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
222
		// ********************* Reading Parameter Input ***************************
223
		param.readParameters(infile.c_str());
David Rohr's avatar
David Rohr committed
224
225

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

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

David Rohr's avatar
David Rohr committed
233
#ifdef WITH_MPI
234
235
236
237
238
239
240
241
242
243
244
245
246
	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
247
		if (DebugOutput >= 2 && mpi_rank == 0) printf("MPI Broadcast of Input Data %f\n", timer.GetCurrentElapsedTime());
248
249
	}
	#endif
250

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

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

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

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

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

284
	return(0);
285
286
}

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

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

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

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

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

322
	return(0);
323
324
325
326
}

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

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

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

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

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

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

355
356
357
				pProbAngle.forAngles = 0.0;
				pProbAngle.ConstAngle = -99999999;
			}
358
		}
359
360
361
362
363
364
365
366
367
368
369
370
371
372
		 if (param.param_device.writeCC)
        	{      int  cc=0;
                	 for (int cent_x = 0; cent_x < param.param_device.NumberPixels -param.param_device.CCdisplace ; cent_x = cent_x + param.param_device.CCdisplace)
                         {
                                 for (int cent_y = 0; cent_y < param.param_device.NumberPixels - param.param_device.CCdisplace ; cent_y = cent_y + param.param_device.CCdisplace)
                                 {
        	                       cout << cc << " " << cent_x << " " << cent_y <<"\n"; 
                                 bioem_Probability_cc& pProbCC = pProb.getProbCC(iRefMap, cc);
                                 pProbCC.forCC = 0.0;
 
                                 cc++;
                                 }
                         }
		}                 
373
	}
David Rohr's avatar
David Rohr committed
374
	// **************************************************************************************
375
376
	deviceStartRun();

David Rohr's avatar
David Rohr committed
377
	// ******************************** MAIN CYCLE ******************************************
David Rohr's avatar
David Rohr committed
378

379
	mycomplex_t* proj_mapsFFT;
380
	myfloat_t* conv_map = NULL;
381
	mycomplex_t* conv_mapFFT;
382
	myfloat_t sumCONV, sumsquareCONV;
383
384

	//allocating fftw_complex vector
385
386
	const int ProjMapSize = (param.FFTMapSize + 64) & ~63;	//Make sure this is properly aligned for fftw..., Actually this should be ensureb by using FFTMapSize, but it is not due to a bug in CUFFT which cannot handle padding properly
	proj_mapsFFT = (mycomplex_t *) myfftw_malloc(sizeof(mycomplex_t) * ProjMapSize * nProjectionsAtOnce);
387
	conv_mapFFT = (mycomplex_t *) myfftw_malloc(sizeof(mycomplex_t) * param.param_device.NumberPixels * param.param_device.NumberFFTPixels1D);
388
	if (!FFTAlgo) conv_map = (myfloat_t*) myfftw_malloc(sizeof(myfloat_t) * param.param_device.NumberPixels * param.param_device.NumberPixels);
389

David Rohr's avatar
David Rohr committed
390
	HighResTimer timer, timer2;
391

392
	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
393
394
395
396

	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
397

398
	for (int iOrientAtOnce = iOrientStart; iOrientAtOnce < iOrientEnd; iOrientAtOnce += nProjectionsAtOnce)
399
	{
David Rohr's avatar
David Rohr committed
400
401
		// ***************************************************************************************
		// ***** Creating Projection for given orientation and transforming to Fourier space *****
David Rohr's avatar
David Rohr committed
402
403
		if (DebugOutput >= 1) timer2.ResetStart();
		if (DebugOutput >= 2) timer.ResetStart();
404
405
406
407
408
409
410
		int iTmpEnd = std::min(iOrientEnd, iOrientAtOnce + nProjectionsAtOnce);
#pragma omp parallel for
		for (int iOrient = iOrientAtOnce; iOrient < iTmpEnd;iOrient++)
		{
			createProjection(iOrient, &proj_mapsFFT[(iOrient - iOrientAtOnce) * ProjMapSize]);
		}
		if (DebugOutput >= 2) printf("\tTime Projection %d: %f (rank %d)\n", iOrientAtOnce, timer.GetCurrentElapsedTime(), mpi_rank);
411

412
		for (int iOrient = iOrientAtOnce; iOrient < iTmpEnd;iOrient++)
413
		{
414
415
416
417
418
419
			mycomplex_t* proj_mapFFT = &proj_mapsFFT[(iOrient - iOrientAtOnce) * ProjMapSize];
			// ***************************************************************************************
			// ***** **** Internal Loop over convolutions **** *****
			for (int iConv = 0; iConv < param.nTotCTFs; iConv++)
			{
				// *** Calculating convolutions of projection map and crosscorrelations ***
420

421
422
423
				if (DebugOutput >= 2) timer.ResetStart();
				createConvolutedProjectionMap(iOrient, iConv, proj_mapFFT, conv_map, conv_mapFFT, sumCONV, sumsquareCONV);
				if (DebugOutput >= 2) printf("\t\tTime Convolution %d %d: %f (rank %d)\n", iOrient, iConv, timer.GetCurrentElapsedTime(), mpi_rank);
424

425
426
427
428
				// ***************************************************************************************
				// *** Comparing each calculated convoluted map with all experimental maps ***
				if (DebugOutput >= 2) timer.ResetStart();
				compareRefMaps(iOrient, iConv, conv_map, conv_mapFFT, sumCONV, sumsquareCONV);
429

430
431
432
433
434
435
436
437
438
439
440
441
442
443
				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;

					printf("\t\tTime 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);
				}
			}
			if (DebugOutput >= 1)
444
			{
445
446
				printf("\tTotal time for projection %d: %f (rank %d)\n", iOrient, timer2.GetCurrentElapsedTime(), mpi_rank);
				timer2.ResetStart();
447
			}
448
449
450
		}
	}
	//deallocating fftw_complex vector
451
	myfftw_free(proj_mapsFFT);
452
	myfftw_free(conv_mapFFT);
453
	if (!FFTAlgo) myfftw_free(conv_map);
David Rohr's avatar
David Rohr committed
454

455
456
	deviceFinishRun();

David Rohr's avatar
David Rohr committed
457
	// ************* Writing Out Probabilities ***************
458

David Rohr's avatar
David Rohr committed
459
	// *** Angular Probability ***
David Rohr's avatar
David Rohr committed
460

David Rohr's avatar
David Rohr committed
461
#ifdef WITH_MPI
462
	if (mpi_size > 1)
David Rohr's avatar
David Rohr committed
463
	{
464
465
		if (DebugOutput >= 1 && mpi_rank == 0) timer.ResetStart();
		//Reduce Constant and summarize probabilities
David Rohr's avatar
David Rohr committed
466
		{
467
468
469
			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
470
471
			for (int i = 0;i < RefMap.ntotRefMap;i++)
			{
472
					tmp1[i] = pProb.getProbMap(i).Constoadd;
David Rohr's avatar
David Rohr committed
473
			}
474
			MPI_Allreduce(tmp1, tmp2, RefMap.ntotRefMap, MY_MPI_FLOAT, MPI_MAX, MPI_COMM_WORLD);
David Rohr's avatar
David Rohr committed
475
476
			for (int i = 0;i < RefMap.ntotRefMap;i++)
			{
477
478
479
480
481
482
483
484
485
486
487
				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
488
				{
489
490
					bioem_Probability_map& pProbMap = pProb.getProbMap(i);
					tmpi1[i] = tmp2[i] <= pProbMap.Constoadd ? mpi_rank : -1;
David Rohr's avatar
David Rohr committed
491
				}
492
493
				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
494
				{
495
					if (tmpi2[i] == -1)
David Rohr's avatar
David Rohr committed
496
					{
497
						if (mpi_rank == 0) printf("Error: Could not find highest probability\n");
David Rohr's avatar
David Rohr committed
498
					}
499
					else if (tmpi2[i] != 0) //Skip if rank 0 already has highest probability
David Rohr's avatar
David Rohr committed
500
					{
501
502
503
504
505
506
507
508
						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
509
510
					}
				}
511
512
				delete[] tmpi1;
				delete[] tmpi2;
David Rohr's avatar
David Rohr committed
513
			}
David Rohr's avatar
David Rohr committed
514

515
			if (mpi_rank == 0)
David Rohr's avatar
David Rohr committed
516
			{
517
518
519
520
521
522
				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
523
			}
David Rohr's avatar
David Rohr committed
524

525
526
527
528
			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
529
		}
530
531
532

		//Angle Reduction and Probability summation for individual angles
		if (param.param_device.writeAngles)
David Rohr's avatar
David Rohr committed
533
		{
534
535
536
537
538
			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
539
			{
540
					tmp1[i] = pProb.getProbMap(i).Constoadd;
David Rohr's avatar
David Rohr committed
541
			}
542
			MPI_Allreduce(tmp1, tmp2, count, MY_MPI_FLOAT, MPI_MAX, MPI_COMM_WORLD);
David Rohr's avatar
David Rohr committed
543
544
545
546
547
			for (int i = 0;i < RefMap.ntotRefMap;i++)
			{
				for (int j = 0;j < param.nTotGridAngles;j++)
				{
					bioem_Probability_angle& pProbAngle = pProb.getProbAngle(i, j);
548
549
550
551
552
553
554
555
556
557
558
559
560
561
					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
562
563
				}
			}
564
565
566
			delete[] tmp1;
			delete[] tmp2;
			delete[] tmp3;
David Rohr's avatar
David Rohr committed
567
		}
568
	}
David Rohr's avatar
David Rohr committed
569
#endif
570

David Rohr's avatar
David Rohr committed
571
	if (mpi_rank == 0)
572
	{
David Rohr's avatar
David Rohr committed
573
574
575
576
577
		ofstream angProbfile;
		if(param.param_device.writeAngles)
		{
			angProbfile.open ("ANG_PROB");
		}
578
		ofstream ccProbfile;
579
580
581
582
		if(param.param_device.writeCC)
		{
			ccProbfile.open ("CROSS_CORRELATION");
		}
David Rohr's avatar
David Rohr committed
583
584
585
586
587
588
589

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

David Rohr's avatar
David Rohr committed
591
			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";
592

David Rohr's avatar
David Rohr committed
593
			outputProbFile << "RefMap " << iRefMap << " Maximizing Param: ";
594

David Rohr's avatar
David Rohr committed
595
596
597
598
599
600
601
602
603
			// *** 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
604
605
606
607
			outputProbFile << pProbMap.max.max_prob_cent_y << " " ;
	                outputProbFile << pProbMap.max.max_prob_norm  << " " ;
        	        outputProbFile << pProbMap.max.max_prob_mu ;
                	outputProbFile << "\n";
608

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

David Rohr's avatar
David Rohr committed
611
			if(param.param_device.writeAngles)
612
			{
David Rohr's avatar
David Rohr committed
613
614
615
				for (int iOrient = 0; iOrient < param.nTotGridAngles; iOrient++)
				{
					bioem_Probability_angle& pProbAngle = pProb.getProbAngle(iRefMap, iOrient);
616

David Rohr's avatar
David Rohr committed
617
618
					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";
				}
619
			}
620
621
622
			if(param.param_device.writeCC)
			{
				int  cc=0;
623
624
625
626
			for (int cent_x = 0; cent_x < param.param_device.NumberPixels -param.param_device.CCdisplace ; cent_x = cent_x + param.param_device.CCdisplace)
                         {
                                 for (int cent_y = 0; cent_y < param.param_device.NumberPixels - param.param_device.CCdisplace ; cent_y = cent_y + param.param_device.CCdisplace)
                                 {
627
						bioem_Probability_cc& pProbCC = pProb.getProbCC(iRefMap, cc);
628
						ccProbfile << " " << iRefMap << " " << cent_x << " " << cent_y << " " << pProbCC.forCC <<"\n";
629
630
				}

631
			}
632
                       }
633
		}
634

David Rohr's avatar
David Rohr committed
635
636
637
638
		if(param.param_device.writeAngles)
		{
			angProbfile.close();
		}
639

640
641
642
643
		if(param.param_device.writeCC)
		{
			ccProbfile.close();
		}
644

David Rohr's avatar
David Rohr committed
645
		outputProbFile.close();
646
	}
647

648
	return(0);
649
650
}

651
int bioem::compareRefMaps(int iOrient, int iConv, const myfloat_t* conv_map, mycomplex_t* localmultFFT, myfloat_t sumC, myfloat_t sumsquareC, const int startMap)
652
{
David Rohr's avatar
David Rohr committed
653
654
	//***************************************************************************************
	//***** BioEM routine for comparing reference maps to convoluted maps *****
655
	if (FFTAlgo)
656
	{
David Rohr's avatar
David Rohr committed
657
		//With FFT Algorithm
658
		#pragma omp parallel for schedule(dynamic, 1)
659
		for (int iRefMap = startMap; iRefMap < RefMap.ntotRefMap; iRefMap ++)
660
		{
661
			const int num = omp_get_thread_num();
662
			calculateCCFFT(iRefMap, iOrient, iConv, sumC, sumsquareC, localmultFFT, param.fft_scratch_complex[num], param.fft_scratch_real[num]);
663
664
665
		}
	}
	else
666
	{
David Rohr's avatar
David Rohr committed
667
		//Without FFT Algorithm
668
		#pragma omp parallel for schedule(dynamic, 1)
669
		for (int iRefMap = startMap; iRefMap < RefMap.ntotRefMap; iRefMap ++)
670
		{
671
			compareRefMapShifted < -1 > (iRefMap, iOrient, iConv, conv_map, pProb, param.param_device, RefMap);
672
673
674
675
676
		}
	}
	return(0);
}

677
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)
678
{
David Rohr's avatar
David Rohr committed
679
680
	//***************************************************************************************
	//***** Calculating cross correlation in FFTALGOrithm *****
Pilar Cossio's avatar
Pilar Cossio committed
681

682
	const mycomplex_t* RefMapFFT = &RefMap.RefMapsFFT[iRefMap * param.FFTMapSize];
683
	for(int i = 0; i < param.param_device.NumberPixels * param.param_device.NumberFFTPixels1D; i++)
684
	{
685
686
		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];
687
688
	}

689
	myfftw_execute_dft_c2r(param.fft_plan_c2r_backward, localCCT, lCC);
690

691
	doRefMapFFT(iRefMap, iOrient, iConv, lCC, sumC, sumsquareC, pProb, param.param_device, RefMap);
692
}
693

694
int bioem::createProjection(int iMap, mycomplex_t* mapFFT)
695
{
David Rohr's avatar
David Rohr committed
696
	// **************************************************************************************
David Rohr's avatar
David Rohr committed
697
698
	// ****  BioEM Create Projection routine in Euler angle predefined grid******************
	// ********************* and turns projection into Fourier space ************************
David Rohr's avatar
David Rohr committed
699
	// **************************************************************************************
700

701
702
	cuda_custom_timeslot("Projection", 0);

703
704
	myfloat3_t RotatedPointsModel[Model.nPointsModel];
	myfloat_t rotmat[3][3];
705
	myfloat_t alpha, gam, beta;
706
	myfloat_t* localproj;
707

708
	localproj = param.fft_scratch_real[omp_get_thread_num()];
709
	memset(localproj, 0, param.param_device.NumberPixels * param.param_device.NumberPixels * sizeof(*localproj));
710

711
712
713
	alpha = param.angles[iMap].pos[0];
	beta = param.angles[iMap].pos[1];
	gam = param.angles[iMap].pos[2];
714

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

David Rohr's avatar
David Rohr committed
717
	// ********** Creat Rotation with pre-defiend grid of orientations**********
718
719
720
721
722
723
724
725
726
727
728
	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++)
729
	{
730
731
732
		RotatedPointsModel[n].pos[0] = 0.0;
		RotatedPointsModel[n].pos[1] = 0.0;
		RotatedPointsModel[n].pos[2] = 0.0;
733
	}
734
	for(int n = 0; n < Model.nPointsModel; n++)
735
	{
736
		for(int k = 0; k < 3; k++)
737
		{
738
			for(int j = 0; j < 3; j++)
739
			{
740
				RotatedPointsModel[n].pos[k] += rotmat[k][j] * Model.points[n].point.pos[j];
741
742
743
744
745
746
			}
		}
	}

	int i, j;

Pilar Cossio's avatar
Pilar Cossio committed
747
	//********** Projection with radius ***************
748
749
750
751
	if(param.doaaradius)
	{
		int irad;
		myfloat_t dist, rad2;
Pilar Cossio's avatar
Pilar Cossio committed
752

753
754
755
756
757
		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);
Pilar Cossio's avatar
Pilar Cossio committed
758
759
			//Getting the radius
			irad=int( Model.points[n].radius / param.pixelSize );
760
			rad2= Model.points[n].radius * Model.points[n].radius;
Pilar Cossio's avatar
Pilar Cossio committed
761

762
763
			//Projecting over the radius
			for(int ii= i - irad; ii < i + irad; ii++)
Pilar Cossio's avatar
Pilar Cossio committed
764
765
766
			{	
				for(int jj = j - irad; jj < j + irad; jj++)
				{
767
768
					dist= ( (myfloat_t) (ii-i)*(ii-i)+(jj-j)*(jj-j) ) *  param.pixelSize ;
					if( dist < rad2 )
Pilar Cossio's avatar
Pilar Cossio committed
769
					{
770
771
						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); 	
Pilar Cossio's avatar
Pilar Cossio committed
772
					}
773
774
				}
			}
Pilar Cossio's avatar
Pilar Cossio committed
775
776
777
		}
	}
	else 
778
	{
779
		// ************ Simple Projection over the Z axis********************
Pilar Cossio's avatar
Pilar Cossio committed
780
781
		for(int n = 0; n < Model.nPointsModel; n++)
		{
782
783
784
			//Getting pixel that represents coordinates & shifting the start at to Numpix/2,Numpix/2 )
			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);
785

786
787
788
789
790
			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;
			}
791

792
793
			localproj[i * param.param_device.NumberPixels + j] += Model.points[n].density / Model.NormDen;
		}
794
795
	}

David Rohr's avatar
David Rohr committed
796
	// **** Output Just to check****
797
#ifdef PILAR_DEBUG
798
	if(iMap == 10)
799
800
801
802
803
804
	{
		ofstream myexamplemap;
		ofstream myexampleRot;
		myexamplemap.open ("MAP_i10");
		myexampleRot.open ("Rot_i10");
		myexamplemap << "ANGLES " << alpha << " " << beta << " " << gam << "\n";
805
		for(int k = 0; k < param.param_device.NumberPixels; k++)
806
		{
807
			for(int j = 0; j < param.param_device.NumberPixels; j++) myexamplemap << "\nMAP " << k << " " << j << " " << localproj[k * param.param_device.NumberPixels + j];
808
809
		}
		myexamplemap << " \n";
810
		for(int n = 0; n < Model.nPointsModel; n++)myexampleRot << "\nCOOR " << RotatedPointsModel[n].pos[0] << " " << RotatedPointsModel[n].pos[1] << " " << RotatedPointsModel[n].pos[2];
811
812
813
		myexamplemap.close();
		myexampleRot.close();
	}
814
#endif
815

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

820
821
	cuda_custom_timeslot_end;

822
823
824
	return(0);
}

825
int bioem::createConvolutedProjectionMap(int iMap, int iConv, mycomplex_t* lproj, myfloat_t* Mapconv, mycomplex_t* localmultFFT, myfloat_t& sumC, myfloat_t& sumsquareC)
826
{
David Rohr's avatar
David Rohr committed
827
828
	// **************************************************************************************
	// ****  BioEM Create Convoluted Projection Map routine, multiplies in Fourier **********
David Rohr's avatar
David Rohr committed
829
830
	// **************** calculated Projection with convoluted precalculated Kernel***********
	// *************** and Backtransforming it to real Space ********************************
David Rohr's avatar
David Rohr committed
831
	// **************************************************************************************
832

833
834
	cuda_custom_timeslot("Convolution", 1);

835
	mycomplex_t* tmp = param.fft_scratch_complex[omp_get_thread_num()];
836

David Rohr's avatar
David Rohr committed
837
	// **** Multiplying FFTmap with corresponding kernel ****
838
	const mycomplex_t* refCTF = &param.refCTF[iConv * param.FFTMapSize];
839
	for(int i = 0; i < param.param_device.NumberPixels * param.param_device.NumberFFTPixels1D; i++)
840
	{
841
842
843
		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";
844
845
	}

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

849
	sumsquareC = 0;
850
	if (FFTAlgo)
851
	{
852
853
		int jloopend = param.param_device.NumberFFTPixels1D;
		if ((param.param_device.NumberPixels & 1) == 0) jloopend--;
854
855
		for(int i = 0; i < param.param_device.NumberPixels; i++)
		{
856
			for (int j = 1;j < jloopend;j++)
857
858
859
860
861
862
			{
				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];
863
864
			if ((param.param_device.NumberPixels & 1) == 0)
			{
865
				k += param.param_device.NumberFFTPixels1D - 1;
866
867
				sumsquareC += localmultFFT[k][0] * localmultFFT[k][0] + localmultFFT[k][1] * localmultFFT[k][1];
			}
868
		}
David Rohr's avatar
David Rohr committed
869

870
871
		myfloat_t norm2 = (myfloat_t) (param.param_device.NumberPixels * param.param_device.NumberPixels);
		sumsquareC = sumsquareC / norm2;
872
	}
873
874
875
876
877
878
879
	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);
880

881
882
883
884
885
886
887
888
889
		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;
	}
David Rohr's avatar
David Rohr committed
890

891
892
	cuda_custom_timeslot_end;

893
	return(0);
894
895
}

896
int bioem::calcross_cor(myfloat_t* localmap, myfloat_t& sum, myfloat_t& sumsquare)
897
{
David Rohr's avatar
David Rohr committed
898
	// *********************** Routine to calculate Cross correlations***********************
899

900
901
	sum = 0.0;
	sumsquare = 0.0;
902
90