diff --git a/EC.jl b/EC.jl
index adfd866..d93367f 100644
--- a/EC.jl
+++ b/EC.jl
@@ -29,12 +29,8 @@ end
 "Train an EC model for a given range of c values.
 If a list is provided for ref_eval, they are used as reference values for picking the closest eigenvalues at each sampling point.
 If a single number is provided for ref_eval, it is used as a reference for the first point, and the previous eigenvalue is used as the reference for each successive point.
-If pseudo_inv_rtol > 0, the GEVP is avoided using Moore-Penrose psuedoinverse, using this value as the relative tolerance for dropping redundant vectors.
-If orthonormalize_threshold > 0, Gram-Schmidt orthonormalization is performed, using this value as the threshold for dropping redundant vectors.
-If both are = 0, GEVP is solved instead."
-function train!(EC::affine_EC, c_vals; ref_eval=-10.0, pseudo_inv_rtol=1e-6, gram_schmidt_threshold=0, CAEC=false, verbose=true, tol=1e-5)
-    @assert pseudo_inv_rtol == 0 || gram_schmidt_threshold == 0 "Cannot perform both psuedoinverse and Gram-Schmidt"
-
+If orthonormalize_threshold > 0, Gram-Schmidt orthonormalization is performed, using this value as the threshold for dropping redundant vectors."
+function train!(EC::affine_EC, c_vals; ref_eval=-10.0, CAEC=false, gram_schmidt_threshold=0, verbose=true, tol=1e-5)
     training_vecs = Vector{ComplexF64}[]
 
     for c in c_vals
@@ -58,66 +54,41 @@ function train!(EC::affine_EC, c_vals; ref_eval=-10.0, pseudo_inv_rtol=1e-6, gra
     end
 
     CAEC && append!(training_vecs, conj.(training_vecs))
-    weights_mat = spdiagm(EC.weights)
 
-    if gram_schmidt_threshold > 0
-        EC.ensemble_size > 0 && generate_resampled_EC_matrices(EC, training_vecs, weights_mat, gram_schmidt_threshold)
-        training_vecs = gram_schmidt!(training_vecs, EC.weights, gram_schmidt_threshold; verbose=verbose)
-    end
+    (EC.H0_EC, EC.H1_EC, EC.N_EC) = get_reduced_matrices(EC, training_vecs, gram_schmidt_threshold; verbose=true)
 
-    EC_basis = hcat(training_vecs...)
-    EC.H0_EC = transpose(EC_basis) * weights_mat * EC.H0 * EC_basis
-    EC.H1_EC = transpose(EC_basis) * weights_mat * EC.H1 * EC_basis
-    EC.N_EC = transpose(EC_basis) * weights_mat * EC_basis
-
-    if pseudo_inv_rtol > 0
-        EC.ensemble_size > 0 && generate_resampled_EC_matrices(EC, pseudo_inv_rtol)
-        inv_N_EC = pinv(EC.N_EC; rtol=pseudo_inv_rtol)
-        EC.H0_EC = inv_N_EC * EC.H0_EC
-        EC.H1_EC = inv_N_EC * EC.H1_EC
-        EC.N_EC = nothing
+    for _ in 1:EC.ensemble_size
+        (H0_EC, H1_EC, N_EC) = get_reduced_matrices(EC, training_vecs, gram_schmidt_threshold, resample(length(training_vecs)); verbose=false)
+        push!(EC.H0_EC_ensemble, H0_EC)
+        push!(EC.H1_EC_ensemble, H1_EC)
+        push!(EC.N_EC_ensemble, N_EC)
     end
 
     EC.trained = true
 end
 
+function get_reduced_matrices(EC::affine_EC, training_vecs, gram_schmidt_threshold, subsample=1:length(training_vecs); verbose=false)
+    vecs = deepcopy(training_vecs[subsample])
+
+    if gram_schmidt_threshold > 0; vecs = gram_schmidt!(sampled_vecs, EC.weights, gram_schmidt_threshold; verbose=verbose); end
+
+    EC_basis = hcat(orth_vecs...)
+    weights_mat = spdiagm(EC.weights)
+    H0_EC = transpose(EC_basis) * weights_mat * EC.H0 * EC_basis
+    H1_EC = transpose(EC_basis) * weights_mat * EC.H1 * EC_basis
+    N_EC = transpose(EC_basis) * weights_mat * EC_basis
+
+    return (H0_EC, H1_EC, N_EC)
+end
+
 resample(n::Int) = rand(1:n, n) |> unique |> sort
 
-"Generate an resampled ensemble of reduced matrices performing Gram-Schmidt orthonormalization for each."
-function generate_resampled_EC_matrices(EC::affine_EC, training_vecs, weights_mat, gram_schmidt_threshold)
-    for _ in 1:EC.ensemble_size
-        sampled_vecs = deepcopy(training_vecs[resample(length(training_vecs))])
-        orth_vecs = gram_schmidt!(sampled_vecs, EC.weights, gram_schmidt_threshold; verbose=false)
-        EC_basis = hcat(orth_vecs...)
-
-        H0_EC = transpose(EC_basis) * weights_mat * EC.H0 * EC_basis
-        H1_EC = transpose(EC_basis) * weights_mat * EC.H1 * EC_basis
-        N_EC = transpose(EC_basis) * weights_mat * EC_basis
-        push!(EC.H0_EC_ensemble, H0_EC)
-        push!(EC.H1_EC_ensemble, H1_EC)
-        push!(EC.N_EC_ensemble, N_EC)
-    end
-end
-
-"Generate an resampled ensemble of reduced matrices performing Moore-Penrose psuedoinverse for each."
-function generate_resampled_EC_matrices(EC::affine_EC, pseudo_inv_rtol)
-    for _ in 1:EC.ensemble_size
-        sample = resample(size(EC.N_EC, 1))
-        new_N_EC = EC.N_EC[sample, sample]
-        new_H0_EC = EC.H0_EC[sample, sample]
-        new_H1_EC = EC.H1_EC[sample, sample]
-
-        inv_N_EC = pinv(new_N_EC; rtol=pseudo_inv_rtol)
-        push!(EC.H0_EC_ensemble, inv_N_EC * new_H0_EC)
-        push!(EC.H1_EC_ensemble, inv_N_EC * new_H1_EC)
-    end
-end
-
 "Extrapolate using a trained EC model for a given range of c values
 If a list is provided for ref_eval, they are used as reference values for picking the closest eigenvalues at each point.
 If a single number is provided for ref_eval, it is used as a reference for the first point, and the previous eigenvalue is used as the reference for each successive point.
-If precalculated_exact_E is provided, ref_eval is ignored."
-function extrapolate!(EC::affine_EC, c_vals; ref_eval=EC.training_E[end], verbose=true, tol=1e-5, precalculated_exact_E=nothing)
+If precalculated_exact_E is provided, ref_eval is ignored.
+If pseudo_inv_rtol > 0, the GEVP is avoided using Moore-Penrose psuedoinverse, using this value as the relative tolerance for dropping redundant vectors."
+function extrapolate!(EC::affine_EC, c_vals; ref_eval=EC.training_E[end], pseudo_inv_rtol=1e-6, verbose=true, tol=1e-5, precalculated_exact_E=nothing)
     @assert EC.trained "EC model must be trained using train() before extrapolation"
 
     for c in c_vals
@@ -144,16 +115,14 @@ function extrapolate!(EC::affine_EC, c_vals; ref_eval=EC.training_E[end], verbos
 
         verbose && println("Extrapolating for c = $c")
 
-        H_EC = EC.H0_EC + c .* EC.H1_EC
-        evals = isnothing(EC.N_EC) ? eigvals(H_EC) : eigvals(H_EC, EC.N_EC)
+        evals = get_extrapolated_evals(EC.H0_EC, EC.H1_EC, EC.N_EC, c, pseudo_inv_rtol)
         push!(EC.extrapolated_E, nearest(evals, current_E))
 
         if EC.ensemble_size > 0
-            E_ensemble = ComplexF64[]
+            E_ensemble = zeros(ComplexF64, EC.ensemble_size)
             for i in 1:EC.ensemble_size
-                H_EC = EC.H0_EC_ensemble[i] + c .* EC.H1_EC_ensemble[i]
-                evals = isempty(EC.N_EC_ensemble) ? eigvals(H_EC) : eigvals(H_EC, EC.N_EC_ensemble[i])
-                push!(E_ensemble, nearest(evals, current_E))
+                evals = get_extrapolated_evals(EC.H0_EC_ensemble[i], EC.H1_EC_ensemble[i], EC.N_EC_ensemble[i], c, pseudo_inv_rtol)
+                E_ensemble[i] = nearest(evals, current_E)
             end
             re_CI = std(real.(E_ensemble))
             im_CI = std(imag.(E_ensemble))
@@ -162,6 +131,18 @@ function extrapolate!(EC::affine_EC, c_vals; ref_eval=EC.training_E[end], verbos
     end
 end
 
+"Solve the GEVP with or without Moore-Penrose psuedoinverse"
+function get_extrapolated_evals(H0_EC, H1_EC, N_EC, c, pseudo_inv_rtol)
+    H_EC = H0_EC + c .* H1_EC
+    if pseudo_inv_rtol > 0
+        inv_N_EC = pinv(N_EC; rtol=pseudo_inv_rtol)
+        H_EC = inv_N_EC * H_EC
+        return eigvals(H_EC)
+    else
+        return eigvals(H_EC, N_EC)
+    end
+end
+
 "Export EC data as CSV file"
 exportCSV(EC::affine_EC, filename) = exportCSV(filename, (EC.training_E, EC.exact_E, EC.extrapolated_E), ("training", "exact", "extrapolated"))