remove loops

bug: optimisation yields two identical components

optim_mixture.jl

@@ -1,40 +1,38 @@
 module OptimMixture
 
+import GLMakie as GLM
 import DelimitedFiles as DF
 import BenchmarkTools as BT
 import UnicodePlots as UP
-import LinearAlgebra: ⋅
+import NLopt
+import LinearAlgebra: norm, mul!, axpy!, axpby!
+
+import Logging
+Logging.global_logger(Logging.ConsoleLogger(stderr, Logging.Info))
 
 const ENERGIES = [2803, 2811, 2819, 2826, 2834, 2842, 2850, 2858, 2866, 2874,
   2882, 2890, 2897, 2905, 2913, 2921, 2929, 2937, 2945, 2953,
   2961, 2969, 2977, 2985, 2993, 3001, 3009, 3018, 3026, 3034,
   3042, 3050]
 
-@kwdef struct Slab
+@kwdef struct Slab{T<:Real}
   energy::Int
-  data::Matrix{Int}
+  data::Matrix{T}
 end
 
-function load_slabs(width::Int=512, height::Int=512)
+function load_slabs(; T::Type=Int, width::Int=512, height::Int=512)
-  slabs = Slab[]()
+  slabs = Slab[]
   i_min, i_max = typemax(Int), 0
   for e in ENERGIES
-    slab = Slab(energy=e, data=DF.readdlm("test_sample/HeLa_F-SRS_512x512_2803cm-1.txt", ',', Int))
+    slab::Slab{T} = Slab{T}(energy=e, data=DF.readdlm("test_sample/HeLa_F-SRS_512x512_$(e)cm-1.txt", ',', T))
     push!(slabs, slab)
 
     i_min, i_max = min(i_min, minimum(slab.data)), max(i_max, maximum(slab.data))
 
-    return slabs, i_min, i_max
   end
-end
+  s = join(size(slabs[1].data), "x")
-
+  @info "Loaded $(length(slabs)) $s slabs"
-function Q_matrix_transpose(X::Matrix{Float64}, Y::Matrix{Float64}, N::Int)
+  return slabs, i_min, i_max
-  return sum(abs2, Y - view(X, 1:N, :) * view(X, (N+1):size(X, 1), :)')
-end
-
-function Q_matrix_reshape(X::Matrix{Float64}, Y::Matrix{Float64}, N::Int)
-  N_plus_M, m = size(X)
-  return sum(abs2, Y - view(X, 1:N, :) * reshape(view(X, (N+1):N_plus_M, :), m, N_plus_M - N))
 end
 
 function E_loop(X::Matrix{Float64}, Y::Matrix{Float64})
@@ -53,8 +51,49 @@ function E_loop(X::Matrix{Float64}, Y::Matrix{Float64})
   return s
 end
 
+function E!(X::Matrix{Float64}, Y::Matrix{Float64}, tmp::Matrix{Float64})
+  N, M = size(Y)
+  Q!(tmp, X, N)
+  # roughly 6ms with N = 512*512, M = 32, m = 2
+  # slighly faster than tmp .= tmp .- Y (7ms)
+  # slighly faster than tmp .-= Y (7ms)
+  axpy!(-1.0, Y, tmp)
+  # roughly 2.5m
+  # slightly faster than norm2(tmp)
+  s = sum(x -> x^2, tmp)
+  return s
+end
+
+function test_E(N::Int, M::Int, m::Int; benchmark::Bool=true)
+  X = rand(N + M, m)
+  Y = rand(N, M)
+
+  tmp = zeros(N, M)
+
+  E_from_loop = E_loop(X, Y)
+  E_from_mul = E!(X, Y, tmp)
+
+  @show abs(E_from_mul - E_from_loop)
+
+  if benchmark
+    @info "Running benchmarks..., this may take a while"
+    show(stdout, "text/plain", BT.@benchmark E_loop($X, $Y))
+    println()
+    show(stdout, "text/plain", BT.@benchmark E!($X, $Y, $tmp))
+  end
+end
+
+
+function Q!(dst::Matrix{Float64}, X::Matrix{Float64}, N::Int)
+  N_plus_M, _ = size(X)
+  # roughly 9ms with N = 512*512, M = 32, m = 2
+  # could be sped up using StrideArray?
+  mul!(dst, view(X, 1:N, :), view(X, N+1:N_plus_M, :)')
+end
+
 function Q_loop!(dst::Matrix{Float64}, X::Matrix{Float64}, N::Int)
   N_plus_M, m = size(X)
+  # roughly 23.5ms with N = 512*512, M = 32, m = 2
   for i in 1:N
     for j in N+1:N_plus_M
       x = 0
@@ -66,58 +105,43 @@ function Q_loop!(dst::Matrix{Float64}, X::Matrix{Float64}, N::Int)
   end
 end
 
+function test_Q(N::Int, M::Int, m::Int; benchmark::Bool=true)
+  X = rand(N + M, m)
+
+  Q_from_loop = zeros(N, M)
+  Q_from_mul = zeros(N, M)
+
+  Q!(Q_from_mul, X, N)
+  Q_loop!(Q_from_loop, X, N)
+
+  @show norm(Q_from_mul - Q_from_loop)
+
+  if benchmark
+    @info "Running benchmarks..., this may take a while"
+    show(stdout, "text/plain", BT.@benchmark Q!($Q_from_mul, $X, $N))
+    println()
+    show(stdout, "text/plain", BT.@benchmark Q_loop!($Q_from_loop, $X, $N))
+  end
+end
+
 function DE_loop!(dst::Matrix{Float64}, X::Matrix{Float64}, Y::Matrix{Float64})
   dst_W = zeros(size(Y))
-  DE_loop!(dst, dst_W, X, Y)
+  DE_loop!!(dst, dst_W, X, Y)
 end
 
-function dot_prod_dual_loop(W::Matrix{Float64}, X::Matrix{Float64}, V::Matrix{Float64})
+function DE!!(dst::Matrix{Float64}, dst_W::Matrix{Float64}, X::Matrix{Float64}, Y::Matrix{Float64})
-  _, m = size(X)
+  N_plus_M, m = size(X)
-  N, M = size(W)
+  N, M = size(Y)
-
+  # dst is the same size as X
-  s = 0
+  # we need storage of size Y
-
+  Q!(dst_W, X, N)
-  for k in 1:m
+  # compute W = 2(Q(X) - Y)
-    for i in 1:N
+  axpby!(-2.0, Y, 2.0, dst_W)
-      x = 0
+  mul!(view(dst, 1:N, :), dst_W, view(X, N+1:N_plus_M, :))
-      for j in 1:M
+  mul!(view(dst, N+1:N_plus_M, :), dst_W', view(X, 1:N, :))
-        @inbounds x += W[i, j] * X[N+j, k]
-      end
-      s += x * V[i, k]
-    end
-    for j in 1:M
-      x = 0
-      for i in 1:N
-        @inbounds x += W[i, j] * X[i, k]
-      end
-      s += x * V[N+j, k]
-    end
-  end
-
-  return s
 end
 
-function dot_prod_primal_loop(W::Matrix{Float64}, X::Matrix{Float64}, V::Matrix{Float64})
+function DE_loop!!(dst::Matrix{Float64}, dst_W::Matrix{Float64}, X::Matrix{Float64}, Y::Matrix{Float64})
-  _, m = size(X)
-  N, M = size(W)
-
-  s = 0
-
-  for i in 1:N
-    for j in 1:M
-      x = 0
-      for k in 1:m
-        @inbounds x += X[i, k] * V[N+j, k] + X[N+j, k] * V[i, k]
-
-      end
-      s += x * W[i, j]
-    end
-  end
-
-  return s
-end
-
-function DE_loop!(dst::Matrix{Float64}, dst_W::Matrix{Float64}, X::Matrix{Float64}, Y::Matrix{Float64})
   _, m = size(X)
   N, M = size(Y)
   # dst is the same size as X
@@ -144,116 +168,162 @@ function DE_loop!(dst::Matrix{Float64}, dst_W::Matrix{Float64}, X::Matrix{Float6
   dst .*= 2
 end
 
-function DQ_V_loop!(dst::Matrix{Float64}, X::Matrix{Float64}, V::Matrix{Float64}, N::Int)
+function check_derivative_E(N::Int, M::Int, m::Int; loop::Bool=false)
-  N_plus_M, m = size(X)
-  for i in 1:N
-    for j in N+1:N_plus_M
-      x = 0
-      for k in 1:m
-        @inbounds x += (V[i, k] * X[j, k] + X[i, k] * V[j, k])
-      end
-      @inbounds dst[i, j-N] = x
-    end
-  end
-end
-
-function Q_loop_transposed(X::Matrix{Float64}, Y::Matrix{Float64}, N::Int)
-  s = 0
-  m, N_plus_M = size(X)
-  for i in 1:N
-    for j in N+1:N_plus_M
-      x = 0
-      for k in 1:m
-        @inbounds x += X[k, i] * X[k, j]
-      end
-      @inbounds s += (Y[j-N, i] - x)^2
-    end
-  end
-  return s
-end
-
-function test_dot_prod(N::Int, M::Int, m::Int)
-  X = rand(N + M, m)
-  W = rand(N, M)
-  V = rand(N + M, m) .- 0.5
-
-  @show dot_prod_primal_loop(W, X, V)
-  @show dot_prod_dual_loop(W, X, V)
-end
-
-function test_DE(N::Int, M::Int, m::Int)
   X = rand(N + M, m)
   Y = rand(N, M)
   V = rand(N + M, m) .- 0.5
 
   EX = E_loop(X, Y)
   grad_EX = zeros(size(X))
-  DE_loop!(grad_EX, X, Y)
+  tmp = zeros(size(Y))
+  if loop
+    DE_loop!!(grad_EX, tmp, X, Y)
+  else
+    DE!!(grad_EX, tmp, X, Y)
+  end
 
   eps = [10.0^k for k in 1:-1:-12]
   errors = zeros(size(eps))
 
   for i in eachindex(eps)
-    EX_V_true = E_loop(X + eps[i] * V, Y)
+    if loop
+      EX_V_true = E_loop(X + eps[i] * V, Y)
+    else
+      EX_V_true = E!(X + eps[i] * V, Y, tmp)
+    end
     EX_V_approx = EX + eps[i] * sum(grad_EX .* V)
-    @show EX_V_true
+    # @show EX_V_true
-    @show EX_V_approx
+    # @show EX_V_approx
     errors[i] = abs(EX_V_true - EX_V_approx)
   end
 
-  @show eps
+  # @show eps
-  @show errors
+  # @show errors
 
   foo = UP.lineplot(eps, errors, xscale=:log10, yscale=:log10)
   UP.lineplot!(foo, eps, eps)
   println(foo)
 end
 
-function test_DQ(N::Int, M::Int, m::Int)
+function test_DE(N::Int, M::Int, m::Int; benchmark::Bool=true)
-  X = rand(N + M, m) * 255
+  X = rand(N + M, m)
-  V = rand(N + M, m)
+  Y = rand(N, M)
 
-  QX = zeros(N, M)
+  tmp_W = zeros(N, M)
-  Q_loop!(QX, X, N)
 
-  dst_true = zeros(N, M)
+  ∇E_loop = zeros(N + M, m)
-  dst_approx = zeros(N, M)
+  ∇E = zeros(N + M, m)
 
-  eps = [10.0^k for k in 1:-1:-8]
+  DE_loop!!(∇E_loop, tmp_W, X, Y)
-  errors = zeros(size(eps))
+  DE!!(∇E, tmp_W, X, Y)
 
-  for i in eachindex(eps)
+  @show norm(∇E_loop - ∇E)
-    Q_loop!(dst_true, X + eps[i] * V, N)
+
-    DQ_V_loop!(dst_approx, X, eps[i] * V, N)
+  if benchmark
-    dst_approx .+= QX
+    @info "Running benchmarks..., this may take a while"
-    errors[i] = sum(abs, dst_approx - dst_true)
+    show(stdout, "text/plain", BT.@benchmark DE_loop!!($∇E_loop, $tmp_W, $X, $Y))
+    println()
+    show(stdout, "text/plain", BT.@benchmark DE!!($∇E, $tmp_W, $X, $Y))
+  end
+end
+
+function benchmark()
+  slabs, _, _ = load_slabs(T=Float64)
+  Y::Matrix{Float64} = hcat([vec(slab.data) for slab in slabs]...)
+  @show size(Y)
+
+  N::Int, M::Int = size(Y)
+  m::Int = 2
+
+  X::Vector{Float64} = fill(0.5, (N + M) * m)
+  ∇E::Vector{Float64} = zeros((N + M) * m)
+  tmp = zeros(size(Y))
+
+  show(stdout, "text/plain", BT.@benchmark DE_loop!!(reshape($∇E, $N + $M, $m), $tmp, reshape($X, $N + $M, $m), reshape($Y, $N, $M)))
+  println()
+  show(stdout, "text/plain", BT.@benchmark E_loop(reshape($X, $N + $M, $m), reshape($Y, $N, $M)))
+end
+
+function plot_result_mixture_1(X::Vector{Float64}, N::Int, M::Int, m::Int)
+  X_ = reshape(X, N + M, m)
+  λ = view(X_, 1:N, :)
+  s = view(X_, N+1:N+M, :)
+
+  @show size(λ)
+  @show size(s)
+
+  fig = GLM.Figure()
+  axes = map(1:m+1) do i
+    aspect = if i < m + 1
+      GLM.DataAspect()
+    else
+      nothing
+    end
+    GLM.Axis(fig[1, i], aspect=aspect)
   end
 
-  @show eps
-  @show errors
 
-  println(UP.lineplot(eps, errors, xscale=:log10, yscale=:log10))
+  n = Int(sqrt(N))
+
+  # Component density
+  GLM.image!(axes[1], reshape(view(λ, :, 1), n, n))
+  GLM.image!(axes[2], reshape(view(λ, :, 2), n, n))
+
+  # Spectra
+  GLM.lines!(axes[3], s[:, 1])
+  GLM.lines!(axes[3], s[:, 2])
+
+  @show norm(λ[:, 1] - λ[:, 2])
+  @show norm(s[:, 1] - s[:, 2])
+
+  display(fig)
 end
 
-function test(N::Int, M::Int, m::Int)
+function optim(algo::Symbol=:LD_MMA)
-  Y = rand(N, M) * 255
+  slabs, _, _ = load_slabs(T=Float64)
-  X = rand(N + M, m) * 255
+  Y::Matrix{Float64} = hcat([vec(slab.data) for slab in slabs]...)
+  @show size(Y)
 
-  X_pre = copy(X)
+  N::Int, M::Int = size(Y)
-  X_pre[N+1:N+M] .= vec(X_pre[N+1:N+M]')
+  m::Int = 2
 
-  r_matrix = Q_matrix_transpose(X, Y, N)
+  X::Vector{Float64} = fill(0.5, (N + M) * m)
-  r_loop = Q_loop(X, Y, N)
+  ∇E::Vector{Float64} = zeros((N + M) * m)
+  tmp = zeros(size(Y))
 
-  @show r_matrix, r_loop, abs(r_matrix - r_loop) / r_loop
+  function objective(X::Vector{Float64}, grad::Vector{Float64})
+    if length(grad) > 0
+      DE!!(reshape(grad, N + M, m), tmp, reshape(X, N + M, m), reshape(Y, N, M))
+    else
+      @debug "No gradient requested!"
+    end
+    # @show "|∇E| = $(norm(grad))"
+    max_Σ_λ = maximum(sum(view(reshape(X, N + M, m), 1:N, :); dims=2))
+    value = E!(reshape(X, N + M, m), reshape(Y, N, M), tmp)
+    @info "E = $value, |∇E| = $(norm(grad)), max_Σ_λ = $max_Σ_λ"
+    return value
+  end
+
+  opt = NLopt.Opt(algo, (N + M) * m)
+  NLopt.lower_bounds!(opt, fill(0.0, (N + M) * m))
+  NLopt.upper_bounds!(opt, vec([fill(1.0, N, m); fill(Inf, M, m)]))
+  NLopt.min_objective!(opt, objective)
+
+  res = NLopt.optimize(opt, X)
+
+  return res
 end
 
+
 end # module
 
-# OptimMixture.test(512 * 512, 20, 2)
+res = OptimMixture.optim(:LD_LBFGS);
-# OptimMixture.test_DQ(512 * 512, 20, 2)
+OptimMixture.plot_result_mixture_1(res[2], 512 * 512, 32, 2)
-OptimMixture.test_DE(512 * 512, 20, 2)
+# OptimMixture.benchmark();
-# OptimMixture.test_DE(20, 2, 2)
+# OptimMixture.test_E(512 * 512, 32, 2)
-# OptimMixture.test_dot_prod(512^2, 20, 2)
+# OptimMixture.test_DE(512 * 512, 32, 2)
+# OptimMixture.test_Q(512 * 512, 32, 2)
+# OptimMixture.check_derivative_E(512 * 512, 32, 2)
+# OptimMixture.check_derivative_E(512 * 512, 32, 2; loop=true)
 
 # vim: ts=2:sw=2:sts=2