From 1ce609836bf45c16131bdec879c6c3fb613385af Mon Sep 17 00:00:00 2001
From: Aaron Fenyes <aaron.fenyes@fareycircles.ooo>
Date: Mon, 15 Jul 2024 22:11:54 -0700
Subject: [PATCH] Implement frozen variables

---
 engine-proto/gram-test/Engine.jl              | 30 +++++++++++++++++--
 engine-proto/gram-test/circles-in-triangle.jl |  3 +-
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/engine-proto/gram-test/Engine.jl b/engine-proto/gram-test/Engine.jl
index 10f4b02..dccb514 100644
--- a/engine-proto/gram-test/Engine.jl
+++ b/engine-proto/gram-test/Engine.jl
@@ -311,7 +311,8 @@ end
 # explicit entry of `gram`. use gradient descent starting from `guess`
 function realize_gram(
   gram::SparseMatrixCSC{T, <:Any},
-  guess::Matrix{T};
+  guess::Matrix{T},
+  frozen = nothing;
   scaled_tol = 1e-30,
   min_efficiency = 0.5,
   init_rate = 1.0,
@@ -336,6 +337,15 @@ function realize_gram(
   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)
@@ -371,7 +381,23 @@ function realize_gram(
     if min_eigval <= 0
       hess -= reg_scale * min_eigval * I
     end
-    base_step = reshape(hess \ reshape(neg_grad, total_dim), dims)
+    
+    # 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
diff --git a/engine-proto/gram-test/circles-in-triangle.jl b/engine-proto/gram-test/circles-in-triangle.jl
index 12975c2..b173ed9 100644
--- a/engine-proto/gram-test/circles-in-triangle.jl
+++ b/engine-proto/gram-test/circles-in-triangle.jl
@@ -66,6 +66,7 @@ guess = hcat(
   Engine.sphere(BigFloat[cos(pi/3), sin(pi/3), 0], BigFloat(1//5)),
   BigFloat[0, 0, 0, 1, 1]
 )
+frozen = [CartesianIndex(j, 9) for j in 4:5]
 #=
 guess = hcat(
   Engine.plane(BigFloat[0, 0, 1], BigFloat(0)),
@@ -86,7 +87,7 @@ L, history = Engine.realize_gram_gradient(gram, guess, scaled_tol = 0.01)
 L_pol, history_pol = Engine.realize_gram_newton(gram, L, rate = 0.3, scaled_tol = 1e-9)
 L_pol2, history_pol2 = Engine.realize_gram_newton(gram, L_pol)
 =#
-L, success, history = Engine.realize_gram(gram, guess)
+L, success, history = Engine.realize_gram(gram, guess, frozen)
 completed_gram = L'*Engine.Q*L
 println("Completed Gram matrix:\n")
 display(completed_gram)