diff --git a/src/poly/solver/private/polySolver.cpp b/src/poly/solver/private/polySolver.cpp
index df88b23..6134560 100644
--- a/src/poly/solver/private/polySolver.cpp
+++ b/src/poly/solver/private/polySolver.cpp
@@ -181,12 +181,24 @@ void PolySolver::solve() const {
     setInitialGuess();
 
     setOperatorEssentialTrueDofs();
+    const auto thetaVec = static_cast<mfem::Vector>(*m_theta); // NOLINT(*-slicing)
+    const auto phiVec = static_cast<mfem::Vector>(*m_phi); // NOLINT(*-slicing)
+
+    // --- Finalize the operator ---
+    // Finalize with the initial state of theta for the initial jacobian calculation
+    m_polytropOperator->finalize(thetaVec);
+
+    if (!m_polytropOperator->isFinalized()) {
+        LOG_ERROR(m_logger, "PolytropeOperator is not finalized. Cannot solve.");
+        throw std::runtime_error("PolytropeOperator is not finalized. Cannot solve.");
+    }
 
     // It's safer to get the offsets directly from the operator after finalization
     const mfem::Array<int>& block_offsets = m_polytropOperator->GetBlockOffsets(); // Assuming a getter exists or accessing member if public/friend
     mfem::BlockVector state_vector(block_offsets);
-    state_vector.GetBlock(0) = static_cast<mfem::Vector>(*m_theta); // NOLINT(*-slicing)
-    state_vector.GetBlock(1) = static_cast<mfem::Vector>(*m_phi); // NOLINT(*-slicing)
+    state_vector.GetBlock(0) = thetaVec; // NOLINT(*-slicing)
+    state_vector.GetBlock(1) = phiVec; // NOLINT(*-slicing)
+
 
     mfem::Vector zero_rhs(block_offsets.Last());
     zero_rhs = 0.0;
@@ -307,17 +319,8 @@ void PolySolver::saveAndViewSolution(const mfem::BlockVector& state_vector) cons
 }
 
 void PolySolver::setOperatorEssentialTrueDofs() const {
-
-    SSE::MFEMArrayPairSet ess_tdof_pair_set = getEssentialTrueDof();
+    const SSE::MFEMArrayPairSet ess_tdof_pair_set = getEssentialTrueDof();
     m_polytropOperator->SetEssentialTrueDofs(ess_tdof_pair_set);
-
-    // -- Finalize the operator --
-    m_polytropOperator->finalize();
-
-    if (!m_polytropOperator->isFinalized()) {
-        LOG_ERROR(m_logger, "PolytropeOperator is not finalized. Cannot solve.");
-        throw std::runtime_error("PolytropeOperator is not finalized. Cannot solve.");
-    }
 }
 
 void PolySolver::LoadSolverUserParams(double &newtonRelTol, double &newtonAbsTol, int &newtonMaxIter, int &newtonPrintLevel, double &gmresRelTol, double &gmresAbsTol, int &gmresMaxIter, int &gmresPrintLevel) const {
diff --git a/src/poly/utils/private/operator.cpp b/src/poly/utils/private/operator.cpp
index 8426d7a..ecbf6ef 100644
--- a/src/poly/utils/private/operator.cpp
+++ b/src/poly/utils/private/operator.cpp
@@ -107,7 +107,7 @@ PolytropeOperator::PolytropeOperator(
     m_f = std::move(f);
 }
 
-void PolytropeOperator::finalize() {
+void PolytropeOperator::finalize(const mfem::Vector &initTheta) {
     if (m_isFinalized) {
         return;
     }
@@ -137,6 +137,9 @@ void PolytropeOperator::finalize() {
     m_jacobian->SetBlock(1, 0, m_negQ_op.get()); // -Q (constant)
     m_jacobian->SetBlock(1, 1, m_Dmat.get()); // D (constant)
 
+    const auto &grad = m_f->GetGradient(initTheta);
+    updatePreconditioner(grad);
+
     m_isFinalized = true;
 }
 
@@ -148,6 +151,9 @@ const mfem::BlockOperator &PolytropeOperator::GetJacobianOperator() const {
 }
 
 mfem::BlockDiagonalPreconditioner& PolytropeOperator::GetPreconditioner() const {
+    if (!m_isFinalized) {
+        MFEM_ABORT("PolytropeOperator::GetPreconditioner called before finalize");
+    }
     return *m_schurPreconditioner;
 }
 
@@ -271,6 +277,9 @@ void PolytropeOperator::updateInverseSchurCompliment() const {
     schurCompliment->Add(-1, *temp_QxdfInvxM); // D - Q * df^{-1} * M
     approxJacobiInvert(*schurCompliment, m_invSchurCompliment, "Schur Compliment");
 
+    if (m_schurPreconditioner == nullptr) {
+        m_schurPreconditioner = std::make_unique<mfem::BlockDiagonalPreconditioner>(m_blockOffsets);
+    }
     m_schurPreconditioner->SetDiagonalBlock(0, m_invNonlinearJacobian.get());
     m_schurPreconditioner->SetDiagonalBlock(1, m_invSchurCompliment.get());
 }
@@ -285,7 +294,7 @@ mfem::Operator& PolytropeOperator::GetGradient(const mfem::Vector &x) const {
     if (!m_isFinalized) {
         MFEM_ABORT("PolytropeOperator::GetGradient called before finalize");
     }
-    // -- Get the gradient of f --
+    // --- Get the gradient of f ---
     mfem::BlockVector x_block(const_cast<mfem::Vector&>(x), m_blockOffsets);
     const mfem::Vector& x_theta = x_block.GetBlock(0);
 
diff --git a/src/poly/utils/public/operator.h b/src/poly/utils/public/operator.h
index 784f232..bb469a6 100644
--- a/src/poly/utils/public/operator.h
+++ b/src/poly/utils/public/operator.h
@@ -48,7 +48,7 @@ public:
 
     bool isFinalized() const { return m_isFinalized; }
 
-    void finalize();
+    void finalize(const mfem::Vector &initTheta);
 
     const mfem::Array<int>& GetBlockOffsets() const { return m_blockOffsets; }
 
@@ -78,7 +78,6 @@ private:
     std::unique_ptr<mfem::ScaledOperator> m_negQ_op;
     mutable std::unique_ptr<mfem::BlockOperator> m_jacobian;
 
-    // TODO I think these need to be calculated in the GetGradient every time since they will always change
     mutable std::unique_ptr<mfem::SparseMatrix> m_invSchurCompliment;
     mutable std::unique_ptr<mfem::SparseMatrix> m_invNonlinearJacobian;