diff --git a/chambollefp.jl b/chambollefp.jl index 34679407c43747268f6ecedb3bd337bd310a9650..5e0274ca026b162ab4f8f13fffa24394d5bc52a8 100644 --- a/chambollefp.jl +++ b/chambollefp.jl @@ -27,7 +27,7 @@ function fpalg(grid, f, A; α, β, τ = 1/8, ϵ = 1e-4) reserr = 1 k = 0 - while reserr > ϵ && k < 3e2 + while reserr > ϵ && k < 30e1 k += 1 #sleep(0.1 / k) (p, reserr) = _fpiter(grid, p, f, A, α, β, τ) diff --git a/chambollefp_dd.jl b/chambollefp_dd.jl index 99be72eb5b0499df2df57092ad51f8826cbd112d..e6de6a8f8b5e2e12d21fcd71134ae1288b24bd2b 100644 --- a/chambollefp_dd.jl +++ b/chambollefp_dd.jl @@ -1,13 +1,102 @@ include("fd.jl") using LinearAlgebra +using SparseArrays import .FDM -grid = FDM.MetaGrid((0:0.1:1, 0:0.1:1), (3, 3), 2) +"partition of unity function" +function theta(grid, idx...) + vidx = FDM.indices(grid, idx...) -f = FDM.GridFunction(x -> norm(x), 1, grid) + slices = Vector{Vector{Float64}}(undef, length(idx)) + for (k, kglobrange) in enumerate(vidx) + krange = Base.OneTo(length(kglobrange)) + slices[k] = min.(krange, reverse(krange), grid.overlap) ./ grid.overlap + # FIXME: dirty hack + if first(kglobrange) == 1 + slices[k][1:grid.overlap] .= 1.0 + end + if last(kglobrange) == size(grid, k) + slices[k][end-grid.overlap:end] .= 1.0 + end + end + + #dst = Array{Float64}(undef, size(grid)...) + dst = zeros(size(grid)) + for (I,J) in zip(CartesianIndices(vidx), CartesianIndices(length.(vidx))) + dst[I] = prod(getindex.(slices, Tuple(J))) + end + dst + + FDM.GridFunction(grid, dst) +end + +function p2u(grid, p, A, f; α, β) + FDM.GridFunction(grid, (A'*A + β*I) \ (α * FDM.divergence(p).data[:] .+ A'*f.data[:])) +end + +# TODO: investigate matrix multiplication performance with GridFunctions +function fpiter(grid, p, f, A, α, β, τ) + x = FDM.GridFunction(grid, (A'*A + β*I) \ (FDM.divergence(p).data[:] .+ A'*f.data[:] ./ α)) + q = FDM.gradient(x) + + #display(x.data) + #FDM.my_plot(x) + #checknan(q) + p_new = FDM.GridFunction(grid, (p .+ τ .* q) ./ (1 .+ τ .* FDM.my_norm(q))) + + reserr = maximum(sum((p_new .- p).^2, dims=ndims(p))) + (p_new, reserr) +end + +function fpalg(grid, f, A; α, β, τ = 1/8, ϵ = 1e-4) + local uimg + p = FDM.GridFunction(x -> 0, ndims(grid), grid) + reserr = 1 + k = 0 + while reserr > ϵ && k < 30e1 + k += 1 + #sleep(0.1 / k) + (p, reserr) = fpiter(grid, p, f, A, α, β, τ) + + u = p2u(grid, p, A, f, α=α, β=β) + + #FDM.my_plot(uimg) + println("[$(lpad(k, 4))] res = $(round(reserr, digits=5))") + println("min/max: $(minimum(uimg)) / $(maximum(uimg))") + println("min/max: $(minimum(p)) / $(maximum(p))") + end + p + #uimg +end + +grid = FDM.MetaGrid((0:0.01:1, 0:0.01:1), (3, 3), 5) + +A = I +#A = spdiagm(0 => ones(length(grid)), 1 => ones(length(grid)-1) ) +#f = A * f + + +f = FDM.GridFunction(x -> float(norm(x .- [0.5, 0.5]) < 0.4), 1, grid) g = FDM.GridFunction(x -> norm(x), 2, grid) + +for part in parts(grid) + gridv = view(grid, part...) + fv = view(f, part...) + fpalg(gridv, fv, A, α=1, β=1, τ=1/100, ϵ=1e-7) +end + +#gridv = view(grid, 1, 1) +#fv = view(f, 1, 1) +#g = view(g, 1, 1) + + +FDM.my_plot(f) + + + + #FDM.plot(f) diff --git a/fd.jl b/fd.jl index 776d4fed0a88921c05dc504bd39822f689ad30c7..43199631bc58b4294c41ab923df4dd0b96c5dd53 100644 --- a/fd.jl +++ b/fd.jl @@ -75,14 +75,37 @@ Base.getindex(grid::Grid, idx...) = SVector(getindex.(grid.domain, idx)) Base.getindex(grid::Grid, I::CartesianIndex) = SVector(getindex.(grid.domain, Tuple(I))) +# GridView + +struct GridView{d} <: AbstractGrid{d} + grid::AbstractGrid{d} + indices::NTuple{d, UnitRange{Int}} +end +GridView(grid, indices::UnitRange{Int}...) = GridView(grid, indices) + +Base.size(grid::GridView) = length.(grid.indices) + +# FIXME: this is terrible! +Base.iterate(grid::GridView) = iterate(grid.grid[I] for I in CartesianIndices(grid.indices)) +Base.iterate(grid::GridView, state) = iterate([grid.grid[I] for I in CartesianIndices(grid.indices)], state) + +Base.getindex(grid::GridView, idx...) = getindex(grid.grid, (first.(grid.indices) .- 1 .+ Base.to_indices(grid, idx))...) + # MetaGrid struct MetaGrid{d} <: AbstractGrid{d} grid::Grid{d} indices::Array{NTuple{d, UnitRange{Int}}, d} + overlap::Int end +""" + MetaGrid(domain::NTuple{d, FDStepRange}, pnum::NTuple{d, Int}, overlap::Int) + +Create a grid on `domain`, partitioned along dimensions according to `pnum` +respecting `overlap`. +""" function MetaGrid(domain::NTuple{d, FDStepRange}, pnum::NTuple{d, Int}, overlap::Int) where d grid = Grid(domain) tsize = size(grid) .+ (pnum .- 1) .* overlap @@ -97,7 +120,7 @@ function MetaGrid(domain::NTuple{d, FDStepRange}, pnum::NTuple{d, Int}, overlap: indices[I] = ntuple(j -> ((I[j] - 1) * psize[j] - (I[j] - 1) * overlap + 1) : ( I[j] * psize[j] - (I[j] - 1) * overlap + overhang(I, j)), d) end - MetaGrid{d}(grid, indices) + MetaGrid{d}(grid, indices, overlap) end Base.size(grid::MetaGrid) = size(grid.grid) @@ -109,9 +132,22 @@ Base.getindex(grid::MetaGrid, idx...) = getindex(grid.grid, idx...) indices(grid::MetaGrid, idx...) = getindex(grid.indices, idx...) +Base.view(grid::MetaGrid, idx...) = GridView(grid, indices(grid, idx...)) + +parts(grid::MetaGrid) = CartesianIndices(grid.indices) + ## GridFunction +abstract type AbstractGridFunction{G,T,d,n} <: AbstractArray{T,1} end + +rangedim(f::AbstractGridFunction{G,T,d,n}) where {G,T,d,n} = n +domaindim(f::AbstractGridFunction{G,T,d}) where {G,T,d} = d + +Base.similar(f::F, ::Type{S}, dims::Int...) where {F<:AbstractGridFunction, S} = + GridFunction(f.grid, Array{S}(undef, length(f))) + + """ GridFunction{G<:AbstractGrid,T,d,n,D} @@ -119,23 +155,31 @@ indices(grid::MetaGrid, idx...) = getindex(grid.indices, idx...) `D == d + 1` """ -struct GridFunction{G<:AbstractGrid,T,d,n,D} <: AbstractArray{T,1} +struct GridFunction{G<:AbstractGrid,T,d,n,A,D} <: AbstractGridFunction{G,T,d,n} grid::G - data::Array{T,D} - function GridFunction{G,T,d,n,D}(grid::G, data::AbstractArray{T,D}) where {G,T,d,n,D} + data::A + function GridFunction{G,T,d,n,A,D}(grid::G, data::A) where {G,T,d,n,D,A<:AbstractArray{T,D}} D == d + 1 || throw(ArgumentError) d == ndims(grid) || throw(ArgumentError) - (size(grid)..., n) == size(data) || throw(ArgumentError) + (size(grid)..., n) == size(data) || throw(ArgumentError("data does not match given size")) new(grid, data) end end -GridFunction(grid::G, data::AbstractArray{T,D}) where {d,G<:AbstractGrid{d},T,D} = - GridFunction{G,T,d,size(data, d+1),d+1}(grid, data) -function GridFunction(grid::G, data::AbstractArray{T,1}) where {G,T} - d = ndims(grid) - n = length(data) ÷ d - data = reshape(data, size(grid)..., n) - GridFunction{G,T,d,n,d+1}(grid, data) +# TODO: check size(data) +#GridFunction(grid::G, data::AbstractArray{T,D}) where {d,G<:AbstractGrid{d},T,D} = +# GridFunction{G,T,d,size(data, d+1),typeof(data),d+1}(grid, data) +function GridFunction(grid::G, data::AbstractArray, ::Val{n}) where {d,G<:AbstractGrid{d}, n} + if length(data) ÷ length(grid) == 1 && n > 1 + data = repeat(data, inner=ntuple(i -> i == d+1 ? n : 1, d+1)) + else + # TODO: check size(data) + data = reshape(data, size(grid)..., n) + end + length(data) ÷ length(grid) == n || throw(ArgumentError("dimensions don't match")) + GridFunction{G,eltype(data),d,n,typeof(data),d+1}(grid, data) +end +function GridFunction(grid::G, data::AbstractArray) where {d,G<:AbstractGrid{d}} + GridFunction(grid, data, Val(length(data) ÷ length(grid))) end function GridFunction(f::Function, n::Int, grid::G) where {G<:AbstractGrid} data = Array{Float64,ndims(grid)+1}(undef, size(grid)..., n) @@ -144,6 +188,10 @@ function GridFunction(f::Function, n::Int, grid::G) where {G<:AbstractGrid} end GridFunction(grid, data) end +function GridFunction(grid::AbstractGrid, ::UndefInitializer, ::Val{n} = Val(1)) where n + data = Array{Float64}(undef, size(grid)..., 1) + GridFunction{typeof(grid), eltype(data), ndims(grid), n, typeof(data), ndims(grid) + 1}(grid, data) +end #Base.show(io::IO, grid::GridFunction) = # print("GridFunction on $(length(grid)) grid points") @@ -157,15 +205,12 @@ Base.setindex!(f::GridFunction, v, idx::Int) = setindex!(f.data, v, idx) #Base.getindex(f::GridFunction, idx...) = getindex(f.data, idx...) #Base.setindex!(f::GridFunction, v, idx...) = setindex!(f.data, v, idx...) -Base.IndexStyle(::Type{<:GridFunction}) = IndexLinear() -#Base.similar(f::GridFunction{G,d,n}, ::Type{S}, size::Integer...) where {G,d,n,S} = -# GridFunction{G,d,n}(f.grid, Array{S}(undef, size)) -# -rangedim(f::GridFunction{G,T,d,n}) where {G,T,d,n} = n -domaindim(f::GridFunction{G,T,d}) where {G,T,d} = d +#Base.IndexStyle(::Type{<:GridFunction}) = IndexLinear() #component(f::GridFunction, k::Integer) = view(reshape(f.data, +Base.view(f::AbstractGridFunction{G}, idx...) where G<:MetaGrid = + GridFunction(GridView(f.grid, indices(f.grid, idx...)), view(f.data, indices(f.grid, idx...)...)) @@ -174,7 +219,8 @@ domaindim(f::GridFunction{G,T,d}) where {G,T,d} = d import Colors: Gray import Plots: plot -plot(img::GridFunction{G,T,d,1}) where {G,T,d} = +# TODO: add clamp option +my_plot(img::AbstractGridFunction{G,T,d,1}) where {G,T,d} = plot(reshape(Gray.((img .- minimum(img)) ./ (maximum(img) - minimum(img))), size(img.grid)), show=true) ## Operations @@ -195,7 +241,8 @@ derivate(f::GridFunction, dir::Integer; fd = FD.forward) = fdfilter(f, _fdkern(f # chambolles algorithm), other choices may be reasonable function gradient(f::GridFunction{G,T,d,1}; fd=FD.forward) where {G,T,d} - A = Array{T,d+1}(undef, (size(f.grid)..., d)) + #A = Array{T,d+1}(undef, (size(f.grid)..., d)) + A = zeros(size(f.grid)..., d) for k = 1:d vidx = ntuple(i -> i <= d ? Colon() : k, d + 1) vidx2 = ntuple(i -> i <= d ? Colon() : 1, d + 1) @@ -216,6 +263,12 @@ function divergence(f::GridFunction{G,T,d}; fd=FD.backward) where {d,G<:Abstract GridFunction(f.grid, A) end +function my_norm(f::GridFunction) + d = domaindim(f) + GridFunction(f.grid, sqrt.(sum(f.data.^2, dims=d+1)), Val(rangedim(f))) +end + + """ diff --git a/filtering.jl b/filtering.jl index 215e591c740658338dd1e52f06892ca1bd98b757..11fa7b2836a0e539a8e05891110162bb931a8765 100644 --- a/filtering.jl +++ b/filtering.jl @@ -55,7 +55,12 @@ function trim_kernel(kernel, bidx) end +""" + imfilter!(dst, img, kern, slice=:inner) +Filters `img` applying `kern` and write to `dst`. `dst` is added to, so you +should initialize with zeros. +""" # TODO: weirdly Base.Slice instead of Base.UnitRange leads to axes that are not # compatible when broadcasting, figure out why! function imfilter!(dst, img, kern, slice=_bidx_to_range.(axes(img), ntuple(i->0, ndims(img)))) @@ -88,7 +93,8 @@ end Filter `img` using a kernel function `kernelf(bidx) = kernel` specifying a kernel for each type of node. """ -imfilter(img, kernelf::Function) = imfilter!(similar(img), img, kernelf) +# need zeros here, because `imfilter!` acts additively +imfilter(img, kernelf::Function) = imfilter!(zeros(size(img)), img, kernelf) """