Polish: reduce motion-refinement memory usage

[hilaolu]

This PR reduces the peak memory and runtime of Bayesian polishing by removing the full corrected detector-movie stack from the motion-refinement path and consolidating particle/frame images into reusable contiguous allocations.

Instead of load everything into RAM at once and index them with offset like

std::vector<std::vector<Image<Complex>>> fullFrame=micrographHandler->loadMovie(...);
fullFrame[i][j];

It now uses an wrapper and index them with wrapper param. The wrapper will load frames only when needed.

ContiguousImageStack<Complex> fullFrame=micrographHandler->loadMovie(...);
fullFrame(i,j);

Benchmark

The benchmark selects one raw micrograph containing 12 particles, which are from K3 superresolutioned dataset EMPIAR-12870. The raw size is 11,520 x 8,184 x 48 (head and tail are dropped).

Run	Threads	Wall time	Peak RSS	Swap
RELION 5.1.1	1	87.35 s	37.48 GiB	0
This PR	1	55.93 s	4.48 GiB	0
This PR	32	44.60 s	4.64 GiB	0

Numerical validation

The multithreaded polishing path is not bit identical for two reasons:

• motion fitting contains order-dependent floating-point reductions;
• detector-defect correction calls the process-global rand() from an OpenMP loop, so scheduling can assign different replacement samples to hot pixels.

With one thread over all 48 production frames:

• all trajectory scalars are exactly equal;
• the trajectory STAR files are bit identical;
• the complete voxel payloads of the FCC cc, w0, and w1 MRC files are bit identical;

Run	Maximum trajectory delta	RMS delta	Correlation
Legacy 1 thread vs this PR 1 thread	0 A	0 A	1
This PR 1 thread vs this PR 32 threads	0.006133 A	0.001591967 A	0.999999822621

The exact single-thread result establishes that the memory-layout and streaming changes preserve the polishing calculation when execution order is fixed.

[biochem-fan]

Thanks for your suggestion.

The benchmark selects one raw micrograph containing 12 particles

This is very low. What happens with datasets with higher particle density?

How does the scaling behaviour compare with the original implementation?

Finally, please declare the use of LLMs if you are using one for code generation and/or analysis.

[hilaolu]

Hi, @biochem-fan, thanks for your suggestions. I use GPT-5.6-sol via Codex for most code generation, all benchmark scripts and numerical validation. The benchmark param matrix is made up myself.

Your claim on particles size is fair, since the smoke micrograph is randomly picked. I rerun the benchmark with selected fiber rich micrographs. I'd like to test how it scales, however scaling legacy polish with MPI procs is impossible in my system since I have only 64 GiB RAM installed. I will test how thread scales.

For each micrograph, particles are extracted with nr_asu=6 and box_size=280A(340 pixels). 40A pass is applied for visualization of fiber like structure.

https://ftp.ebi.ac.uk/empiar/world_availability/12870/data/RAW_Tif/Peter_08032022_206-5_00010_X%201Y%200-1.tif (202 particles extracted)

https://ftp.ebi.ac.uk/empiar/world_availability/12870/data/RAW_Tif/Peter_08032022_191-14_00013_X%200Y%200-1.tif (242 particles extracted)

https://ftp.ebi.ac.uk/empiar/world_availability/12870/data/RAW_Tif/Peter_08032022_170-17_00027_X%201Y%201-3.tif (208 particles extracted)

Benchmark

Particles	Legacy wall	This PR wall	Runtime change	Legacy RSS	This PR RSS	RSS change
242	184.00 s	133.95 s	-27.20%	37.48 GiB	17.54 GiB	-53.19%
208	168.41 s	122.74 s	-27.12%	37.48 GiB	15.59 GiB	-58.39%
202	163.77 s	114.79 s	-29.91%	37.48 GiB	15.25 GiB	-59.31%

[biochem-fan]

The wrapper will load frames only when needed.

Where does this deferring happen? Can you provide more details?
I don't want to read large changes made by an AI.

Where are your movies stored? SSD?

How does the access pattern change with your patch? In the original code, a full movie is read by the main thread before the multi-thread section starts, right? Or was it already parallel? I don't remember well. My concern is that if each thread reads (part of) movies, it generates many non-sequential disk reads and makes HDD-based systems very slow.

I'd like to test how it scales, however scaling legacy polish with MPI procs is impossible in my system since I have only 64 GiB RAM installed

It is a pity. The code does not scale to many threads / MPI because the number of particles per movie is limited. MPI scaling behaviour is also important.

[hilaolu]

Where are your movies stored? SSD?

They are in HDD, however this doesn't matter since the tests are run multiple times on same micrographs, they are in memory cache already.

[hilaolu]

In src/jaz/single_particle/micrograph_handler.cpp,

	BufferedImage<float> muGraph;
	
	...

	muGraph = MovieLoader::readEER<float>(
		movieFn, gainRefToUse, defectMaskToUse,
		frame0, fc,
		eer_upsampling, eer_grouping,
		nr_omp_threads);


	//the copy assignment creates temporary deep copy of muGraph

	...

	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);

	//muGraph is no longer used from now on

Here, the muGraph is a full frame float32 movie, each pixel takes 4 bytes. The muGraph is actually consumed by extractMovieStackFS, which extract and return out[particle][frames]. You may notice that the movies are actually extracted frame by frame.

src/jaz/single_particle/spa_extraction.h

template <typename T>
std::vector<std::vector<Image<Complex>>> SpaExtraction::extractMovieStackFS(
		const MetaDataTable& mdt, 
		const RawImage<T>& movie,
		int boxSize,
		double outPs, double coordsPs, double moviePs, double dataPs,
		const std::vector<std::vector<gravis::d2Vector>>* offsets_in,
		std::vector<std::vector<gravis::d2Vector>>* offsets_out,
		int num_threads)
{
    ...

	const int w0 = movie.xdim;
	const int h0 = movie.ydim;
	const int fc = movie.zdim;

	...

    
	#pragma omp parallel for num_threads(num_threads)
	for (long int f = 0; f < fc; f++)
	{
	    ...

		for (long p = 0; p < pc; p++)
		{


    	    ...
         
			out[p][f] = Image<Complex>(sqMg,sqMg);
			
    	    ...

			for (long int y = 0; y < sqMg; y++)
			for (long int x = 0; x < sqMg; x++)
			{

                ...

				DIRECT_NZYX_ELEM(aux0[t].data, 0, 0, y, x) = movie(xx,yy,f);

    			if (outPs == moviePs)
    			{
    				fts[t].FourierTransform(aux0[t](), out[p][f]());
    			}
    			else
    			{
    				fts[t].FourierTransform(aux0[t](), aux1[t]());
    				out[p][f] = FilterHelper::cropCorner2D(aux1[t], boxSize/2+1, boxSize);
    			}

                ...

				
			}
		}
	}

	return out;
}

[hilaolu]

Where does this deferring happen? Can you provide more details?

I made a mistake that memory saving/defer loading/streaming doesn't yield from the wrapper. The wrapper actually accounts for runtime optimization. In that case, particles or other are stored in a continuous memory block, reduce overhead of allocating/reclaiming on heap. The memory usage optimization and the wrapper are technically orthogonal to each other. The wrapper touchs too many things and should be separated.

It is a pity. The code does not scale to many threads / MPI because the number of particles per movie is limited. MPI scaling behaviour is also important.

It does scale in my system. For legacy, MPI mprocs scaling is limited since ~40GiB per process exceeds 64 GiB RAM easily, regardless of particle count of the micrograph. In this PR, the most significant memory footprint cost shifts to particles count and box size. I can run with -np 6 --j 32 through the entire ~4000 micrographs without any OOM, since ~200 particles per micrograph is an extreme case and it is very unlikely that the parallel 6 mprocs all have such great number of particles.

[biochem-fan]

I am getting confused. If deferred reading is not implemented here, where does the memory saving come from?

the copy assignment creates temporary deep copy of muGraph
muGraph is no longer used from now on

Is the core problem the copy assignment and life time of muGraph? Instad of creating a new class, cannot you just improve this by for example using move semantics or references and releasing it immediately after extractMovieStackFS?

You may notice that the movies are actually extracted frame by frame.

I meant movie frame loading inMovieLoader, not particle extraction.