Mercurial > repos > PredictPDPS / changeset

--- /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
+
+