inpaint.jl

"""
    inpaint(A)

Inpaints `missing` values if `A` contains some.

Inspired by MATLAB's `inpaint_nans`'s (by John d'Errico: [link](https://www.mathworks.com/matlabcentral/fileexchange/4551-inpaint_nans)).

# Example
```jldoctest
julia> A = convert(Array{Union{Float64, Missing},2}, (1:3)*(1:4)') ; A[1:2, 1:2] .= missing ; A
3×4 Array{Union{Missing, Float64},2}:
  missing   missing  3.0   4.0
  missing   missing  6.0   8.0
 3.0       6.0       9.0  12.0

julia> inpaint(A)
3×4 Array{Union{Missing, Float64},2}:
 1.0  2.0  3.0   4.0
 2.0  4.0  6.0   8.0
 3.0  6.0  9.0  12.0
```
"""
inpaint(A::AbstractArray{Union{Missing, T}}; method=1, cycledims=Int64[]) where {T<:AbstractFloat} = inpaint(ismissing, A, method=method, cycledims=cycledims)

"""
    inpaint(A)

Inpaints `NaN` values if the elements of `A` are all float (i.e., `eltype(A) <: AbstractFloat`).

# Example

```jldoctest
julia> A = float((1:3)*(1:4)') ; A[1:2, 1:2] .= NaN ; A
3×4 Array{Float64,2}:
 NaN    NaN    3.0   4.0
 NaN    NaN    6.0   8.0
   3.0    6.0  9.0  12.0

julia> inpaint(A)
3×4 Array{Float64,2}:
 1.0  2.0  3.0   4.0
 2.0  4.0  6.0   8.0
 3.0  6.0  9.0  12.0
```
"""
inpaint(A::AbstractArray{T}; method=1, cycledims=Int64[]) where {T<:AbstractFloat} = inpaint(isnan, A, method=method, cycledims=cycledims)

"""
    inpaint(A, missing)

Inpaints `missing` values.
Should be the same as `inpaint(A)`.
"""
inpaint(A::AbstractArray, ::Missing; method=1, cycledims=Int64[]) = inpaint(ismissing, A, method=method, cycledims=cycledims)

"""
    inpaint(A, value_to_fill)

Inpaints `value_to_fill` values (can be `NaN`).
Useful in the case data was generated by filling missing values with an
arbitrary chosen one.

# Examples

```jldoctest
julia> A = float((1:3)*(1:4)') ; A[1:2, 1:2] .= 999 ; A
3×4 Array{Float64,2}:
 999.0  999.0  3.0   4.0
 999.0  999.0  6.0   8.0
   3.0    6.0  9.0  12.0

julia> inpaint(A, 999)
3×4 Array{Float64,2}:
 1.0  2.0  3.0   4.0
 2.0  4.0  6.0   8.0
 3.0  6.0  9.0  12.0
```

```jldoctest
julia> A = float((1:3)*(1:4)') ; A[1:2, 1:2] .= NaN ; A
3×4 Array{Float64,2}:
 NaN    NaN    3.0   4.0
 NaN    NaN    6.0   8.0
   3.0    6.0  9.0  12.0

julia> inpaint(A, NaN)
3×4 Array{Float64,2}:
 1.0  2.0  3.0   4.0
 2.0  4.0  6.0   8.0
 3.0  6.0  9.0  12.0
```
"""
function inpaint(A::AbstractArray{T}, value_to_fill; method=1, cycledims=Int64[]) where T<:AbstractFloat
    return if isnan(value_to_fill)
        inpaint(isnan, A, method=method, cycledims=cycledims)
    else
        inpaint(x -> x == value_to_fill, A, method=method, cycledims=cycledims)
    end
end

"""
    inpaint(f, A)

Inpaints values of `A` for which `f(A) == true`.

# Examples

```jldoctest
julia> A = float((1:3)*(1:4)') ; A[1:2, 1:2] .= 999 ; A
3×4 Array{Float64,2}:
 999.0  999.0  3.0   4.0
 999.0  999.0  6.0   8.0
   3.0    6.0  9.0  12.0

julia> inpaint(x -> x == 999, A)
3×4 Array{Float64,2}:
 1.0  2.0  3.0   4.0
 2.0  4.0  6.0   8.0
 3.0  6.0  9.0  12.0
```

```jldoctest
julia> A = float((1:3)*(1:4)') ; A[1:2, [1, end]] .= NaN ; A
3×4 Array{Float64,2}:
 NaN    2.0  3.0  NaN
 NaN    4.0  6.0  NaN
   3.0  6.0  9.0   12.0

julia> inpaint(isnan, A)
3×4 Array{Float64,2}:
 1.0  2.0  3.0   4.0
 2.0  4.0  6.0   8.0
 3.0  6.0  9.0  12.0
```

    inpaint(f, A; method=1, cycledims=[])

With optional arguments, you can chose the inpainting method and if dimensions are cyclic.

```jldoctest
julia> A = float((1:3)*(1:4)') ; A[1:2, [1, end]] .= NaN ; A
3×4 Array{Float64,2}:
 NaN    2.0  3.0  NaN
 NaN    4.0  6.0  NaN
   3.0  6.0  9.0   12.0

julia> inpaint(A, cycledims=[2])
3×4 Array{Float64,2}:
 2.16475  2.0  3.0   2.83525
 3.76245  4.0  6.0   6.23755
 3.0      6.0  9.0  12.0
```
"""
function inpaint(f, A; method=1, cycledims=Int64[])
    return @match method begin
        1 => inpaint_method1(f, A, cycledims=cycledims) # Default
        _ => error("Method $method not available yet. Suggest it on Github!")
    end
end

"""
    list_neighbors(A, idx, neighbors)

Lists all the neighbors of the indices in `idx` in Array `A`.
Neighbors already in `idx` are accepted.
The argument `neighnors` contains a list of the neighbors about the
origin coordinate `(0, 0, ...)`.
In other words, it is a `Vector` of `CartesianIndex` such that
the direct neighbors of index `i` are given by `i + n for n in neighbors`.
Inspired by MATLAB's `inpaint_nans`'s `identify_neighbors` (by John d'Errico).
See https://www.mathworks.com/matlabcentral/fileexchange/4551-inpaint_nans.
"""
function list_neighbors(R, idx, neighbors)
    out = [i + n for i in idx for n in neighbors if i + n ∈ R]
    out = sort(unique([out; idx]))
end
list_neighbors(R, idx::Vector{Int64}, neighbors) = list_neighbors(R, CartesianIndex.(idx), neighbors)

"""
    inpaint_method1(f, A::Array, cycledims=Int64[])

Inpaints values in `A` that `f` gives `true` on by solving a simple diffusion PDE.
Default method for `inpaint`.
The partial differential equation (PDE) is defined by the standard Laplacian, `Δ = ∇^2`.
Inspired by MATLAB's `inpaint_nans`'s method `0` for matrices (by John d'Errico).
See https://www.mathworks.com/matlabcentral/fileexchange/4551-inpaint_nans.
The discrete 1D stencil used for `Δ` looks like
```
┌───┐ ┌───┐ ┌───┐
│ 1 ├─┤-2 ├─┤ 1 │
└───┘ └───┘ └───┘
```
and is applied where it fits at the borders in all dimensions.
However, the user can supply a list of dimensions that should be considered cyclic.
In this case, the sentil will be used also at the borders and "jump" to the other side.
This is particularly useful for, e.g., world maps with longitudes spanning the entire globe.
"""
function inpaint_method1(f, A::Array; cycledims=Int64[])
    I_unknown = findall(@. f(A))
    I_known = findall(@. !f(A))

    # direct neighbors
    Ns = [CartesianIndex(ntuple(x -> x == d ? 1 : 0, ndims(A))) for d in 1:ndims(A)]
    neighbors = [Ns; -Ns]
    # Add neighbors on opposite border if dimensions are cyclic
    N′s = [CartesianIndex(ntuple(x -> x == d ? size(A,d) - 1 : 0, ndims(A))) for d in 1:ndims(A)]
    for d in cycledims # Add neighbors in cyclic case
        neighbors = push!(neighbors, N′s[d], -N′s[d])
    end

    # "working" list: the indices of all unkowns and their neighbors
    R = CartesianIndices(size(A))
    Iwork = list_neighbors(R, I_unknown, neighbors)
    iwork = LinearIndices(size(A))[Iwork]

    # Preallocate the Laplacian (not the fastest way but easier-to-read code)
    nA = length(A)
    Δ = sparse([], [], Vector{Float64}(), nA, nA)

    # Fill the Laplacian within the borders
    for d in 1:ndims(A)
        N = Ns[d]
        # R′ is the range of indices without the borders in dimension `d`
        R′ = CartesianIndices(size(A) .- 2 .* N.I)
        R′ = [r + N for r in R′]
        # Convert to linear indices to build the Laplacian
        u = vec(LinearIndices(size(A))[R′])
        n = LinearIndices(size(A))[first(R) + N] - LinearIndices(size(A))[first(R)]
        # Build the Laplacian (not the fastest way but easier-to-read code)
        Δ += sparse(u, u     , -2.0, nA, nA)
        Δ += sparse(u, u .- n, +1.0, nA, nA)
        Δ += sparse(u, u .+ n, +1.0, nA, nA)
    end

    # Add Laplacian along border if cyclic along dimension `d`
    for d in cycledims
        N = Ns[d]   # usual neighbor
        N′ = N′s[d] # neighbor on the opposite border
        R′₊ = [r for r in R if r + N ∉ R] # one border
        R′₋ = [r for r in R if r - N ∉ R] # the opposite border
        # Convert to linear indices to build the Laplacian
        n = LinearIndices(size(A))[first(R) + N] - LinearIndices(size(A))[first(R)]
        n′ = LinearIndices(size(A))[first(R) + N′] - LinearIndices(size(A))[first(R)]
        u₊ = LinearIndices(size(A))[R′₊]
        u₋ = LinearIndices(size(A))[R′₋]
        # Add to the Laplacian
        Δ += sparse(u₊, u₊      , -2.0, nA, nA)
        Δ += sparse(u₊, u₊ .- n , +1.0, nA, nA)
        Δ += sparse(u₊, u₊ .- n′, +1.0, nA, nA)
        Δ += sparse(u₋, u₋      , -2.0, nA, nA)
        Δ += sparse(u₋, u₋ .+ n , +1.0, nA, nA)
        Δ += sparse(u₋, u₋ .+ n′, +1.0, nA, nA)
    end

    # Use linear indices to access the sparse Laplacian
    i_unknown = LinearIndices(size(A))[I_unknown]
    i_known = LinearIndices(size(A))[I_known]

    # Place the knowns to right hand side
    rhs = -Δ[iwork, i_known] * A[i_known]

    # Solve for the unknowns
    B = zeros(size(A))
    B[i_known] .= A[i_known]
    B[i_unknown] .= Δ[iwork, i_unknown] \ rhs
    return B
end

export inpaint