differentiable_solver.hpp

Go to the documentation of this file.

// Copyright (c) Lawrence Livermore National Security, LLC and
// other Smith Project Developers. See the top-level LICENSE file for
// details.
//
// SPDX-License-Identifier: (BSD-3-Clause)
 
#pragma once
 
#include <memory>
#include <functional>
 
namespace mfem {
class Solver;
class Vector;
class HypreParMatrix;
class BlockOperator;
}  // namespace mfem
 
namespace smith {
 
class EquationSolver;
class BoundaryConditionManager;
class FiniteElementState;
class FiniteElementDual;
class Mesh;
struct NonlinearSolverOptions;
struct LinearSolverOptions;
 
class DifferentiableSolver {
 public:
  virtual ~DifferentiableSolver() {}
 
  virtual void completeSetup(const smith::FiniteElementState& u) = 0;
 
  virtual std::shared_ptr<smith::FiniteElementState> solve(
      const smith::FiniteElementState& u_guess, std::function<mfem::Vector(const smith::FiniteElementState&)> equation,
      std::function<std::unique_ptr<mfem::HypreParMatrix>(const smith::FiniteElementState&)> jacobian) const = 0;
 
  virtual std::shared_ptr<smith::FiniteElementState> solveAdjoint(
      const smith::FiniteElementDual& u_bar, std::unique_ptr<mfem::HypreParMatrix> jacobian_transposed) const = 0;
 
  virtual void clearMemory() const {}
};
 
class LinearDifferentiableSolver : public DifferentiableSolver {
 public:
  LinearDifferentiableSolver(std::unique_ptr<mfem::Solver> s, std::unique_ptr<mfem::Solver> p);
 
  void completeSetup(const smith::FiniteElementState& u) override;
 
  std::shared_ptr<smith::FiniteElementState> solve(
      const smith::FiniteElementState& u_guess, std::function<mfem::Vector(const smith::FiniteElementState&)> equation,
      std::function<std::unique_ptr<mfem::HypreParMatrix>(const smith::FiniteElementState&)> jacobian) const override;
 
  std::shared_ptr<smith::FiniteElementState> solveAdjoint(
      const smith::FiniteElementDual& u_bar, std::unique_ptr<mfem::HypreParMatrix> jacobian_transposed) const override;
 
  mutable std::unique_ptr<mfem::Solver> mfem_solver;          
  mutable std::unique_ptr<mfem::Solver> mfem_preconditioner;  
};
 
class NonlinearDifferentiableSolver : public DifferentiableSolver {
 public:
  NonlinearDifferentiableSolver(std::unique_ptr<EquationSolver> s);
 
  void completeSetup(const smith::FiniteElementState& u) override;
 
  std::shared_ptr<smith::FiniteElementState> solve(
      const smith::FiniteElementState& u_guess, std::function<mfem::Vector(const smith::FiniteElementState&)> equation,
      std::function<std::unique_ptr<mfem::HypreParMatrix>(const smith::FiniteElementState&)> jacobian) const override;
 
  std::shared_ptr<smith::FiniteElementState> solveAdjoint(
      const smith::FiniteElementDual& u_bar, std::unique_ptr<mfem::HypreParMatrix> jacobian_transposed) const override;
 
  void clearMemory() const override;
 
  mutable std::unique_ptr<mfem::HypreParMatrix> J_;  
  mutable std::unique_ptr<EquationSolver>
      nonlinear_solver_;  
};
 
class DifferentiableBlockSolver {
 public:
  virtual ~DifferentiableBlockSolver() {}
 
  using FieldT = FiniteElementState;                        
  using FieldPtr = std::shared_ptr<FieldT>;                 
  using FieldD = FiniteElementDual;                         
  using DualPtr = std::shared_ptr<FieldD>;                  
  using MatrixPtr = std::unique_ptr<mfem::HypreParMatrix>;  
 
  virtual void completeSetup(const std::vector<FieldT>& us) = 0;
 
  virtual std::vector<FieldPtr> solve(
      const std::vector<FieldPtr>& u_guesses,
      std::function<std::vector<mfem::Vector>(const std::vector<FieldPtr>&)> residuals,
      std::function<std::vector<std::vector<MatrixPtr>>(const std::vector<FieldPtr>&)> jacobians) const = 0;
 
  virtual std::vector<FieldPtr> solveAdjoint(const std::vector<DualPtr>& u_bars,
                                             std::vector<std::vector<MatrixPtr>>& jacobian_transposed) const = 0;
 
  virtual void clearMemory() const {}
};
 
class LinearDifferentiableBlockSolver : public DifferentiableBlockSolver {
 public:
  LinearDifferentiableBlockSolver(std::unique_ptr<mfem::Solver> s, std::unique_ptr<mfem::Solver> p);
 
  void completeSetup(const std::vector<FieldT>& us) override;
 
  std::vector<FieldPtr> solve(
      const std::vector<FieldPtr>& u_guesses,
      std::function<std::vector<mfem::Vector>(const std::vector<FieldPtr>&)> residuals,
      std::function<std::vector<std::vector<MatrixPtr>>(const std::vector<FieldPtr>&)> jacobians) const override;
 
  std::vector<FieldPtr> solveAdjoint(const std::vector<DualPtr>& u_bars,
                                     std::vector<std::vector<MatrixPtr>>& jacobian_transposed) const override;
 
  mutable std::unique_ptr<mfem::Solver> mfem_solver;          
  mutable std::unique_ptr<mfem::Solver> mfem_preconditioner;  
};
 
class NonlinearDifferentiableBlockSolver : public DifferentiableBlockSolver {
 public:
  NonlinearDifferentiableBlockSolver(std::unique_ptr<EquationSolver> s);
 
  void completeSetup(const std::vector<FieldT>& us) override;
 
  std::vector<FieldPtr> solve(
      const std::vector<FieldPtr>& u_guesses,
      std::function<std::vector<mfem::Vector>(const std::vector<FieldPtr>&)> residuals,
      std::function<std::vector<std::vector<MatrixPtr>>(const std::vector<FieldPtr>&)> jacobians) const override;
 
  std::vector<FieldPtr> solveAdjoint(const std::vector<DualPtr>& u_bars,
                                     std::vector<std::vector<MatrixPtr>>& jacobian_transposed) const override;
 
  mutable std::unique_ptr<mfem::BlockOperator>
      block_jac_;  
  mutable std::vector<std::vector<MatrixPtr>>
      matrix_of_jacs_;  
 
  mutable std::unique_ptr<EquationSolver>
      nonlinear_solver_;  
};
 
std::shared_ptr<LinearDifferentiableSolver> buildDifferentiableLinearSolver(LinearSolverOptions linear_opts,
                                                                            const smith::Mesh& mesh);
 
std::shared_ptr<NonlinearDifferentiableSolver> buildDifferentiableNonlinearSolver(NonlinearSolverOptions nonlinear_opts,
                                                                                  LinearSolverOptions linear_opts,
                                                                                  const smith::Mesh& mesh);
 
std::shared_ptr<NonlinearDifferentiableBlockSolver> buildDifferentiableNonlinearBlockSolver(
    NonlinearSolverOptions nonlinear_opts, LinearSolverOptions linear_opts, const smith::Mesh& mesh);
 
}  // namespace smith