TV/TSV regularizer is not equivalent to source

[kadri-nizam]

While profiling for bottlenecks in the code, I happened to inspect the implementation of the TV regularizer and noticed that it is slightly different from the way it is implemented in the cited source.

I don't think the difference will affect the results very much, but I thought I'd point it out. Feel free to close this if it's not an issue.

Comparison

Here's an MWE that illustrates this. First I'll define the variables we need.

import torch

torch.manual_seed(2023)
sky_cube = torch.rand((8, 10, 10))
epsilon = 1e-6

nchan = sky_cube.size(0)
npix = sky_cube.size(-1)
img_dim = npix * npix

horizontal_pad = torch.zeros(nchan, 1, npix)
vertical_pad = torch.zeros(nchan, npix, 1)

MPoL Implementation

After refactoring the code to use built-in torch functions to facilitate easier comparison, the current MPoL implementation is:

diff_mm = torch.diff(sky_cube[:, :, 0:-1], dim=1)
diff_ll = torch.diff(sky_cube[:, 0:-1, :], dim=2)
loss_old = torch.sqrt(diff_ll**2 + diff_mm**2 + epsilon).sum()

eht-imaging Implementation

The biggest difference lies in the way that the finite difference is taken. Here, they add a zero-padding to the end of the tensor in the dimension that they are taking the difference.

diff_mm = torch.diff(sky_cube, dim=1, append=horizontal_pad)
diff_ll = torch.diff(sky_cube, dim=2, append=vertical_pad)
# eht-imaging returns the negative loss, but that's just a minor detail
loss_new = torch.sqrt(diff_ll**2 + diff_mm** 2 + epsilon).sum()

Loss Values

The EHT implementation results in a larger loss compared to the implementation of MPoL.

>>> loss_new       # eht-imaging
tensor(443.1116)
>>> loss_old       # MPoL
tensor(341.5393)

[iancze]

Thanks for catching this. I guess the difference between our implementations is how edges should be handled. In the EHT implementation, I think this would be equivalent to saying there is a penalty if the edge of the image isn't close to 0.

We should discuss if that behavior in the TV/TSV regularizers is important to us (and possibly change). We've mostly been working with images that have edge values close to 0, so we've sort of side-stepped this issue, I think. It probably makes sense to adopt the EHT definition (and cite their codebase in this routine, specifically).

[iancze]

In fe9c786 I implemented your updates using torch.diff, but kept the calculation the same as we have for consistency. This is still technically wrong, but may not matter.

The definition of TV in EHT IV is

So from a given reference pixel, it's a right pixel minus current and a bottom pixel minus current. There's a question of what to do when l=L and m=M and l+1 and m+1 would be off the image entirely. You could pad the image with zeros, but then you're assuming that as a reference value. I.e., a constant-valued image should in spirit have a TV score of 0. But under this definition there would be a penalty between the last constant-value and the padded 0. In practice this is a good thing, since you always want to be imaging all of your flux in a dataset anyway, so your image boundaries should be 0, otherwise you have bigger problems.

Whereas in our implementation we only compute TV between valid pixels in the array. This is fine, but results in diff_ll and diff_mm of different shapes by 1. But because we are adding the differences in the l and m directions inside the sum, we need diff_ll and diff_mm to have the same shape, so we delete the extra row of differences in each dimension.

One other potential solution is to calculate diff_mm and diff_ll as we do, then append a row of zeros (i.e., diff is zero) to augment each array to make them the same shape. This is technically more correct as regards the definition since it treats the extra image edge as extrapolated nearest-neighbor (i.e., diff = 0), but as I previously commented, the push towards 0 at edges is probably diserable anyway.

@kadri-nizam
I see eht-imaging is GPLv3, whereas we are MIT. I don't understand enough about open source licensees to know what's permissible in this situation w.r.t. using your solution referenced from the eht-imaging source. I'd prefer not to have to relicense MPoL under GPLv3. But I am happy to cite the eht-imaging software for the implementation, if that's sufficient.