diff -r befb8d5125cd -r 98b79c837a30 src/PET/AlgorithmZeroDual.jl --- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/PET/AlgorithmZeroDual.jl Sun Apr 21 13:43:18 2024 +0300 @@ -0,0 +1,181 @@ +#################################################################### +# Predictive online PDPS for optical flow with known velocity field +#################################################################### + +__precompile__() + +module AlgorithmZeroDual + +identifier = "pdps_known_zerodual" + +using Printf + +using AlgTools.Util +import AlgTools.Iterate +using ImageTools.Gradient +using ImageTools.Translate + +using ..Radon +using ImageTransformations +using Images, CoordinateTransformations, Rotations, OffsetArrays +using ImageCore, Interpolations + +using ..OpticalFlow: ImageSize, + Image, + petpdflow! + +######################### +# Iterate initialisation +######################### + +function init_rest(x::Image) + imdim=size(x) + y = zeros(2, imdim...) + Δx = copy(x) + Δy = copy(y) + x̄ = copy(x) + radonx = copy(x) + return x, y, Δx, Δy, x̄, radonx +end + +function init_iterates(xinit::Image) + return init_rest(copy(xinit)) +end + +function init_iterates(dim::ImageSize) + return init_rest(zeros(dim...)) +end + +######################### +# PETscan related +######################### +function petvalue(x, b, c) + tmp = similar(b) + radon!(tmp, x) + return sum(@. tmp - b*log(tmp+c)) +end + +function petgrad!(res, x, b, c, S) + tmp = similar(b) + radon!(tmp, x) + @. tmp = S .- b/(tmp+c) + backproject!(res, S.*tmp) +end + +function proj_nonneg!(y) + @inbounds @simd for i=1:length(y) + if y[i] < 0 + y[i] = 0 + end + end + return y +end + +############ +# Algorithm +############ + +function solve( :: Type{DisplacementT}; + dim :: ImageSize, + iterate = AlgTools.simple_iterate, + params::NamedTuple) where DisplacementT + + ################################ + # Extract and set up parameters + ################################ + α, ρ = params.α, params.ρ + R_K² = ∇₂_norm₂₂_est² + γ = 1 + L = params.L + τ₀, σ₀ = params.τ₀, params.σ₀ + τ = τ₀/L + σ = σ₀*(1-τ₀)/(R_K²*τ) + + println("Step length parameters: τ=$(τ), σ=$(σ)") + + λ = params.λ + c = params.c*ones(params.radondims...) + + + ###################### + # Initialise iterates + ###################### + + x, y, Δx, Δy, x̄, r∇ = init_iterates(dim) + + if params.L_experiment + oldpetgradx = zeros(size(x)...) + petgradx = zeros(size(x)) + oldx = ones(size(x)) + end + + #################### + # Run the algorithm + #################### + + v = iterate(params) do verbose :: Function, + b :: Image, # noisy_sinogram + v_known :: DisplacementT, + theta_known :: DisplacementT, + b_true :: Image, + S :: Image + + ################### + # Prediction steps + ################### + + petpdflow!(x, Δx, y, Δy, v_known, theta_known, false) + y .= zeros(size(y)...) + + if params.L_experiment + @. oldx = x + end + + ############ + # PDPS step + ############ + + ∇₂ᵀ!(Δx, y) # primal step: + @. x̄ = x # | save old x for over-relax + petgrad!(r∇, x, b, c, S) # | Calculate gradient of fidelity term + + @. x = x-(τ*λ)*r∇-τ*Δx # | + proj_nonneg!(x) # | non-negativity constaint prox + @. x̄ = 2x - x̄ # over-relax: x̄ = 2x-x_old + ∇₂!(Δy, x̄) # dual step: + @. y = y + σ*Δy # | + proj_norm₂₁ball!(y, α) # | prox + + ##################### + # L update if needed + ##################### + if params.L_experiment + petgrad!(petgradx, x, b, c, S) + petgrad!(oldpetgradx, oldx, b, c, S) + if norm₂(x-oldx)>1e-12 + L = max(0.9*norm₂(petgradx - oldpetgradx)/norm₂(x-oldx),L) + println("Step length parameters: L=$(L)") + τ = τ₀/L + σ = σ₀*(1-τ₀)/(R_K²*τ) + end + end + + ################################ + # Give function value if needed + ################################ + + v = verbose() do + ∇₂!(Δy, x) + value = λ*petvalue(x, b, c) + params.α*norm₂₁(Δy) + value, x, [NaN, NaN], nothing, τ, σ + end + + v + end + + return x, y, v +end + +end # Module + +