451 lines
13 KiB
Julia
451 lines
13 KiB
Julia
module Engine
|
|
|
|
using LinearAlgebra
|
|
using GenericLinearAlgebra
|
|
using SparseArrays
|
|
using Random
|
|
using Optim
|
|
|
|
export
|
|
rand_on_shell, Q, DescentHistory,
|
|
realize_gram_gradient, realize_gram_newton, realize_gram_optim, realize_gram
|
|
|
|
# === guessing ===
|
|
|
|
sconh(t, u) = 0.5*(exp(t) + u*exp(-t))
|
|
|
|
function rand_on_sphere(rng::AbstractRNG, ::Type{T}, n) where T
|
|
out = randn(rng, T, n)
|
|
tries_left = 2
|
|
while dot(out, out) < 1e-6 && tries_left > 0
|
|
out = randn(rng, T, n)
|
|
tries_left -= 1
|
|
end
|
|
normalize(out)
|
|
end
|
|
|
|
##[TO DO] write a test to confirm that the outputs are on the correct shells
|
|
function rand_on_shell(rng::AbstractRNG, shell::T) where T <: Number
|
|
space_part = rand_on_sphere(rng, T, 4)
|
|
rapidity = randn(rng, T)
|
|
sig = sign(shell)
|
|
nullmix * [sconh(rapidity, sig)*space_part; sconh(rapidity, -sig)]
|
|
end
|
|
|
|
rand_on_shell(rng::AbstractRNG, shells::Array{T}) where T <: Number =
|
|
hcat([rand_on_shell(rng, sh) for sh in shells]...)
|
|
|
|
rand_on_shell(shells::Array{<:Number}) = rand_on_shell(Random.default_rng(), shells)
|
|
|
|
# === elements ===
|
|
|
|
point(pos) = [pos; 0.5; 0.5 * dot(pos, pos)]
|
|
|
|
plane(normal, offset) = [-normal; 0; -offset]
|
|
|
|
function sphere(center, radius)
|
|
dist_sq = dot(center, center)
|
|
[
|
|
center / radius;
|
|
0.5 / radius;
|
|
0.5 * (dist_sq / radius - radius)
|
|
]
|
|
end
|
|
|
|
# === Gram matrix realization ===
|
|
|
|
# basis changes
|
|
nullmix = [Matrix{Int64}(I, 3, 3) zeros(Int64, 3, 2); zeros(Int64, 2, 3) [-1 1; 1 1]//2]
|
|
unmix = [Matrix{Int64}(I, 3, 3) zeros(Int64, 3, 2); zeros(Int64, 2, 3) [-1 1; 1 1]]
|
|
|
|
# the Lorentz form
|
|
## [old] Q = diagm([1, 1, 1, 1, -1])
|
|
Q = [Matrix{Int64}(I, 3, 3) zeros(Int64, 3, 2); zeros(Int64, 2, 3) [0 -2; -2 0]]
|
|
|
|
# project a matrix onto the subspace of matrices whose entries vanish at the
|
|
# given indices
|
|
function proj_to_entries(mat, indices)
|
|
result = zeros(size(mat))
|
|
for (j, k) in indices
|
|
result[j, k] = mat[j, k]
|
|
end
|
|
result
|
|
end
|
|
|
|
# the difference between the matrices `target` and `attempt`, projected onto the
|
|
# subspace of matrices whose entries vanish at each empty index of `target`
|
|
function proj_diff(target::SparseMatrixCSC{T, <:Any}, attempt::Matrix{T}) where T
|
|
J, K, values = findnz(target)
|
|
result = zeros(size(target))
|
|
for (j, k, val) in zip(J, K, values)
|
|
result[j, k] = val - attempt[j, k]
|
|
end
|
|
result
|
|
end
|
|
|
|
# a type for keeping track of gradient descent history
|
|
struct DescentHistory{T}
|
|
scaled_loss::Array{T}
|
|
neg_grad::Array{Matrix{T}}
|
|
base_step::Array{Matrix{T}}
|
|
hess::Array{Hermitian{T, Matrix{T}}}
|
|
slope::Array{T}
|
|
stepsize::Array{T}
|
|
positive::Array{Bool}
|
|
backoff_steps::Array{Int64}
|
|
last_line_L::Array{Matrix{T}}
|
|
last_line_loss::Array{T}
|
|
|
|
function DescentHistory{T}(
|
|
scaled_loss = Array{T}(undef, 0),
|
|
neg_grad = Array{Matrix{T}}(undef, 0),
|
|
hess = Array{Hermitian{T, Matrix{T}}}(undef, 0),
|
|
base_step = Array{Matrix{T}}(undef, 0),
|
|
slope = Array{T}(undef, 0),
|
|
stepsize = Array{T}(undef, 0),
|
|
positive = Bool[],
|
|
backoff_steps = Int64[],
|
|
last_line_L = Array{Matrix{T}}(undef, 0),
|
|
last_line_loss = Array{T}(undef, 0)
|
|
) where T
|
|
new(scaled_loss, neg_grad, hess, base_step, slope, stepsize, positive, backoff_steps, last_line_L, last_line_loss)
|
|
end
|
|
end
|
|
|
|
# seek a matrix `L` for which `L'QL` matches the sparse matrix `gram` at every
|
|
# explicit entry of `gram`. use gradient descent starting from `guess`
|
|
function realize_gram_gradient(
|
|
gram::SparseMatrixCSC{T, <:Any},
|
|
guess::Matrix{T};
|
|
scaled_tol = 1e-30,
|
|
min_efficiency = 0.5,
|
|
init_stepsize = 1.0,
|
|
backoff = 0.9,
|
|
max_descent_steps = 600,
|
|
max_backoff_steps = 110
|
|
) where T <: Number
|
|
# start history
|
|
history = DescentHistory{T}()
|
|
|
|
# scale tolerance
|
|
scale_adjustment = sqrt(T(nnz(gram)))
|
|
tol = scale_adjustment * scaled_tol
|
|
|
|
# initialize variables
|
|
stepsize = init_stepsize
|
|
L = copy(guess)
|
|
|
|
# do gradient descent
|
|
Δ_proj = proj_diff(gram, L'*Q*L)
|
|
loss = dot(Δ_proj, Δ_proj)
|
|
for _ in 1:max_descent_steps
|
|
# stop if the loss is tolerably low
|
|
if loss < tol
|
|
break
|
|
end
|
|
|
|
# find negative gradient of loss function
|
|
neg_grad = 4*Q*L*Δ_proj
|
|
slope = norm(neg_grad)
|
|
dir = neg_grad / slope
|
|
|
|
# store current position, loss, and slope
|
|
L_last = L
|
|
loss_last = loss
|
|
push!(history.scaled_loss, loss / scale_adjustment)
|
|
push!(history.neg_grad, neg_grad)
|
|
push!(history.slope, slope)
|
|
|
|
# find a good step size using backtracking line search
|
|
push!(history.stepsize, 0)
|
|
push!(history.backoff_steps, max_backoff_steps)
|
|
empty!(history.last_line_L)
|
|
empty!(history.last_line_loss)
|
|
for backoff_steps in 0:max_backoff_steps
|
|
history.stepsize[end] = stepsize
|
|
L = L_last + stepsize * dir
|
|
Δ_proj = proj_diff(gram, L'*Q*L)
|
|
loss = dot(Δ_proj, Δ_proj)
|
|
improvement = loss_last - loss
|
|
push!(history.last_line_L, L)
|
|
push!(history.last_line_loss, loss / scale_adjustment)
|
|
if improvement >= min_efficiency * stepsize * slope
|
|
history.backoff_steps[end] = backoff_steps
|
|
break
|
|
end
|
|
stepsize *= backoff
|
|
end
|
|
|
|
# [DEBUG] if we've hit a wall, quit
|
|
if history.backoff_steps[end] == max_backoff_steps
|
|
break
|
|
end
|
|
end
|
|
|
|
# return the factorization and its history
|
|
push!(history.scaled_loss, loss / scale_adjustment)
|
|
L, history
|
|
end
|
|
|
|
function basis_matrix(::Type{T}, j, k, dims) where T
|
|
result = zeros(T, dims)
|
|
result[j, k] = one(T)
|
|
result
|
|
end
|
|
|
|
# seek a matrix `L` for which `L'QL` matches the sparse matrix `gram` at every
|
|
# explicit entry of `gram`. use Newton's method starting from `guess`
|
|
function realize_gram_newton(
|
|
gram::SparseMatrixCSC{T, <:Any},
|
|
guess::Matrix{T};
|
|
scaled_tol = 1e-30,
|
|
rate = 1,
|
|
max_steps = 100
|
|
) where T <: Number
|
|
# start history
|
|
history = DescentHistory{T}()
|
|
|
|
# find the dimension of the search space
|
|
dims = size(guess)
|
|
element_dim, construction_dim = dims
|
|
total_dim = element_dim * construction_dim
|
|
|
|
# list the constrained entries of the gram matrix
|
|
J, K, _ = findnz(gram)
|
|
constrained = zip(J, K)
|
|
|
|
# scale the tolerance
|
|
scale_adjustment = sqrt(T(length(constrained)))
|
|
tol = scale_adjustment * scaled_tol
|
|
|
|
# use Newton's method
|
|
L = copy(guess)
|
|
for step in 0:max_steps
|
|
# evaluate the loss function
|
|
Δ_proj = proj_diff(gram, L'*Q*L)
|
|
loss = dot(Δ_proj, Δ_proj)
|
|
|
|
# store the current loss
|
|
push!(history.scaled_loss, loss / scale_adjustment)
|
|
|
|
# stop if the loss is tolerably low
|
|
if loss < tol || step > max_steps
|
|
break
|
|
end
|
|
|
|
# find the negative gradient of loss function
|
|
neg_grad = 4*Q*L*Δ_proj
|
|
|
|
# find the negative Hessian of the loss function
|
|
hess = Matrix{T}(undef, total_dim, total_dim)
|
|
indices = [(j, k) for k in 1:construction_dim for j in 1:element_dim]
|
|
for (j, k) in indices
|
|
basis_mat = basis_matrix(T, j, k, dims)
|
|
neg_dΔ = basis_mat'*Q*L + L'*Q*basis_mat
|
|
neg_dΔ_proj = proj_to_entries(neg_dΔ, constrained)
|
|
deriv_grad = 4*Q*(-basis_mat*Δ_proj + L*neg_dΔ_proj)
|
|
hess[:, (k-1)*element_dim + j] = reshape(deriv_grad, total_dim)
|
|
end
|
|
hess = Hermitian(hess)
|
|
push!(history.hess, hess)
|
|
|
|
# compute the Newton step
|
|
step = hess \ reshape(neg_grad, total_dim)
|
|
L += rate * reshape(step, dims)
|
|
end
|
|
|
|
# return the factorization and its history
|
|
L, history
|
|
end
|
|
|
|
LinearAlgebra.eigen!(A::Symmetric{BigFloat, Matrix{BigFloat}}; sortby::Nothing) =
|
|
eigen!(Hermitian(A))
|
|
|
|
function convertnz(type, mat)
|
|
J, K, values = findnz(mat)
|
|
sparse(J, K, type.(values))
|
|
end
|
|
|
|
function realize_gram_optim(
|
|
gram::SparseMatrixCSC{T, <:Any},
|
|
guess::Matrix{T}
|
|
) where T <: Number
|
|
# find the dimension of the search space
|
|
dims = size(guess)
|
|
element_dim, construction_dim = dims
|
|
total_dim = element_dim * construction_dim
|
|
|
|
# list the constrained entries of the gram matrix
|
|
J, K, _ = findnz(gram)
|
|
constrained = zip(J, K)
|
|
|
|
# scale the loss function
|
|
scale_adjustment = length(constrained)
|
|
|
|
function loss(L_vec)
|
|
L = reshape(L_vec, dims)
|
|
Δ_proj = proj_diff(gram, L'*Q*L)
|
|
dot(Δ_proj, Δ_proj) / scale_adjustment
|
|
end
|
|
|
|
function loss_grad!(storage, L_vec)
|
|
L = reshape(L_vec, dims)
|
|
Δ_proj = proj_diff(gram, L'*Q*L)
|
|
storage .= reshape(-4*Q*L*Δ_proj, total_dim) / scale_adjustment
|
|
end
|
|
|
|
function loss_hess!(storage, L_vec)
|
|
L = reshape(L_vec, dims)
|
|
Δ_proj = proj_diff(gram, L'*Q*L)
|
|
indices = [(j, k) for k in 1:construction_dim for j in 1:element_dim]
|
|
for (j, k) in indices
|
|
basis_mat = basis_matrix(T, j, k, dims)
|
|
neg_dΔ = basis_mat'*Q*L + L'*Q*basis_mat
|
|
neg_dΔ_proj = proj_to_entries(neg_dΔ, constrained)
|
|
deriv_grad = 4*Q*(-basis_mat*Δ_proj + L*neg_dΔ_proj) / scale_adjustment
|
|
storage[:, (k-1)*element_dim + j] = reshape(deriv_grad, total_dim)
|
|
end
|
|
end
|
|
|
|
optimize(
|
|
loss, loss_grad!, loss_hess!,
|
|
reshape(guess, total_dim),
|
|
Newton()
|
|
)
|
|
end
|
|
|
|
# seek a matrix `L` for which `L'QL` matches the sparse matrix `gram` at every
|
|
# explicit entry of `gram`. use gradient descent starting from `guess`
|
|
function realize_gram(
|
|
gram::SparseMatrixCSC{T, <:Any},
|
|
guess::Matrix{T},
|
|
frozen = nothing;
|
|
scaled_tol = 1e-30,
|
|
min_efficiency = 0.5,
|
|
init_rate = 1.0,
|
|
backoff = 0.9,
|
|
reg_scale = 1.1,
|
|
max_descent_steps = 200,
|
|
max_backoff_steps = 110
|
|
) where T <: Number
|
|
# start history
|
|
history = DescentHistory{T}()
|
|
|
|
# find the dimension of the search space
|
|
dims = size(guess)
|
|
element_dim, construction_dim = dims
|
|
total_dim = element_dim * construction_dim
|
|
|
|
# list the constrained entries of the gram matrix
|
|
J, K, _ = findnz(gram)
|
|
constrained = zip(J, K)
|
|
|
|
# scale the tolerance
|
|
scale_adjustment = sqrt(T(length(constrained)))
|
|
tol = scale_adjustment * scaled_tol
|
|
|
|
# list the un-frozen indices
|
|
has_frozen = !isnothing(frozen)
|
|
if has_frozen
|
|
is_unfrozen = fill(true, size(guess))
|
|
is_unfrozen[frozen] .= false
|
|
unfrozen = findall(is_unfrozen)
|
|
unfrozen_stacked = reshape(is_unfrozen, total_dim)
|
|
end
|
|
|
|
# initialize variables
|
|
grad_rate = init_rate
|
|
L = copy(guess)
|
|
|
|
# use Newton's method with backtracking and gradient descent backup
|
|
Δ_proj = proj_diff(gram, L'*Q*L)
|
|
loss = dot(Δ_proj, Δ_proj)
|
|
for step in 1:max_descent_steps
|
|
# stop if the loss is tolerably low
|
|
if loss < tol
|
|
break
|
|
end
|
|
|
|
# find the negative gradient of loss function
|
|
neg_grad = 4*Q*L*Δ_proj
|
|
|
|
# find the negative Hessian of the loss function
|
|
hess = Matrix{T}(undef, total_dim, total_dim)
|
|
indices = [(j, k) for k in 1:construction_dim for j in 1:element_dim]
|
|
for (j, k) in indices
|
|
basis_mat = basis_matrix(T, j, k, dims)
|
|
neg_dΔ = basis_mat'*Q*L + L'*Q*basis_mat
|
|
neg_dΔ_proj = proj_to_entries(neg_dΔ, constrained)
|
|
deriv_grad = 4*Q*(-basis_mat*Δ_proj + L*neg_dΔ_proj)
|
|
hess[:, (k-1)*element_dim + j] = reshape(deriv_grad, total_dim)
|
|
end
|
|
hess = Hermitian(hess)
|
|
push!(history.hess, hess)
|
|
|
|
# regularize the Hessian
|
|
min_eigval = minimum(eigvals(hess))
|
|
push!(history.positive, min_eigval > 0)
|
|
if min_eigval <= 0
|
|
hess -= reg_scale * min_eigval * I
|
|
end
|
|
|
|
# compute the Newton step
|
|
neg_grad_stacked = reshape(neg_grad, total_dim)
|
|
if has_frozen
|
|
hess = hess[unfrozen_stacked, unfrozen_stacked]
|
|
neg_grad_compressed = neg_grad_stacked[unfrozen_stacked]
|
|
else
|
|
neg_grad_compressed = neg_grad_stacked
|
|
end
|
|
base_step_compressed = hess \ neg_grad_compressed
|
|
if has_frozen
|
|
base_step_stacked = zeros(total_dim)
|
|
base_step_stacked[unfrozen_stacked] .= base_step_compressed
|
|
else
|
|
base_step_stacked = base_step_compressed
|
|
end
|
|
base_step = reshape(base_step_stacked, dims)
|
|
push!(history.base_step, base_step)
|
|
|
|
# store the current position, loss, and slope
|
|
L_last = L
|
|
loss_last = loss
|
|
push!(history.scaled_loss, loss / scale_adjustment)
|
|
push!(history.neg_grad, neg_grad)
|
|
push!(history.slope, norm(neg_grad))
|
|
|
|
# find a good step size using backtracking line search
|
|
push!(history.stepsize, 0)
|
|
push!(history.backoff_steps, max_backoff_steps)
|
|
empty!(history.last_line_L)
|
|
empty!(history.last_line_loss)
|
|
rate = one(T)
|
|
step_success = false
|
|
for backoff_steps in 0:max_backoff_steps
|
|
history.stepsize[end] = rate
|
|
L = L_last + rate * base_step
|
|
Δ_proj = proj_diff(gram, L'*Q*L)
|
|
loss = dot(Δ_proj, Δ_proj)
|
|
improvement = loss_last - loss
|
|
push!(history.last_line_L, L)
|
|
push!(history.last_line_loss, loss / scale_adjustment)
|
|
if improvement >= min_efficiency * rate * dot(neg_grad, base_step)
|
|
history.backoff_steps[end] = backoff_steps
|
|
step_success = true
|
|
break
|
|
end
|
|
rate *= backoff
|
|
end
|
|
|
|
# if we've hit a wall, quit
|
|
if !step_success
|
|
return L_last, false, history
|
|
end
|
|
end
|
|
|
|
# return the factorization and its history
|
|
push!(history.scaled_loss, loss / scale_adjustment)
|
|
L, true, history
|
|
end
|
|
|
|
end |