Release memory after Bayesian optimization

[sezelt]

Previously the BO code retained a reference to the ptycho reconstruction, which in some cases would include a copy of the dataset. This ensures that these resources are released. Additionally it puts pbar.close() into a try/finally clause so the progressbar is always terminated correctly.

[gvarnavi]

Thanks @sezelt - looks good to me!

Quick-note re: how BO plays w/ the new device/storage memory management:

The current default behavior for clear_fft_cache=True is to clear the cupy FFT plans at the end of preprocess/reconstruct. This might make less sense for BO since we're running reconstruct over and over again.

The user can always specify clear_fft_cache=False in init_args to disable this, but perhaps it might make more sense to add this to _set_optimizer_defaults and simply call ptycho.clear_device_mem(ptycho._device,True) during the finally clause.

[sezelt]

I'm curious, how much time does it actually take to generate the plan, and how much memory does it use?

[gvarnavi]

I think it's actually not insignificant, since it's of the order of (max_batch_size,Qx,Qy).

You can get a feel on how slow it is by seeing how much slower the first ptycho iteration is from subsequent iterations with clear_fft_cache=True, and test the memory usage with the notebook I attached here.

[sezelt]

I added clear_fft_cache=False to the default optimizer args. Calling ptycho.clear_device_mem is trickier because the ptycho object is wrapped by the optimizer function, so it's not in scope at the level of optimizer.optimize. I think the better solution (and one which potentially will fix other GPU memory usage issues I've been having in my notebooks recently even without using the optimizer) is for the ptycho object to clear the plans on its destruction.

[sezelt]

There also seems to be something odd happening with the black style check, since it wants to reformat 21 files unrelated to this PR. Perhaps the version of black on the GH action has changed?

[sezelt]

So it does seem a new version of black was released, and that has led to new rules in some places. I've applied the new formatting in this PR (the only functional changes are the ones originally described).

[sezelt]

@gvarnavi I made the change you asked for, but haven't tested everything. Is the current approach to your liking?

[gvarnavi]

lgtm, thanks @sezelt!