dyna3/engine-proto/gram-test/Engine.jl
Aaron Fenyes 74c7f64b0c Correct sign of normal in plane utility
Clarify the relevant notes too.
2024-07-18 00:03:12 -07:00

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