Polish: reduce motion-refinement memory usage

src/jaz/image/buffered_image.h

@@ -1,6 +1,8 @@
 #ifndef JAZ_VECTOR_IMAGE_H
 #define JAZ_VECTOR_IMAGE_H
 
+#include <utility>
+
 #include "raw_image.h"
 
 
@@ -15,6 +17,7 @@ class BufferedImage : public RawImage<T>
 		BufferedImage();
 		BufferedImage(size_t xdim, size_t ydim = 1, size_t zdim = 1);
 		BufferedImage(const BufferedImage& vi);
+		BufferedImage(BufferedImage&& vi) noexcept;
 		BufferedImage(const RawImage<T>& vi);
 		BufferedImage(const Image<T>& vi);
 		BufferedImage(std::string filename);
@@ -37,6 +40,7 @@ class BufferedImage : public RawImage<T>
 		RawImage<T> getRef();
 
 		BufferedImage& operator = (const BufferedImage& other);
+		BufferedImage& operator = (BufferedImage&& other) noexcept;
 };
 
 
@@ -71,6 +75,27 @@ BufferedImage<T>& BufferedImage<T>::operator = (const BufferedImage<T>& other)
 	return *this;
 }
 
+template <class T>
+BufferedImage<T>& BufferedImage<T>::operator = (BufferedImage<T>&& other) noexcept
+{
+	if (this != &other)
+	{
+		this->xdim = other.xdim;
+		this->ydim = other.ydim;
+		this->zdim = other.zdim;
+		dataVec = std::move(other.dataVec);
+		this->data = dataVec.data();
+
+		other.xdim = 0;
+		other.ydim = 0;
+		other.zdim = 0;
+		other.dataVec.clear();
+		other.data = 0;
+	}
+
+	return *this;
+}
+
 template <class T>
 BufferedImage<T>::BufferedImage()
 	: RawImage<T>()
@@ -92,6 +117,20 @@ BufferedImage<T>::BufferedImage(const BufferedImage<T>& vi)
 	RawImage<T>::data = &(dataVec[0]);
 }
 
+template <class T>
+BufferedImage<T>::BufferedImage(BufferedImage<T>&& vi) noexcept
+	:   RawImage<T>(vi.xdim, vi.ydim, vi.zdim, 0),
+	  dataVec(std::move(vi.dataVec))
+{
+	RawImage<T>::data = dataVec.data();
+
+	vi.xdim = 0;
+	vi.ydim = 0;
+	vi.zdim = 0;
+	vi.dataVec.clear();
+	vi.data = 0;
+}
+
 template <class T>
 BufferedImage<T>::BufferedImage(const RawImage<T>& vi)
 	:   RawImage<T>(vi)

src/jaz/image/contiguous_image_stack.h

@@ -0,0 +1,93 @@
+#ifndef JAZ_CONTIGUOUS_IMAGE_STACK_H
+#define JAZ_CONTIGUOUS_IMAGE_STACK_H
+
+#include <src/image.h>
+#include <src/jaz/image/raw_image.h>
+
+#include <vector>
+
+
+template <typename T>
+class ContiguousImageStack
+{
+	public:
+
+		ContiguousImageStack()
+			: particle_count(0), frame_count(0), width(0), height(0)
+		{}
+
+		ContiguousImageStack(const ContiguousImageStack&) = delete;
+		ContiguousImageStack& operator = (const ContiguousImageStack&) = delete;
+
+		void resize(long int particleCount, long int frameCount,
+				long int width, long int height)
+		{
+			images.clear();
+
+			particle_count = particleCount;
+			frame_count = frameCount;
+			this->width = width;
+			this->height = height;
+
+			if (particleCount <= 0 || frameCount <= 0 || width <= 0 || height <= 0)
+			{
+				storage.clear();
+				particle_count = frame_count = this->width = this->height = 0;
+				return;
+			}
+
+			const long int imageCount = particleCount * frameCount;
+			storage.data.setDimensions(width, height, 1, imageCount);
+			storage.data.coreAllocateReuse();
+
+			images.resize(imageCount);
+
+			const long int imageSize = width * height;
+			for (long int i = 0; i < imageCount; i++)
+			{
+				images[i].data.setDimensions(width, height, 1, 1);
+				images[i].data.data = storage.data.data + i * imageSize;
+				images[i].data.destroyData = false;
+			}
+		}
+
+		long int particleCount() const { return particle_count; }
+		long int frameCount() const { return frame_count; }
+		long int xdim() const { return width; }
+		long int ydim() const { return height; }
+
+		Image<T>& operator () (long int particle, long int frame)
+		{
+			return images[particle * frame_count + frame];
+		}
+
+		const Image<T>& operator () (long int particle, long int frame) const
+		{
+			return images[particle * frame_count + frame];
+		}
+
+		RawImage<T> getParticleRef(long int particle)
+		{
+			return RawImage<T>(
+				width, height, frame_count,
+				storage.data.data + particle * frame_count * width * height);
+		}
+
+		const RawImage<T> getParticleRef(long int particle) const
+		{
+			return RawImage<T>(
+				width, height, frame_count,
+				storage.data.data + particle * frame_count * width * height);
+		}
+
+		const T* data() const { return storage.data.data; }
+		T* data() { return storage.data.data; }
+
+	private:
+
+		long int particle_count, frame_count, width, height;
+		Image<T> storage;
+		std::vector<Image<T>> images;
+};
+
+#endif

src/jaz/single_particle/img_proc/filter_helper.cpp

@@ -407,6 +407,27 @@ Image<Complex> FilterHelper::cropCorner2D(const Image<Complex>& img, int w, int
     return out;
 }
 
+void FilterHelper::cropCorner2D(const Image<Complex>& img, Image<Complex>& out)
+{
+	const int w = out.data.xdim;
+	const int h = out.data.ydim;
+	const int w1 = img.data.xdim;
+	const int h1 = img.data.ydim;
+
+	for (int y = 0; y < h1; y++)
+	for (int x = 0; x < w1; x++)
+	{
+		const int x1 = x;
+		const int y1 = y < h1 / 2? y : y - h1;
+
+		if (x1 < w && y1 < h / 2 && y1 >= -h / 2)
+		{
+			const int y0 = y1 < 0? y1 + h : y1;
+			DIRECT_A2D_ELEM(out.data, y0, x1) = DIRECT_A2D_ELEM(img.data, y, x);
+		}
+	}
+}
+
 Image<RFLOAT> FilterHelper::zeroOutsideCorner2D(Image<RFLOAT> &img, double radius)
 {
     const int w = img.data.xdim;

src/jaz/single_particle/img_proc/filter_helper.h

@@ -65,6 +65,7 @@ class FilterHelper
         static Image<RFLOAT> padCorner2D(const Image<RFLOAT> &img, int w, int h);
         static Image<RFLOAT> cropCorner2D(const Image<RFLOAT> &img, int w, int h);
         static Image<Complex> cropCorner2D(const Image<Complex> &img, int w, int h);
+        static void cropCorner2D(const Image<Complex>& img, Image<Complex>& out);
         static Image<RFLOAT> zeroOutsideCorner2D(Image<RFLOAT>& img, double radius);
         static void GaussianEnvelopeCorner2D(Image<RFLOAT>& img, double sigma);
 		static Image<RFLOAT> raisedCosEnvCorner2D(Image<RFLOAT>& img, double radIn, double radOut);

src/jaz/single_particle/micrograph_handler.cpp

@@ -351,9 +351,10 @@ void MicrographHandler::validatePixelSize(RFLOAT angpix) const
 	}
 }
 
-std::vector<std::vector<Image<Complex>>> MicrographHandler::loadMovie(
+void MicrographHandler::loadMovie(
 		const MetaDataTable &mdt, int s,
 		double angpix, std::vector<ParFourierTransformer>& fts,
+		ContiguousImageStack<Complex>& movie,
 		const std::vector<std::vector<gravis::d2Vector>>* offsets_in,
 		std::vector<std::vector<gravis::d2Vector>>* offsets_out,
 		double data_angpix,
@@ -430,8 +431,6 @@ std::vector<std::vector<Image<Complex>>> MicrographHandler::loadMovie(
 				DIRECT_MULTIDIM_ELEM(defectMask, n) = true;
 	}
 
-	BufferedImage<float> muGraph;
-
 	RawImage<RFLOAT> gainRef_new(lastGainRef);
 	RawImage<bool> defectMask_new(defectMask);
 
@@ -448,7 +447,19 @@ std::vector<std::vector<Image<Complex>>> MicrographHandler::loadMovie(
 
 	const int frame0 = returnSingleFrame? single_frame_relative_index : firstFrame;
 	const int fc = returnSingleFrame? 1 : lastFrame - firstFrame + 1;
-
+
+	std::vector<Image<RFLOAT>> extractionAuxReal;
+	std::vector<Image<Complex>> extractionAuxFourier;
+
+	auto extractFrame = [&](const RawImage<float>& frame, long int f)
+	{
+		SpaExtraction::extractMovieFrameFS(
+			mdt, frame, f, fc, s,
+			angpix, coords_angpix, movie_angpix, data_angpix,
+			offsets_in, offsets_out, movie, fts,
+			extractionAuxReal, extractionAuxFourier);
+	};
+
 	if (isEER)			
 	{
 		if (eer_upsampling < 0)
@@ -461,47 +472,39 @@ std::vector<std::vector<Image<Complex>>> MicrographHandler::loadMovie(
 			eer_grouping = micrograph.getEERGrouping();
 		}
 
-		muGraph = MovieLoader::readEER<float>(
+		MovieLoader::readEERFrames<float>(
 			movieFn, gainRefToUse, defectMaskToUse,
 			frame0, fc,
 			eer_upsampling, eer_grouping,
-			nr_omp_threads);
+			nr_omp_threads, extractFrame);
 	}
 	else
 	{
-		muGraph = MovieLoader::readDense<float>(
+		MovieLoader::readDenseFrames<float>(
 			movieFn, gainRefToUse, defectMaskToUse,
 			frame0, fc,
 			hotCutoff,
-			nr_omp_threads);
+			nr_omp_threads, extractFrame);
 	}
 
-	std::vector<std::vector<Image<Complex>>> movie = SpaExtraction::extractMovieStackFS(
-			mdt, muGraph, s,
-			angpix, coords_angpix, movie_angpix, data_angpix,
-			offsets_in, offsets_out, 
-			nr_omp_threads);
-
-	const int pc = movie.size();
+	const int pc = movie.particleCount();
 
 	if (!returnSingleFrame)
 	{
 		#pragma omp parallel for num_threads(nr_omp_threads)
 		for (int p = 0; p < pc; p++)
 		{
-			RFLOAT scale2 = StackHelper::computePower(movie[p], false);
+			RFLOAT scale2 = StackHelper::computePower(movie.getParticleRef(p), false);
 
 			for (int f = 0; f < fc; f++)
 			{
 				// NOTE: sqrt(fc) is a legacy scaling factor.
 				// It probably shouldn't be there.
 				//                   --JZ
-				RawImage<Complex>(movie[p][f]) /= s * sqrt(scale2 / fc); 
+				RawImage<Complex>(movie(p, f)) /= s * sqrt(scale2 / fc);
 			}
 		}
 	}
-
-	return movie;
 }
 
 void MicrographHandler::loadInitialTracks(

src/jaz/single_particle/micrograph_handler.h

@@ -26,6 +26,7 @@
 
 #include <src/jaz/gravis/t2Vector.h>
 #include <src/jaz/single_particle/parallel_ft.h>
+#include <src/jaz/image/contiguous_image_stack.h>
 
 #include <src/micrograph_model.h>
 #include <src/image.h>
@@ -76,11 +77,11 @@ class MicrographHandler
 		const std::vector<MetaDataTable>& mdts,
 		int fc, int verb);
 
-	// load a movie and extract all particles
-	// returns a per-particle vector of per-frame images of size (s/2+1) x s
-	std::vector<std::vector<Image<Complex>>> loadMovie(
+	// load a movie and extract all particles into contiguous per-frame storage
+	void loadMovie(
 		const MetaDataTable& mdt, int s, double angpix, 
 		std::vector<ParFourierTransformer>& fts,
+		ContiguousImageStack<Complex>& movie,
 		const std::vector<std::vector<gravis::d2Vector>>* offsets_in = 0,
 		std::vector<std::vector<gravis::d2Vector>>* offsets_out = 0,
 		double data_angpix = -1,

src/jaz/single_particle/motion/frame_recombiner.cpp

@@ -317,6 +317,7 @@ void FrameRecombiner::process(const std::vector<MetaDataTable>& mdts, long g_sta
 
 		BufferedImage<Complex> sumStack(sh_out[ogmg], s_out[ogmg], pc);
 		sumStack.fill(Complex(0,0));
+		ContiguousImageStack<Complex> fullFrame;
 
 		// This is an ugly code duplication.
 		// Frame-by-frame reading (the normal route) is more memory efficient but
@@ -325,16 +326,16 @@ void FrameRecombiner::process(const std::vector<MetaDataTable>& mdts, long g_sta
 		// CompressedMRCReader (and probably EERRenderer).
 		if (readAtOnce)
 		{
-			std::vector<std::vector<Image<Complex>>> fullFrame = micrographHandler->loadMovie(
-						mdtOut, s_out[ogmg], angpix_out[ogmg], fts,
+			micrographHandler->loadMovie(
+						mdtOut, s_out[ogmg], angpix_out[ogmg], fts, fullFrame,
 						&priorShift, &shift, data_angpix[ogmg]);
 
 			for (int f = 0; f < fc; f++)
 			{
 				#pragma omp parallel for num_threads(nr_omp_threads)
 				for (int p = 0; p < pc; p++)
 				{
-					RawImage<Complex> obs(fullFrame[p][f]);
+					RawImage<Complex> obs(fullFrame(p, f));
 
 					Translation::shiftInFourierSpace2D(obs, -shift[p][f].x, -shift[p][f].y);
 
@@ -357,14 +358,14 @@ void FrameRecombiner::process(const std::vector<MetaDataTable>& mdts, long g_sta
 					priorShift_f[p] = {priorShift[p][f]};
 				}
 
-				std::vector<std::vector<Image<Complex>>> fullFrame = micrographHandler->loadMovie(
-							mdtOut, s_out[ogmg], angpix_out[ogmg], fts,
+				micrographHandler->loadMovie(
+							mdtOut, s_out[ogmg], angpix_out[ogmg], fts, fullFrame,
 							&priorShift_f, &shift, data_angpix[ogmg], f);
 
 				#pragma omp parallel for num_threads(nr_omp_threads)
 				for (int p = 0; p < pc; p++)
 				{
-					RawImage<Complex> obs(fullFrame[p][0]);
+					RawImage<Complex> obs(fullFrame(p, 0));
 
 					Translation::shiftInFourierSpace2D(obs, -shift[p][0].x, -shift[p][0].y);