Line data    Source code

       1              : // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
       2              : // SPDX-License-Identifier: Apache-2.0
       3              : 
       4              : #include "jacobi.hpp"
       5              : 
       6              : #include <mfem/general/forall.hpp>
       7              : 
       8              : namespace palace
       9              : {
      10              : 
      11              : namespace
      12              : {
      13              : 
      14            0 : double GetLambdaMax(MPI_Comm comm, const Operator &A, const Vector &dinv)
      15              : {
      16              :   // Assumes A SPD (diag(A) > 0) to use Hermitian eigenvalue solver.
      17              :   DiagonalOperator Dinv(dinv);
      18            0 :   ProductOperator DinvA(Dinv, A);
      19            0 :   return linalg::SpectralNorm(comm, DinvA, true);
      20              : }
      21              : 
      22            0 : double GetLambdaMax(MPI_Comm comm, const ComplexOperator &A, const ComplexVector &dinv)
      23              : {
      24              :   // Assumes A SPD (diag(A) > 0) to use Hermitian eigenvalue solver.
      25              :   ComplexDiagonalOperator Dinv(dinv);
      26            0 :   ComplexProductOperator DinvA(Dinv, A);
      27            0 :   return linalg::SpectralNorm(comm, DinvA, A.IsReal());
      28              : }
      29              : 
      30              : template <bool Transpose = false>
      31            0 : inline void Apply(const Vector &dinv, const Vector &x, Vector &y)
      32              : {
      33            0 :   const bool use_dev = dinv.UseDevice() || x.UseDevice() || y.UseDevice();
      34              :   const int N = dinv.Size();
      35            0 :   const auto *DI = dinv.Read(use_dev);
      36            0 :   const auto *X = x.Read(use_dev);
      37            0 :   auto *Y = y.Write(use_dev);
      38            0 :   mfem::forall_switch(use_dev, N, [=] MFEM_HOST_DEVICE(int i) { Y[i] = DI[i] * X[i]; });
      39            0 : }
      40              : 
      41              : template <bool Transpose = false>
      42            0 : inline void Apply(const ComplexVector &dinv, const ComplexVector &x, ComplexVector &y)
      43              : {
      44            0 :   const bool use_dev = dinv.UseDevice() || x.UseDevice() || y.UseDevice();
      45              :   const int N = dinv.Size();
      46            0 :   const auto *DIR = dinv.Real().Read(use_dev);
      47            0 :   const auto *DII = dinv.Imag().Read(use_dev);
      48            0 :   const auto *XR = x.Real().Read(use_dev);
      49            0 :   const auto *XI = x.Imag().Read(use_dev);
      50            0 :   auto *YR = y.Real().Write(use_dev);
      51            0 :   auto *YI = y.Imag().Write(use_dev);
      52              :   if constexpr (!Transpose)
      53              :   {
      54              :     mfem::forall_switch(use_dev, N,
      55            0 :                         [=] MFEM_HOST_DEVICE(int i)
      56              :                         {
      57            0 :                           YR[i] = DIR[i] * XR[i] - DII[i] * XI[i];
      58            0 :                           YI[i] = DII[i] * XR[i] + DIR[i] * XI[i];
      59              :                         });
      60              :   }
      61              :   else
      62              :   {
      63              :     mfem::forall_switch(use_dev, N,
      64              :                         [=] MFEM_HOST_DEVICE(int i)
      65              :                         {
      66              :                           YR[i] = DIR[i] * XR[i] + DII[i] * XI[i];
      67              :                           YI[i] = -DII[i] * XR[i] + DIR[i] * XI[i];
      68              :                         });
      69              :   }
      70            0 : }
      71              : 
      72              : }  // namespace
      73              : 
      74              : template <typename OperType>
      75            0 : void JacobiSmoother<OperType>::SetOperator(const OperType &op)
      76              : {
      77            0 :   dinv.SetSize(op.Height());
      78            0 :   dinv.UseDevice(true);
      79            0 :   op.AssembleDiagonal(dinv);
      80            0 :   dinv.Reciprocal();
      81              : 
      82              :   // Damping factor. If the given damping is zero, estimate the spectral radius-minimizing
      83              :   // damping factor.
      84            0 :   if (omega == 0.0)
      85              :   {
      86            0 :     auto lambda_max = GetLambdaMax(comm, op, dinv);
      87            0 :     auto lambda_min = (sf_max - 1.0) * lambda_max;
      88            0 :     omega = 2.0 / (lambda_min + lambda_max);
      89              :   }
      90            0 :   if (omega != 1.0)
      91              :   {
      92            0 :     dinv *= omega;
      93              :   }
      94              : 
      95            0 :   this->height = op.Height();
      96            0 :   this->width = op.Width();
      97            0 : }
      98              : 
      99              : template <typename OperType>
     100            0 : void JacobiSmoother<OperType>::Mult(const VecType &x, VecType &y) const
     101              : {
     102              :   MFEM_ASSERT(!this->initial_guess, "JacobiSmoother does not use initial guess!");
     103            0 :   Apply(dinv, x, y);
     104            0 : }
     105              : 
     106              : template class JacobiSmoother<Operator>;
     107              : template class JacobiSmoother<ComplexOperator>;
     108              : 
     109              : }  // namespace palace