Cov2d noise_var=0 appears to have a bug (crash)

[garrettwrong]

There appears to be a discrepancy between 0 and 1 for the variable noise_var in cov2d family of code. This is similar to an issue I found in the legacy code, but I was expecting it to be handled in ASPIRE-Python because it is specifically mentioned in the docstrings... However, trying to use noise_var=0 with this code (for using clean images in development) results in a crash.

It would appear most of this cov2d code uses a default noise_var=1, I'm guessing to superficially circumvent this crashing in the shrinker codes.

To repro easiest thing is to alter the cov2d tutorial:

(aspire_cls_avg) ➜  ASPIRE-Python git:(cls_avg_cov2d) ✗ diff tutorials/examples/cov2d_bugcheck.py tutorials/examples/cov2d_simulation.py 
41c41
< noise_var = 0 #1.3957e-4
---
> noise_var = 1.3957e-4

bt:

(aspire_cls_avg) ➜  ASPIRE-Python git:(cls_avg_cov2d) ✗ python tutorials/examples/cov2d_bugcheck.py 
2021-05-25 11:50:27,700 INFO This script illustrates 2D covariance Wiener filtering functionality in ASPIRE package.
2021-05-25 11:50:27,700 INFO Simulation running in <class 'numpy.float32'> precision.
2021-05-25 11:50:27,700 INFO Initialize simulation object and CTF filters.
2021-05-25 11:50:27,700 INFO Load 3D map and downsample 3D map to desired grids of 64 x 64 x 64.
2021-05-25 11:50:28,162 INFO Use downsampled map to creat simulation object.
2021-05-25 11:50:28,181 INFO Appending a NoiseAdder to generation pipeline
2021-05-25 11:50:28,181 INFO Expanding 2D image in a frequency-domain Fourier–Bessel basis using the fast method.
2021-05-25 11:50:28,647 INFO Apply CTF filters to clean images.
2021-05-25 11:50:29,057 INFO Trying NFFT backend finufft
2021-05-25 11:50:29,060 INFO NFFT backend finufft usable.
2021-05-25 11:50:29,060 INFO Trying NFFT backend cufinufft
2021-05-25 11:50:29,060 INFO NFFT backend cufinufft not usable:
	No module named 'pycuda'
2021-05-25 11:50:29,060 INFO Trying NFFT backend pynfft
2021-05-25 11:50:29,061 INFO NFFT backend pynfft not usable:
	No module named 'pynfft'
2021-05-25 11:50:29,061 INFO Selected NFFT backend = finufft.
tutorials/examples/cov2d_bugcheck.py:97: RuntimeWarning: divide by zero encountered in double_scalars
  sn_ratio = power_clean / noise_var
2021-05-25 11:50:32,228 INFO Signal to noise ratio is inf.
2021-05-25 11:50:37,565 INFO Applying forward transformations in pipeline
2021-05-25 11:50:37,565 INFO All forward transformations applied
2021-05-25 11:50:37,565 INFO Loaded 1024 images
2021-05-25 11:50:37,565 INFO Get coefficients of clean and noisy images in FFB basis.
2021-05-25 11:50:40,000 INFO Get 2D covariance matrices of clean and noisy images using FB coefficients.
/home/garrett/work/ClassAveraging/ASPIRE-Python/src/aspire/covariance/covar2d.py:32: RuntimeWarning: divide by zero encountered in true_divide
  covar = covar_in / noise_var
2021-05-25 11:50:43,297 CRITICAL Array must not contain infs or NaNs
Traceback:
  File "tutorials/examples/cov2d_bugcheck.py", line 154, in <module>
    covar_est_opt=covar_opt,
  File "/home/garrett/work/ClassAveraging/ASPIRE-Python/src/aspire/covariance/covar2d.py", line 305, in get_covar
    covar_est_opt["shrinker"],
  File "/home/garrett/work/ClassAveraging/ASPIRE-Python/src/aspire/covariance/covar2d.py", line 372, in shrink_covar_backward
    b_ell = shrink_covar(b_ell, noise_var, p / n, shrinker)
  File "/home/garrett/work/ClassAveraging/ASPIRE-Python/src/aspire/covariance/covar2d.py", line 34, in shrink_covar
    lambs, eig_vec = eig(make_symmat(covar))
  File "/home/garrett/anaconda3/envs/aspire_cls_avg/lib/python3.6/site-packages/numpy/linalg/linalg.py", line 1305, in eig
    _assertFinite(a)
  File "/home/garrett/anaconda3/envs/aspire_cls_avg/lib/python3.6/site-packages/numpy/linalg/linalg.py", line 218, in _assertFinite
    raise LinAlgError("Array must not contain infs or NaNs")

Error in sys.excepthook:
Traceback (most recent call last):
  File "/home/garrett/work/ClassAveraging/ASPIRE-Python/src/aspire/exceptions.py", line 87, in handle_exception
    raise exc_value
  File "tutorials/examples/cov2d_bugcheck.py", line 154, in <module>
    covar_est_opt=covar_opt,
  File "/home/garrett/work/ClassAveraging/ASPIRE-Python/src/aspire/covariance/covar2d.py", line 305, in get_covar
    covar_est_opt["shrinker"],
  File "/home/garrett/work/ClassAveraging/ASPIRE-Python/src/aspire/covariance/covar2d.py", line 372, in shrink_covar_backward
    b_ell = shrink_covar(b_ell, noise_var, p / n, shrinker)
  File "/home/garrett/work/ClassAveraging/ASPIRE-Python/src/aspire/covariance/covar2d.py", line 34, in shrink_covar
    lambs, eig_vec = eig(make_symmat(covar))
  File "/home/garrett/anaconda3/envs/aspire_cls_avg/lib/python3.6/site-packages/numpy/linalg/linalg.py", line 1305, in eig
    _assertFinite(a)
  File "/home/garrett/anaconda3/envs/aspire_cls_avg/lib/python3.6/site-packages/numpy/linalg/linalg.py", line 218, in _assertFinite
    raise LinAlgError("Array must not contain infs or NaNs")
numpy.linalg.LinAlgError: Array must not contain infs or NaNs

Original exception was:
Traceback (most recent call last):
  File "/home/garrett/work/ClassAveraging/ASPIRE-Python/src/aspire/exceptions.py", line 87, in handle_exception
    raise exc_value
  File "tutorials/examples/cov2d_bugcheck.py", line 154, in <module>
    covar_est_opt=covar_opt,
  File "/home/garrett/work/ClassAveraging/ASPIRE-Python/src/aspire/covariance/covar2d.py", line 305, in get_covar
    covar_est_opt["shrinker"],
  File "/home/garrett/work/ClassAveraging/ASPIRE-Python/src/aspire/covariance/covar2d.py", line 372, in shrink_covar_backward
    b_ell = shrink_covar(b_ell, noise_var, p / n, shrinker)
  File "/home/garrett/work/ClassAveraging/ASPIRE-Python/src/aspire/covariance/covar2d.py", line 34, in shrink_covar
    lambs, eig_vec = eig(make_symmat(covar))
  File "/home/garrett/anaconda3/envs/aspire_cls_avg/lib/python3.6/site-packages/numpy/linalg/linalg.py", line 1305, in eig
    _assertFinite(a)
  File "/home/garrett/anaconda3/envs/aspire_cls_avg/lib/python3.6/site-packages/numpy/linalg/linalg.py", line 218, in _assertFinite
    raise LinAlgError("Array must not contain infs or NaNs")
numpy.linalg.LinAlgError: Array must not contain infs or NaNs

Fixing this bug will look like:

• Reverse engineering what the code and value intent probably was.
• Patching the code.
• Patching the docstrings.
• Testing the case above.
• Confirming running clean images through the cov2d tool yields a sane result wrt the known clean images and opening a new issue if that is a larger problem.

[janden]

I trace it to here:

ASPIRE-Python/src/aspire/covariance/covar2d.py

Line 32 in 4b25c09

covar = covar_in / noise_var

Basically, the covariance matrix is normalized by noise_var to simplify some of the calculations below. This will involve removing that division and sprinkling in a few noise_var * in the rest of that function.

[janden]

So there's two ways to solve this: treat noise_var = 0 as a special case or rewrite the formulas so that they are correct in this special case. The former is obviously easier, but the latter may help with some instability issues when the noise is very small.

I've been able to fix up (I think) shrink_covar to deal gracefully with noise_var = 0, but get_cwf_coeffs is trickier. Might be simpler to do a bypass there for now.

[garrettwrong]

Yes that is one problem. I've tried a few things (including scaling the noise_var through lamba (constant just comes through right)). I am finding the tutorial example (which is more/less like what I would try to do in other codes...) comes down later with Singular Matrix in get_cwf_coeffs. The unit test does not and I don't know why yet...

[janden]

I've tried a few things (including scaling the noise_var through lamba (constant just comes through right)).

Right. I think I've gotten that to work. Will push a proposed fix.

comes down later with Singular Matrix in get_cwf_covar.

If I had to guess, it's related to the fact that when forcing the covariance matrix to be PSD, we may set certain eigenvalues to zero. Now this means that when we try to invert it in get_cwf_coeffs, it's singular. If there's noise, this isn't a problem, but since noise_var = 0, we tickle this instability.

One solution is to project all the images onto the space of covariance eigenvectors with non-zero eigenvalues and do all our work there (that's a pain). Another is to just return the original coefficients when noise_var = 0 (Wiener filtering is a no-op if there is no noise). I suggest the latter here.

[garrettwrong]

Okay cool, thanks for you help. I'll bring your patch down (closed the other PR, will do cleanup separate later).

[garrettwrong]

Closed crash issue with #421. Some improvement still on the table #424.