petsc.h

// Copyright (C) 2004-2018 Johan Hoffman, Johan Jansson, Anders Logg and
// Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#ifdef HAS_PETSC
 
#include "Vector.h"
#include <cassert>
#include <cstdint>
#include <format>
#include <functional>
#include <optional>
#include <petscksp.h>
#include <petscmat.h>
#include <petscoptions.h>
#include <petscvec.h>
#include <span>
#include <string>
#include <string_view>
#include <vector>
 
namespace dolfinx::common
{
class IndexMap;
} // namespace dolfinx::common
 
namespace dolfinx::la
{
class SparsityPattern;
enum class Norm : std::int8_t;

namespace petsc
{
void error(PetscErrorCode error_code, std::string_view filename,
           std::string_view petsc_function);

std::vector<Vec>
create_vectors(MPI_Comm comm,
               const std::vector<std::span<const PetscScalar>>& x);

Vec create_vector(const common::IndexMap& map, int bs);

Vec create_vector(MPI_Comm comm, std::array<std::int64_t, 2> range,
                  std::span<const std::int64_t> ghosts, int bs);

Vec create_vector_wrap(const common::IndexMap& map, int bs,
                       std::span<const PetscScalar> x);

template <class V>
Vec create_vector_wrap(const la::Vector<V>& x)
{
  assert(x.index_map());
  return create_vector_wrap(*x.index_map(), x.bs(), x.array());
}

std::vector<IS> create_index_sets(
    const std::vector<
        std::pair<std::reference_wrapper<const common::IndexMap>, int>>& maps);

std::vector<std::vector<PetscScalar>> get_local_vectors(
    const Vec x,
    const std::vector<
        std::pair<std::reference_wrapper<const common::IndexMap>, int>>& maps);

void scatter_local_vectors(
    Vec x, const std::vector<std::span<const PetscScalar>>& x_b,
    const std::vector<
        std::pair<std::reference_wrapper<const common::IndexMap>, int>>& maps);

Mat create_matrix(MPI_Comm comm, const SparsityPattern& sp,
                  std::optional<std::string_view> type = std::nullopt);

MatNullSpace create_nullspace(MPI_Comm comm, std::span<const Vec> basis);

namespace options
{
void set(std::string option);

template <typename T>
  requires requires(const T& value) { std::format("{}", value); }
void set(std::string option, const T& value)
{
  if (option[0] != '-')
    option = '-' + option;
 
  PetscErrorCode ierr;
  ierr = PetscOptionsSetValue(nullptr, option.c_str(),
                              std::format("{}", value).c_str());
  if (ierr != 0)
    petsc::error(ierr, __FILE__, "PetscOptionsSetValue");
}

void clear(std::string option);

void clear();
} // namespace options

class Vector
{
public:
  Vector(const common::IndexMap& map, int bs);
 
  // Delete copy constructor to avoid accidental copying of 'heavy' data
  Vector(const Vector& x) = delete;

  Vector(Vector&& x) noexcept;

  Vector(Vec x, bool inc_ref_count);

  virtual ~Vector();
 
  // Assignment operator (disabled)
  Vector& operator=(const Vector& x) = delete;

  Vector& operator=(Vector&& x) noexcept;

  Vector copy() const;

  std::int64_t size() const;

  std::int32_t local_size() const;

  std::array<std::int64_t, 2> local_range() const;

  MPI_Comm comm() const;

  void set_options_prefix(std::string_view options_prefix);

  std::string get_options_prefix() const;

  void set_from_options();

  Vec vec() const;
 
private:
  // PETSc Vec pointer
  Vec _x;
};

class Operator
{
public:
  Operator(Mat A, bool inc_ref_count);
 
  // Copy constructor (deleted)
  Operator(const Operator& A) = delete;

  Operator(Operator&& A) noexcept;

  virtual ~Operator();

  Operator& operator=(const Operator& A) = delete;

  Operator& operator=(Operator&& A) noexcept;

  std::array<std::int64_t, 2> size() const;

  Vec create_vector(std::size_t dim) const;

  Mat mat() const;
 
protected:
  // PETSc Mat pointer
  Mat _matA;
};

class Matrix : public Operator
{
public:
  static auto set_fn(Mat A, InsertMode mode)
  {
    return [A, mode, cache = std::vector<PetscInt>()](
               std::span<const std::int32_t> rows,
               std::span<const std::int32_t> cols,
               std::span<const PetscScalar> vals) mutable -> int
    {
      PetscErrorCode ierr;
#ifdef PETSC_USE_64BIT_INDICES
      cache.resize(rows.size() + cols.size());
      std::ranges::copy(rows, cache.begin());
      std::ranges::copy(cols, std::next(cache.begin(), rows.size()));
      const PetscInt* _rows = cache.data();
      const PetscInt* _cols = cache.data() + rows.size();
      ierr = MatSetValuesLocal(A, rows.size(), _rows, cols.size(), _cols,
                               vals.data(), mode);
#else
      ierr = MatSetValuesLocal(A, rows.size(), rows.data(), cols.size(),
                               cols.data(), vals.data(), mode);
#endif
 
#ifndef NDEBUG
      if (ierr != 0)
        petsc::error(ierr, __FILE__, "MatSetValuesLocal");
#endif
      return ierr;
    };
  }

  static auto set_block_fn(Mat A, InsertMode mode)
  {
    return [A, mode, cache = std::vector<PetscInt>()](
               std::span<const std::int32_t> rows,
               std::span<const std::int32_t> cols,
               std::span<const PetscScalar> vals) mutable -> int
    {
      PetscErrorCode ierr;
#ifdef PETSC_USE_64BIT_INDICES
      cache.resize(rows.size() + cols.size());
      std::ranges::copy(rows, cache.begin());
      std::ranges::copy(cols, std::next(cache.begin(), rows.size()));
      const PetscInt* _rows = cache.data();
      const PetscInt* _cols = cache.data() + rows.size();
      ierr = MatSetValuesBlockedLocal(A, rows.size(), _rows, cols.size(), _cols,
                                      vals.data(), mode);
#else
      ierr = MatSetValuesBlockedLocal(A, rows.size(), rows.data(), cols.size(),
                                      cols.data(), vals.data(), mode);
#endif
 
#ifndef NDEBUG
      if (ierr != 0)
        petsc::error(ierr, __FILE__, "MatSetValuesBlockedLocal");
#endif
      return ierr;
    };
  }

  static auto set_block_expand_fn(Mat A, int bs0, int bs1, InsertMode mode)
  {
    return [A, bs0, bs1, mode, cache0 = std::vector<PetscInt>(),
            cache1 = std::vector<PetscInt>()](
               std::span<const std::int32_t> rows,
               std::span<const std::int32_t> cols,
               std::span<const PetscScalar> vals) mutable -> int
    {
      PetscErrorCode ierr;
      cache0.resize(bs0 * rows.size());
      cache1.resize(bs1 * cols.size());
      for (std::size_t i = 0; i < rows.size(); ++i)
        for (int k = 0; k < bs0; ++k)
          cache0[bs0 * i + k] = bs0 * rows[i] + k;
 
      for (std::size_t i = 0; i < cols.size(); ++i)
        for (int k = 0; k < bs1; ++k)
          cache1[bs1 * i + k] = bs1 * cols[i] + k;
 
      ierr = MatSetValuesLocal(A, cache0.size(), cache0.data(), cache1.size(),
                               cache1.data(), vals.data(), mode);
#ifndef NDEBUG
      if (ierr != 0)
        petsc::error(ierr, __FILE__, "MatSetValuesLocal");
#endif
      return ierr;
    };
  }

  Matrix(MPI_Comm comm, const SparsityPattern& sp,
         std::optional<std::string_view> type = std::nullopt);

  Matrix(Mat A, bool inc_ref_count);
 
  // Copy constructor (deleted)
  Matrix(const Matrix& A) = delete;

  Matrix(Matrix&& A) = default;

  ~Matrix() = default;

  Matrix& operator=(const Matrix& A) = delete;

  Matrix& operator=(Matrix&& A) = default;

  enum class AssemblyType : std::int8_t
  {
    FINAL,
    FLUSH
  };

  void apply(AssemblyType type);

  double norm(Norm norm_type) const;
 
  //--- Special PETSc Functions ---

  void set_options_prefix(std::string_view options_prefix);

  std::string get_options_prefix() const;

  void set_from_options();
};

class KrylovSolver
{
public:
  explicit KrylovSolver(MPI_Comm comm);

  KrylovSolver(KSP ksp, bool inc_ref_count);
 
  // Copy constructor (deleted)
  KrylovSolver(const KrylovSolver& solver) = delete;

  KrylovSolver(KrylovSolver&& solver) noexcept;

  ~KrylovSolver();
 
  // Assignment operator (deleted)
  KrylovSolver& operator=(const KrylovSolver&) = delete;

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

  void set_operator(const Mat A);

  void set_operators(const Mat A, const Mat P);

  int solve(Vec x, const Vec b, bool transpose = false) const;

  void set_options_prefix(std::string_view options_prefix);

  std::string get_options_prefix() const;

  void set_from_options() const;

  KSP ksp() const;
 
private:
  // PETSc solver pointer
  KSP _ksp;
};
} // namespace petsc
} // namespace dolfinx::la
 
#endif