utils.h

// Copyright (C) 2018-2022 Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <dolfinx/common/types.h>
#include <format>
#include <ranges>
#include <span>
#include <stdexcept>
 
namespace dolfinx::la
{
enum class Norm : std::int8_t
{
  l1,
  l2,
  linf,
  frobenius
};

template <class U, class T>
concept MatSet
    = std::invocable<U, std::span<const std::int32_t>,
                     std::span<const std::int32_t>, std::span<const T>>;
template <dolfinx::scalar T>
void set_diagonal(MatSet<T> auto&& set_fn,
                  const common::LocalIndexRange auto& rows, T diagonal = T(1))
{
  std::span<const T, 1> diag_span(&diagonal, 1);
  for (std::size_t i = 0; i < std::ranges::size(rows); ++i)
  {
    std::span<const std::int32_t, 1> row(std::ranges::data(rows) + i, 1);
    set_fn(row, row, diag_span);
  }
}

template <dolfinx::scalar T>
void set_diagonal(MatSet<T> auto&& set_fn,
                  const common::LocalIndexRange auto& rows,
                  std::span<const T> diagonals)
{
  if (diagonals.size() != std::ranges::size(rows))
  {
    throw std::invalid_argument(
        std::format("Number of diagonal values ({}) does not match number "
                    "of rows ({}).",
                    diagonals.size(), std::ranges::size(rows)));
  }
 
  for (std::size_t i = 0; i < diagonals.size(); ++i)
  {
    set_diagonal(
        set_fn,
        std::span<const std::int32_t, 1>(std::ranges::data(rows) + i, 1),
        diagonals[i]);
  }
}
} // namespace dolfinx::la