Ndim fft

[herjy]

All right, I am mildly fed up with this, so I will submit it for review. The critical part, which is the ndimensional fft for Observation works marvelously well. The same trick does not work at all for LowResObservation. If anything it makes everything slower. Here is a recap of what you'll find:

• On the simple example from the notebook, the execution of blend().fit drops from 88 ms to 18 ms. I will do more tests to see hoe the various dimensions of the arrays affect this, but it looks nice!
• Problem: It does not pass the tests and I am not entirely sure why.
• make_operator is now a method of LowResObsevation, which makes more sense.
• I considered switching resampling.match_patches to a method of that class too, but its use in match and match_psfs made me think twice.
• I timed the setup time for LowRes and the run times for blend.fit in the notebooks. For comparison, the same thing is done on the current implementation in a branch I called 'speed_test' which I will use to investigate further the gain that the new implementation gives us.
  On the bad news side (yes this was all good news):
• The more I try to use Ndim_fft in LowRes, the slower it gets. I tried an implementation full-on ndim fft (it was beautiful, not a single for loop), but after running (building the resampling operator) for MANY minutes I lost patience and gave up. You will find a hybrid implementation that still uses scipy.
• The prototypical code for ndim fft in `make_operator is left commented. I could never make it work for reasons that are still unclear.

Since I already spent days on trying to make LowRes faster without success, Idon't want to delay this any further, I suggest that you guys review the part in Observation. If I can find why LowRes is slow, I'll update everything. Writing this makes me think that it could be a problem of fftpack_shapes but I am not convinced.
Good luck.

[fred3m]

Looks good to me, I didn't see anything that Peter didn't already comment on.

[pmelchior]

Besides the three open comments, it would be nice to have a simple comparison of the time for obs.match(frame) (with PSF matching) and of obs.render(model) in three configurations of (C,Ny,Nx)

• 1, 1000,1000
• 5, 100, 100
• 100, 100, 100

The model can be simple zeros.

[herjy]

Sorry, I saw your latest comment just now. I had similar stuffs done, so I'll do that quickly tomorrow. In the mean time, here is a new commit with Ndim_fft in place where it can be done without losing anything
I also added 5bands to the data, so that we have nicer colour images ;)

[herjy]

On the performance side, here is what this PR changes. In this plot, the y-axis is the time gained in percents by using this implementation compared to the master branch at the time of the PR for three different tasks making use of rfftn(axes = (1, 2)): psf convolution observation.render(), observation setup observation().match(frame), fitting blend.fit().

[pmelchior]

OK, this looks good in the regime where we need it. The drop off at the right is probably because of memory thrashing / cache misses.

[herjy]

Finally, I figured it out! So: for some values of the psf size, the fast shape of the fft returns an even value, which autograd.fft cannot deal with. To have autograd run smooth and clean, we tested the fast shape and made it uneven, but only along one axis, which, apparently, for the 2-d fft was ok, but not when going Ndim (this is still a little weird to me). So now I make both axes uneven whenever things fast shapes are even and we the Fourrier operations run nominally.

[fred3m]

Weird. It might be worth creating an issue with autograd or even fixing this ourselves upstream because I doubt that we're the only ones to get dinged by this.

[pmelchior]

Is this a problem from autograd or from numpy FFTs?

[fred3m]

It's a problem with autograd. Numpy.rfftn has no problem with even shapes, in fact I think that it prefers them and this code (without the check for an even last dimension) is the same basic algorithm that scipy.signal.fftconvolve uses.

But with autograd they raised an Exception when using rfftn with an even numbered last dimension, hence the fix that is in master (https://github.com/fred3m/scarlet/blob/master/scarlet/observation.py#L190-L194). But I had never run across a problem with PSF sizes, so I didn't notice that I needed to use the same fix when we calculated the optimal shape for PSF matching.

[pmelchior]

So, it looks like we are in a good state now. @herjy please confirm that we can merge this PR.

[herjy]

To answer Fred's last comment, It is somewhat strange that the problem did not arise earlier, but did now in the ndim. It could be related to the ndim implementation and the use of axes, but looking at the code for fft, I don't see it.
Anyway, I now have a much better grasp at this part of the code, and I think we are ready to merge.

[pmelchior]

I just spotted this: cb8bd40 deletes the file docs/changes.rst. That's our change log and should not be delete. I'll restore it but future future reference: git rm needs to be used with caution!