xgcm / xrft

Fourier transforms on xarray data structures
https://xrft.readthedocs.io/en/latest/
MIT License
153 stars 44 forks source link

Proposal: Output length parameter for FTs #146

Closed bennyrowland closed 8 months ago

bennyrowland commented 3 years ago

Big fan of what you guys have already done with xrft. What do people think about the idea of adding an extra parameter to the FT functions to specify the length/shape of the outputs along each transformed dimension? If the requested output length is larger than the current dimension size, it should be zero filled before FT, if the requested length is smaller, then the post FT output should be truncated. Both of these operations are very common in MRI reconstruction. There are obviously a few corner cases - e.g. for an IFT where there is an imbalance in the number of points either side of 0, points could be added asymmetrically to balance them, making the IFT possible.

The best way to define such a parameter is not 100% clear to me, some possible options: 1) output_length: an iterable of the same length as dim (or a scalar if dim is a string) defining the shape in the same order as the dimensions. Quite compact but risks the order of dims getting out of sync between the two parameters. 2) output_length as a dict with the dim names as keys and lengths as values. Can be constructed from the above paradigm as dict(zip(dim, output_length)) but feels a bit more unwieldy and gives more duplication as you repeat all the dim names in the output_length argument as well. Not that different to the way xarray arrays are defined with coords though. 3) allow passing the same dictionary as above to the dim argument - the keys are the dims to use and the values are the output sizes of each dimension. Obviously this would be an additional alternative to the dim parameter, not replacing the existing version. This is the most compact option but also represents a bigger change to existing usage, offering the possibility for confusion and possible bugs in implementation.

I am happy to have a go at putting together a PR if there is enthusiasm for the idea.

roxyboy commented 3 years ago

Is this in a way proposing an automated method for "zero padding"?

bennyrowland commented 3 years ago

It does either of two things - zero padding an input (to increase apparent resolution) and truncating an output before returning (which removes oversampling). Both of these operations are endemic in MRI, sometimes we will actually do both, zero filling to double the size of an input, then truncating the output to the original size which halves the field of view while maintaining the number of pixels. As proposed, only one or the other is possible with this keyword, I had considered the possibility of suggesting a separate zero padding parameter to allow both to be specified, but that does add some extra complexity, and adding a separate zero padding method for example would not be difficult (I have already implemented an xarray accessor providing this method in my own code). The truncated output is probably what I would rather have in the dft() and idft() functions, because although it is easy to implement, there is enough book keeping to make it a pain. The zero padding is then an obvious bonus to include.

bennyrowland commented 3 years ago

@roxyboy did you have any further thoughts about this proposal?

roxyboy commented 3 years ago

I'm honestly ambivalent about this... I personally probably won't use this functionality but if you think it would benefit others and can implement it in a way that would pass the pre-existing tests along with a new test that will check for the zero padding and truncation, I'd welcome a PR :) Maybe @rabernat will have more to say.

rabernat commented 3 years ago

I agree with what @roxyboy said. I'd be happy to accept a PR for this that:

santisoler commented 3 years ago

I'm very interested in seeing an implementation of padding in xrft (as mentioned in #154). Just a minor thought. Wouldn't be better to have a separate function to apply padding rather than having it included into the xrft.xrft.dft function?

roxyboy commented 3 years ago

I'm very interested in seeing an implementation of padding in xrft (as mentioned in #154). Just a minor thought. Wouldn't be better to have a separate function to apply padding rather than having it included into the xrft.xrft.dft function?

Yes, if we are to implement this, there should be a separate function (maybe something like _padding(da, dim, N)) with N being a list of number of grids to pad per dimension, which xrft.fft should then call.

roxyboy commented 3 years ago

@santisoler Would you like to give it a go for a PR?

santisoler commented 3 years ago

I'm very interested in seeing an implementation of padding in xrft (as mentioned in #154). Just a minor thought. Wouldn't be better to have a separate function to apply padding rather than having it included into the xrft.xrft.dft function?

Yes, if we are to implement this, there should be a separate function (maybe something like _padding(da, dim, N)) with N being a list of number of grids to pad per dimension, which xrft.fft should then call.

What I meant was having a public function for the padding and leaving the step to the user, instead of including the padding process inside fft. So a typical workflow would be like:

# Compute FFT of padded grid
grid_padded = xrft.pad(grid, ...)
grid_ft = xrft.fft.fft(grid_padded)

# Apply filters, plot, etc
...

# Compute inverse FFT and unpad the grid
grid_padded = xrft.fft.ifft(grid_ft)
grid = xrft.unpad(grid_padded, ...)

What do you think?

santisoler commented 3 years ago

@santisoler Would you like to give it a go for a PR?

Sure! I can tackle this. I would like to have some common thoughts about the design of the user interface before I jump into the actual implementation. Let me know if you like the idea I posted on the previous comment, and if not what would you change.

roxyboy commented 3 years ago

What I meant was having a public function for the padding and leaving the step to the user, instead of including the padding process inside fft. So a typical workflow would be like:

# Compute FFT of padded grid
grid_padded = xrft.pad(grid, ...)
grid_ft = xrft.fft.fft(grid_padded)

# Apply filters, plot, etc
...

# Compute inverse FFT and unpad the grid
grid_padded = xrft.fft.ifft(grid_ft)
grid = xrft.unpad(grid_padded, ...)

What do you think?

Sounds good to me. We already have isotropize and detrend as public functions so maybe you can create a new file like padding.py as part of xrft?

roxyboy commented 3 years ago

But, along with the public function, I still think it would be beneficial to have an argument within fft(pad=True), which would call the public padding function. This is how we do it for detrending for example.

bennyrowland commented 3 years ago

@santisoler nice to see someone else with enthusiasm for this. I think that your proposal would work, and as @roxyboy says it would be easy to then add that functionality inside the FT functions. I agree it is nice to have it there as it can help to convey more intent, but not too bothered as it is easy to wrap the two functions up anyway. However, pad/unpad are going to have more arguments than just a global True/False so they might get messy inside the function call.

I can see multiple ways to specify pad() and unpad(), do you give the number of points to add, or the target size? Either way, I would probably argue for adding points asymmetrically if the zero of the dimension coords is not in the middle of the coords (in order to symmetrize it for FT operations), at least as a kwarg. For unpad() I would probably like the ability to specify an output size N and take the central N points by coord, but again you could also specify the number of points to remove (or an amount from each end). Perhaps a pad_by() and a pad_to() would be appropriate? I am also less sure about the name unpad(), it conveys the conjugacy with pad() which is good, but personally I would probably often call it independently of pad() and it would then make less sense. One possible alternative could be trim() (trim_to()/trim_by()?) and I am sure there are others. Not too bothered, I merely raised the point for discussion.

santisoler commented 3 years ago

Regarding including the padding inside fft, I agree with @bennyrowland. I think pad will end up having many more arguments than just True/False, and they could scale in the future if we implement some fancy method of padding (e.g., some authors have filled the padded region with white noise based on the spectrum of the original data -citation needed-). This is why I would implement it as a separate function, at least for now. Specially because fft currently has quite a bunch of arguments and I don't want to clutter it.

I can see multiple ways to specify pad() and unpad(), do you give the number of points to add, or the target size?

For simplicity I would build a new xrft.pad function that wraps numpy.pad but also extends the coordinates arrays. This is very similar to what xarray is doing with xarray.DataArray.pad, although we would shrink the scope of the function to regular grids (evenly spaced coordinates). That being said, I would like to pass the number of points that will be used at each side of the array. But since we are aware of the coordinates of the array, I would love this function to be able to get the actual width (in the same units of the coordinates) of the desired padding at each side of the array.

Either way, I would probably argue for adding points asymmetrically if the zero of the dimension coords is not in the middle of the coords (in order to symmetrize it for FT operations), at least as a kwarg.

I need to disagree on this, mainly because on the type of arrays that we'll be applying padding on Harmonica. The coordinates of these arrays are actual geographic coordinates, therefore having an array centered around zero would be just a coincidence, and there shouldn't be any difference in applying the pad function to any other array that isn't centered around zero. Nevertheless, you could still apply asymmetrical padding if you want, you would just need to specify different pad_width for that dimension. Does it make sense?

For unpad() I would probably like the ability to specify an output size N and take the central N points by coord, but again you could also specify the number of points to remove (or an amount from each end). Perhaps a pad_by() and a pad_to() would be appropriate?

I would like that the unpad (or trim, read below) function to take the same arguments than pad, so you could safely use the same arguments for padding and trimming a grid and get the same set of coordinates, like this:

# Compute FFT of padded grid
path_width = {
    "x": (10, 10),
    "y": (5, 5),
}  # define amount of padding
grid_padded = xrft.pad(grid, pad_width)
grid_ft = xrft.fft.fft(grid_padded)

# Apply filters, plot, etc
...

# Compute inverse FFT and unpad the grid
grid_padded = xrft.fft.ifft(grid_ft)
grid = xrft.unpad(grid_padded, pad_width)

I am also less sure about the name unpad(), it conveys the conjugacy with pad() which is good, but personally I would probably often call it independently of pad() and it would then make less sense. One possible alternative could be trim() (trim_to()/trim_by()?) and I am sure there are others. Not too bothered, I merely raised the point for discussion.

I like both unpad and trim. If we expect users to use this function besides the padding-unpadding process, I'm more prone to call it trim. Nevertheless, with xarray is fairly simple to trim a DataArray by slicing it, and I wouldn't like people to get used to using a single function of xrft just to slice a portion of an array. That makes me more prone to call it unpad. I'm not entirely sure yet and open to suggestions.

bennyrowland commented 3 years ago

If you want to simply match the xarray pad() behaviour, but handling extrapolating coordinates better, then it might be more sensible to approach that directly in xarray. @shoyer for example has previously suggested that having a special case for xarray.pad() where coordinates are equally spaced would be appropriate, and it would then reach a larger audience.

My use case (in MRI) is obviously very different to yours. My data comes in the frequency domain and will generally receive a single IFT, rather than an FT/processing/IFT sequence. To do the IFT the data does have to be centred on 0 (as enforced by xrft) and my data is often undersampled and therefore asymmetric. Thus code for me might look more like:

raw_data # is 192 x 256
padded_data = xrft.pad_to(raw_data, freq_x=256, freq_y=256, fill_value=0, make_sym=True)
full_image = xrft.idft(padded_data)
final_image = xrft.trim_to(full_image, x=128, y=128)

In this case, I am more interested in specifying the size of the arrays after each operation rather than the number of points to add each side, particularly to balance out the asymmetry. Automatically balancing the asymmetry is more relevant in pad_to style than pad itself, but I do think it is a useful convenience function. It is not too difficult to calculate the necessary values on either side, but having the function makes it nice and concise and the intent is clearer. Also note that my pad and trim steps are not paired the way that yours are, I am not trying to reverse steps to get back to the initial configuration, the two operations are performing totally different functions (one in frequency domain and one in space domain). That is mainly why I am less keen on unpad from my side, but I see it works for your example. unpad_to also doesn't sound quite right. Of course, my code could also look like:

raw_data # is 192 x 256
padded_data = xrft.pad_to(raw_data, freq_x=256, freq_y=256, fill_value=0, make_sym=True)
full_image = xrft.idft(padded_data, output_size={x:128, y:128})

if we include the argument in the FT functions, and then trim_to becomes a less compelling addition, since (I think) trim_to would generally only be called immediately after the IFT in this way anyway.

I see no reason though why there should not be both of these styles of functions available in xrft, in fact a natural way to implement e.g. pad_to is to use pad under the hood with a bit of extra calculation. I could also imagine having either extra functions or a kwarg which make the padding/trimming use coordinate units instead of indices.

shoyer commented 3 years ago

Generic padding/trimming as separate functionality seems fine, too, but I'll just note that output length padding/trimming is already built into all the numpy.fft functions. So this might also be implemented just by passing on the length information into the underlying FFT.

roxyboy commented 3 years ago

Generic padding/trimming as separate functionality seems fine, too, but I'll just note that output length padding/trimming is already built into all the numpy.fft functions. So this might also be implemented just by passing on the length information into the underlying FFT.

If the padding/trimming functionality in numpy already exists, we should definitely take advantage of that by calling it in xrft rather than us reinventing the wheel from scratch.

santisoler commented 3 years ago

If you want to simply match the xarray pad() behaviour, but handling extrapolating coordinates better, then it might be more sensible to approach that directly in xarray. @shoyer for example has previously suggested that having a special case for xarray.pad() where coordinates are equally spaced would be appropriate, and it would then reach a larger audience.

You're right about this. Maybe it would be better to extend xarray.DataArray.pad to support padding coordinates if they are evenly spaced. @shoyer are you interested in incorporating this feature? I've already started drafting some code for that, I can quickly move it to xarray and open a PR.

@bennyrowland Now I understand better your needs. I think we can have pad_to and trim_to functions that wrap the pad and unpad functions but takes the entire width of the array instead of the "extra" points. Once we have a pad function, I think their implementation is very straight forward.

I could also imagine having either extra functions or a kwarg which make the padding/trimming use coordinate units instead of indices.

Absolutely. Since we are now aware of the coordinates of the array, I think we must support passing either number of points or actual coordinates widths. Keeping number of extra points is very well suited for your needs. For our needs in Harmonica, we really want to pad through coordinates units: pad the array with 10km at each side of the grid.

Generic padding/trimming as separate functionality seems fine, too, but I'll just note that output length padding/trimming is already built into all the numpy.fft functions. So this might also be implemented just by passing on the length information into the underlying FFT.

You're right. So maybe we can support xrft.xrft.fft to take some arguments to apply padding before computing the fft. Maybe we could make fft to apply a pad full of zeros so we can lower the number of arguments, and in case a user might want to apply fancier padding, they would need to use the actual pad function. Does it make sense?

If the padding/trimming functionality in numpy already exists, we should definitely take advantage of that by calling it in xrft rather than us reinventing the wheel from scratch.

You're right @roxyboy. My idea was always to use numpy.pad for applying the padding both on the array and also on the coordinates. I'm not planning to reimplement any of these features, our functions would need to connect the xarray-friendly interface with the numpy functions.