code: more work

FD/src/FD.jl

@@ -5,6 +5,13 @@ using StaticArrays
 using OffsetArrays
 #using LinearMaps
 #using IterativeSolvers
+import Colors: Gray
+import Plots: plot
+
+
+export Grid, MetaGrid, GridView, GridFunction
+export grid
+
 
 include("filtering.jl")
 
@@ -54,6 +61,7 @@ Base.ndims(grid::AbstractGrid{d}) where d = d
 Base.axes(grid::AbstractGrid) = Base.OneTo.(size(grid))
 Base.length(grid::AbstractGrid) = prod(size(grid))
 Base.IndexStyle(::Type{<:AbstractGrid}) = IndexCartesian()
+Base.parent(grid::AbstractGrid) = grid
 
 # aka `StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}`
 const FDStepRange = typeof(0:0.1:1)
@@ -91,6 +99,9 @@ Base.iterate(grid::GridView, state) = iterate([grid.grid[I] for I in CartesianIn
 
 Base.getindex(grid::GridView, idx...) = getindex(grid.grid, (first.(grid.indices) .- 1 .+ Base.to_indices(grid, idx))...)
 
+Base.parent(grid::GridView) = grid.grid
+
+
 # MetaGrid
 
 
@@ -146,6 +157,7 @@ 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)))
+Base.parent(f::AbstractGridFunction) = f
 
 
 """
@@ -169,10 +181,9 @@ end
 #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
+  if length(data) == length(grid) && n > 1
     data = repeat(data, inner=ntuple(i -> i == d+1 ? n : 1, d+1))
-  else
-    # TODO: check size(data)
+  elseif size(data) != (size(grid)..., n)
     data = reshape(data, size(grid)..., n)
   end
   length(data) ÷ length(grid) == n || throw(ArgumentError("dimensions don't match"))
@@ -208,16 +219,80 @@ Base.setindex!(f::GridFunction, v, idx::Int) = setindex!(f.data, v, idx)
 #Base.IndexStyle(::Type{<:GridFunction}) = IndexLinear()
 
 #component(f::GridFunction, k::Integer) = view(reshape(f.data,
+grid(f::GridFunction) = f.grid
 
+# TODO: call is different, so maybe make thin not part of Base
 Base.view(f::AbstractGridFunction{G}, idx...) where G<:MetaGrid =
     GridFunction(GridView(f.grid, indices(f.grid, idx...)), view(f.data, indices(f.grid, idx...)..., :))
 
-grid(f::GridFunction) = f.grid
+Base.parent(f::AbstractGridFunction{<:GridView}) =
+    GridFunction(parent(grid(f)), parent(f.data))
+
 
 
+function _check_shape(a::GridFunction, b::GridFunction)
+  #TODO: why does fv.grid == gv.grid fail?
+  a.grid === b.grid || throw(ArgumentError("grids need to match"))
+  rangedim(a) == rangedim(b) || throw(ArgumentError("range dimensions need to match"))
+end
+
 Base.:*(a::Number, b::GridFunction) = GridFunction(b.grid, a * b.data, Val(rangedim(b)))
+Base.:-(a::GridFunction) = GridFunction(a.grid, -a.data, Val(rangedim(a)))
+function Base.:+(a::GridFunction, b::GridFunction)
+  _check_shape(a, b)
+  GridFunction(a.grid, a.data .+ b.data, Val(rangedim(a)))
+end
+function Base.:-(a::GridFunction, b::GridFunction)
+  _check_shape(a, b)
+  GridFunction(a.grid, a.data .- b.data, Val(rangedim(a)))
+end
+
+
+function Base.:+(a::GridFunction, b::GridFunction{<:GridView})
+  a.grid == parent(b.grid) || throw(ArgumentError("grids need to match"))
+  data = copy(a.data)
+  datav = view(data, b.grid.indices..., :)
+  datav .+= b.data
+  GridFunction(a.grid, data, Val(rangedim(a)))
+end
+Base.:+(a::GridFunction{<:GridView}, b::GridFunction) = b + a
+
+# TODO: cleanup! use metaprogramming
+function Base.:+(a::GridFunction{<:GridView}, b::GridFunction{<:GridView})
+  _check_shape(a, b)
+  GridFunction(a.grid, a.data .+ b.data, Val(rangedim(a)))
+end
+# TODO: cleanup! use metaprogramming
+function Base.:-(a::GridFunction{<:GridView}, b::GridFunction{<:GridView})
+  _check_shape(a, b)
+  GridFunction(a.grid, a.data .- b.data, Val(rangedim(a)))
+end
+
+function Base.:-(a::GridFunction, b::GridFunction{<:GridView})
+  a.grid == parent(b.grid) || throw(ArgumentError("grids need to match"))
+  data = copy(a.data)
+  datav = view(data, b.grid.indices..., :)
+  datav .-= b.data
+  GridFunction(a.grid, data, Val(rangedim(a)))
+end
+Base.:-(a::GridFunction{<:GridView}, b::GridFunction) = b - a
+
+
+
+function my_plot(img::FD.GridFunction{G,T,d,1}) where {G,T,d}
+  print("Plotting ... ")
 
+  data = reshape(copy(img.data), size(img.grid))
 
+  min = minimum(data)
+  max = maximum(data)
+  @show min, max
+
+  #@. data = (data - min) / (max - min)
+  @. data = clamp(data, 0, 1)
+
+  plot(Gray.(data), show=true)
+end
 
 
 ## Operations
@@ -260,6 +335,34 @@ function divergence(f::GridFunction{G,T,d}; fd=FDiff.backward) where {d,G<:Abstr
   GridFunction(f.grid, A)
 end
 
+function _kernel_zero_right_bnd!(kernel, bidx, k)
+  if bidx[k] == 1
+    kernel[ntuple(i -> 0, ndims(kernel))...] = 0
+  end
+  kernel
+end
+
+function gradient_c(f::GridFunction{G,T,d,1}; fd=FDiff.forward) where {G,T,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)
+    imfilter!(view(A, vidx...), f.data[vidx2...], bidx -> trim_kernel(_fdkern(fd, ndims(f.grid), k), bidx))
+  end
+  GridFunction(f.grid, A)
+end
+
+function divergence_z(f::GridFunction{G,T,d}; fd=FDiff.backward) where {d,G<:AbstractGrid{d},T}
+  A = zeros(size(f.grid)..., 1)
+  for k = 1:d
+    vidx = ntuple(i -> i <= d ? Colon() : 1, d + 1)
+    vidx2 = ntuple(i -> i <= d ? Colon() : k, d + 1)
+    imfilter!(view(A, vidx...), f.data[vidx2...], bidx -> _kernel_zero_right_bnd!(trim_kernel(_fdkern(fd, ndims(f.grid), k), bidx), bidx, k))
+  end
+  GridFunction(f.grid, A)
+end
+
 function pw_norm(f::GridFunction)
   d = domaindim(f)
   GridFunction(f.grid, sqrt.(sum(f.data.^2, dims=d+1)), Val(1))

chambollefp.jl

@@ -10,13 +10,14 @@ using FD
 
 imf(f, grid) = [f(x) for x in grid]
 
-imshow(grid, img) = plot(Gray.((reshape(float.(img), size(grid)) .- minimum(float.(img))) / (maximum(float.(img)) - minimum(float.(img))) ), show=true)
+imshow(grid, img) = FD.my_plot(FD.GridFunction(grid, img))
+#plot(Gray.((reshape(float.(img), size(grid)) .- minimum(float.(img))) / (maximum(float.(img)) - minimum(float.(img))) ), show=true)
 
 function _fpiter(grid, p, f, A, α, β, τ)
-  x = (A'*A + β*I) \ (FD.divergence(grid, p) .+ A'*f ./ α)
+  x = (A'*A + β*I) \ (FD.divergence(grid, p) .+ A'*f)
   q = FD.gradient(grid, x)
 
-  p_new = (p .+ τ .* q) ./ (1 .+ τ .* FD.my_norm(grid, q))
+  p_new = (p .+ τ .* q) ./ (1 .+ (τ/α) .* FD.my_norm(grid, q))
 
   reserr = maximum(sum((p_new .- p).^2, dims=ndims(p)))
   (p_new, reserr)
@@ -27,17 +28,16 @@ function fpalg(grid, f, A; α, β, τ = 1/8, ϵ = 1e-4)
 
   reserr = 1
   k = 0
-  while reserr > ϵ && k < 30e1
+  while reserr > ϵ && k < 300e1
     k += 1
     #sleep(0.1 / k)
     (p, reserr) = _fpiter(grid, p, f, A, α, β, τ)
 
-    uimg = (A'*A + β*I) \ (α * FD.divergence(grid, p) .+ A'*f)
+    uimg = (A'*A + β*I) \ (FD.divergence(grid, p) .+ A'*f)
 
     imshow(grid, uimg)
     println("[$(lpad(k, 4))] res = $(round(reserr, digits=5))")
-    println("min/max: $(minimum(uimg)) / $(maximum(uimg))")
-    println("min/max: $(minimum(p)) / $(maximum(p))")
+    println("min/max: $(minimum(FD.my_norm(FD.GridFunction(grid,p)))) / $(maximum(FD.my_norm(FD.GridFunction(grid,p))))")
   end
 end
 
@@ -54,15 +54,15 @@ f = reshape([float(norm(x .- [0.5, 0.5]) < 0.4) for x in grid], length(grid))
 #f += clamp.(0.2 * rand(eltype(f), size(f)), 0, 1)
 
 
-#A = I
-A = spdiagm(0 => ones(length(grid)), 1 => ones(length(grid)-1) )
+A = I
+#A = spdiagm(0 => ones(length(grid)), 1 => ones(length(grid)-1) )
 
 f = A * f
 
 #imshow(grid, f)
 #sleep(5)
 
-img = fpalg(grid, f, A, α=1, β=1, τ=1/100, ϵ=1e-7)
+img = fpalg(grid, f, A, α=10000, β=0, τ=1/8, ϵ=1e-7)
 
 
 #A = Array{Float64}(undef, (length(grid), length(grid)))

chambollefp_dd.jl

@@ -8,9 +8,6 @@ import FD
 import Colors: Gray
 import Plots: plot
 
-# TODO: add clamp option
-my_plot(img::FD.AbstractGridFunction{G,T,d,1}) where {G,T,d} =
-    plot(reshape(Gray.((img .- minimum(img)) ./ (maximum(img) - minimum(img))), size(img.grid)), show=true)
 
 
 "partition of unity function"
@@ -31,115 +28,174 @@ function theta(grid, idx...)
   FD.GridFunction(grid, dst)
 end
 
+
 function p2u(grid, p, A, f; α, β)
-  FD.GridFunction(grid, (A'*A + β*I) \ (α .* FD.divergence(p).data[:] .+ A'*f.data[:]))
+  FD.GridFunction(grid, (A'*A + β*I) \ (FD.divergence_z(p).data[:] .+ A'*f.data[:]))
 end
 
+
 # TODO: investigate matrix multiplication performance with GridFunctions
-function fpiter(grid, p, w, A, α, β, τ)
-  x = FD.GridFunction(grid, (A'*A + β*I) \ (FD.divergence(p).data[:] .+ w ./ α))
-  q = FD.gradient(x)
+function fpiter(grid, p, wg, A, α, β, τ)
+  ## local variant
+  #w = FD.GridFunction(grid, view(wg.data, grid.indices..., :))
+  #x = FD.GridFunction(grid, (A'*A + β*I) \ (FD.divergence(p) + w).data[:])
+  #q = FD.gradient(x)
+
+  ## global variant
+  x = FD.GridFunction(FD.parent(grid), (A'*A + β*I) \ (FD.divergence(p) + wg).data[:])
+  #FD.my_plot(FD.divergence(p) + wg)
+
+  q = FD.GridFunction(grid, view(FD.gradient(x).data, grid.indices..., :))
+  q2 = FD.gradient(FD.GridFunction(grid, view(x.data, grid.indices..., :)))
+
+  #display(x.data[10:13,51:54,1])
+  #display(q.data[10:13,52:53,2])
+  #display(q2.data[10:13,52:53,2])
+  #sleep(0.1)
+
+  #q = FD.gradient(FD.GridFunction(grid, view(x.data, grid.indices..., :)))
+
 
-  #display(x.data)
+  #FD.my_plot(FD.pw_norm(q))
   #FD.my_plot(x)
-  #checknan(q)
-  p_new = FD.GridFunction(grid, FD.pwmul(FD.GridFunction(grid, 1 ./ (1 .+ τ .* FD.pw_norm(q))), FD.GridFunction(grid, p .+ τ .* q)))
+  #FD.my_plot(α)
+
+  p_new = FD.GridFunction(grid, FD.pwmul(FD.GridFunction(grid, α ./ (α .+ τ .* FD.pw_norm(q))),
+                                         FD.GridFunction(grid, p .+ τ .* q)))
+
+
 
   reserr = maximum(sum((p_new .- p).^2, dims=ndims(p)))
   (p_new, reserr)
 end
 
-function fpalg(grid, w, A; α, β, τ = 1/8, ϵ = 1e-4)
+
+function fpalg(grid, w, A, p = FD.GridFunction(x -> 0, ndims(grid), grid);
+               α, β, τ = 1/8, ϵ = 1e-4)
   local uimg
-  p = FD.GridFunction(x -> 0, ndims(grid), grid)
   reserr = 1
   k = 0
-  while reserr > ϵ && k < 30e1
+  while reserr > ϵ && k < 10e1
     k += 1
     (p, reserr) = fpiter(grid, p, w, A, α, β, τ)
 
-    #FD.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))")
+    #u = p2u(grid, p, A, f, α=α, β=β)
+    #@show k
+    #FD.my_plot(u)
   end
-  p
-  #uimg
+  return p
 end
 
-grid = FD.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 = FD.GridFunction(x -> float(norm(x .- [0.5, 0.5]) < 0.4), 1, grid)
-g = FD.GridFunction(x -> norm(x), 2, grid)
-
-α=1
-β=1
-τ=1/100
-ϵ=1e-7
-ρ=0.25
-
-p = FD.GridFunction(x -> 0, 2, grid)
 
 function dditer(grid, p, A, f; α, β, τ, ϵ, ρ)
-  p .= (1-ρ) .* p
+  #pzero = FD.GridFunction(x -> 0, ndims(grid), grid)
+
+  p̂ = (1-ρ) * p
 
+  i = 0
   for part in FD.parts(grid)
     gridv = view(grid, Tuple(part)...)
     fv = view(f, Tuple(part)...)
     #p.data .= 0
+    p̂v = view(p̂, Tuple(part)...)
     pv = view(p, Tuple(part)...)
+    #pzerov = view(pzero, Tuple(part)...)
 
     θ = theta(grid, Tuple(part)...)
     θv = view(θ, Tuple(part)...)
-    w = FD.GridFunction(grid, A' * f.data[:] - FD.divergence(FD.pwmul(FD.GridFunction(grid, 1 .- θ), p)))
+
+    #y = FD.divergence(FD.pwmul(FD.GridFunction(grid, 1 .- θ), p))
+    #FD.my_plot(y)
+    #sleep(1)
+
+    w = FD.GridFunction(grid, A' * f.data[:]) + FD.divergence(FD.pwmul(FD.GridFunction(grid, 1 .- θ), p))
     wv = view(w, Tuple(part)...)
 
+    #FD.my_plot(w)
+    #sleep(1)
+
     #wv = A' * fv
 
-    pv .+= ρ .* fpalg(gridv, wv, A, α=α*θv, β=β, τ=τ, ϵ=ϵ)
+    #u = p2u(grid, p̂, A, f, α=α, β=β)
+    #uv = view(u, Tuple(part)...)
 
-    #u = p2u(grid, p, A, f, α=α*θ, β=β)
-    u = p2u(grid, p, A, f, α=α, β=β)
-    uv = view(u, Tuple(part)...)
+    #FD.my_plot(uv)
 
-    my_plot(θv)
     #sleep(1)
-    #my_plot(FD.pw_norm(pv))
-    @show maximum(pv), minimum(pv)
-    #sleep(1)
-    #my_plot(θv)
-    return
-  end
-end
 
-dditer(grid, p, A, f, α=α, β=β, τ=τ, ϵ=ϵ, ρ=ρ)
+    p̂v .+= ρ .* fpalg(gridv, w, A, α=α*θv, β=β, τ=τ, ϵ=ϵ)
+    #p̂v .= fpalg(gridv, wv, A, α=α*θv, β=β, τ=τ, ϵ=ϵ)
 
-#for n = 1:10
-#  println("iteration $n")
-#  dditer(grid, p, A, f, α=α, β=β, τ=τ, ϵ=ϵ, ρ=ρ)
-#end
+  #FD.my_plot(FD.pw_norm(p))
+  #sleep(1)
+  #FD.my_plot(FD.pw_norm(f))
+  #sleep(1)
+  #FD.my_plot(FD.pw_norm(p̂))
+  #display(p̂.data[10:13,52:54,:])
+    #break
+    #FD.my_plot(FD.pw_norm(p̂))
+    #sleep(1)
+    #FD.my_plot(θ)
+    #display(p̂v.data)
+    #break
+    #sleep(1)
 
+    #u = p2u(grid, p, A, f, α=α*θ, β=β)
+    u = p2u(grid, p̂, A, f, α=α, β=β)
+    uv = view(u, Tuple(part)...)
 
+    #FD.my_plot(θv)
 
+    #sleep(1)
+    #FD.my_plot(FD.pw_norm(p̂v))
+    #@show maximum(u), minimum(u)
+    #@show maximum(p̂v), minimum(p̂v)
+    #sleep(1)
+    #FD.my_plot(FD.pw_norm(p̂))
+    #FD.my_plot(u)
+    #FD.my_plot(p̂v)
+    #y = FD.GridFunction(gridv, FD.gradient(θv).data[:,:,1])
+    #display(y.data)
+    #display(θv.data)
+    #FD.my_plot(FD.GridFunction(gridv, FD.gradient(θv).data[:,:,2]))
 
-#gridv = view(grid, 1, 1)
-#fv = view(f, 1, 1)
-#g = view(g, 1, 1)
+    #display(FD.gradient(θv).data)
+    #break
 
+    #@show size(θv.data), size(p̂v.data)
 
-#FD.my_plot(f)
+    #FD.my_plot(FD.pw_norm(p̂))
 
+    i += 1
+    #break
+    #i > 3 && break
+  end
+  p .= p̂
 
+  u = p2u(grid, p, A, f, α=α, β=β)
+  #FD.my_plot(u)
+  FD.my_plot(FD.pw_norm(p))
+end
 
 
-#FD.plot(f)
+grid = FD.MetaGrid((0:0.01:1, 0:0.01:1), (1, 2), 5)
 
+A = I
+#A = spdiagm(0 => ones(length(grid)), 1 => ones(length(grid)-1) )
+#f = A * f
 
+f = FD.GridFunction(x -> float(norm(x .- [0.5, 0.5]) < 0.4), 1, grid)
+g = FD.GridFunction(x -> norm(x), 2, grid)
 
+α=100000
+β=0
+τ=1/8
+ϵ=1e-9
+ρ=0.5
 
+p = FD.GridFunction(x -> 0, 2, grid)
 
+for n = 1:1
+  println("iteration $n")
+  dditer(grid, p, A, f, α=α, β=β, τ=τ, ϵ=ϵ, ρ=ρ)
+end