From c39ce0b8bf284a4994d5ae023a113fa97a33606c Mon Sep 17 00:00:00 2001
From: Michael Tupek <tupek2@llnl.gov>
Date: Tue, 6 Aug 2024 12:40:46 -0700
Subject: [PATCH] Add a method to base_physics to reset adjoint state back to
 the end of time. Also fix divide by zeros in output.

---
 src/serac/numerics/equation_solver.cpp        |  6 ++--
 src/serac/physics/base_physics.hpp            | 12 +++++++-
 .../tests/quasistatic_solid_adjoint.cpp       | 28 ++++++-------------
 3 files changed, 23 insertions(+), 23 deletions(-)

diff --git a/src/serac/numerics/equation_solver.cpp b/src/serac/numerics/equation_solver.cpp
index f3e68bbf6..7d02ad618 100644
--- a/src/serac/numerics/equation_solver.cpp
+++ b/src/serac/numerics/equation_solver.cpp
@@ -97,7 +97,7 @@ class NewtonSolver : public mfem::NewtonSolver {
       if (print_options.iterations) {
         mfem::out << "Newton iteration " << std::setw(3) << it << " : ||r|| = " << std::setw(13) << norm;
         if (it > 0) {
-          mfem::out << ", ||r||/||r_0|| = " << std::setw(13) << norm / initial_norm;
+          mfem::out << ", ||r||/||r_0|| = " << std::setw(13) << (initial_norm != 0.0 ? norm / initial_norm : norm);
         }
         mfem::out << '\n';
       }
@@ -383,7 +383,7 @@ class TrustRegion : public mfem::NewtonSolver {
     for (cgIter = 1; cgIter <= settings.maxCgIterations; ++cgIter) {
       hess_vec_func(d, Hd);
       const double curvature = Dot(d, Hd);
-      const double alphaCg   = rPr / curvature;
+      const double alphaCg   = curvature != 0.0 ? rPr / curvature : 0.0;
 
       auto& zPred = Hd;  // re-use Hd, this is where bugs come from
       add(z, alphaCg, d, zPred);
@@ -490,7 +490,7 @@ class TrustRegion : public mfem::NewtonSolver {
       if (print_options.iterations) {
         mfem::out << "Newton iteration " << std::setw(3) << it << " : ||r|| = " << std::setw(13) << norm;
         if (it > 0) {
-          mfem::out << ", ||r||/||r_0|| = " << std::setw(13) << norm / initial_norm;
+          mfem::out << ", ||r||/||r_0|| = " << std::setw(13) << (initial_norm != 0.0 ? norm / initial_norm : norm);
         }
         mfem::out << '\n';
       }
diff --git a/src/serac/physics/base_physics.hpp b/src/serac/physics/base_physics.hpp
index 84ca38d9f..c9edd0912 100644
--- a/src/serac/physics/base_physics.hpp
+++ b/src/serac/physics/base_physics.hpp
@@ -126,13 +126,23 @@ class BasePhysics {
   virtual const std::vector<double>& timesteps() const;
 
   /**
-   * @brief Base method to reset physics states to zero.  This does not reset design parameters or shape.
+   * @brief Base method to reset physics states to the initial time.  This does not reset design parameters or shape.
    *
    * @param[in] cycle The simulation cycle (i.e. timestep iteration) to intialize the physics module to
    * @param[in] time The simulation time to initialize the physics module to
    */
   virtual void resetStates(int cycle = 0, double time = 0.0) = 0;
 
+  /**
+   * @brief Base method to reset physics states back to the end of time to start adjoint calculations again.  This does not reset design parameters or shape.
+   *
+   */
+  virtual void resetAdjointStates()
+  {
+    time_ = max_time_;
+    cycle_ = max_cycle_;
+  }
+
   /**
    * @brief Complete the setup and allocate the necessary data structures
    *
diff --git a/src/serac/physics/tests/quasistatic_solid_adjoint.cpp b/src/serac/physics/tests/quasistatic_solid_adjoint.cpp
index 8e071a5d5..bf13b2868 100644
--- a/src/serac/physics/tests/quasistatic_solid_adjoint.cpp
+++ b/src/serac/physics/tests/quasistatic_solid_adjoint.cpp
@@ -74,16 +74,6 @@ double forwardPass(serac::BasePhysics* solid, qoiType* qoi, mfem::ParMesh* /*mes
 {
   solid->resetStates();
 
-  // set up plotting
-  /*
-  mfem::ParGridFunction uGF(const_cast<mfem::ParFiniteElementSpace*>(&solid->state("displacement").space()));
-  solid->state("displacement").fillGridFunction(uGF);
-  mfem::VisItDataCollection visit_dc(saveName, meshPtr);
-  visit_dc.RegisterField("u", &uGF);
-  visit_dc.SetCycle(0);
-  visit_dc.Save();
-  */
-
   double                           qoiValue = 0.0;
   const serac::FiniteElementState& E        = solid->parameter("E");
   const serac::FiniteElementState& v        = solid->parameter("v");
@@ -94,13 +84,6 @@ double forwardPass(serac::BasePhysics* solid, qoiType* qoi, mfem::ParMesh* /*mes
     // solve
     solid->advanceTimestep(timeStep);
 
-    // plot
-    /*
-    u.fillGridFunction(uGF);
-    visit_dc.SetCycle(i + 1);
-    visit_dc.Save();
-    */
-
     // accumulate
     double curr = (*qoi)(solid->time(), E, v, u);
     qoiValue += timeStep * 0.5 * (prev + curr);  // trapezoid
@@ -169,10 +152,10 @@ TEST(quasistatic, finiteDifference)
                                                          .print_level    = 1};
   auto seracSolid = ::std::make_unique<solidType>(nonlinear_options, serac::solid_mechanics::direct_linear_options,
                                                   ::serac::solid_mechanics::default_quasistatic_options,
-                                                  ::serac::GeometricNonlinearities::Off, physics_prefix, mesh_tag,
+                                                  ::serac::GeometricNonlinearities::On, physics_prefix, mesh_tag,
                                                   std::vector<std::string>{"E", "v"});
 
-  using materialType = ParameterizedLinearIsotropicSolid;
+  using materialType = ParameterizedNeoHookeanSolid;
   materialType material;
   seracSolid->setMaterial(::serac::DependsOn<0, 1>{}, material);
 
@@ -222,6 +205,13 @@ TEST(quasistatic, finiteDifference)
   double Ederiv, vderiv;
   adjointPass(seracSolid.get(), qoi.get(), nTimeSteps, timeStep, *param_space, Ederiv, vderiv);
 
+  seracSolid->resetAdjointStates();
+
+  double Ederiv2, vderiv2;
+  adjointPass(seracSolid.get(), qoi.get(), nTimeSteps, timeStep, *param_space, Ederiv2, vderiv2);
+  EXPECT_EQ(Ederiv, Ederiv2);
+  EXPECT_EQ(vderiv, vderiv2);
+
   // FINITE DIFFERENCE GRADIENT
   double h = 1e-7;