diff -r 4bbb246f4cac -r e4ad8f7ce671 src/PET/ImGenerate.jl
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/src/PET/ImGenerate.jl	Fri Apr 19 17:00:37 2024 +0300
@@ -0,0 +1,282 @@
+###################
+# Image generation
+###################
+
+module ImGenerate
+
+using ColorTypes: Gray
+import TestImages
+# We don't really *directly* depend on QuartzImageIO. The import here is
+# merely a workaround to suppress warnings when  loading TestImages.
+# Something is broken in those packages.
+import QuartzImageIO
+
+using AlgTools.Util
+using AlgTools.Comms
+using ImageTools.Translate
+
+using ..OpticalFlow: Image, DisplacementConstant, DisplacementFull
+using ..Radon
+
+# Added for reproducibility
+import StableRNGs: StableRNG, Random
+const rng = StableRNG(9182737465)
+
+# Added for PET
+import PoissonRandom: pois_rand
+import Random: shuffle
+import Images: center, warp
+import CoordinateTransformations: recenter
+import Rotations: RotMatrix
+import Interpolations: Flat
+
+
+##############
+# Our exports
+##############
+
+export ImGen,
+       OnlineData,
+       imgen_square,
+       imgen_shake,
+       PetOnlineData,
+       imgen_shepplogan_radon
+
+##################
+# Data structures
+##################
+
+struct ImGen
+    f :: Function
+    dim :: Tuple{Int64,Int64}
+    Λ :: Float64
+    dynrange :: Float64
+    name :: String
+end
+
+struct OnlineData{DisplacementT}
+    b_true :: Image
+    b_noisy :: Image
+    v :: DisplacementT
+    v_true :: DisplacementT
+    v_cumul_true :: DisplacementT
+end
+
+struct PetOnlineData{DisplacementT}
+    b_true :: Image
+    sinogram_true :: Image
+    sinogram_noisy :: Image
+    v :: DisplacementT
+    v_true :: DisplacementT
+    v_cumul_true :: DisplacementT
+    theta :: DisplacementT              # theta = thetaknown, theta_cumul
+    S :: Image
+end
+
+###################
+# Shake generation
+###################
+
+function make_const_v(displ, sz)
+    v = zeros(2, sz...)
+    v[1, :, :] .= displ[1]
+    v[2, :, :] .= displ[2]
+    return v
+end
+
+function shake(params)
+    if !haskey(params, :shaketype) || params.shaketype == :gaussian
+        return () -> params.shake.*randn(rng,2)
+    elseif params.shaketype == :disk
+        return () -> begin
+            θ = 2π*rand(rng,Float64)
+            r = params.shake*√(rand(rng,Float64))
+            return [r*cos(θ), r*sin(θ)]
+        end
+    elseif params.shaketype == :circle
+        return () -> begin
+            θ = 2π*rand(rng,Float64)
+            r = params.shake
+            return [r*cos(θ), r*sin(θ)]
+        end        
+    else
+        error("Unknown shaketype $(params.shaketype)")
+    end 
+end
+
+pixelwise = (shakefn, sz) -> () -> make_const_u(shakefn(), sz)
+
+
+function rotatebytheta(params)
+    r = params.rotation_factor*(2*rand(rng)-1)
+    return r
+end
+
+function generate_radonmask(params)
+    imdim = params.radondims
+    sino_sparse = params.sino_sparsity
+    numone = Int64(round(sino_sparse*imdim[1]*imdim[2]))
+    numzero = imdim[1]*imdim[2]-numone
+    A = shuffle(rng,reshape([ones(numone); zeros(numzero)],(imdim[1],imdim[2])))
+    return A
+end
+
+################
+# Moving square
+################
+
+function generate_square(sz,
+                         :: Type{DisplacementT},
+                         datachannel :: Channel{OnlineData{DisplacementT}},
+                         params) where DisplacementT
+
+    if false
+        v₀ = make_const_v(0.1.*(-1, 1), sz)
+        nextv = () -> v₀
+    elseif DisplacementT == DisplacementFull
+        nextv = pixelwise(shake(params), sz)
+    elseif DisplacementT == DisplacementConstant
+        nextv = shake(params)
+    else
+        @error "Invalid DisplacementT"
+    end
+
+    # Constant linear displacement everywhere has Jacobian determinant one
+    # (modulo the boundaries which we ignore here)
+    m = round(Int, sz[1]/5)
+    b_orig = zeros(sz...)
+    b_orig[sz[1].-(2*m:3*m), 2*m:3*m] .= 1
+
+    v_true = nextv()
+    v_cumul = copy(v_true)
+
+    while true
+        # Flow original data and add noise
+        b_true = zeros(sz...)
+        translate_image!(b_true, b_orig, v_cumul; threads=true)
+        b = b_true .+ params.noise_level.*randn(rng,sz...)
+        v = v_true.*(1.0 .+ params.shake_noise_level.*randn(rng,size(v_true)...))
+        # Pass true data to iteration routine
+        data = OnlineData{DisplacementT}(b_true, b, v, v_true, v_cumul)
+        if !put_unless_closed!(datachannel, data)
+            return
+        end
+        # Next step shake
+        v_true = nextv()
+        v_cumul .+= v_true
+    end
+end
+
+function imgen_square(sz)
+    return ImGen(curry(generate_square, sz), sz, 1, 1, "square$(sz[1])x$(sz[2])")
+end
+
+################
+# Shake a photo
+################
+
+function generate_shake_image(im, sz,
+                              :: Type{DisplacementConstant},
+                              datachannel :: Channel{OnlineData{DisplacementConstant}},
+                              params :: NamedTuple)
+    
+    # Restart the seed to enable comparison across predictors                       
+    Random.seed!(rng,9182737465)   
+
+    nextv = shake(params)
+    v_true = nextv()
+    v_cumul = copy(v_true)
+
+    while true
+        # Extract subwindow of original image and add noise
+        b_true = zeros(sz...)
+        extract_subimage!(b_true, im, v_cumul; threads=true)
+        b = b_true .+ params.noise_level.*randn(rng,sz...)
+        v = v_true.*(1.0 .+ params.shake_noise_level.*randn(rng,size(v_true)...))
+        # Pass data to iteration routine
+        data = OnlineData{DisplacementConstant}(b_true, b, v, v_true, v_cumul)
+        if !put_unless_closed!(datachannel, data)
+            return
+        end
+        # Next step shake
+        v_true = nextv()
+        v_cumul .+= v_true
+    end
+end
+
+function imgen_shake(imname, sz)
+    im = Float64.(Gray.(TestImages.testimage(imname)))
+    dynrange = maximum(im)
+    return ImGen(curry(generate_shake_image, im, sz), sz, 1, dynrange,
+                 "$(imname)$(sz[1])x$(sz[2])")
+end
+
+
+
+########################################################################
+# PETscan
+########################################################################
+function generate_radon(im, sz,
+    :: Type{DisplacementConstant},
+    datachannel :: Channel{PetOnlineData{DisplacementConstant}},
+    params :: NamedTuple)
+
+    # Restart the seed to enable comparison across predictors                       
+    Random.seed!(rng,9182737465)  
+
+    nextv = shake(params)
+    v_true = nextv()
+    v_cumul = copy(v_true)
+
+    S_true = generate_radonmask(params)
+
+    theta_true = rotatebytheta(params)
+    theta_cumul = copy(theta_true)
+
+    while true
+        # Define the transformation matrix
+        center_point = (sz[1]/2 + v_true[1], sz[2]/2 + v_true[2])
+        tform = recenter(RotMatrix(theta_cumul), center_point)
+
+        # Apply the transformation to the image using warp
+        b_true = copy(warp(im, tform, axes(im), fillvalue=Flat()))
+
+        v = v_true.*(1.0 .+ params.shake_noise_level.*randn(rng,size(v_true)...))
+        theta = theta_true*(1.0 + params.rotation_noise_level.*randn(rng))
+
+        # Generate the true and noisy sinograms
+        sinogram_true = zeros(params.radondims...)
+        sinogram_true .*= params.scale
+        radon!(sinogram_true, b_true)
+        sinogram_noisy = copy(sinogram_true)
+
+        for i=1:params.radondims[1], j=1:params.radondims[2]
+            sinogram_noisy[i, j] += pois_rand(rng,params.noise_level)
+        end
+
+        # Pass data to iteration routine
+        data = PetOnlineData{DisplacementConstant}(b_true, sinogram_true, sinogram_noisy, v, v_true, v_cumul, [theta, theta_cumul], S_true)
+        if !put_unless_closed!(datachannel, data)
+        return
+        end
+
+        # Next step shake
+        v_true = nextv()
+        v_cumul .+= v_true
+        # Next theta
+        theta_true = rotatebytheta(params)
+        theta_cumul += theta_true
+        # Next sinogram mask 
+        S_true = generate_radonmask(params)
+    end
+end
+
+
+function imgen_shepplogan_radon(origsize,sz)
+    im = convert(Array{Float64},TestImages.shepp_logan(origsize, highContrast=true))
+    dynrange = maximum(im)
+    return ImGen(curry(generate_radon, im, sz), sz, 1, dynrange, "shepplogan$(sz[1])x$(sz[2])")
+end
+
+
+end # Module