improve op add_to_diagonal for eigen sparse matrix

include/pressio/ops/eigen/ops_add_to_diagonal.hpp

@@ -49,7 +49,33 @@
 #ifndef OPS_EIGEN_OPS_ADD_TO_DIAGONAL_HPP_
 #define OPS_EIGEN_OPS_ADD_TO_DIAGONAL_HPP_
 
-namespace pressio{ namespace ops{
+namespace pressio{
+
+namespace impl{
+
+template<class T>
+bool _eigen_matrix_has_diagonal_elements(T & M) {
+  using value_type = typename ::pressio::Traits<T>::scalar_type;
+
+  /*
+    https://eigen.tuxfamily.org/dox-devel/classEigen_1_1SparseMatrix.html#ae6f4db56c76e4d374eb8507fbef5bdb5
+
+    If the diagonal entries are written, then all diagonal entries must already exist, 
+    otherwise an assertion will be raised.
+    Since we do not catch the assertion, the program terminates which is what we want.
+    Note that we need to add zero here so that if M is satisfies the condition, 
+    then its values are not changed.
+  */
+  const volatile value_type zero = value_type(0);
+  auto d = M.diagonal();
+  for (int i=0; i<M.rows(); i++) {
+    d(i) += zero;
+  }
+  return true;
+}
+} // end namespace impl
+
+namespace ops{
 
 /*
   constrained via is_convertible because the impl is using
@@ -63,10 +89,8 @@ ::pressio::mpl::enable_if_t<
   >
 add_to_diagonal(T & o, const ScalarType & value)
 {
-  auto identity(o);
-  identity.setIdentity();
-  identity.coeffs() *= value;
-  o += identity;
+  assert(::pressio::impl::_eigen_matrix_has_diagonal_elements(o));
+  o.diagonal().array() += value;
 }
 
 /*

tests/functional_small/ode_steppers/ode_all_implicit_schemes_check_app_called_with_correct_time.cc

@@ -25,6 +25,10 @@ struct MyApp1
 
   jacobian_type createJacobian() const{
     jacobian_type J(3,3);
+    // ensure that the diagonal elements exist
+    for (int i=0; i<J.innerSize(); i++) {
+      J.coeffRef(i, i) = 0;
+    }
     return J;
   }
 
@@ -90,6 +94,10 @@ struct MyApp2
 
   jacobian_type createJacobian() const{
     jacobian_type J(3,3);
+    // ensure that the diagonal elements exist
+    for (int i=0; i<J.innerSize(); i++) {
+      J.coeffRef(i, i) = 0;
+    }
     return J;
   }
 

tests/functional_small/ode_steppers/ode_bdf1_correctness_with_step_recovery.cc

@@ -19,6 +19,10 @@ struct MyApp
 
   jacobian_type createJacobian() const{
     jacobian_type J(3,3);
+    // ensure that the diagonal elements exist
+    for (int i=0; i<J.innerSize(); i++) {
+      J.coeffRef(i, i) = 0;
+    }
     return J;
   }
 

tests/functional_small/ode_steppers/ode_bdf1_strong_condition_correctness.cc

@@ -13,7 +13,15 @@ struct MyApp
   state_type createState() const{ return state_type(3); }
 
   rhs_type createRhs() const{ rhs_type f(3); return f; }
-  jacobian_type createJacobian() const{ jacobian_type J(3,3); return J; }
+
+  jacobian_type createJacobian() const{
+    jacobian_type J(3,3);
+    // ensure that the diagonal elements exist
+    for (int i=0; i<J.innerSize(); i++) {
+      J.coeffRef(i, i) = 0;
+    }
+    return J;
+  }
 
   void rhsAndJacobian(const state_type & /*unused*/,
 		      const independent_variable_type& evaltime,

tests/functional_small/ode_steppers/ode_bdf2_strong_condition_correctness.cc

@@ -13,7 +13,15 @@ struct MyApp
   state_type createState() const{ return state_type(3); }
 
   rhs_type createRhs() const{ rhs_type f(3); return f; }
-  jacobian_type createJacobian() const{ jacobian_type J(3,3); return J; }
+
+  jacobian_type createJacobian() const{
+    jacobian_type J(3,3);
+    // ensure that the diagonal elements exist
+    for (int i=0; i<J.innerSize(); i++) {
+      J.coeffRef(i, i) = 0;
+    }
+    return J;
+  }
 
   void rhsAndJacobian(const state_type & /*unused*/,
 		      const independent_variable_type& evaltime,

tests/functional_small/ode_steppers/ode_crank_nicolson_simple_correctness_eigen.cc

@@ -12,8 +12,17 @@ struct MyApp
 
 public:
   state_type createState() const{ return state_type(3); }
+
   rhs_type createRhs() const{ return rhs_type(3); }
-  jacobian_type createJacobian() const{return jacobian_type(3,3);}
+
+  jacobian_type createJacobian() const{
+    jacobian_type J(3,3);
+    // ensure that the diagonal elements exist
+    for (int i=0; i<J.innerSize(); i++) {
+      J.coeffRef(i, i) = 0;
+    }
+    return J;
+  }
 
   void rhsAndJacobian(const state_type & y,
 		      const independent_variable_type& evaltime,