Line data    Source code

       1           0 : // Distributed under the MIT License.
       2             : // See LICENSE.txt for details.
       3             : 
       4             : #pragma once
       5             : 
       6             : #include <algorithm>
       7             : #include <cstddef>
       8             : #include <fstream>
       9             : #include <string>
      10             : #include <tuple>
      11             : #include <vector>
      12             : 
      13             : #include "DataStructures/DataBox/DataBox.hpp"
      14             : #include "DataStructures/DynamicMatrix.hpp"
      15             : #include "DataStructures/DynamicVector.hpp"
      16             : #include "NumericalAlgorithms/Convergence/HasConverged.hpp"
      17             : #include "NumericalAlgorithms/LinearSolver/BuildMatrix.hpp"
      18             : #include "NumericalAlgorithms/LinearSolver/LinearSolver.hpp"
      19             : #include "Options/Auto.hpp"
      20             : #include "Options/String.hpp"
      21             : #include "Parallel/Tags/ArrayIndex.hpp"
      22             : #include "Utilities/ErrorHandling/Error.hpp"
      23             : #include "Utilities/Gsl.hpp"
      24             : #include "Utilities/MakeWithValue.hpp"
      25             : #include "Utilities/NoSuchType.hpp"
      26             : #include "Utilities/Serialization/CharmPupable.hpp"
      27             : #include "Utilities/Serialization/PupStlCpp17.hpp"
      28             : #include "Utilities/TMPL.hpp"
      29             : 
      30             : namespace LinearSolver::Serial {
      31             : 
      32             : /// \cond
      33             : template <typename LinearSolverRegistrars>
      34             : struct ExplicitInverse;
      35             : /// \endcond
      36             : 
      37           1 : namespace Registrars {
      38             : /// Registers the `LinearSolver::Serial::ExplicitInverse` linear solver
      39           1 : using ExplicitInverse = Registration::Registrar<Serial::ExplicitInverse>;
      40             : }  // namespace Registrars
      41             : 
      42             : /*!
      43             :  * \brief Linear solver that builds a matrix representation of the linear
      44             :  * operator and inverts it directly
      45             :  *
      46             :  * This solver first constructs an explicit matrix representation by "sniffing
      47             :  * out" the operator, i.e. feeding it with unit vectors, and then directly
      48             :  * inverts the matrix. The result is an operator that solves the linear problem
      49             :  * in a single step. This means that each element has a large initialization
      50             :  * cost, but all successive solves converge immediately.
      51             :  *
      52             :  * \par Advice on using this linear solver:
      53             :  *
      54             :  * - This solver is entirely agnostic to the structure of the linear operator.
      55             :  *   It is usually better to implement a linear solver that is specialized for
      56             :  *   your linear operator to take advantage of its properties. For example, if
      57             :  *   the operator has a tensor-product structure, the linear solver might take
      58             :  *   advantage of that. Only use this solver if no alternatives are available
      59             :  *   and if you have verified that it speeds up your solves.
      60             :  * - Since this linear solver stores the full inverse operator matrix it can
      61             :  *   have significant memory demands. For example, an operator representing a 3D
      62             :  *   first-order Elasticity system (9 variables) discretized on 12 grid points
      63             :  *   per dimension requires ca. 2GB of memory (per element) to store the matrix,
      64             :  *   scaling quadratically with the number of variables and with a power of 6
      65             :  *   with the number of grid points per dimension. Therefore, make sure to
      66             :  *   distribute the elements on a sufficient number of nodes to meet the memory
      67             :  *   requirements.
      68             :  * - This linear solver can be `reset()` when the operator changes (e.g. in each
      69             :  *   nonlinear-solver iteration). However, when using this solver as
      70             :  *   preconditioner it can be advantageous to avoid the reset and the
      71             :  *   corresponding cost of re-building the matrix and its inverse if the
      72             :  *   operator only changes "a little". In that case the preconditioner solves
      73             :  *   subdomain problems only approximately, but possibly still sufficiently to
      74             :  *   provide effective preconditioning.
      75             :  */
      76             : template <typename LinearSolverRegistrars =
      77             :               tmpl::list<Registrars::ExplicitInverse>>
      78           1 : class ExplicitInverse : public LinearSolver<LinearSolverRegistrars> {
      79             :  private:
      80           0 :   using Base = LinearSolver<LinearSolverRegistrars>;
      81             : 
      82             :  public:
      83           0 :   struct WriteMatrixToFile {
      84           0 :     using type = Options::Auto<std::string, Options::AutoLabel::None>;
      85           0 :     static constexpr Options::String help =
      86             :         "Write the matrix representation of the linear operator to a "
      87             :         "space-delimited CSV file with this name. A '.txt' extension will be "
      88             :         "added. Also a suffix with the element ID will be added if this linear "
      89             :         "solver runs on an array element, so one file per element will be "
      90             :         "written.";
      91             :   };
      92             : 
      93           0 :   using options = tmpl::list<WriteMatrixToFile>;
      94           0 :   static constexpr Options::String help =
      95             :       "Build a matrix representation of the linear operator and invert it "
      96             :       "directly. This means that the first solve has a large initialization "
      97             :       "cost, but all subsequent solves converge immediately.";
      98             : 
      99           0 :   ExplicitInverse(const ExplicitInverse& /*rhs*/) = default;
     100           0 :   ExplicitInverse& operator=(const ExplicitInverse& /*rhs*/) = default;
     101           0 :   ExplicitInverse(ExplicitInverse&& /*rhs*/) = default;
     102           0 :   ExplicitInverse& operator=(ExplicitInverse&& /*rhs*/) = default;
     103           0 :   ~ExplicitInverse() = default;
     104             : 
     105           0 :   explicit ExplicitInverse(
     106             :       std::optional<std::string> matrix_filename = std::nullopt)
     107             :       : matrix_filename_(std::move(matrix_filename)) {}
     108             : 
     109             :   /// \cond
     110             :   explicit ExplicitInverse(CkMigrateMessage* m) : Base(m) {}
     111             :   using PUP::able::register_constructor;
     112             :   WRAPPED_PUPable_decl_template(ExplicitInverse);  // NOLINT
     113             :   /// \endcond
     114             : 
     115             :   /*!
     116             :    * \brief Solve the equation \f$Ax=b\f$ by explicitly constructing the
     117             :    * operator matrix \f$A\f$ and its inverse. The first solve is computationally
     118             :    * expensive and successive solves are cheap.
     119             :    *
     120             :    * Building a matrix representation of the `linear_operator` requires
     121             :    * iterating over the `SourceType` (which is also the type returned by the
     122             :    * `linear_operator`) in a consistent way. This can be non-trivial for
     123             :    * heterogeneous data structures because it requires they define a data
     124             :    * ordering. Specifically, the `SourceType` must have a `size()` function as
     125             :    * well as `begin()` and `end()` iterators that point into the data. If the
     126             :    * iterators have a `reset()` function it is used to avoid repeatedly
     127             :    * re-creating the `begin()` iterator. The `reset()` function must not
     128             :    * invalidate the `end()` iterator.
     129             :    */
     130             :   template <typename LinearOperator, typename VarsType, typename SourceType,
     131             :             typename... OperatorArgs>
     132           1 :   Convergence::HasConverged solve(
     133             :       gsl::not_null<VarsType*> solution, const LinearOperator& linear_operator,
     134             :       const SourceType& source,
     135             :       const std::tuple<OperatorArgs...>& operator_args = std::tuple{}) const;
     136             : 
     137             :   /// Flags the operator to require re-initialization. No memory is released.
     138             :   /// Call this function to rebuild the solver when the operator changed.
     139           1 :   void reset() override { size_ = std::numeric_limits<size_t>::max(); }
     140             : 
     141             :   /// Size of the operator. The stored matrix will have `size^2` entries.
     142           1 :   size_t size() const { return size_; }
     143             : 
     144             :   /// The matrix representation of the solver. This matrix approximates the
     145             :   /// inverse of the subdomain operator.
     146             :   const blaze::DynamicMatrix<double, blaze::columnMajor>&
     147           1 :   matrix_representation() const {
     148             :     return inverse_;
     149             :   }
     150             : 
     151             :   // NOLINTNEXTLINE(google-runtime-references)
     152           0 :   void pup(PUP::er& p) override {
     153             :     p | matrix_filename_;
     154             :     p | size_;
     155             :     p | inverse_;
     156             :     if (p.isUnpacking() and size_ != std::numeric_limits<size_t>::max()) {
     157             :       source_workspace_.resize(size_);
     158             :       solution_workspace_.resize(size_);
     159             :     }
     160             :   }
     161             : 
     162           0 :   std::unique_ptr<Base> get_clone() const override {
     163             :     return std::make_unique<ExplicitInverse>(*this);
     164             :   }
     165             : 
     166             :  private:
     167           0 :   std::optional<std::string> matrix_filename_{};
     168             :   // Caches for successive solves of the same operator
     169             :   // NOLINTNEXTLINE(spectre-mutable)
     170           0 :   mutable size_t size_ = std::numeric_limits<size_t>::max();
     171             :   // We currently store the matrix representation in a dense matrix because
     172             :   // Blaze doesn't support the inversion of sparse matrices (yet).
     173             :   // NOLINTNEXTLINE(spectre-mutable)
     174           0 :   mutable blaze::DynamicMatrix<double, blaze::columnMajor> inverse_{};
     175             : 
     176             :   // Buffers to avoid re-allocating memory for applying the operator
     177             :   // NOLINTNEXTLINE(spectre-mutable)
     178           0 :   mutable blaze::DynamicVector<double> source_workspace_{};
     179             :   // NOLINTNEXTLINE(spectre-mutable)
     180           0 :   mutable blaze::DynamicVector<double> solution_workspace_{};
     181             : };
     182             : 
     183             : template <typename LinearSolverRegistrars>
     184             : template <typename LinearOperator, typename VarsType, typename SourceType,
     185             :           typename... OperatorArgs>
     186           0 : Convergence::HasConverged ExplicitInverse<LinearSolverRegistrars>::solve(
     187             :     const gsl::not_null<VarsType*> solution,
     188             :     const LinearOperator& linear_operator, const SourceType& source,
     189             :     const std::tuple<OperatorArgs...>& operator_args) const {
     190             :   if (UNLIKELY(size_ == std::numeric_limits<size_t>::max())) {
     191             :     const auto& used_for_size = source;
     192             :     size_ = used_for_size.size();
     193             :     source_workspace_.resize(size_);
     194             :     solution_workspace_.resize(size_);
     195             :     inverse_.resize(size_, size_);
     196             :     // Construct explicit matrix representation by "sniffing out" the operator,
     197             :     // i.e. feeding it unit vectors
     198             :     auto operand_buffer = make_with_value<VarsType>(used_for_size, 0.);
     199             :     auto result_buffer = make_with_value<SourceType>(used_for_size, 0.);
     200             :     build_matrix(make_not_null(&inverse_), make_not_null(&operand_buffer),
     201             :                  make_not_null(&result_buffer), linear_operator, operator_args);
     202             :     // Write to file before inverting
     203             :     if (UNLIKELY(matrix_filename_.has_value())) {
     204             :       const auto filename_suffix =
     205             :           [&operator_args]() -> std::optional<std::string> {
     206             :         using DataBoxType =
     207             :             std::decay_t<tmpl::front<tmpl::list<OperatorArgs..., NoSuchType>>>;
     208             :         if constexpr (tt::is_a_v<db::DataBox, DataBoxType>) {
     209             :           if constexpr (db::tag_is_retrievable_v<Parallel::Tags::ArrayIndex,
     210             :                                                  DataBoxType>) {
     211             :             const auto& box = std::get<0>(operator_args);
     212             :             return "_" + get_output(db::get<Parallel::Tags::ArrayIndex>(box));
     213             :           } else {
     214             :             (void)operator_args;
     215             :             return std::nullopt;
     216             :           }
     217             :         } else {
     218             :           (void)operator_args;
     219             :           return std::nullopt;
     220             :         }
     221             :       }();
     222             :       std::ofstream matrix_file(matrix_filename_.value() +
     223             :                                 filename_suffix.value_or("") + ".txt");
     224             :       write_csv(matrix_file, inverse_, " ");
     225             :     }
     226             :     // Directly invert the matrix
     227             :     try {
     228             :       blaze::invert(inverse_);
     229             :     } catch (const std::invalid_argument& e) {
     230             :       ERROR("Could not invert subdomain matrix (size " << size_
     231             :                                                        << "): " << e.what());
     232             :     }
     233             :   }
     234             :   // Copy source into contiguous workspace. In cases where the source and
     235             :   // solution data are already stored contiguously we might avoid the copy and
     236             :   // the associated workspace memory. However, compared to the cost of building
     237             :   // and storing the matrix this is likely insignificant.
     238             :   std::copy(source.begin(), source.end(), source_workspace_.begin());
     239             :   // Apply inverse
     240             :   solution_workspace_ = inverse_ * source_workspace_;
     241             :   // Reconstruct solution data from contiguous workspace
     242             :   std::copy(solution_workspace_.begin(), solution_workspace_.end(),
     243             :             solution->begin());
     244             :   return {0, 0};
     245             : }
     246             : 
     247             : /// \cond
     248             : template <typename LinearSolverRegistrars>
     249             : // NOLINTNEXTLINE
     250             : PUP::able::PUP_ID ExplicitInverse<LinearSolverRegistrars>::my_PUP_ID = 0;
     251             : /// \endcond
     252             : 
     253             : }  // namespace LinearSolver::Serial