From f649f24053a15d46a5162ddbdc1e24464bdf314a Mon Sep 17 00:00:00 2001
From: Logan Mondal Bhamidipaty
 <76822456+FlyingWorkshop@users.noreply.github.com>
Date: Fri, 3 May 2024 19:33:35 -0700
Subject: [PATCH] shape changes

---
 src/explicit.jl |  6 +++---
 src/implicit.jl |  6 +++---
 src/utils.jl    | 17 ++++++++---------
 3 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/src/explicit.jl b/src/explicit.jl
index df52cdd..b7cb17f 100644
--- a/src/explicit.jl
+++ b/src/explicit.jl
@@ -60,13 +60,13 @@ function fit!(
     X;
     mu=epca.mu,
     maxoutdim=1, 
-    maxiter=1000,
+    maxiter=10,
     verbose=false,
     steps_per_print=10,
     epsilon=eps(),
 )
     L = _make_loss(epca, X, epsilon, mu)
-    A =  _fit!(epca, X, maxoutdim, L, verbose, steps_per_print, maxiter)
+    A =  _fit!(epca, X, maxoutdim, L, maxiter, verbose, steps_per_print)
     return A
 end
 
@@ -75,7 +75,7 @@ function compress(
     epca::ExplicitEPCA, 
     X;
     mu=epca.mu,
-    maxiter=100,
+    maxiter=10,
     verbose=false,
     steps_per_print=10,
     epsilon=eps()
diff --git a/src/implicit.jl b/src/implicit.jl
index 8dd7021..9f3f17a 100644
--- a/src/implicit.jl
+++ b/src/implicit.jl
@@ -68,14 +68,14 @@ function fit!(
     X,
     mu;
     maxoutdim=1, 
-    maxiter=100,
+    maxiter=10,
     verbose=false,
     steps_per_print=10,
     epsilon=eps(),
     tol=eps()
 )
     L = _make_loss(epca, X, mu, epsilon; tol=tol)
-    A =  _fit!(epca, X, maxoutdim, L, verbose, steps_per_print, maxiter)
+    A =  _fit!(epca, X, maxoutdim, L, maxiter, verbose, steps_per_print)
     return A
 end
 
@@ -83,7 +83,7 @@ function compress(
     epca::ImplicitEPCA, 
     X;
     mu=1,  # NOTE: mu = 1 may not be valid for all link functions. 
-    maxiter=100,
+    maxiter=10,
     verbose=false,
     steps_per_print=10,
     epsilon=eps(),
diff --git a/src/utils.jl b/src/utils.jl
index 45452e3..fb5790d 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -1,10 +1,10 @@
-function _single_fit_iter(L::Function, V, A, verbose, i, steps_per_print)
+function _single_fit_iter(L::Function, V, A, verbose::Bool, i::Integer, steps_per_print::Integer, maxiter::Integer)
     V = Optim.minimizer(optimize(V_hat->L(A * V_hat), V))
-    A = _single_compress_iter(L, V, A, verbose, i, steps_per_print)
+    A = _single_compress_iter(L, V, A, verbose, i, steps_per_print, maxiter)
     return V, A
 end
 
-function _single_compress_iter(L::Function, V, A, verbose, i, steps_per_print)
+function _single_compress_iter(L::Function, V, A, verbose::Bool, i::Integer, steps_per_print::Integer, maxiter::Integer)
     result = optimize(A_hat->L(A_hat * V), A)
     A = Optim.minimizer(result)
     if verbose && (i % steps_per_print == 0 || i == 1)
@@ -14,24 +14,23 @@ function _single_compress_iter(L::Function, V, A, verbose, i, steps_per_print)
     return A
 end
 
-function _fit!(epca::EPCA, X, maxoutdim, L, verbose, steps_per_print, maxiter)
+function _fit!(epca::EPCA, X, maxoutdim::Integer, L::Function, maxiter::Integer, verbose::Bool, steps_per_print::Integer)
     n, d = size(X)
     A = ones(n, maxoutdim)
     V = ismissing(epca.V) ? ones(maxoutdim, d) : epca.V
     for i in 1:maxiter
-        V, A = _single_fit_iter(L, V, A, verbose, i, steps_per_print)
+        V, A = _single_fit_iter(L, V, A, verbose, i, steps_per_print, maxiter)
     end
     epca.V = V
     return A
 end
 
-function _compress(epca::EPCA, X, L, maxiter, verbose, steps_per_print)
-    V = epca.V
+function _compress(epca::EPCA, X, L::Function, maxiter::Integer, verbose::Bool, steps_per_print::Integer)
     n, _ = size(X)
-    outdim = size(V)[1]
+    outdim = size(epca.V)[1]
     A = ones(n, outdim)
     for i in 1:maxiter
-        _single_compress_iter(L, V, A, verbose, i, steps_per_print)
+        A = _single_compress_iter(L, epca.V, A, verbose, i, steps_per_print, maxiter)
     end
     return A
 end