Zygote, Flux, and Interpolations

Specific Domains Machine Learning
1

Hey everyone!

I’m working on radio telescope imaging and am trying to do image reconstruction with an iterative, maximum entropy method. Many papers suggest using conjugate gradient, and have intricate derivations of gradients. I thought it would be neat to do the same thing, but pushing the model through some sort of Autodiff.

In this process, I am needing to take samples of an interpolated FFT and calculate a mean squared error to measured radio telescope data.

Something like this

function vis_res(image, vis_data, uv)
    # Generate interpolated visibilities from the image
    N = length(vis_data)
    image_fft = fft(image) |> fftshift
    freqs = fftfreq(size(image)[1]) |> fftshift
    interpolation = LinearInterpolation((freqs, freqs), image_fft)
    vis_interp = [interpolation(uv[:,i]...) for i ∈ 1:N]
    # Calculate visibility residuals
    return (abs.(vis_interp .- vis_data)).^2
end

However, Flux/Zygote doesn’t seem to be happy with the indexing of the interpolation. Throwing the error:

ERROR: ArgumentError: unable to check bounds for indices of type Interpolations.WeightedAdjIndex{2, Float64}

Any help in this regard would be greatly appreciated.

2 Likes
2

Funny, I just ran into the same error. I’m trying to differentiate through a likelihood function that uses an interpolation using Zygote.

I expect the incompatibility is within the Interpolations library rather than the indexing you show. My likelihood function also calls a manual bi-linear interpolation function I wrote, and it has no issue with that as far as I can tell.

3

You might also want to try testing ForwardDiff unless you have to use Zygote for some reason.

1 Like
4

Yeah I dug in deep to this - and it actually does work on the current master branch of Interpolations, as the release hasn’t been bumped in a few months. There is still a strange indexing problem, stemming from strangeness in the dimensionality of the gradient. I solved my specific example by just providing an explicit gradient.

5

It would help to have a full stacktrace as well. As-is (having no knowledge of how interpolations works), I have no idea which line is even causing the error.

6

Looks like Interpolations didn’t gain AD support until after the latest release, so that makes sense Edit: saw you commented already on a linked issue :). Worth reporting an issue if you’re getting incorrect gradient values.

7

For those finding your way here from Google, the following function gives a bilinear interpolation over the axes of a matrix and is easily auto differentiable with Zygote and others. You should generally prefer the battle-tested Interpolations.jl implementations where you can, but this one works in a pinch.


function bilininterp(data::AbstractMatrix, x, y)
    x1 = floor(Int,x)
    x2 = ceil(Int,x)
    y1 = floor(Int,y)
    y2 = ceil(Int,y)

    # Handle case of perfect grid alignment
    if x1 == x2
        x2 = x1 + 1
    end
    if y1 == y2
        y2 = y1 + 1
    end

    # Handle boundary conditions
    if x1 == 0
        x1 += 1
        x2 += 1
    end
    if y1 == 0
        y1 += 1
        y2 += 1
    end
    if x1 == size(data,1)
        x1 -= 1
        x2 -= 1
    end
    if y1 == size(data,2)
        y1 -= 1
        y2 -= 1
    end

    mat = [
        1 x1 y1 (x1 * y1)
        1 x1 y2 (x1 * y2)
        1 x2 y1 (x2 * y1)
        1 x2 y2 (x2 * y2)
    ]
    # Get surrounding pixels if finite (fallback to image median otherwise)
    col = [
        data[x1, y1]
        data[x1, y2]
        data[x2, y1]
        data[x2, y2]
    ]
    # Solve the system of equations without inverting
    coeffs = mat \ col
    interp =
        coeffs[1] + coeffs[2] * x + coeffs[3] * y + coeffs[4] * x * y
    return interp
end
3 Likes