diff --git a/src/poly/utils/private/operator.cpp b/src/poly/utils/private/operator.cpp
index ca70062..8426d7a 100644
--- a/src/poly/utils/private/operator.cpp
+++ b/src/poly/utils/private/operator.cpp
@@ -26,6 +26,8 @@
 #include "linalg/vector.hpp"
 #include <memory>
 
+#include "linalg/tensor.hpp"
+
 void writeDenseMatrixToCSV(const std::string &filename, int precision, const mfem::DenseMatrix *mat) {
     if (!mat) {
         throw std::runtime_error("The operator is not a SparseMatrix.");
@@ -211,27 +213,39 @@ void PolytropeOperator::Mult(const mfem::Vector &x, mfem::Vector &y) const {
 
 }
 
-// TODO: I was *very* stupid and I accidently deleted a lot of the code
+// TODO: I was *very* stupid and I accidentally deleted a lot of the code
 //       which I had written to find and use the preconditioner. This needs
 //       to be reimplemented. Once that is working you can get back to
 //       Trying to understand the multimodal hump in the residuals vector.
 //       There is a jupyter notebook about this. I was thinking that this was
-//       perhapse related to the non consistent application of boundary conditions.
+//       perhaps related to the non consistent application of boundary conditions.
 void approxJacobiInvert(const mfem::SparseMatrix& mat, std::unique_ptr<mfem::SparseMatrix>& invMat, const std::string& name="matrix") {
     // PERF: This likely can be made much more efficient and will probably be called in tight loops, a good
     // PERF:   place for some easy optimization might be here.
+
+    // Confirm that mat is a square matrix
+    MFEM_ASSERT(mat.Height() == mat.Width(), "Matrix " + name + " is not square, cannot invert.");
+
     mfem::Vector diag;
     mat.GetDiag(diag);
 
     // Invert the diagonal
     for (int i = 0; i < diag.Size(); i++) {
-        if (diag(i) == 0.0) {
-            MFEM_ABORT("Diagonal element (" + std::to_string(i) +") in " + name + " is zero, cannot invert.");
-        }
+        MFEM_ASSERT(diag(i) != 0, "Diagonal element (" + std::to_string(i) +") in " + name + " is zero, cannot invert.");
         diag(i) = 1.0 / diag(i);
     }
 
-    invMat = std::make_unique<mfem::SparseMatrix>(diag);
+    // If the matrix is already inverted, just set the diagonal to avoid reallocation
+    if (invMat != nullptr) {
+        MFEM_ASSERT(invMat->Height() == invMat->Width(), "invMat (result matrix) is not square, cannot invert " + name + " into it.");
+        MFEM_ASSERT(invMat->Height() == mat.Height(), "Incompatible matrix sizes for inversion of " + name + ", expected " + std::to_string(mat.Height()) + " but got " + std::to_string(invMat->Height()));
+        for (int i = 0; i < diag.Size(); i++) {
+            MFEM_ASSERT(diag(i) != 0, "Diagonal element (" + std::to_string(i) +") in " + name + " is zero, resulting matrix would be singular.");
+            invMat->Elem(i, i) = diag(i);
+        }
+    } else {
+        invMat = std::make_unique<mfem::SparseMatrix>(diag);
+    }
 }
 
 void PolytropeOperator::updateInverseNonlinearJacobian(const mfem::Operator &grad) const {
@@ -261,6 +275,12 @@ void PolytropeOperator::updateInverseSchurCompliment() const {
     m_schurPreconditioner->SetDiagonalBlock(1, m_invSchurCompliment.get());
 }
 
+void PolytropeOperator::updatePreconditioner(const mfem::Operator &grad) const {
+    updateInverseNonlinearJacobian(grad);
+    updateInverseSchurCompliment();
+}
+
+
 mfem::Operator& PolytropeOperator::GetGradient(const mfem::Vector &x) const {
     if (!m_isFinalized) {
         MFEM_ABORT("PolytropeOperator::GetGradient called before finalize");
@@ -270,8 +290,7 @@ mfem::Operator& PolytropeOperator::GetGradient(const mfem::Vector &x) const {
     const mfem::Vector& x_theta = x_block.GetBlock(0);
 
     auto &grad = m_f->GetGradient(x_theta);
-    updateInverseNonlinearJacobian(grad);
-    updateInverseSchurCompliment();
+    updatePreconditioner(grad);
 
     m_jacobian->SetBlock(0, 0, &grad);
     // The other blocks are set up in finalize since they are constant. Only J00 depends on the current state of theta