State-dependent custom operators for block preconditioning

RELEASE-NOTES.md

@@ -30,6 +30,8 @@ The Smith project release numbers follow [Semantic Versioning](http://semver.org
 - Added composable solid-mechanics and thermo-mechanics examples, tutorials, and regression tests covering coupled
   sensitivities, finite-difference checks, field parameters, and solves.
 - Added axisymmetric solid mechanics materials and loads for 2D `(r, z)` meshes.
+- Added diagonal, triangular, and Schur block preconditioners for linearized block systems.
+- Added support for custom block operators that are rebuilt from the current nonlinear state at each Newton iteration.
 
 ### Removed
 

src/smith/differentiable_numerics/CMakeLists.txt

@@ -10,6 +10,7 @@ set(differentiable_numerics_sources
     differentiable_physics.cpp
     lumped_mass_explicit_newmark_state_advancer.cpp
     nonlinear_block_solver.cpp
+    weak_form_block_operator.cpp
     system_solver.cpp
     system_base.cpp
     nonlinear_solve.cpp
@@ -24,6 +25,7 @@ set(differentiable_numerics_headers
     state_advancer.hpp
     reaction.hpp
     nonlinear_block_solver.hpp
+    weak_form_block_operator.hpp
     system_solver.hpp
     differentiable_physics.hpp
     timestep_estimator.hpp

src/smith/differentiable_numerics/nonlinear_block_solver.cpp

@@ -58,7 +58,7 @@ NonlinearBlockSolver::NonlinearBlockSolver(std::unique_ptr<EquationSolver> s, MP
 
 std::shared_ptr<NonlinearBlockSolver> NonlinearBlockSolver::cloneFresh() const
 {
-  if (!retained_nonlinear_options_ || !retained_linear_options_) {
+  if (state_dependent_solver_ || !retained_nonlinear_options_ || !retained_linear_options_) {
     return nullptr;
   }
 
@@ -199,6 +199,9 @@ std::vector<NonlinearBlockSolverBase::FieldPtr> NonlinearBlockSolver::solve(
         for (int row_i = 0; row_i < num_rows; ++row_i) {
           *u_guesses[static_cast<size_t>(row_i)] = u->GetBlock(row_i);
         }
+        if (state_dependent_solver_) {
+          state_dependent_solver_->updateForState(*u, block_offsets);
+        }
         matrix_of_jacs_ = jacobian_funcs(u_guesses);
         if (num_rows == 1) {
           auto& J = matrix_of_jacs_[0][0];
@@ -296,4 +299,21 @@ std::shared_ptr<NonlinearBlockSolver> buildNonlinearBlockSolver(NonlinearSolverO
                                                 nonlinear_opts.relative_tol, nonlinear_opts, linear_opts);
 }
 
+std::shared_ptr<NonlinearBlockSolver> buildNonlinearBlockSolver(NonlinearSolverOptions nonlinear_opts,
+                                                                LinearSolverOptions linear_opts,
+                                                                const smith::Mesh& mesh,
+                                                                std::unique_ptr<mfem::Solver> preconditioner)
+{
+  auto solid_solver =
+      std::make_unique<EquationSolver>(nonlinear_opts, linear_opts, std::move(preconditioner), mesh.getComm());
+  auto* state_dependent_preconditioner = dynamic_cast<StateDependentSolver*>(&solid_solver->preconditioner());
+  auto nonlinear_block_solver =
+      std::make_shared<NonlinearBlockSolver>(std::move(solid_solver), mesh.getComm(), nonlinear_opts.absolute_tol,
+                                             nonlinear_opts.relative_tol, std::nullopt, linear_opts);
+  if (state_dependent_preconditioner) {
+    nonlinear_block_solver->setStateDependentSolver(state_dependent_preconditioner);
+  }
+  return nonlinear_block_solver;
+}
+
 }  // namespace smith

src/smith/differentiable_numerics/nonlinear_block_solver.hpp

@@ -15,13 +15,16 @@
 #include <memory>
 #include <functional>
 #include <optional>
+#include <utility>
 #include <vector>
 #include <mpi.h>
 
 #include "smith/numerics/solver_config.hpp"
 #include "smith/numerics/nonlinear_convergence.hpp"
 
 namespace mfem {
+template <typename T>
+class Array;
 class Solver;
 class Vector;
 class HypreParMatrix;
@@ -35,6 +38,7 @@ class BoundaryConditionManager;
 class FiniteElementState;
 class FiniteElementDual;
 class Mesh;
+class StateDependentSolver;
 struct NonlinearSolverOptions;
 struct LinearSolverOptions;
 
@@ -140,7 +144,10 @@ class NonlinearBlockSolver : public NonlinearBlockSolverBase {
   /// @brief Set the inner tolerance multiplier.
   void setInnerToleranceMultiplier(double multiplier) override { inner_tol_multiplier_ = multiplier; }
 
-  /// @brief Build a fresh solver instance from retained config.
+  /// @brief Register a solver refreshed with the current nonlinear state before Jacobian assembly.
+  void setStateDependentSolver(StateDependentSolver* solver) { state_dependent_solver_ = solver; }
+
+  /// @brief Build a fresh solver instance from retained config, or nullptr when this solver has injected state.
   std::shared_ptr<NonlinearBlockSolver> cloneFresh() const;
 
   mutable std::unique_ptr<mfem::BlockOperator>
@@ -158,6 +165,7 @@ class NonlinearBlockSolver : public NonlinearBlockSolverBase {
   double inner_tol_multiplier_ = 1.0;  ///< multiplier for tolerances during inner solves
   std::optional<NonlinearSolverOptions> retained_nonlinear_options_ = std::nullopt;  ///< retained nonlinear config
   std::optional<LinearSolverOptions> retained_linear_options_ = std::nullopt;        ///< retained linear config
+  StateDependentSolver* state_dependent_solver_ = nullptr;  ///< optional solver refreshed during Jacobian evaluation
 };
 
 /// @brief Create an equation-backed nonlinear block solver.
@@ -168,4 +176,14 @@ std::shared_ptr<NonlinearBlockSolver> buildNonlinearBlockSolver(NonlinearSolverO
                                                                 LinearSolverOptions linear_opts,
                                                                 const smith::Mesh& mesh);
 
+/// @brief Create an equation-backed nonlinear block solver with a custom preconditioner.
+/// @param nonlinear_opts nonlinear options struct
+/// @param linear_opts linear options struct
+/// @param mesh mesh
+/// @param preconditioner custom preconditioner attached to the linear solver
+std::shared_ptr<NonlinearBlockSolver> buildNonlinearBlockSolver(NonlinearSolverOptions nonlinear_opts,
+                                                                LinearSolverOptions linear_opts,
+                                                                const smith::Mesh& mesh,
+                                                                std::unique_ptr<mfem::Solver> preconditioner);
+
 }  // namespace smith

src/smith/differentiable_numerics/tests/CMakeLists.txt

@@ -19,7 +19,10 @@ set(differentiable_numerics_test_source
     test_multiphysics_time_integrator.cpp
     test_thermo_mechanics_with_internal_vars.cpp
     test_mixed_poisson.cpp
+    test_nonlinear_mixed_diffusion.cpp
     test_amg_vdim_block_preconditioner.cpp
+    test_state_dependent_preconditioner.cpp
+    test_weak_form_block_operator.cpp
     )
 
 smith_add_tests( SOURCES       ${differentiable_numerics_test_source}

src/smith/differentiable_numerics/tests/test_mixed_poisson.cpp

@@ -1,5 +1,12 @@
 #include <gtest/gtest.h>
 
+#include <cmath>
+#include <memory>
+#include <set>
+#include <string>
+#include <utility>
+#include <vector>
+
 #include "smith/infrastructure/application_manager.hpp"
 #include "smith/numerics/equation_solver.hpp"
 #include "smith/numerics/solver_config.hpp"
@@ -13,6 +20,7 @@
 #include "smith/differentiable_numerics/nonlinear_block_solver.hpp"
 #include "smith/differentiable_numerics/nonlinear_solve.hpp"
 #include "smith/differentiable_numerics/paraview_writer.hpp"
+#include "smith/differentiable_numerics/weak_form_block_operator.hpp"
 #include "smith/numerics/block_preconditioner.hpp"
 
 #include "gretl/data_store.hpp"
@@ -112,6 +120,8 @@ TEST_P(BlockPreconditionerTest, BlockSolve)
       "constitutive_eqn", mesh, space(flux), spaces({flux, potential}));
   smith::FunctionalWeakForm<2, Space, smith::Parameters<VectorSpace, Space>> bal_form(
       "balance_eqn", mesh, space(potential), spaces({flux, potential}));
+  smith::FunctionalWeakForm<2, Space, smith::Parameters<Space>> potential_diffusion_form(
+      "potential_diffusion", mesh, space(potential), spaces({potential}));
 
   con_form.addBodyIntegral(DependsOn<0, 1>{}, mesh->entireBodyName(),
                            [](auto /* t */, auto /* x */, auto SIGMA, auto U) {
@@ -132,6 +142,11 @@ TEST_P(BlockPreconditionerTest, BlockSolve)
     auto f = 2.0 * pi * pi * sin(pi * x[0]) * sin(pi * x[1]);
     return smith::tuple{-f + div_sigma, smith::zero{}};
   });
+  potential_diffusion_form.addBodyIntegral(DependsOn<0>{}, mesh->entireBodyName(),
+                                           [](auto /* time_info */, auto /* x */, auto U) {
+                                             auto grad_u = get<DERIVATIVE>(U);
+                                             return smith::tuple{smith::zero{}, grad_u};
+                                           });
   // u_exact = sin(M_PI * x(0)) * sin(M_PI * x(1));
   // sigma_exact
   // pi * cos(pi * x(0)) * sin(pi * x(1))
@@ -177,6 +192,8 @@ TEST_P(BlockPreconditionerTest, BlockSolve)
   auto time = graph->create_state<double, double>(0.0);
   auto dt = graph->create_state<double, double>(0.025);
   size_t cycle = 0;
+  const auto time_info = smith::TimeInfo(time.get(), dt.get(), cycle);
+  std::unique_ptr<mfem::Solver> diffusion_precond;
   std::vector<smith::FieldState> params;
   auto& flux_params = params;
   auto& potential_params = params;
@@ -203,7 +220,16 @@ TEST_P(BlockPreconditionerTest, BlockSolve)
     case BlockPrecondType::SchurFullCustom:
       linear_options.preconditioner = smith::Preconditioner::BlockSchur;
       linear_options.block_schur_type = smith::BlockSchurType::Full;
-      /// linear_options.schur_approx_type = smith::BlockSchurType::Custom;
+      linear_options.schur_approx_type = smith::SchurApproxType::Custom;
+      std::vector<smith::BlockProviderOverride> overrides;
+      overrides.push_back(smith::makeWeakFormBlockProviderOverride(1, potential_diffusion_form, shape_disp, {potential},
+                                                                   {1.0}, time_info, potential_bc_manager.get()));
+
+      auto solvers =
+          smith::buildBlockPreconditionerSubSolvers(linear_options.sub_block_linear_solver_options, mesh->getComm());
+
+      diffusion_precond = std::make_unique<smith::BlockSchurPreconditioner>(
+          std::move(solvers), linear_options.block_schur_type, linear_options.schur_approx_type, std::move(overrides));
       break;
   }
 
@@ -214,11 +240,14 @@ TEST_P(BlockPreconditionerTest, BlockSolve)
   nonlin_opts.max_iterations = 1;
   nonlin_opts.print_level = linear_options.print_level;
 
-  auto nonlinear_block_solver = smith::buildNonlinearBlockSolver(nonlin_opts, linear_options, *mesh);
+  auto nonlinear_block_solver =
+      diffusion_precond
+          ? smith::buildNonlinearBlockSolver(nonlin_opts, linear_options, *mesh, std::move(diffusion_precond))
+          : smith::buildNonlinearBlockSolver(nonlin_opts, linear_options, *mesh);
 
   auto sols = block_solve({&con_form, &bal_form}, {{0, 1}, {0, 1}}, shape_disp, {con_arguments, bal_arguments},
-                          {flux_params, potential_params}, smith::TimeInfo(time.get(), dt.get(), cycle),
-                          nonlinear_block_solver.get(), {flux_bc_manager.get(), potential_bc_manager.get()});
+                          {flux_params, potential_params}, time_info, nonlinear_block_solver.get(),
+                          {flux_bc_manager.get(), potential_bc_manager.get()});
 
   auto pv_writer = smith::createParaviewWriter(*mesh, sols, physics_name);
   pv_writer.write(0, 0.0, sols);