Use Newton's method for polishing

engine-proto/gram-test/Engine.jl

@@ -51,11 +51,21 @@ end
 # the Lorentz form
 Q = diagm([1, 1, 1, 1, -1])
 
+# 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)...)
+  result = zeros(size(target))
   for (j, k, val) in zip(J, K, values)
     result[j, k] = val - attempt[j, k]
   end
@@ -87,7 +97,7 @@ 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(
+function realize_gram_gradient(
   gram::SparseMatrixCSC{T, <:Any},
   guess::Matrix{T};
   scaled_tol = 1e-30,
@@ -111,7 +121,7 @@ function realize_gram(
   # do gradient descent
   Δ_proj = proj_diff(gram, L'*Q*L)
   loss = dot(Δ_proj, Δ_proj)
-  for step in 1:max_descent_steps
+  for _ in 1:max_descent_steps
     # stop if the loss is tolerably low
     if loss < tol
       break
@@ -160,4 +170,73 @@ function realize_gram(
   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
+    
+    # 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
+
 end