Add freqkernel and spacekernel (#100)

Author: timholy

src/ImageFiltering.jl

@@ -14,7 +14,8 @@ export Kernel, KernelFactors,
     Algorithm,
     imfilter, imfilter!,
     mapwindow, mapwindow!,
-    imgradients, padarray, centered, kernelfactors, reflect
+    imgradients, padarray, centered, kernelfactors, reflect,
+    freqkernel, spacekernel
 
 FixedColorant{T<:Normed} = Colorant{T}
 StaticOffsetArray{T,N,A<:StaticArray} = OffsetArray{T,N,A}

src/utils.jl

@@ -1,5 +1,5 @@
 """
-    centered(kernel) -> shiftedkernel
+    shiftedkernel = centered(kernel)
 
 Shift the origin-of-coordinates to the center of `kernel`. The
 center-element of `kernel` will be accessed by `shiftedkernel[0, 0,
@@ -13,6 +13,71 @@ See also: [`imfilter`](@ref).
 """
 centered(A::AbstractArray) = OffsetArray(A, map(n->-((n+1)>>1), size(A)))
 
+"""
+    kernfft = freqkernel([T::Type], kern, sz=size(kern); rfft=false)
+
+Return a frequency-space representation of `kern`.
+This embeds `kern` in an array of size `sz`, in a manner that implicitly imposes periodic boundary
+conditions, and then returns the fourier transform. This is sometimes called the optical transfer
+function, and known in some frameworks as `psf2otf`. If `rfft` is `true`, the fft for real-valued
+arrays (`rfft`) is returned instead and the first dimension size will be approximately half of `sz[1]`.
+
+`kern` should be zero-centered, i.e., `kern[0, 0]` should reference the center of your kernel.
+See [`centered`](@ref). Optionally specify the numeric type `T` (which must be one of the
+types supported by FFTW, either `Float32` or `Float64`).
+
+The inverse of `freqkernel` is [`spacekernel`](@ref).
+"""
+function freqkernel(::Type{T}, kern::AbstractArray, sz::Dims=size(kern); rfft=false) where T<:Union{Float32,Float64}
+    wrapindex(i, s) = i<1 ? i+s : i
+    all(size(kern) .<= sz) || throw(DimensionMismatch("kernel size $(size(kern)) is bigger than supplied size $sz"))
+    kernw = zeros(T, sz...)
+    for I in CartesianIndices(kern)
+        J = CartesianIndex(map(wrapindex, Tuple(I), sz))
+        kernw[J] = kern[I]
+    end
+    return rfft ? FFTW.rfft(kernw) : fft(kernw)
+end
+freqkernel(kern::AbstractArray{T}, args...; rfft=false) where T =
+    freqkernel(ffteltype(T), kern, args...; rfft=rfft)
+
+"""
+    kern = spacekernel(kernfft, axs; rfftsz=0)
+
+Return a real-space representation of `kernfft`, a frequency-space representation of a kernel.
+This performs an inverse fourier transform, implicitly imposes periodic boundary conditions,
+and then trims & truncates axes of the output to `axs`. By default `kernfft` is assumed to have
+been generated by `fft`; if it was instead generated by `rfft`, specify the original size
+of the first dimension. (If `kernfft` was generated by [`freqkernel`](@ref), this is just `sz[1]`.)
+
+The inverse of `spacekernel` is [`freqkernel`](@ref).
+"""
+function spacekernel(kernfft::AbstractArray, axs::Indices; rfftsz=0)
+    wrapindex(i, s) = i<1 ? i+s : i
+    kernw = rfftsz > 0 ? irfft(kernfft, rfftsz) : ifft(kernfft)
+    kern = zeros(eltype(kernw), axs...)
+    sz = size(kernw)
+    for I in CartesianIndices(kern)
+        J = CartesianIndex(map(wrapindex, Tuple(I), sz))
+        kern[I] = kernw[J]
+    end
+    return kern
+end
+function spacekernel(kernfft::AbstractArray; rfftsz=0)
+    sz = size(kernfft)
+    if rfftsz > 0
+        sz = (rfftsz, Base.tail(sz)...)
+    end
+    upper = map(s->s>>1, sz)
+    axs = map((u, s)-> u-s+1:u, upper, sz)
+    return spacekernel(kernfft, axs; rfftsz=rfftsz)
+end
+
+ffteltype(::Type{T}) where T<:Union{Float32,Float64} = T
+ffteltype(::Type{Float16}) = Float32
+ffteltype(::Type{T}) where T<:Normed = ffteltype(FixedPointNumbers.floattype(T))
+ffteltype(::Type{T}) where T = Float64
+
 dummyind(::Base.OneTo) = Base.OneTo(1)
 dummyind(::AbstractUnitRange) = 0:0
 

test/basic.jl

@@ -1,4 +1,4 @@
-using ImageFiltering, OffsetArrays, Logging, ImageMetadata, Test
+using ImageFiltering, OffsetArrays, Logging, ImageMetadata, FixedPointNumbers, Test
 import AxisArrays
 using AxisArrays: AxisArray, Axis
 
@@ -88,4 +88,26 @@ end
     check_range_axes(axs[1].val, -1, 1)
 end
 
-nothing
+@testset "freqkernel/spacekernel" begin
+    k = centered(reshape([1,-1], 2, 1))
+    kfft = freqkernel(k, (31, 31))
+    @test size(kfft) == (31, 31)
+    k2 = real.(spacekernel(kfft, axes(k)))
+    @test k2 ≈ k
+    kfft = freqkernel(k, (31, 31); rfft=true)
+    @test size(kfft) == (16, 31)
+    k2 = spacekernel(kfft, axes(k); rfftsz=31)
+    @test k2 ≈ k
+
+    for T in (Float64, Float32, Float16, N0f16)
+        k = T.(Kernel.gaussian(3))
+        kfft = freqkernel(k)
+        @test size(kfft) == size(k)
+        k2 = real.(spacekernel(kfft))
+        @test T.(k2) ≈ k
+
+        kfft = freqkernel(k; rfft=true)
+        k2 = spacekernel(kfft; rfftsz=size(k, 1))
+        @test T.(k2) ≈ k
+    end
+end