remove loops
bug: optimisation yields two identical components
This commit is contained in:
parent
b8af1fa23c
commit
11f955d265
1 changed files with 206 additions and 136 deletions
342
optim_mixture.jl
342
optim_mixture.jl
|
|
@ -1,41 +1,39 @@
|
|||
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)
|
||||
slabs = Slab[]()
|
||||
function load_slabs(; T::Type=Int, width::Int=512, height::Int=512)
|
||||
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))
|
||||
|
||||
end
|
||||
s = join(size(slabs[1].data), "x")
|
||||
@info "Loaded $(length(slabs)) $s slabs"
|
||||
return slabs, i_min, i_max
|
||||
end
|
||||
end
|
||||
|
||||
function Q_matrix_transpose(X::Matrix{Float64}, Y::Matrix{Float64}, N::Int)
|
||||
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})
|
||||
s = 0
|
||||
|
|
@ -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})
|
||||
_, m = size(X)
|
||||
N, M = size(W)
|
||||
|
||||
s = 0
|
||||
|
||||
for k in 1:m
|
||||
for i in 1:N
|
||||
x = 0
|
||||
for j in 1:M
|
||||
@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
|
||||
function DE!!(dst::Matrix{Float64}, dst_W::Matrix{Float64}, X::Matrix{Float64}, Y::Matrix{Float64})
|
||||
N_plus_M, m = size(X)
|
||||
N, M = size(Y)
|
||||
# dst is the same size as X
|
||||
# we need storage of size Y
|
||||
Q!(dst_W, X, N)
|
||||
# compute W = 2(Q(X) - Y)
|
||||
axpby!(-2.0, Y, 2.0, dst_W)
|
||||
mul!(view(dst, 1:N, :), dst_W, view(X, N+1:N_plus_M, :))
|
||||
mul!(view(dst, N+1:N_plus_M, :), dst_W', view(X, 1:N, :))
|
||||
end
|
||||
|
||||
return s
|
||||
end
|
||||
|
||||
function dot_prod_primal_loop(W::Matrix{Float64}, X::Matrix{Float64}, V::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})
|
||||
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)
|
||||
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)
|
||||
function check_derivative_E(N::Int, M::Int, m::Int; loop::Bool=false)
|
||||
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)
|
||||
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_approx
|
||||
# @show EX_V_true
|
||||
# @show EX_V_approx
|
||||
errors[i] = abs(EX_V_true - EX_V_approx)
|
||||
end
|
||||
|
||||
@show eps
|
||||
@show errors
|
||||
# @show eps
|
||||
# @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)
|
||||
X = rand(N + M, m) * 255
|
||||
V = rand(N + M, m)
|
||||
function test_DE(N::Int, M::Int, m::Int; benchmark::Bool=true)
|
||||
X = rand(N + M, m)
|
||||
Y = rand(N, M)
|
||||
|
||||
QX = zeros(N, M)
|
||||
Q_loop!(QX, X, N)
|
||||
tmp_W = zeros(N, M)
|
||||
|
||||
dst_true = zeros(N, M)
|
||||
dst_approx = zeros(N, M)
|
||||
∇E_loop = zeros(N + M, m)
|
||||
∇E = zeros(N + M, m)
|
||||
|
||||
eps = [10.0^k for k in 1:-1:-8]
|
||||
errors = zeros(size(eps))
|
||||
DE_loop!!(∇E_loop, tmp_W, X, Y)
|
||||
DE!!(∇E, tmp_W, X, Y)
|
||||
|
||||
for i in eachindex(eps)
|
||||
Q_loop!(dst_true, X + eps[i] * V, N)
|
||||
DQ_V_loop!(dst_approx, X, eps[i] * V, N)
|
||||
dst_approx .+= QX
|
||||
errors[i] = sum(abs, dst_approx - dst_true)
|
||||
@show norm(∇E_loop - ∇E)
|
||||
|
||||
if benchmark
|
||||
@info "Running benchmarks..., this may take a while"
|
||||
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
|
||||
|
||||
@show eps
|
||||
@show errors
|
||||
function benchmark()
|
||||
slabs, _, _ = load_slabs(T=Float64)
|
||||
Y::Matrix{Float64} = hcat([vec(slab.data) for slab in slabs]...)
|
||||
@show size(Y)
|
||||
|
||||
println(UP.lineplot(eps, errors, xscale=:log10, yscale=:log10))
|
||||
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 test(N::Int, M::Int, m::Int)
|
||||
Y = rand(N, M) * 255
|
||||
X = rand(N + M, m) * 255
|
||||
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, :)
|
||||
|
||||
X_pre = copy(X)
|
||||
X_pre[N+1:N+M] .= vec(X_pre[N+1:N+M]')
|
||||
@show size(λ)
|
||||
@show size(s)
|
||||
|
||||
r_matrix = Q_matrix_transpose(X, Y, N)
|
||||
r_loop = Q_loop(X, Y, N)
|
||||
|
||||
@show r_matrix, r_loop, abs(r_matrix - r_loop) / r_loop
|
||||
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
|
||||
|
||||
|
||||
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 optim(algo::Symbol=:LD_MMA)
|
||||
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))
|
||||
|
||||
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)
|
||||
# OptimMixture.test_DQ(512 * 512, 20, 2)
|
||||
OptimMixture.test_DE(512 * 512, 20, 2)
|
||||
# OptimMixture.test_DE(20, 2, 2)
|
||||
# OptimMixture.test_dot_prod(512^2, 20, 2)
|
||||
res = OptimMixture.optim(:LD_LBFGS);
|
||||
OptimMixture.plot_result_mixture_1(res[2], 512 * 512, 32, 2)
|
||||
# OptimMixture.benchmark();
|
||||
# OptimMixture.test_E(512 * 512, 32, 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
|
||||
|
|
|
|||
Loading…
Reference in a new issue