Add Denoising Problem

Project.toml

@@ -28,19 +28,25 @@ julia = "^1.6.0"
 [extras]
 ADNLPModels = "54578032-b7ea-4c30-94aa-7cbd1cce6c9a"
 DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
+FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 ProximalOperators = "a725b495-10eb-56fe-b38b-717eba820537"
 QuadraticModels = "f468eda6-eac5-11e8-05a5-ff9e497bcd19"
 ShiftedProximalOperators = "d4fd37fa-580c-4e43-9b30-361c21aae263"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Wavelets = "29a6e085-ba6d-5f35-a997-948ac2efa89a"
 
 [targets]
 test = [
   "ADNLPModels",
   "DifferentialEquations",
+  "FFTW",
+  "Images",
   "MLDatasets",
   "ProximalOperators",
   "QuadraticModels",
   "ShiftedProximalOperators",
   "Test",
+  "Wavelets",
 ]

src/RegularizedProblems.jl

@@ -42,6 +42,13 @@ function __init__()
   @require QuadraticModels = "f468eda6-eac5-11e8-05a5-ff9e497bcd19" begin
     include("qp_rand_model.jl")
   end
+  @require FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" begin
+    @require Wavelets = "29a6e085-ba6d-5f35-a997-948ac2efa89a" begin
+      @require Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0" begin
+        include("denoising_model.jl")
+      end
+    end
+  end
 end
 
 end

src/denoising_data.jl

@@ -0,0 +1,162 @@
+export generate_uniform_blur, generate_gaussian_blur
+
+"""
+Implementation of the denoising problem described in 
+
+Stella, L., Themelis, A. & Patrinos, P. 
+Forward–backward quasi-Newton methods for nonsmooth optimization problems. 
+Comput Optim Appl 67, 443–487 (2017). https://doi.org/10.1007/s10589-017-9912-y
+
+and adapted from the original implementation in Python by the authors of the paper:
+
+Chouzenoux, E., Martin, S. & Pesquet, JC. 
+A Local MM Subspace Method for Solving Constrained Variational Problems in Image Recovery. 
+J Math Imaging Vis 65, 253–276 (2023). https://doi.org/10.1007/s10851-022-01112-z
+"""
+
+# Function to unpad an array
+function unpad(x, n_p, m_p, n)
+  a = n_p - n + 1
+  x = reshape_array(x, (n_p, m_p))
+  unpadded_x = x[a:end, a:end]
+  return unpadded_x
+end
+
+# Function to pad an array
+function pad(z, s1, s2)
+  x = cat(zeros(size(z, 1), s1), z, dims = 2)
+  x = cat(x, zeros(size(z, 1), s2), dims = 2)
+  x = cat(x, zeros(s2, size(x, 2)), dims = 1)
+  x = cat(zeros(s1, size(x, 2)), x, dims = 1)
+  return x
+end
+
+# Function to pad a specific array to match suitable dimensions
+function pad_x(z, a)
+  t = cat(zeros(size(z, 1), a), z, dims = 2)
+  t = cat(zeros(a, size(z, 2) + a), t, dims = 1)
+  return t
+end
+
+# Function to create a meshgrid for the gaussian kernel
+function meshgrid(x_range, y_range)
+  m = length(x_range)
+  n = length(y_range)
+  x = repeat(x_range, inner = (1, n))
+  y = repeat(y_range', outer = (m, 1))
+  return x, y
+end
+
+# Function to generate a Gaussian kernel
+function my_gaussian_kernel(kernel_size, kernel_sigma)
+  x, y = meshgrid((-kernel_size):kernel_size, (-kernel_size):kernel_size)
+  normal = 1 / (2 * pi * kernel_sigma^2)
+  kernel = exp.(-((x .^ 2 .+ y .^ 2) / (2.0 * kernel_sigma^2))) * normal
+  kernel ./= sum(kernel)
+  return kernel
+end
+
+# Function to generate a uniform kernel
+function uniform_kernel(kernel_size)
+  kernel = ones(2 * kernel_size + 1, 2 * kernel_size + 1)
+  kernel = kernel / sum(kernel)
+  return kernel
+end
+
+# Main function to generate H, H_T, W, and W_T functions for gaussian kernel
+function generate_gaussian_blur(shape, shape_p, KERNEL_SIZE, KERNEL_SIGMA = 1.5)
+  (n, m) = shape
+  (n_p, m_p) = shape_p
+  a = n_p - n
+
+  kernel_h = my_gaussian_kernel(KERNEL_SIZE, KERNEL_SIGMA)
+
+  sz = (n_p - (2 * KERNEL_SIZE + 1), m_p - (2 * KERNEL_SIZE + 1))
+  kernel_h = pad(kernel_h, div(sz[1], 2), div(sz[1], 2) + 1)
+  kernel_h = FFTW.ifftshift(kernel_h)
+  fft_h = FFTW.fft(kernel_h)
+
+  # Function H: Applies a linear transformation H which models the blur to an input x
+  function H(x)
+    x = reshape_array(x, (n, m))
+    x = pad_x(x, a)
+    fft_x = FFTW.fft(x)
+    x_new = real(FFTW.ifft(fft_x .* fft_h))
+    return unpad(x_new, n_p, m_p, n)[:]
+  end
+
+  # Function H_T: Applies the transpose of the linear transformation H to an input x
+  function H_T(x)
+    x = reshape_array(x, (n, m))
+    x = pad_x(x, a)
+    fft_x = FFTW.fft(x)
+    x_new = real(FFTW.ifft(fft_x .* conj.(fft_h)))
+    return unpad(x_new, n_p, m_p, n)[:]
+  end
+
+  wt = Wavelets.wavelet(Wavelets.WT.haar)
+
+  # ----- Discrete Wavelet Transform (DWT) -----
+
+  # Function W: Applies the DWT to an input x
+  function W(x)
+    return Wavelets.dwt(reshape_array(x, (n, m)), wt, 4)[:]
+  end
+
+  # Function W_T: Applies the inverse DWT to an input x
+  function W_T(x)
+    return Wavelets.idwt(reshape_array(x, (n, m)), wt, 4)[:]
+  end
+
+  # Return the generated functions
+  return H, H_T, W, W_T
+end
+
+# Main function to generate H, H_T, W, and W_T functions for gaussian kernel
+function generate_uniform_blur(shape, shape_p, KERNEL_SIZE)
+  (n, m) = shape
+  (n_p, m_p) = shape_p
+  a = n_p - n
+
+  kernel_h = uniform_kernel(KERNEL_SIZE)
+
+  sz = (n_p - (2 * KERNEL_SIZE + 1), m_p - (2 * KERNEL_SIZE + 1))
+  kernel_h = pad(kernel_h, div(sz[1], 2), div(sz[1], 2) + 1)
+  kernel_h = FFTW.ifftshift(kernel_h)
+  fft_h = FFTW.fft(kernel_h)
+
+  # Function H: Applies a linear transformation H which models the blur to an input x
+  function H(x)
+    x = reshape_array(x, (n, m))
+    x = pad_x(x, a)
+    fft_x = FFTW.fft(x)
+    x_new = real(FFTW.ifft(fft_x .* fft_h))
+    return unpad(x_new, n_p, m_p, n)[:]
+  end
+
+  # Function H_T: Applies the transpose of the linear transformation H to an input x
+  function H_T(x)
+    x = reshape_array(x, (n, m))
+    x = pad_x(x, a)
+    fft_x = FFTW.fft(x)
+    x_new = real(FFTW.ifft(fft_x .* conj.(fft_h)))
+    return unpad(x_new, n_p, m_p, n)[:]
+  end
+
+  wt = Wavelets.wavelet(Wavelets.WT.haar)
+
+  # ----- Discrete Wavelet Transform (DWT) -----
+
+  # Function W: Applies the DWT to an input x
+  function W(x)
+    return Wavelets.dwt(reshape_array(x, (n, m)), wt, 4)[:]
+  end
+
+  # Function W_T: Applies the inverse DWT to an input x
+  function W_T(x)
+    return Wavelets.idwt(reshape_array(x, (n, m)), wt, 4)[:]
+  end
+
+  # Return the generated functions
+  return H, H_T, W, W_T
+end

src/denoising_model.jl

@@ -0,0 +1,34 @@
+export denoising_model
+
+include("denoising_data.jl")
+
+function denoising_model(shape, shape_p, KERNEL_SIZE, KERNEL_SIGMA = 1.5)
+  sigma = 10^-3
+  data_path = joinpath(@__DIR__, "..", "images/cameraman.png")
+  cameraman_image = Images.load(data_path)
+  x_t = vec(Float64.(cameraman_image))
+  H, H_T, W, W_T = generate_gaussian_blur(shape, shape_p, KERNEL_SIZE, KERNEL_SIGMA)
+  (n, m) = shape
+  b = H(x_t) + randn(n * m) * sigma
+  y = similar(b)
+  z = similar(b)
+
+  function obj(x)
+    y .= W_T(x)
+    y .= H(y)
+    z .= log.((y .- b) .^ 2 .+ 1)
+    return sum(z)
+  end
+
+  function grad!(g, x)
+    y .= H(W_T(x))
+    z .= 1 ./ ((y .- b) .^ 2 .+ 1)
+    @. z = 2 * z * (y - b)
+    g .= W(H_T(z))
+    return g
+  end
+
+  x0 = W(b)
+
+  FirstOrderModel(obj, grad!, x0, name = "denoing_model"), x_t
+end

test/runtests.jl

@@ -1,5 +1,6 @@
 using LinearAlgebra, Test
 using ADNLPModels, DifferentialEquations, ManualNLPModels, MLDatasets, NLPModels, QuadraticModels
+using Images, FFTW, Wavelets
 using RegularizedProblems
 
 function test_well_defined(model, nls_model, sol)
@@ -165,4 +166,17 @@ end
   @test all(model.meta.x0 .== 0)
 end
 
+@testset "denoising_model" begin
+  n, m = 256, 256
+  n_p, m_p = 260, 260
+  kernel_size = 9
+  model, sol = denoising_model((n, m), (n_p, m_p), kernel_size)
+  @test typeof(model) <: FirstOrderModel
+  @test typeof(sol) == typeof(model.meta.x0)
+  @test model.meta.nvar == n * m
+  x = model.meta.x0
+  @test typeof(obj(model, x)) <: Float64
+  @test typeof(grad(model, x)) <: Vector{Float64}
+end
+
 include("rmodel_tests.jl")