Slow birefringence reconstruction for z-stacks about waveorder

[talonchandler]

Thanks for writing this up @edyoshikun. I agree that we should expect ~linear reconstruction times with increasing FOV size since this is a voxel-by-voxel reconstruction.

We'll need to do some debugging to figure out if this bottleneck is caused by:

• saturating our compute resources? (I don't think we're filling RAM, but we may be filling a cache or other resource)
• an accidental bottleneck somewhere in the birefringence code...an accidental non-linear step?

A "trivial" (but more work intensive) temporary workaround might be to split the large stacks into pieces across CPUs/jobs.

[edyoshikun]

Ok, so thanks to @ziw-liu suggestion to run the individual z-planes separately. I used the multiprocess library to allocate 61 processes and got them down from 547s to 22s processing time. This is a short-term fix.

I used the cProfiler and snakeviz visualizer and this points out to Pol_recon and the Stokes_recon function.

Start Reconstructor
Initializing Reconstructor...
Finished Initializing Reconstructor (0.00 min)
/hpc/projects/comp_micro/projects/zebrafish-infection/2012_12_09_zebrafish_GFPisg15_dtomato_mpeg/2023_01_06_141529/multiposition_intoto_1
No existing directory found. Creating new directory at /hpc/projects/comp_micro/projects/zebrafish-infection/2012_12_09_zebrafish_GFPisg15_dtomato_mpeg/2023_01_06_141529/multiposition_intoto_1
Creating new zarr store at /hpc/projects/comp_micro/projects/zebrafish-infection/2012_12_09_zebrafish_GFPisg15_dtomato_mpeg/2023_01_06_141529/multiposition_intoto_1/20230106_birefringence_pos_0.zarr
Begin stokes calculation
Stokes elapsed time:9.410654783248901
Shape of background corrected data Stokes: (5, 5, 2048, 2048)
Begin birefringence calculation
Bire elapsed time:2.1207518577575684
Total elapsed time:11.531463384628296
         70 function calls in 9.411 seconds

   Ordered by: cumulative time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.002    0.002    9.411    9.411 reconstructions.py:259(reconstruct_qlipp_stokes)
        1    8.549    8.549    8.549    8.549 waveorder_reconstructor.py:1049(Stokes_recon)
        1    0.834    0.834    0.834    0.834 waveorder_reconstructor.py:1113(Stokes_transform)
        5    0.000    0.000    0.026    0.005 {built-in method numpy.core._multiarray_umath.implement_array_function}
        1    0.000    0.000    0.025    0.025 <__array_function__ internals>:177(copy)
        1    0.000    0.000    0.025    0.025 function_base.py:870(copy)
        1    0.025    0.025    0.025    0.025 {built-in method numpy.array}`
```

[mattersoflight]

Thanks for logging this @edyoshikun. For deeper debugging and proper fix, you can profile both reconstruct_qlipp_stokes and reconstruct_qlipp_birefringence with snakeviz, which uses cProfile module of python.

[mattersoflight]

Hi @edyoshikun , @talonchandler All the operations we apply to compute birefringence are straightforward with linear complexity (inverse instrument matrix, background correction, square roots).

My observations/ questions:

• It makes sense that Stokes Elapsed Time increases linearly.
• It doesn't make sense that Birefringence Elapsed Time increases exponentially.
• Does reconstructor = initialize_reconstructor(pipeline="birefringence", **reconstructor_args) compute the phase OTF and operate on it anyway? If it does, that is a bug that can explain the exponential increase in compute time.
• If you set "mode": "3D" parameter to 2D, it should not affect the birefringence. The parameter currently affects only phase reconstruction as far as I remember.

[edyoshikun]

• Yes, it doesn't make sense that by increasing the number of slices the code is exponentially slower.
• The birefringence calculations are just linear as you mention and pixel to pixel, and I've checked the flags in recorder and waveorder related to the OTF calculations.
• I had the dictionary of arguments such as mode:3D, pad_z,z_step,n_slice which are not needed for birefringence but left them in place as in theory they don't change the reconstruction pipeline and are used for QLIPP. I can take them out and take a closer look if that makes things different.

[talonchandler]

@edyoshikun you can try out recOrder #322 or you can replace
birefringence = reconstruct_qlipp_birefringence(stokes, reconstructor)
with
birefringence = reconstructor.Polarization_Recon(stokes) to skip all of the transposes and copies.