SNESSolver.h

// Copyright (C) 2026 Jack S. Hale
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#ifdef HAS_PETSC
 
#include <exception>
#include <functional>
#include <mpi.h>
#include <petscmat.h>
#include <petscsnes.h>
#include <petscsystypes.h>
#include <petscvec.h>
#include <string>
#include <string_view>
 
namespace dolfinx::nls::petsc
{
class SNESSolver
{
public:
  explicit SNESSolver(MPI_Comm comm);

  SNESSolver(SNES snes, bool inc_ref_count);
 
  // Copy constructor (deleted)
  SNESSolver(const SNESSolver& solver) = delete;

  SNESSolver(SNESSolver&& solver) noexcept;

  ~SNESSolver();
 
  // Copy assignment (deleted)
  SNESSolver& operator=(const SNESSolver& solver) = delete;

  SNESSolver& operator=(SNESSolver&& solver) noexcept;

  void set_F(std::function<void(const Vec x, Vec b)> F, Vec b_layout);

  void set_J(std::function<void(const Vec x, Mat Jmat, Mat Pmat)> J,
             Mat J_layout, Mat P_layout = nullptr);

  void set_update(std::function<void(PetscInt step)> update);

  [[nodiscard(
      "check the converged reason - positive on convergence, negative on "
      "divergence")]] SNESConvergedReason
  solve(Vec x);

  void set_options_prefix(std::string_view options_prefix);

  std::string get_options_prefix() const;

  void set_from_options() const;

  SNES snes() const;
 
private:
  // Callbacks passed to SNESSetFunction/SNESSetJacobian. The context
  // pointer is the owning SNESSolver.
  static PetscErrorCode residual(SNES snes, Vec x, Vec b, void* ctx);
  static PetscErrorCode jacobian(SNES snes, Vec x, Mat Jmat, Mat Pmat,
                                 void* ctx);
 
  // Callback passed to SNESSetUpdate. It takes no context argument, so
  // the owning SNESSolver is recovered from the context that set_F
  // registered on the SNES object.
  static PetscErrorCode update_step(SNES snes, PetscInt step);
 
  // Register the callbacks that have been set, and attach this as the
  // context that the callbacks recover the solver from
  void set_callbacks();
 
  // Function for computing the residual vector
  std::function<void(const Vec x, Vec b)> _fnF;
 
  // Function for computing the Jacobian and preconditioner matrices
  std::function<void(const Vec x, Mat Jmat, Mat Pmat)> _fnJ;
 
  // Function called before each nonlinear iteration
  std::function<void(PetscInt step)> _fnupdate;
 
  // Exception thrown by a callback during a solve, re-thrown by solve
  std::exception_ptr _exception;
 
  // Residual vector
  Vec _b = nullptr;
 
  // Jacobian matrix
  Mat _matJ = nullptr;
 
  // Preconditioner matrix
  Mat _matP = nullptr;
 
  // PETSc solver pointer
  SNES _snes;
};
} // namespace dolfinx::nls::petsc
 
#endif