Implement frozen variables

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